前言
Diffusion Transformer（DiT）是近年来扩散模型领域的重要创新。虽然最初主要应用于图像生成，但随着技术演进，其在视频生成和机器人动作预测等任务中也展现出巨大潜力。本文将深入解析 DiT 的核心架构及其在多个前沿场景中的应用。
第一部分 Diffusion Transformer (DiT)：将扩散过程中的 U-Net 换成 ViT
1.1 什么是 DiT
1.1.1 核心概念
在 Vision Transformer (ViT) 出现之前，图像领域的扩散模型噪声估计器通常采用卷积架构的 U-Net。随着 Transformer 架构的崛起，研究者开始探索将其引入扩散过程。
2022 年 12 月，William Peebles 和 Saining Xie 在论文《Scalable Diffusion Models with Transformers》中提出了 DiT 结构。其核心思想是将 DDPM 中的卷积 U-Net 骨干替换为基于 Patch 操作的 Transformer。
- 架构融合：DiT = DDPM + ViT，结合了视觉 Transformer 和扩散模型的优势。
- 骨干替换：用 Transformer 替代了常见的 U-Net 作为潜在扩散模型的骨干网络。
1.1.2 条件策略
DiT 架构通过多种方式融入条件信息（Conditioning），主要包括以下四种策略：
- adaLN-Zero block：借鉴 ResNets 的初始化经验，将残差块初始化为恒等函数。作者对 adaLN DiT 块进行了修改，除了回归 $\gamma$ 和 $\beta$，还回归在 DiT 块内任何残差连接之前应用的缩放参数 $\alpha$。实验表明自适应层归一化效果最好。
- 交叉注意力块：将时间步 $t$ 和条件 $c$ 的嵌入连接成序列，与图像 token 分开处理。在多头自注意力之后增加一个交叉注意层，类似于 LDM 的设计。这会增加约 15% 的计算开销。
- 上下文条件化 (In-context conditioning)：将 $t$ 和 $c$ 的向量嵌入作为额外 token 追加到输入序列中，像 ViT 中的 cls token 一样对待。这种方法无需修改标准 ViT 块，计算开销可忽略。
- 自适应层归一化 (Adaptive layer norm, adaLN)：用自适应层归一化替换 transformer 块中的标准层归一化。不是直接学习维度方向的缩放和偏移参数，而是从 $t$ 和类别标签 $c$ 的嵌入向量之和中回归它们。
1.2 DiT 在视频生成领域中的应用
1.2.1 改造方案
要将 DiT 用于视频生成，需要在原有基础上进行两项关键改造：
- 输入线性化：将噪音 Patch 线性化后，并入 Text Prompt Condition 和 Time Step Condition，共同作为 Transformer 的输入。
- 引入时间维度：Transformer 内部需包含三个子模块：
- Local Spatial Attention：负责收集视频帧的空间信息。
- Causal Time Attention：负责收集历史时间信息，确保第 $i$ 帧能看到之前的 $k$ 帧内容，维持时间一致性。
- MLP 模块：负责对时间和空间信息进行非线性融合。
为了支持不同长宽比和分辨率的视频，需要设计特定的 Attention Mask 矩阵。例如，假设 Batch 内序列最大长度为 8，可设置 8*8 的 0/1 Attention Mask。只有对角线正方形子 Block 位置为 1，其他为 0。这样既能保证同一帧内的 Patch 相互可见，又能防止跨帧的信息泄露。
1.2.2 典型应用
CogVideo 是 DiT 在视频生成领域的代表性应用之一，展示了该架构在处理长序列视频时的能力。
1.3 相关工作
1.3.1 U-ViT：早期的 Transformer 尝试
早在 2022 年 9 月，清华朱军团队便发布了《All are Worth Words: A ViT Backbone for Diffusion Models》，提出了 U-ViT 架构。该工作同样尝试用 Transformer 替代 U-Net，并引入了长跳跃连接（Long Skip Connections）以辅助低频特征传递，加速训练收敛。
