论文题目Scalable Diffusion Models with Transformers中文翻译使用 Transformer 的可扩展扩散模型会议ICCV 2023ICCV 2023 官方论文页官方 PyTorch 代码摘要本文探索了一类基于 Transformer 架构的新型扩散模型。作者训练图像的潜空间扩散模型用作用于潜变量 Patch 的 Transformer 替代扩散模型中常用的 U-Net Backbone并从前向传播计算复杂度 Gflops 的角度研究所提出 Diffusion TransformerDiT的可扩展性。实验发现无论通过增加 Transformer 的深度和宽度还是增加输入 Token 数量只要 DiT 的 Gflops 增大FID 都会持续降低。除了具有良好的可扩展性之外最大的 DiT-XL/2 模型在 ImageNet 512×512 和 256×256 类条件图像生成基准上超过此前的扩散模型并在 ImageNet 256×256 上获得 2.27 的 FID。一、研究背景与核心问题扩散模型早期有一个非常稳定的架构习惯从 DDPM 到 ADM再到 Latent Diffusion大部分图像扩散模型都会使用卷积 U-Net 作为去噪网络。这并不是因为理论上规定 diffusion 必须搭配 U-Net而更多是一种历史路径依赖。U-Net 最初是卷积架构随后扩散模型又在低分辨率特征层中加入 Spatial Self-Attention但整体骨架并没有发生根本变化。作者因此提出一个很直接的问题扩散模型真的必须使用 U-Net 吗与此同时Transformer 已经在 NLP 和视觉任务中表现出了非常明显的 Scaling 特性那么第二个问题自然出现如果把 diffusion 的 U-Net 换成纯 Transformer它是否也能获得类似的 Scaling 能力这正是 DiT 的核心研究目标。论文明确指出他们并不是想重新设计整个扩散过程而是希望把注意力集中在Backbone Architecture上然后研究之间究竟存在什么关系。1.1 Figure 2 已经提前给出了论文最重要的结论论文 Figure 2DiT 的 Gflops 与 ImageNet 生成质量Figure 2 左边展示不同 DiT 模型的计算量与 FID。最明显的趋势就是右侧又进一步把最佳 DiT-XL/2 与 ADM、LDM 等 U-Net 扩散模型进行比较。作者最终想证明的并不仅仅是Transformer 也能做 diffusion。真正重要的是Transformer 不但能替代 U-Net而且表现出了非常稳定的规模扩展规律。论文把这种新 Backbone 称为二、DiT 整体框架本质上是 Latent Diffusion Vision Transformer理解 DiT 之前需要先明确一个容易产生误解的地方DiT 并不是直接在 RGB Pixel 上运行一个巨大 Transformer。它采用的仍然是 Latent Diffusion 框架。整个生成流程可以写成论文直接使用 Stable Diffusion 中预训练好的 VAE。对于的 RGB 图像经过 8 倍下采样以后得到的 Latent Representation。真正的 diffusion 就是在这个的 Z-space 中完成。因此 DiT 更准确的结构其实是而不是完全抛弃 Stable Diffusion / LDM 的潜空间思路。2.1 Figure 3DiT 的完整结构论文 Figure 3Diffusion Transformer ArchitectureFigure 3 是全文最核心的方法图。DiT 所做的事情与 ViT 处理图像的思想非常相似先把二维空间特征切成 Patch再把 Patch 变成 Token让 Transformer 处理。不同之处在于ViT 最终输出分类结果而 DiT 最终需要预测 diffusion 的NoiseCovariance。Figure 3 右边还给出了三种条件注入方式In-Context ConditioningCross-AttentionAdaptive LayerNorm。最终实验发现adaLN-Zero效果最好。三、DiT 的核心设计Patchify、条件注入与 adaLN-Zero3.1 Patchify把 Latent 变成 Transformer Token论文 Figure 4DiT Patchify设输入 Latent 的空间尺寸为Patch Size 为那么最终 Token 数量为论文测试这里有一个非常重要的关系如果 Patch Size 减半那么 Token 数量将变成原来的4T于是 Transformer 的计算量也会显著增加。这也是 DiT 后面 Scaling 实验的关键。因为作者不仅可以增加模型深度/宽度还可以减小Patch Size、增加Token数量两种方式都会提高 Gflops。3.2 条件信息应该怎么送进 TransformerDiffusion 不只是处理图像 Token还需要知道当前以及类别因此必须解决怎样把 (t) 和 (c) 注入 Transformer作者比较了几种方案。In-Context Conditioning优点是简单几乎不用改标准 Transformer。Cross-Attention这个方案和后来的很多 Text-to-Image 模型比较接近。但论文发现Cross-Attention 会增加大约15% Gflops。Adaptive Layer NormalizationadaLN 则换了一个思路。不是增加额外 Attention而是由tc产生 LayerNorm 的即因此条件信息直接控制 Transformer Block 内部特征的 Scale 和 Shift。这种方案计算量非常低。3.3 adaLN-ZeroDiT 中最重要的设计细节DiT 最终使用的是adaLN-Zero它在 adaLN 基础上再增加残差分支的缩放参数并把输出这些 (\alpha) 的 MLP 初始化为零。因此网络初始化时每一个 DiT Block 在训练开始时近似Identity Function之后再逐渐学习如何修改特征。这一设计看起来只是初始化策略但后面的 Figure 5 表明它对训练影响非常明显。四、DiT 怎么 Scale模型大小与 Token 数同时控制计算量DiT 没有只训练一个模型。作者构造了一个完整 Design Space。Table 1四种 Transformer 规模同时 Patch Size842所以最终形成12个 DiT 模型。因此 DiT 的 Scaling 有两条路径路径一扩大 Transformer也就是提高DepthWidthAttention Heads。路径二增加 TokenPatch 越小进而值得注意的是减小 Patch Size 基本不会明显增加参数量。因此这篇论文能够区分到底是参数数量重要还是实际计算量重要这也是后面 Scaling 实验最关键的设计。五、实验结果与 Scaling 分析5.1 实验设置DiT 在ImageNet Class-Conditional Generation上进行测试包含和两个分辨率。主要指标FIDsFIDInception ScorePrecisionRecall。扩散过程使用的线性 Noise Schedule。模型采用 AdamW学习率Batch Size 256 并使用 EMA0.9999论文的大部分 Scaling 实验使用250 DDPM Sampling Steps计算 FID-50K。5.2 Figure 5为什么最后选择 adaLN-Zero论文 Figure 5Conditioning Strategy Ablation作者使用 DiT-XL/2 比较In-ContextCross-AttentionadaLNadaLN-Zero。结果非常清楚在整个训练过程中基本都保持最低 FID。在 400K Step 时adaLN-Zero 的 FID 接近 In-Context 方法的一半。同时Cross-Attention约 137.6 GflopsadaLN / adaLN-Zero约 118.6 Gflops。所以 adaLN-Zero 并不是单纯准确率更高还具有更好的计算效率。这组实验说明如何把条件信息注入 Transformer并不是一个无关紧要的实现细节而会明显影响 diffusion 的训练质量。5.3 Figure 6Transformer 越大FID 越低论文 Figure 6Model Size / Patch Size ScalingFigure 6 是论文最关键的 Scaling 实验证据之一。上排固定 Patch Size比较结果几乎始终表现为下排固定模型规模再比较同样发现而 Patch Size 改变并不会显著提高参数量。因此论文得到一个非常重要的观察DiT 性能提升并不能简单解释成“参数更多”真正与性能高度相关的是 Forward Pass Compute。5.4 Figure 8Gflops 与 FID 的相关系数达到 -0.93论文 Figure 8Transformer Gflops vs. FID作者把所有 DiT 模型Gflops和画到一张图上。得到相关系数-0.93也就是说Gflops和之间具有非常强的相关关系。而且非常有意思的是不同架构只要 Gflops 接近例如和最终 FID 也比较接近。所以这篇论文真正提出的经验规律可以概括为DiT质量主要随计算量Scale。而不是只要参数越多就越好5.5 Figure 9大模型反而更“省算力”论文 Figure 9Training Compute vs. FID有人可能会自然产生一个疑问大模型单次训练更贵那我训练一个小模型更久不就行了吗Figure 9 专门回答这个问题。作者估算训练总计算量结果发现小模型即使训练得更久达到一定阶段以后也会比大模型更不具计算效率。比如 XL/2 在训练达到一定计算预算后会超过 XL/4。换句话说大模型不仅上限更高而且在高计算预算下利用Compute更有效。这其实与今天大模型 Scaling 的思路非常相似。5.6 Figure 7Scaling 不只是 FID 数字变化论文 Figure 712 个 DiT 模型使用相同 Noise 的生成结果作者给所有 12 个模型相同Initial NoiseSampling NoiseClass Label。从 Figure 7 可以明显看到物体结构越来越稳定类别形态越来越准确背景更加合理局部细节逐渐清晰小模型中的畸变逐渐减少。因此 Figure 7 可以看作 Figure 6、Figure 8 的视觉版本。Figure 1 则展示了最终 DiT-XL/2 在 256×256 和 512×512 ImageNet 上的高质量生成样例。六、主实验DiT-XL/2 到底有多强6.1 Table 2ImageNet 256×256说明 Latent Space Transformer 在这一实验设置中具有很高的计算效率。值得注意的是Table 2 中 StyleGAN-XL 的 FID 为 2.30而 DiT 为 2.27所以论文在该基准下也略低于 StyleGAN-XL。6.2 Table 3ImageNet 512×512在更高分辨率下DiT-XL/2 处理DiT 在 512×512 上超过此前列出的diffusion models但 Table 3 中 StyleGAN-XL 的 FID 仍然更低。论文正文称其在该设置下取得最低 FID 的表述与表中 StyleGAN-XL2.41 存在值得注意的比较口径问题博客中更稳妥的结论是DiT在上显著超过此前U-Net扩散模型另外计算量方面ADM1983 GflopsADM-U2813 GflopsDiT-XL/2524.6 Gflops。这也说明 DiT 在高分辨率 Latent Diffusion 中具有比较明显的计算优势。6.3 Figure 10小模型多采样几步能不能追上大模型论文 Figure 10Model Compute vs. Sampling Compute也就是与其使用一个容量不足的小模型不断增加采样步数不如先把 Backbone 做大。这一观察对后来大型生成模型的发展非常有启发性。七、总结与思考如果把 DiT 压缩成一句话DiT 证明了扩散模型并不依赖 U-Net只要把 Latent 切成 Token标准 Transformer 就可以成为 diffusion backbone而且随着模型深度、宽度和 Token 数增加生成质量会表现出非常稳定的 Scaling 特性。这篇论文真正有影响力的地方因此并不只是“把 U-Net 换成 Transformer”。它真正建立的是一种新的研究范式把扩散模型的 Backbone 从高度任务定制的卷积 U-Net转向一种可以直接利用 Transformer Scaling 规律的通用架构。这也是为什么 DiT 后来会如此重要。再往后看PixArt、Stable Diffusion 3 的 MMDiT以及大量新一代图像和视频生成模型都开始越来越明显地向 Transformer Backbone 靠拢。不过需要特别注意原始 DiT 本身并不是一个 Text-to-Image 模型。所以从历史位置来看我认为理解 DiT 最准确的方式不是“它是一个新的 Stable Diffusion。”而是它证明 Transformer 可以成为扩散生成模型的通用可扩展 Backbone并为后续 SD3 等大型 Diffusion Transformer 打下了架构基础。