Diffusion Transformer (DiT)：架构演进与应用
什么是 DiT
在 Vision Transformer (ViT) 出现之前，图像生成领域基本由 CNN 主导，扩散模型中的噪声估计器也普遍采用卷积架构的 U-Net。随着 ViT 展现出强大的建模能力，研究者自然开始思考能否将其引入扩散过程。
2022 年 12 月，William Peebles 和 Saining Xie 发表了论文《Scalable Diffusion Models with Transformers》，正式提出了 DiT 架构。其核心思想非常直接：将 DDPM 中的卷积 U-Net 骨干替换为 Transformer 架构。简单来说，DiT = DDPM + ViT。
条件策略与架构细节
DiT 的核心在于如何有效地加入条件信息（Conditioning）。常见的策略有三种：
- adaLN-Zero block：借鉴 ResNets 的经验，将残差块初始化为恒等函数有助于训练。作者对 adaLN DiT 块进行了修改，除了回归 γ 和 β，还回归缩放参数 α，效果显著优于其他变体。
- 交叉注意力块 (Cross-Attention)：将时间步 t 和类别 c 的嵌入连接成序列，作为额外的 token 输入。这会增加约 15% 的 Gflops 开销，但能灵活处理多模态条件。
- In-context conditioning：将 t 和 c 的向量嵌入作为额外 token 追加到输入序列中，类似 ViT 中的 cls token。这种方法无需修改标准 Transformer 块，计算开销可忽略。
在推理过程中，潜在输入被分解为 patch，经过多个 DiT blocks 处理。最终输出不仅包含噪声预测，还有协方差矩阵。通过 T 个采样步骤，即可从纯噪声还原出降噪后的 latent 表示。
从图像到视频：DiT 的改造
DiT 原生支持图像生成，若要应用于视频生成，主要需进行两项改造：
- 时空 Token 化：将噪音 Patch 线性化后，并入 Text Prompt Condition 和 Time Step Condition，共同作为 Transformer 的输入。
- 引入时间维度：Transformer 内部增加 Causal Time Attention 模块。该模块负责收集历史时间信息，确保第 i 帧能看到之前的 k 帧内容，从而维护视频生成的时间一致性。同时，Local Spatial Attention 负责收集单帧内的空间信息。
为了支持不同长宽比和分辨率的视频，可以采用 0/1 Attention Mask 矩阵。例如，设置一个 8*8 的掩码，仅在对角线子 Block 位置设为 1，这样每帧的 Patch 只能看到同帧内的其他 Patch，避免了跨帧污染。这种机制有效解决了 NaViT 导致的分辨率不一致问题。
相关工作对比
U-ViT：清华朱军团队的早期探索
早在 2022 年 9 月，清华朱军团队就发布了《All are Worth Words: A ViT Backbone for Diffusion Models》。他们提出的 U-ViT 同样尝试用 Transformer 替代 U-Net，比 DiT 早了两个月左右。
U-ViT 保留了 U-Net 的跳跃连接结构，这对像素级预测任务至关重要，因为低级特征需要捷径来辅助训练。实验表明，Patch Size 为 2*2 时效果最佳，且两者都验证了 Scale 特性。不过 U-ViT 在小数据集上的实验更为丰富，并引入了长连接技术加速收敛。
Simple Diffusion：Google Research 的端到端方案
2023 年 1 月，Google Brain Team 提出了 Simple Diffusion。其架构在低层次用 MLP 块（Transformer）取代卷积层，高层保留残差连接。研究指出，加噪应根据图像大小调整，且仅需将 U-Net 缩放到 16×16 分辨率即可改善性能。配合蒸馏技术，该模型能在 0.4 秒内生成一张高分辨率图像。
U-DiT：北大与华为的混合架构
2024 年 11 月，北大和华为提出 U-DiT，试图重新结合 U-Net 和 Transformer。研究发现，全局自注意力运算消耗巨大，而 U-Net 主干特征图多为低频信号。因此，他们提出了下采样自注意力机制，在自注意力前先将特征图 2 倍下采样，再并行计算四个分支的注意力，最后融合。相比传统全局自注意力，算力降低 3/4，效果却更优。
机器人动作预测：PAD 框架
DiT 不仅在视觉生成领域表现优异，在具身智能中也展现了潜力。2024 年 11 月，清华大学及上海人工智能实验室等机构提出了 PAD（Prediction with Action）框架。
核心思路
传统的扩散策略通常只关注动作去噪，而 PAD 利用 DiT 架构无缝合并所有模态输入，通过联合去噪同时预测未来图像和动作。这种设计基于一个假设：图像预测能力可以显著增强机器人策略学习，因为它们共享相同的物理法则。
相比于两阶段方法（先生成目标图像再训练策略），PAD 在统一去噪过程中整合了图像生成和机器人动作，避免了中间接口带来的信息损失。
模型架构
PAD 将所有模态映射到潜在空间：
- RGB 图像：通过冻结的 VAE 编码器处理，分词为 token 序列。
- 自然语言指令：通过 CLIP 编码器生成文本嵌入。
- 机器人姿态：使用 MLP 编码并线性转换为 token。
- 深度图像：可选模态，下采样后分词。
这些 token 序列被连接在一起，经历联合的潜在去噪过程。DiT 骨干网络能够适应缺失的模态，例如在仅有视频数据的数据集中，只需将 token 序列扩展到组合长度，并在自注意力块中引入掩码排除填充 token。
训练与实验
训练分为两个阶段：先在互联网 RGB 视频数据集（如 BridgeData-v2）上预训练以建立图像预测先验，再在机器人演示数据上进行微调。损失函数由三部分组成，分别对应图像、动作和深度模态的去噪误差，并通过超参数平衡。
实验结果显示，PAD 在 Metaworld 和现实世界任务中均取得了良好效果。相比 RT-1、RT-2 以及 GR-1 等方法，PAD 预测的未来图像更精确，这有助于提升复杂操作任务的成功率。此外，引入深度模态进一步提升了操控性能，证明了多模态联合预测的有效性。
总的来说，DiT 及其变体正在重塑扩散模型的架构范式，从单纯的图像生成扩展到视频理解、机器人控制等更广泛的具身智能场景。
