【论文精读】Diffusion Transformer（DiT）

摘要

       以往的Diffusion模型主要以卷积UNet作为主干网络，本文发现U-Net的归纳偏差对Diffusion模型的性能并不是至关重要的，可以用Transformer等取代。通过利用其他视觉方向成熟的Transformer方案和训练方法，Diffusion模型可以从这些架构中获益，且能保留UNet可扩展性、鲁棒性和效率等特性。

       故本文提出一种基于Transformer的扩散模型，称为Diffusion Transformer（DiTs），DiT遵循ViT的技术方法。具体有：

• 用Transformer架构取代了以往在VAE的潜空间训练的Latent Diffusion Models（LDMs）框架中常用的U-Net主干
• 实验发现，增加Transformer的深度/宽度或增加输入token的数量会提高DiT的Gflops，但同时FID会降低。这证明DiTs是一个可扩展的架构，其网络复杂度（Gflops）和样本质量（FID）之间有很强的相关性
• 在基于class conditional的512×512和256×256 ImageNet基准上，最大的DiT-XL/2（118.6 Gflops）模型性能优于了先前所有的Diffusion模型，在256×256 ImageNet基准上实现了最先进的2.27的FID
  

       上图为在512×512和256×256 ImageNet上训练的DiT-XL/2模型的生成示例。

框架

Preliminaries

Diffusion formulation

       高斯扩散模型（DDPMs）假设了一个前向noise过程，该过程会将noise逐步添加真实数据$x_0:q(x_t|x_0)=N(x_t;\sqrt{{\tilde{\alpha }}_t}{x}_0,(1-{\tilde{\alpha }}_t)\mathrm{I})$，其中常数${\tilde{\alpha }}_t$是超参数。通过对其应用重参化技巧，可以采样$x_t=\sqrt{{\tilde{\alpha }}_t}{x}_0+\sqrt{1-{\tilde{\alpha }}_t}{ϵ}_t$，其中${ϵ}_t\sim N(0,\mathrm{I})$。

       Diffusion模型被训练来学习正向加噪过程的反向过程，即$p_{\theta }(x_{t-1}|x_t)=N({\mu }_{\theta }(x_t),{\Sigma }_{\theta }(x_t))$，其中神经网络用于预测$p_{\theta }$的统计数据。反向过程模型用 x 0