Sean's Blog

DiT

11分钟阅读

下面按“多义辨析 → 在生成模型里的 DiT (Diffusion Transformer) 核心思想 → 架构细节 → 与传统 U-Net 扩散的差异 → 关键技术点（AdaLN、Patch 潜空间等）→ 训练与推理流程 → 性能与可扩展性 → 典型变体与后续工作 → 适用/局限”结构说明。若你指的是别的领域（如 Department of IT, Data-in-Transit 等），请再告诉我。

一、可能含义快速辨析

在当前机器学习/视觉生成语境里，“DiT”最常指 Diffusion Transformer：用 Vision Transformer (ViT) 风格 Transformer 作为扩散模型的骨干。本文聚焦这一含义。

二、DiT 是什么（一句话）

DiT（Diffusion Transformer）= “把扩散模型的去噪网络由卷积 U-Net 换成纯 Transformer（ViT 风格）并在 VAE 潜空间中以图像 patch token 形式处理，从而获得更好的可扩展性与高端算力下更优生成质量。”

三、为什么提出 DiT

传统扩散（DDPM/Stable Diffusion）常用多尺度 U-Net：

  • 优点：局部感受野、归纳偏置、训练早期稳定。
  • 瓶颈：纵向扩展（更深更宽）时多尺度跳连与大卷积开销增加，跨尺度全局交互有限；难与大规模 Transformer 生态（LayerScale、Flash Attention、分布式并行）直接复用。 Transformer 在图像理解中已证明更好的“可扩展性曲线”（scaling law）。因此 DiT 探索：能否在扩散去噪阶段直接用 ViT Backbone，在保持/提升 FID 的同时让模型在更大参数规模下效率更高。

四、整体框架（高层）

  1. 使用 VAE/Autoencoder 将原图 x 压缩到潜空间 z（例如 Stable Diffusion 的做法）降低分辨率与计算。
  2. 在扩散过程中对潜表示 z_t (被加噪后的潜变量) 进行 Transformer 去噪：预测噪声 ε_θ(z_t, t, y) 或直接预测 x_0 depending on parameterization。
  3. Token 化：把 z_t 切分为固定大小 patch（例如 2×2 latent patch）展平后线性映射成 token 序列。
  4. 加上可学习或正弦/二维位置编码。
  5. 利用时间步 t 与（可选）类别/条件 y 通过嵌入 + 调制（Adaptive LayerNorm，AdaLN）注入条件。
  6. 若是类条件模型，可以用 Classifier-Free Guidance 进行推理增强。

五、关键模块细节

(1) Patch Embedding：

  • 把潜空间张量 (C,H,W) 分成若干 patch；每个 patch Flatten → 线性层 → d 维 token。
  • 位置编码可为 2D 可学习矩阵拆成行/列或直接二维展开。

(2) Timestep & 条件编码：

  • 时间步 t → 映射为向量 e_t（通常先正弦时间编码，再 MLP）。
  • 条件（类别、文本 pooled embedding 等）→ e_c。
  • 合并得到条件向量 c（例如拼接后 MLP 或逐元素加）。

(3) AdaLN (Adaptive LayerNorm)：

  • 对每层的 LayerNorm 输出进行“仿射调制”： norm(h) = (h - μ)/σ AdaLN: γ(c) ⊙ norm(h) + β(c) 其中 γ(c), β(c) 由条件向量通过小 MLP 生成。
  • 作用：向每个 Transformer Block 注入时间步与条件信息，替代 U-Net 中大量的 FiLM / AdaGN。

(4) Transformer Block：

  • 标准多头自注意力 (MHA) + MLP（前馈）结构。
  • 有时使用 Pre-Norm + DropPath + LayerScale 改善稳定性。
  • 可集成 Flash Attention 加速。

(5) 输出层：

  • 将 Transformer 最后一层 token 序列 reshape 回潜空间布局，通过线性映射输出预测噪声 ε̂ 或残差。
  • 若使用预测的噪声形式，则 DDPM 采样步骤中： z_{t-1} = 1/√α_t ( z_t - (β_t /√(1- \bar α_t)) ε̂ ) + 采样噪声项 （可省略具体公式细节，除非你需要推导）。

六、与卷积 U-Net 扩散的主要差异

方面U-NetDiT
表示粒度像素特征金字塔Patch token 序列（全局注意力）
条件注入多尺度特征层的 AdaGN/Concat全层 AdaLN 调制
感受野需多层卷积 + 上/下采样自注意力天然全局
可扩展性参数增大后 FID 改善趋缓更线性/可预测的 scaling
并行友好卷积 + 跳连，不同分辨率不易均衡单序列注意力易做张量并行
内存复杂度与空间分辨率线性注意力 O(N^2)（依赖 patch 数）

通过在 VAE 潜空间操作 + 小 patch，DiT 控制 token 数量，缓解注意力 O(N^2) 的成本。

七、训练策略要点

  • 数据：大规模图像（类条件：ImageNet；文本条件：配合文本编码器 CLIP/文本 Transformer）。
  • 目标：最小化噪声预测 MSE（或 v- 预测形式提升稳定性）。
  • Classifier-Free Guidance：随机丢弃条件（设为空）一定比例（例如 10–20%），推理时用 ε_guided = ε_uncond + w (ε_cond - ε_uncond)，w>1。
  • 归一化与优化：AdamW + 余弦或多段学习率调度，LayerNorm 前置。
  • 数据增广：随机裁剪、翻转；若是高分辨图像可先下采样再由 VAE 编码。

八、推理/采样

  1. 文本或类别条件 → 得到 embedding e_c。
  2. 初始化 z_T ~ N(0,I)（潜空间）。
  3. 迭代 t = T..1：用 DiT 预测 ε̂_t；应用采样公式（DDIM 或 DDPM 变体）更新 z_{t-1}。
  4. 最终 z_0 通过 VAE 解码 → 输出图像。 可使用：
  • 减步数技术：DDIM、DPM-Solver、Euler Ancestral 等。
  • Guidance w 调节细节 vs. 多样性。

九、性能与扩展性

  • 论文显示在相同或更少计算预算下，较同规模 U-Net 扩散有更好的 FID/IS（特别是大模型区间）。
  • Scaling Law：参数翻倍 → 质量改进更可预测，有利于规划训练预算。
  • 可与更大语义编码器（如 T5、LLaMA 文本嵌入）耦合实现多模态生成。

十、常见变体与后续发展

  • 文本条件 DiT：把文本序列 embedding 和时间/类别统一调制；或将文本 token 与图像 patch token 拼接做 cross/self 注意力。
  • Hybrid DiT：前几层局部注意力 + 后几层全局注意力降低成本。
  • Sparse / Window Attention DiT：减少 O(N^2) 成本，以获更高分辨率。
  • Masked Pretraining：先做 MAE/VQ 形式预训练，再微调用于扩散去噪。
  • Video DiT：在 token 维度加时间维（space-time patches），或分解时空注意力。
  • Latent Upsampler DiT：先生成低分辨潜图，再用上采样 DiT 扩散放大。

十一、典型优势

  1. 更强扩展性：借助成熟的 Transformer 工程栈（并行、混合并行、Flash Attention）。
  2. 全局建模：长距依赖及一致性（结构、纹理重复等）更自然。
  3. 模块化：更易插入额外条件（布局、深度、风格 token）。
  4. 与多模态扩散生态兼容：文本、控制信号在同一 token 框架中整合。

十二、局限与挑战

  1. 注意力平方复杂度：超高分辨率仍昂贵；需稀疏/分块技巧。
  2. 归纳偏置不足：小数据或低算力下可能不如卷积稳定/样本效率。
  3. 几何细节：高频局部结构（微纹理）在纯全局注意力下早期训练收敛略慢；常借助 VAE/感知损失补强。
  4. 推理延迟：仍需多步采样；虽可用加速 ODE/SDE Solver，但不如单步生成（GAN）快。

十三、简化伪代码

(仅示意)

# z_t: latent noisy image
tokens = PatchEmbed(z_t) + PosEncoding
cond = CondMLP(TimeEmbed(t), Class/TextEmbed(y))
for block in TransformerBlocks:
    tokens = block(tokens, cond)  # 内部对每个 LayerNorm 做 AdaLN 调制
pred_noise = Head(tokens).reshape_like(z_t)
return pred_noise

十四、与 Stable Diffusion 的关系

Stable Diffusion ≈ (VAE 编码) + U-Net 扩散 + 文本条件 cross-attn。 DiT 可被视为：把 U-Net 替换为 patch Transformer，并用 AdaLN 而非 cross-attn（或结合），同时仍可在 VAE 潜空间运行。日后许多高端模型（如部分开源大尺寸图像/视频模型）采用此思路或其变体。

十五、快速记忆口诀 “DiT = 先 VAE 降维，再把潜图切成 patch → Transformer + AdaLN 注入时间与条件 → 预测噪声 → 多步去噪采样。”

十六、一句话总结 DiT 是利用 Vision Transformer 作为扩散去噪网络的设计，使生成模型在大规模参数下具有更好的扩展曲线与全局建模能力，逐步成为高端多模态图像/视频扩散体系的核心骨干之一。

如果你想：

  • 看更详细的 AdaLN 数学形式
  • 了解与最新多模态（文本/控制）扩展怎么融合
  • 对比 DiT 与最近的 U-ViT、PixArt-Alpha、SDXL 内部差异 告诉我，我再继续展开。

目录

Graph View