Diffusion Transformer (DiT)

扩散模型在图像生成上长期以卷积式 U-Net 为主干,但论文Scalable Diffusion Models with Transformers提出以 Transformer 作为扩散模型的主干(Diffusion Transformer,简称 DiT),并系统研究其“可扩展性”:当通过增大网络深/宽或增加输入 token 数使前向计算量(以 Gflops 计)增加时,生成质量(以 FID 计)是否稳定提升。作者给出明确证据:更高的前向 Gflops 与更低的 FID 强负相关(相关系数约 −0.93),同时在 ImageNet 256×256 与 512×512 上取得当时的 SOTA,其中 256×256 的 FID 降至 2.27。论文还展示了在潜空间(LDM 框架)中以 ViT 风格 patch 化的方式替换 U-Net 的可行性与效率优势。

论文作者为William Peebles, Saining Xie。

一、方法总览:Latent-space DiT
作者采用“潜空间扩散”范式:先用预训练 VAE 将图像压缩为较小的空间表示 z(如 256×256×3 图像被编码为 32×32×4 的 latent),扩散模型在 z 空间训练与采样,最后用 VAE 解码回像素。DiT 以 ViT 为蓝本:将 z 按 p×p patch 线性嵌入为 token 序列(长度 T=(I/p)²),叠加正弦/余弦位置编码后,送入一串 Transformer blocks 进行噪声与对角协方差预测,最终线性投影并按原空间排布还原至 p×p×2C(对应噪声与协方差)。p 越小,token 越多,Gflops 增长更快,但参数量基本不变,从而可将“算力”与“参数量”解耦到一定程度。

二、条件注入与 DiT Block 设计
论文系统比较了四种将条件(扩散时间步 t、类别标签 c 等)注入 Transformer 的方式:
1)In-context conditioning:将 t、c 的向量直接作为额外 tokens 拼到序列中,几乎不增加 Gflops。
2)Cross-Attention:在自注意力后加入多头交叉注意力,用条件序列查询图像 tokens,约带来 15% 的额外 Gflops。
3)Adaptive LayerNorm(adaLN):用(t+c) 的嵌入回归层归一化中的 γ、β,对所有 tokens 施加同一条件仿射变换,额外 Gflops 极小。
4)adaLN-Zero:在 adaLN 的基础上,引入残差前的逐维缩放 α,并将其 MLP 初始化为全零,使每个 Transformer block 初始为“恒等”映射(identity-like initialization),训练更稳更快。实验显示,adaLN-Zero 在全程训练中 FID 最优且计算最省。

三、模型规格与缩放维度
作者沿用 ViT 的 S、B、L、XL 四档配置(层数 N、隐层宽 d、注意力头数随之增长),并与 patch 大小 p∈{2,4,8} 组合,形成 12 个模型(如“DiT-XL/2”表示 XL 配置且 p=2)。两条主要缩放轴:
(A)增大模型深宽(提高每层维度与层数);
(B)减小 patch 尺寸、增加 token 数(主要增加 Gflops 而几乎不变参数量)。两轴均显著降低 FID。

四、可扩展性关键发现
1)“Gflops 而非参数量”驱动质量提升:在固定参数量(固定配置)下,仅通过减小 p 提高 token 数、增加前向 Gflops,也能大幅降 FID;不同配置只要 Gflops 相近,其 FID 往往接近(如 S/2 与 B/4)。
2)强相关性:在 400k 步训练后,各模型的 FID-50K 与 Transformer 前向 Gflops 呈显著负相关(≈−0.93)。
3)更大模型的“训练算力利用率”更高:以“总训练算力 = 模型 Gflops × batch × steps ×3”近似衡量时,小模型即便延长训练也会在单位算力效率上落后于大模型;同为 XL 配置时,p=2 的 XL/2 在约 10¹⁰ Gflops 后就超过 p=4 的 XL/4。

五、与 U-Net 架构的对比与 SOTA 结果
在 ImageNet 256×256(类别条件)上,DiT-XL/2(在采样时使用 classifier-free guidance)将 FID-50K 推进到 2.27,优于同为潜空间的 LDM 与像素空间的 ADM 系列;在 512×512 上也刷新扩散模型最佳(文中报告 3.04),同时其前向 Gflops 相比 ADM/ADM-U 更低,显示出“高质量 + 计算效率”的双重优势。值得注意的是,作者强调 DiT 的 compute-efficiency:同等或更低 Gflops 下优于 U-Net 系列。

六、训练与评测设置
1)训练设置:在 JAX + TPU v3-Pods 上训练,优化器 AdamW,学习率 1e-4,批量 256,无权重衰减与预热,仅水平翻转增广;对所有尺寸与 p 取相同超参;维护 0.9999 的 EMA。
2)扩散超参:t_max=1000 的线性噪声日程,预测噪声(ϵ-parametrization),并学习对角协方差 Σθ。
3)潜空间:使用 Stable Diffusion 的预训练 VAE(编码下采样 8 倍),论文还对不同 VAE 解码器做了对照,指标相近。
4)评测:主指标为 FID-50K(默认 250 DDPM 采样步),另报告 IS、sFID、Precision/Recall 并呈现与 Gflops 的关联趋势曲线。

七、采样计算 vs 模型计算
扩散模型可通过增加采样步数在推理时“加算力”。作者对比了固定训练步(400k)下,不同模型在不同采样步(16→1000)时的 FID-10K:结论是“小模型即使把采样算力堆到比大模型更高,仍难以弥合差距”。例如 L/2 用 1000 步采样其单图采样算力是 XL/2(128 步)的 ~5 倍,但 XL/2 的 FID 仍更好。模型规模带来的“能力”不能单靠增加采样步在测试时补齐。

八、为何 adaLN-Zero 有效
在 Transformer 中以“条件驱动的自适应层归一化”替换标准 LN,可在不引入昂贵 cross-attn 的前提下将条件强耦合到每一层、每个通道的尺度与偏置上;再加上“零初始化的残差缩放 α”让每个 block 刚开始更像“恒等映射”,缓解深层网络早期优化不稳定,实践中在等算力下显著优于 in-context 与 cross-attn。这个设计既稳定训练、又省 Gflops,体现了针对扩散去噪任务的“薄而稳”的条件注入优势。

九、工程与实践要点
1)两条扩展轴:优先增加“有效算力”而非盲目扩参;在预算有限时,减小 patch 增 token 往往更高效。
2)潜空间优先:在相似感知质量下,潜空间扩散相对于像素空间显著省算。
3)指导策略:classifier-free guidance 依然是质量/多样性的关键旋钮;文中还观察到对潜通道子集施加 guidance 也能取得接近效果。
4)统一架构红利:用标准化的 ViT 设计与训练 recipe(无额外正则/预热亦稳),便于跨任务迁移与持续扩容。

十、局限与开放问题
1)算力敏感:尽管 DiT 相对高效,但最优结果仍依赖大量训练算力与长时训练;如何进一步提升数据与算力利用率仍是问题。
2)文本条件/多模态泛化:本文以类别条件与潜空间为主,如何在文本到图像、视频、多模态条件下复用 DiT 的可扩展性,需要系统验证。
3)推理速度:即使模型更强,扩散采样本身仍慢,如何与蒸馏、加速采样器(如 DDIM、EDM 系列)协同以保持质量与速度平衡,是工程落地关键。

十一、结论与启示
论文以严格对照证明:扩散模型并不依赖 U-Net 的卷积归纳偏置,标准化的 ViT 架构在潜空间同样能稳定训练并具备一流可扩展性;“提升前向 Gflops”是质量提升的核心旋钮,而不是“单纯扩参数”。结合 adaLN-Zero 等轻量条件注入机制,DiT 在 ImageNet 基准上以更优算力效率刷新成绩,为“跨任务统一主干(Transformer)”提供了强力证据,也为文本到图像等更广泛生成任务的主干替换与规模化训练奠定了方法论基础。

十二、给研发者的落地建议(结合论文观察)
1)若已有 LDM 训练/推理基础,优先尝试将 U-Net 主干替换为 ViT-style DiT,并使用 adaLN-Zero。
2)在固定显存下,优先通过更小的 patch(更多 token)来增加“有效算力”;若带宽/显存吃紧,再折中增加层宽/深。
3)保证评测一致性(如统一 FID 评估管线、固定采样步等),并用“总训练算力”与“单图采样算力”两维来度量“性价比”,避免单看参数或 FLOPs。


Diffusion Transformer(DiT) on GitHub: https://github.com/facebookresearch/DiT

发表评论

您的邮箱地址不会被公开。 必填项已用 * 标注