[论文解读] All are Worth Words: A ViT Backbone for Diffusion Models

Fan Bao, Nie, Shen+5 more|arXiv (Cornell University)|Sep 25, 2022

Generative Adversarial Networks and Image Synthesis被引用 7

一句话总结

本文提出U-ViT，一种基于视觉变换器（ViT）的扩散模型主干网络，将时间、条件和噪声图像块视为标记（tokens），并在浅层与深层之间采用长跳跃连接。U-ViT在ImageNet 256×256类别条件生成任务上实现了2.29的SOTA FID分数，在MS-COCO文本到图像生成任务上实现了5.48的SOTA FID分数，表明采用跳跃连接的ViT主干网络可超越无需下采样/上采样操作的CNN-based U-Net，实现更优性能。

ABSTRACT

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.

研究动机与目标

• 探究视觉变换器（ViT）是否能有效替代扩散模型中基于CNN的U-Net主干网络。
• 评估U-Net架构中下采样和上采样操作在基于扩散的图像生成任务中的必要性。
• 确定长跳跃连接在基于ViT的扩散模型主干网络中的影响。
• 通过一种简单的ViT主干架构，在无条件、类别条件和文本到图像生成任务中实现SOTA性能。

提出的方法

• U-ViT将所有输入——噪声图像块、时间步嵌入和条件嵌入——视为标记，实现统一的基于标记的输入表示。
• 该模型在浅层与深层之间采用长跳跃连接，以保留对扩散模型像素级重建至关重要的低级特征。
• 采用标准的ViT变换器编码器结构，包含分类标记和位置嵌入，随后接一个带有学习位置编码的解码器头。
• 在最终输出前可选地添加一个3×3卷积块，以提升视觉质量。
• 该架构被应用于像素空间和潜在空间的扩散模型中，训练采用噪声预测目标。
• 在多个基准上进行评估：CIFAR-10、ImageNet 256×256、ImageNet 512×512和MS-COCO，使用FID及其他指标。

实验结果

研究问题

• RQ1 在不损失性能的前提下，视觉变换器主干是否能替代扩散模型中的基于CNN的U-Net？
• RQ2 U-Net中的下采样和上采样机制是否对扩散模型中高质量图像生成至关重要？
• RQ3 与标准ViT架构相比，长跳跃连接在基于ViT的扩散模型主干网络中有多关键？
• RQ4 基于ViT的架构是否能在不依赖大规模外部数据集的前提下，实现在文本到图像和类别条件图像生成任务中的SOTA性能？

主要发现

• U-ViT在类别条件生成的ImageNet 256×256任务上实现了2.29的SOTA FID分数，优于所有先前方法，且未使用大规模外部数据集。
• 在MS-COCO文本到图像生成任务中，U-ViT-S（Deep）实现了5.48的新SOTA FID分数，表现出更优的语义对齐和图像质量。
• 在无条件和条件图像生成任务中，U-ViT在参数量和FLOPs与U-Net相近的情况下，性能均优于基于CNN的U-Net。
• 消融研究证实，长跳跃连接对性能至关重要，而U-Net中的下采样/上采样操作并非高质量生成的必要条件。
• 与U-Net相比，U-ViT生成的样本在语义上更准确，尤其在文本到图像生成中，能更准确地捕捉提示词中的物体（如蝙蝠和球）。
• 该模型在不同采样步数和指标（包括Inception Score、精确率和召回率）下均保持强大性能，表明其生成过程稳健且质量高。