[论文解读] Exploring Vision Transformers as Diffusion Learners

He Cao, Jianan Wang+5 more|arXiv (Cornell University)|Dec 28, 2022

Generative Adversarial Networks and Image Synthesis被引用 7

一句话总结

本文提出 ASCEND，一种基于视觉 Transformer 的扩散模型主干网络，采用非对称编码器-解码器架构，在图像生成任务中显著提升性能。通过用强大的 Swin Transformer 编码器和轻量级卷积解码器替代 U-Net，ASCEND 在 CIFAR-10、CelebA、LSUN 和 CUB-Bird 上取得具有竞争力的结果，并首次实现单阶段、端到端的文本到图像扩散模型训练，分辨率超过 64×64，为端到端高分辨率生成树立了新基准。

ABSTRACT

Score-based diffusion models have captured widespread attention and funded fast progress of recent vision generative tasks. In this paper, we focus on diffusion model backbone which has been much neglected before. We systematically explore vision Transformers as diffusion learners for various generative tasks. With our improvements the performance of vanilla ViT-based backbone (IU-ViT) is boosted to be on par with traditional U-Net-based methods. We further provide a hypothesis on the implication of disentangling the generative backbone as an encoder-decoder structure and show proof-of-concept experiments verifying the effectiveness of a stronger encoder for generative tasks with ASymmetriC ENcoder Decoder (ASCEND). Our improvements achieve competitive results on CIFAR-10, CelebA, LSUN, CUB Bird and large-resolution text-to-image tasks. To the best of our knowledge, we are the first to successfully train a single diffusion model on text-to-image task beyond 64x64 resolution. We hope this will motivate people to rethink the modeling choices and the training pipelines for diffusion-based generative models.

研究动机与目标

• 系统评估视觉 Transformer 作为扩散模型主干网络的性能，尤其与 U-Net 进行对比。
• 通过改进 U-ViT 架构（得到 IU-ViT），缩小原始 ViT 与 U-Net 在低分辨率图像生成中的性能差距。
• 提出并验证一种非对称编码器-解码器架构（ASCEND），用于可扩展的、高分辨率的、多模态的扩散生成。
• 通过展示基于视觉 Transformer 主干网络的端到端、单阶段高分辨率文本到图像生成，挑战多阶段训练流水线的现状。
• 激励研究社区重新思考扩散生成建模中的主干网络设计与训练范式。

提出的方法

• 提出 IU-ViT，即对 U-ViT 的改进版本，通过增强跳跃连接和残差下采样/上采样机制，替代原有的补丁合并/展开操作。
• 设计 ASCEND，一种采用 Swin Transformer 编码器和卷积解码器的非对称编码器-解码器主干网络，用于分层特征重建。
• 采用残差下采样和上采样机制，而非补丁合并/展开，以改善特征层次结构和梯度流动。
• 集成预训练的 T5-XXL 编码器，用于文本条件化在文本到图像生成任务中的应用。
• 在 LAION-Aesthetics 等大规模数据集上，对扩散模型进行端到端训练，实现 128×128 和 256×256 分辨率的图像生成。
• 在 CIFAR-10 上进行消融实验，验证架构选择，包括跳跃连接、下采样方法以及编码器/解码器配置。

实验结果

研究问题

• RQ1 视觉 Transformer 能否有效用作扩散模型主干网络？其在低分辨率和高分辨率图像生成中与 U-Net 的表现如何比较？
• RQ2 为使基于原始 ViT 的主干网络在扩散建模中与 U-Net 竞争，需要哪些架构改进？
• RQ3 非对称编码器-解码器结构是否能提升扩散图像生成中的性能与可扩展性？
• RQ4 基于视觉 Transformer 主干网络的单阶段、端到端扩散模型能否在 64×64 以上分辨率生成高质量图像？
• RQ5 编码器架构的选择（如 Swin Transformer）如何影响高分辨率和文本到图像生成中的性能表现？

主要发现

• ASCEND 在 CIFAR-10 上实现了 2.98 的 Fréchet Inception Distance（FID），与经过精细调优的 U-Net 主干网络性能相当。
• 在 ASCEND 中，将补丁合并/展开替换为残差下采样/上采样，使 FID 从 12.81 降低至 2.98，证明了架构改进的显著成效。
• 减少 ASCEND 中的跳跃连接会使 FID 上升 3.54 个点，证实其在特征重建中的关键作用。
• ASCEND 使用参数量为 590M 的模型，成功生成了高质量的 128×128 文本到图像样本，标志着首次在 64×64 以上分辨率实现单阶段训练的成功。
• ASCEND 在高分辨率生成中优于 IU-ViT 和 U-Net，展现出更快的收敛速度和在训练后期更优的样本质量。
• 消融实验确认，使用 SwinBlocks 作为编码器、卷积解码器的配置，相比对称或其它替代配置，能取得更优结果。