All are Worth Words: A ViT Backbone for Diffusion Models

Vision transformers (ViT) have shown promise in various vision tasks while the U-Net based on a convolutional neural network (CNN) remains dominant in diffusion models. We design a simple and general ViT-based architecture (named U-ViT) for image generation with diffusion models. U-ViT is characterized by treating all inputs including the time, condition and noisy image patches as tokens and employing long skip connections between shallow and deep layers. We evaluate U-ViT in unconditional and classconditional image generation, as well as text-to-image generation tasks, where U-ViT is comparable if not superior to a CNN-based U-Net of a similar size. In particular, latent diffusion models with U-ViT achieve…

• NN
  National Natural Science Foundation of China

  Award: 62061136001,61620106010,62076145,U19B2034,U1811461,U19A2081,6197222
• TU
  Tsinghua University
• BN
  Beijing Nova Program
• FR
  Fundamental Research Funds for the Central Universities