[论文解读] DiffiT: Diffusion Vision Transformers for Image Generation

Ali Hatamizadeh, Jiaming Song+3 more|arXiv (Cornell University)|Dec 4, 2023

Generative Adversarial Networks and Image Synthesis被引用 7

一句话总结

本文提出 DiffiT，一种基于视觉 Transformer 的扩散模型，用分层的 Transformer 编码器-解码器替代传统的 U-Net 架构。它提出了一种时间依赖的自注意力（TMSA）机制，可动态适应不同去噪阶段，实现在潜在空间中 ImageNet-256 的 SOTA FID 得分为 1.73，ImageNet-512 的得分为 2.67。

ABSTRACT

Diffusion models with their powerful expressivity and high sample quality have achieved State-Of-The-Art (SOTA) performance in the generative domain. The pioneering Vision Transformer (ViT) has also demonstrated strong modeling capabilities and scalability, especially for recognition tasks. In this paper, we study the effectiveness of ViTs in diffusion-based generative learning and propose a new model denoted as Diffusion Vision Transformers (DiffiT). Specifically, we propose a methodology for finegrained control of the denoising process and introduce the Time-dependant Multihead Self Attention (TMSA) mechanism. DiffiT is surprisingly effective in generating high-fidelity images with significantly better parameter efficiency. We also propose latent and image space DiffiT models and show SOTA performance on a variety of class-conditional and unconditional synthesis tasks at different resolutions. The Latent DiffiT model achieves a new SOTA FID score of 1.73 on ImageNet256 dataset while having 19.85%, 16.88% less parameters than other Transformer-based diffusion models such as MDT and DiT,respectively. Code: https://github.com/NVlabs/DiffiT

研究动机与目标

• 探究视觉 Transformer 在基于扩散的图像生成中的有效性，超越卷积 U-Net 架构。
• 解决现有扩散模型中缺乏细粒度时间建模的问题，即注意力机制无法适应不同去噪阶段。
• 通过引入分层的 Transformer 架构，统一去噪网络的设计模式，实现图像空间与潜在空间生成的一体化。
• 在多个基准测试中，通过 FID 得分衡量，实现图像生成质量的 SOTA 性能。

提出的方法

• 提出一种新颖的时间依赖自注意力（TMSA）模块，将时间步信息直接整合到自注意力的查询、键和值投影中，实现在不同去噪阶段的动态适应。
• 设计一种分层的、U 形的 Transformer 编码器-解码器架构（DiffiT），保持特征层次结构，并在多尺度上实现有效的特征优化。
• 引入潜在 DiffiT，一种高分辨率图像生成模型，将 TMSA 模块应用于潜在空间，实现高效且高保真度的图像合成。
• 通过可学习投影将时间标记注入自注意力机制，使模型能够根据当前去噪步骤调节注意力行为。
• 结合相对位置偏置与时间标记，在保留空间归纳偏差的同时实现时间适应性。
• 避免使用自适应层归一化（AdaLN）进行条件控制；相反，将时间依赖性直接嵌入注意力机制中。

Figure 2 : Overview of the image-space DiffiT model. Downsample and Upsample denote convolutional downsampling and upsampling layers, respectively. Please see the supplementary materials for more information regarding the DiffiT architecture.

实验结果

研究问题

• RQ1 视觉 Transformer 是否能有效替代扩散模型中的 U-Net 架构，同时保持或提升生成样本质量？
• RQ2 如何使扩散模型中的自注意力机制适应不同的去噪时间步，以更好地捕捉随时间演化的图像结构？
• RQ3 将时间依赖行为直接集成到注意力机制中，而非通过简单的位置嵌入或归一化层，其影响是什么？
• RQ4 具有分层设计的统一 Transformer 架构是否在图像空间与潜在空间生成中，均优于现有的 U-Net 和 DiT 风格模型？
• RQ5 所提出的时变自注意力机制在 CIFAR10、FFHQ-64 和 ImageNet-256/512 等多样化基准上的 FID 得分提升程度如何？

主要发现

• 当用标准自注意力替换时，所提出的时变自注意力（TMSA）模块在 VE 和 VP 设置下，分别将 CIFAR10 的 FID 降低 0.28 和 0.25。
• 潜在 DiffiT 在 ImageNet-256 数据集上实现了新的 SOTA FID 得分为 1.73，优于先前方法（包括 DDPM++ 和 MDT）。
• 消融实验证实，TMSA 显著优于其他替代方案：移除 TMSA 中的时间依赖性后，FID 上升至 3.97；而仅在 MLP 层中使用时间标记时，FID 为 3.81。
• 定性注意力图可视化显示，TMSA 能够捕捉细粒度的物体细节，并在不同去噪步骤中实现动态适应，而无 TMSA 的模型则不具备此能力。
• 采用无分类器引导与幂余弦调度时，ImageNet-256 的 FID 降低至 1.73，最优引导尺度为 4.6；引导尺度过高反而导致性能下降。
• 完整 DiffiT 架构（包含编码器与解码器模块）在 CIFAR10 上实现 FID 为 1.95，证明了完整的 U 形设计对高保真生成的重要性。

Figure 3 : The DiffiT Transformer block applies linear projection to spatial and time-embedding tokens before combining them together to form query, key, and value vectors for each token. These vectors are then used to compute multi-head self-attention activations, followed by two linear layers.