Diffusion Transformer (DiT) 详解
一、DiT 核心原理
1.1 什么是 DiT
在 Vision Transformer (ViT) 出现之前，图像领域基本由 CNN 主导，包括扩散模型中的噪声估计器（U-Net）也主要采用卷积架构。随着 ViT 的横空出世，研究者自然开始思考：能否用 Transformer 架构来替代扩散过程中的 U-Net？
2022 年 12 月，William Peebles（当时在 UC Berkeley）和 Saining Xie（当时在纽约大学）通过论文《Scalable Diffusion Models with Transformers》提出了 DiT 结构。
其核心思想可以概括为两点：
- 结合优势：DiT = DDPM + ViT，结合了视觉 Transformer 和扩散模型的优点。
- 架构替换：将 DDPM 中常用的卷积架构 U-Net 骨干替换为基于 Patch 的 Transformer。
本质上，这是利用 Transformer 强大的全局建模能力来处理潜在空间（Latent Space）中的去噪任务。
1.2 DiT 的条件策略
DiT 架构支持多种条件输入方式（Conditioning），主要包括以下三种策略：
1. adaLN-Zero Block
受 ResNets 工作的启发，作者发现将残差块初始化为恒等函数有益训练。在 DiT 中，除了回归缩放参数 γ 和 β，还额外回归一个缩放参数 α，用于在残差连接前对特征进行自适应层归一化。实验表明，这种自适应层归一化效果最好。
2. 交叉注意力块 (Cross-Attention)
将时间步 t 和类别 c 的嵌入连接成一个长度为二的序列，与图像 token 序列分开。Transformer 块被修改为在多头自注意力之后，增加一个额外的多头交叉注意力层。这类似于 Latent Diffusion Models (LDM) 的设计。不过，这种方式会增加约 15% 的计算开销（GFLOPs）。
3. 上下文条件化 (In-context Conditioning)
将时间步 t 和条件 c 的向量嵌入作为两个额外的 token 追加到输入序列中，与图像 token 无异地对待。这类似于 ViT 中的 cls token，允许使用标准 Transformer 块而无需修改架构。这种方法几乎不增加计算量。
1.3 相关工作对比
U-ViT (清华团队)
早在 2022 年 9 月，清华朱军团队就发布了《All are Worth Words: A ViT Backbone for Diffusion Models》，提出了 U-ViT。它同样尝试用 Transformer 替代 U-Net，并引入了长跳跃连接（Long Skip Connections）以保留低级特征，加速收敛。相比 DiT，U-ViT 在小数据集和图文数据集上做了更广泛的实验。
Simple Diffusion (Google Research)
2023 年 1 月，Google Brain Team 提出了 Simple Diffusion，本质上也是用 MLP 块（Transformer）取代低层次的卷积层，并在高层保留残差连接。该工作强调了分辨率调整、Dropout 策略以及端到端训练的重要性。
U-DiT (北大 & 华为)
2024 年 11 月，北大和华为提出 U-DiT，试图重新拾起 U-Net 架构并与 Transformer 结合。他们发现全局自注意力在 U-Net 主干中存在冗余，因此提出了下采样自注意力机制，在降低算力的同时保持性能。
二、DiT 在视频生成中的应用
DiT 最初主要用于图像生成，若要应用于视频生成，需要进行两项关键改造：
- 时空 Token 化：将噪音 Patch 线性化后，并入 Text Prompt Condition 和 Time Step Condition，一起作为 Transformer 的输入。
- 引入时间维度：在 Transformer 内部加入 Causal Time Attention 模块，负责收集历史时间信息，确保生成视频的时间一致性。
此外，为了支持不同长宽比和分辨率的视频，需要设计特定的 Attention Mask 矩阵（如 NaViT 方案）。通过设置 0/1 Attention Mask，可以保证某帧内的 Patch 只能看到本帧内容，避免跨帧污染，从而灵活处理变长视频序列。
