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 class-conditional 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 record-breaking FID scores of 2.29 in class-conditional image generation on ImageNet 256x256, and 5.48 in text-to-image generation on MS-COCO, among methods without accessing large external datasets during the training of generative models. Our results suggest that, for diffusion-based image modeling, the long skip connection is crucial while the down-sampling and up-sampling operators in CNN-based U-Net are not always necessary. We believe that U-ViT can provide insights for future research on backbones in diffusion models and benefit generative modeling on large scale cross-modality datasets.

翻译：视觉Transformer（ViT）在多种视觉任务中展现出潜力，但基于卷积神经网络（CNN）的U-Net在扩散模型中仍占主导地位。我们设计了一个简单且通用的基于ViT的架构（命名为U-ViT），用于扩散模型的图像生成。U-ViT的特点是将所有输入（包括时间步、条件和噪声图像块）视为词元，并在浅层和深层之间采用长跳跃连接。我们在无条件图像生成、类别条件图像生成以及文本到图像生成任务中评估了U-ViT，结果表明U-ViT与同等规模的基于CNN的U-Net性能相当甚至更优。特别地，在生成模型训练期间不访问大型外部数据集的方法中，采用U-ViT的潜在扩散模型在ImageNet 256x256上的类别条件图像生成任务中实现了破纪录的FID分数2.29，在MS-COCO上的文本到图像生成任务中实现了5.48。我们的结果表明，对于基于扩散的图像建模，长跳跃连接至关重要，而基于CNN的U-Net中的下采样和上采样算子并非始终必要。我们相信U-ViT能为扩散模型骨干网络的未来研究提供见解，并有益于大规模跨模态数据集的生成建模。