Diffusion Transformer (DiT)：从 U-Net 到 ViT 的架构演进
前言
Diffusion Transformer (DiT) 是近年来扩散模型领域的重要创新。除了广泛应用于视频生成，它在机器人动作预测等具身智能领域的应用也日益增多。本文将独立解析 DiT 的核心架构及其在多个前沿场景中的落地实践。
第一部分 Diffusion Transformer (DiT)：将扩散过程中的 U-Net 换成 ViT
1.1 什么是 DiT
1.1.1 DiT：在 VAE 框架之下扩散去噪中的卷积架构换成 Transformer 架构
在 Vision Transformer (ViT) 出现之前，图像领域基本由 CNN 主导，包括扩散模型中用于噪声估计的 U-Net 也是基于卷积架构。随着 ViT 的横空出世，研究者自然开始思考能否用 Transformer 替代噪声估计器。
2022 年 12 月，William Peebles（当时在 UC Berkeley）和 Saining Xie（当时在纽约大学）通过论文《Scalable Diffusion Models with Transformers》提出了 DiT 结构。
- 核心结合：DiT 结合了视觉 Transformer 和扩散模型的优点，即 DiT = DDPM + ViT。
- 架构替换：它将 DDPM 中常用的卷积架构 U-Net 换成了 Transformer 架构。具体来说，是在潜在空间中对图像 Patch 进行操作。
1.1.2 DiT 的三种条件策略
DiT 的架构设计允许通过多种方式加入条件信息（Conditioning），主要包括以下三种策略：
- 右侧 DiT Blocks 细节：作者试验了标准 Transformer 块的变体，通过自适应层归一化、交叉注意力和额外输入 Token 来加入条件。其中，自适应层归一化效果最好。
- 左侧训练流程：潜在输入被分解成 Patch，经过几个 DiT Blocks 处理。本质上是噪声图片减去预测的噪声以实现逐步复原。
以一张 256x256x3 的图片为例，切 Patch 后投影得到每个 Patch 的 Token，形成 32x32x4 的 Noised Latent（加噪后的潜在表示）。结合当前的 Timestep t 和 Label y 作为输入，经过 N 个 DiT Block（基于 Transformer）通过 MLP 输出，从而得到噪声预测以及对应的协方差矩阵。最后经过 T 步采样，得到降噪后的潜在表示。
接下来详细介绍上图右侧的四种不同条件策略（即如何加入 Conditioning）：
-
adaLN-Zero block 借鉴 ResNets 的经验，将残差块初始化为恒等函数有益。例如，Goyal 等人发现零初始化批量归一化尺度因子可加速训练。DiT 对 adaLN 块做了类似修改，除了回归 γ 和 β，还回归在残差连接前应用的缩放参数 α。
-
交叉注意力块 将时间 t 和条件 c 的嵌入连接成长度为二的序列，与图像 Token 序列分开。Transformer 块在多头自注意力之后，包含一个额外的多头交叉注意力层。这增加了约 15% 的计算开销。
-
上下文条件化 (In-context conditioning) 将 t 和 c 的向量嵌入作为两个额外的 Token 追加到输入序列中，像处理图像 Token 一样对待它们。这类似于 ViT 中的 cls Token，允许使用标准 ViT 块而无需修改。这种方法几乎不增加计算量。
-
自适应层归一化 (Adaptive layer norm, adaLN) block 鉴于自适应归一化层在 GANs 和扩散模型中的广泛使用，DiT 用自适应层归一化替换了 Transformer 块中的标准层归一化。它不是直接学习维度方向的缩放和偏移参数，而是从 t 和类别标签 c 的嵌入向量之和中回归它们。
1.2 DiT 在视频生成领域中的应用
1.2.1 如何改造 DiT 以支持视频生成
DiT 原生用于图像生成，若用于视频生成，需进行两项关键改造：
- 输入线性化：将噪音 Patch 线性化后，并入 Text Prompt Condition 和 Time Step Condition，一起作为 Transformer 的输入。
