Diffusion Transformer (DiT) 架构深度解析:从 U-Net 替换到 AdaLN-Zero 条件调制的下一代扩散模型主干网络演进

Diffusion Transformer (DiT) 架构深度解析:从 U-Net 替换到 AdaLN-Zero 条件调制的下一代扩散模型主干网络演进

• 核心痛点:扩散模型的 U-Net 主干长期受限于卷积归纳偏置与跨尺度特征工程复杂度,导致规模化(Scaling)规律难以显现、训练效率低下,Stable Diffusion 3 / Sora / PixArt-Σ 等新一代模型的架构跃迁需要从原理层面讲清楚 DiT 如何用 Transformer 取代 U-Net
• 适配人群:具备 PyTorch / Transformer / DDPM 基础的 AI 算法工程师、生成式模型研究员、AIGC 应用架构师、扩散模型训练与推理优化工程师
• 收获能力:深入理解 DiT 的 Patchify 流程、AdaLN-Zero 调制机制、In-Context Conditioning 与 Cross-Attention 变体的对比、扩散 Transformer 的规模化规律、MMDiT 多模态扩展、Spacetime Patch 在视频生成中的应用以及 DiT 训练与推理的工程实践要点

• 技术背景与演进逻辑

  • U-Net 在扩散模型中的历史定位
    • 扩散模型(DDPM、DDIM、Score-based)自 2015 年起以 U-Net 作为去噪主干,其优势源于局部归纳偏置、对像素网格的天然适配和跨尺度 skip connection
    • 局限:在 ImageNet 256×256 高分辨率训练上,模型容量与 FID 改善呈现边际递减;U-Net 的复杂结构使得训练/调参/扩展困难
    • 结论:U-Net 难以稳定吸收 LLM 时代揭示的 Scaling Law,推动研究者寻求更通用、可扩展的主干
  • Transformer 在视觉领域的成熟
    • ViT(2020)证明纯 Transformer 视觉主干在大规模数据下可超越 CNN
    • MAE(2022)、BeiT(2021)进一步证明 Mask 预训练 + Transformer 在视觉表征上的有效性
    • 启示:既然 ViT 已经接管视觉识别,那么扩散模型的视觉主干也理应可被 Transformer 接管
  • DiT 的诞生动机
    • William Peebles 与 Saining Xie(ICCV 2023)首次提出以纯 Transformer 作为 Latent Diffusion 的去噪主干
    • 关键实验:S, B, L, XL 四档 DiT 在 ImageNet 256×256 上 FID 分别达到 9.62, 5.91, 4.64, 4.17,验证 Transformer 在扩散任务上同样服从 Gflops 与 FID 的负相关缩放关系
    • 结论:DiT 不仅替换 U-Net,还首次在扩散模型上观察到清晰的 Compute-FID Scaling Law
  • 演进时间线(text 树表达)

    时间线 ├── 2015 -> DDPM: U-Net 作为扩散模型主流主干奠基 ├── 2020 -> DDIM/LDM: 加速采样与潜空间扩散降低算力门槛 ├── 2020 -> ViT: 视觉 Transformer 证明可作为通用主干 ├── 2022.12 -> DiT 论文首次提交: Peebles & Xie 提出 Latent DiT ├── 2023.06 -> MDT / U-ViT: 掩码扩散 Transformer 进一步提升 FID ├── 2023.10 -> PixArt-α: DiT 高效训练范式,首次证明 0.6B 参数即可匹敌 SDXL ├── 2024.02 -> Sora: OpenAI 把 DiT 扩展到 Spacetime Patch 视频生成 ├── 2024.04 -> SD3 MMDiT: Stability AI 用双流 Transformer 取代 UNet+CLIP ├── 2024.04 -> PixArt-Σ: DiT 弱到强训练,直接支持 4K 文生图 ├── 2025 -> DiT 衍生生态: WALT / MAR / DiffiT / SiT 等多种变体竞相出现 └── 2026 -> DiT 主导地位: 几乎所有新一代扩散模型均采用 Transformer 主干

  • 总结
    • 技术迭代必然性:LLM 时代揭示的 Scaling Law 必然要向扩散模型渗透,DiT 是这一必然性的开路者
    • 核心价值:为扩散模型提供可扩展、可预测、可工程化的主干,降低了高分辨率生成模型的训练成本与门槛

• 核心原理深度解析

  • Latent Diffusion 回顾
    • 动机:像素空间扩散在高分辨率下代价过高,LDM(Latent Diffusion Model)先用一个 VAE 将图像编码到潜空间z i n m a t h b b R h × w × c z in mathbb{R}^{h × w × c}zinmathbbRh×w×c,在潜空间做扩散训练
    • 流程:Encoder(x xx) ->z zz-> 扩散前向加噪z t z_tzt​-> DiT 去噪z h e t a ( z t , t , y ) z_{ heta}(z_t, t, y)zheta​(zt​,t,y)-> Decoder ->h a t x hat{x}hatx
    • 关键参数:Stable Diffusion VAE 给出 8× 空间压缩率,256×256 图像对应 32×32 潜空间,这就是 DiT 的 token 序列起点
  • Patchify 与 Token 序列构造
    • 动机:Transformer 接受 1D token 序列,需要把 2D 潜空间切分成 patch
    • 操作:Patch Size 取 2 或 4,使用一个卷积核(Patch Embed 层)把z i n m a t h b b R H × W × C z in mathbb{R}^{H × W × C}zinmathbbRH×W×C转成T = ( H / p ) c d o t ( W / p ) T = (H/p) cdot (W/p)T=(H/p)cdot(W/p)个 token,每个 token 维度为D DD
    • 位置编码:加入 2D sin-cos 可学习位置编码,使得 Transformer 感知空间结构
    • 推导流程:z zz-> Patch Embed ->T TTtokens + 位置编码 ->N NN个 DiT Block -> 输出T TTtokens -> Final Layer ->h a t z hat{z}hatz-> VAE Decoder ->h a t x hat{x}hatx
  • DiT Block 结构
    • 整体结构:沿用 ViT 的残差块范式,但用 Adaptive Layer Norm(AdaLN)替代标准 LayerNorm,并将时间步t tt与类别条件y yy通过 MLP 调制每个 block
    • 残差路径:h a t x i = x i + m a t h r m A d a L N α i , β i , γ i ( t , y ) c d o t m a t h r m M H S A ( x i ) hat{x}_i = x_i + mathrm{AdaLN}_{α_i, β_i, γ_i}(t, y) cdot mathrm{MHSA}(x_i)hatxi​=xi​+mathrmAdaLNαi​,βi​,γi​​(t,y)cdotmathrmMHSA(xi​)-> 再经过m a t h r m A d a L N c d o t m a t h r m M L P mathrm{AdaLN} cdot mathrm{MLP}mathrmAdaLNcdotmathrmMLP-> 残差
    • 核心洞察:条件信息不参与注意力主路径,而是通过归一化参数α , β , γ α, β, γα,β,γ在通道维度上调节激活,这与 LayerNorm 设计哲学高度契合
  • AdaLN-Zero 调制机制
    • AdaLN 基础:对每个 token 沿通道维度做m a t h r m a d a L N ( h ) = γ ( t , y ) c d o t m a t h r m L a y e r N o r m ( h ) + β ( t , y ) mathrm{adaLN}(h) = γ(t, y) cdot mathrm{LayerNorm}(h) + β(t, y)mathrmadaLN(h)=γ(t,y)cdotmathrmLayerNorm(h)+β(t,y),其中γ , β γ, βγ,β由t , y t, yt,y经小型 MLP 输出
    • Zero 初始化:在 AdaLN-Zero 中,MLP 输出额外一个残差缩放因子a l p h a alphaalpha,初始化为m a t h b f 0 mathbf{0}mathbf0,等价于让整个 block 在训练起步时表现为 Identity,显著提升训练稳定性
    • 优势:相比 Cross-Attention 条件化,AdaLN-Zero 计算代价低(无 KV 投影),且实验显示 FID 优于 Cross-Attention 与 In-Context Conditioning