前言
扩散模型近年来在视觉生成领域取得了巨大突破。其中，Diffusion Transformer（DiT）通过将传统 U-Net 架构替换为 Vision Transformer（ViT），显著提升了模型的扩展性和性能。除了视频生成，DiT 在机器人动作预测等具身智能任务中也展现出强大潜力。本文深入解析 DiT 的核心原理、架构变体及其在视频与机器人领域的典型应用。
DiT 核心架构与条件策略
1. 什么是 DiT
在 ViT 出现之前，图像生成领域的噪声估计器主要依赖卷积架构的 U-Net。随着 Transformer 在视觉领域的成功，研究者开始探索将其引入扩散过程。2022 年 12 月，William Peebles 和 Saining Xie 在论文《Scalable Diffusion Models with Transformers》中提出了 DiT 结构。
DiT 的本质是将 DDPM 中的卷积 U-Net 骨干替换为 Transformer 架构。其核心优势在于结合了视觉 Transformer 的全局建模能力与扩散模型的生成能力。在 VAE 框架下，输入潜在空间被分解为 Patch，通过多个 DiT Blocks 处理，最终输出噪声预测及协方差矩阵。
2. 条件化策略
为了让模型能够根据特定条件（如时间步 $t$、类别标签 $c$）进行去噪，DiT 设计了多种条件融合方式，效果最好的是自适应层归一化（adaLN-Zero）。
- Adaptive Layer Norm (adaLN)：不直接学习缩放参数 $eta$ 和 $ ext{scale}$，而是从时间步 $t$ 和类别标签 $c$ 的嵌入向量之和回归得到。这种设计允许模型动态调整激活函数的尺度，加速训练收敛。
- 交叉注意力块：将 $t$ 和 $c$ 的嵌入连接成序列，作为额外的 token 输入到多头自注意力之后。这种方式增加了约 15% 的计算开销，但提供了灵活的条件控制。
- 上下文条件化 (In-context Conditioning)：将 $t$ 和 $c$ 的嵌入作为额外 token 追加到输入序列中，类似 ViT 中的 cls token。这种方法无需修改标准 Transformer 块，计算开销可忽略。
DiT 在视频生成中的应用
将 DiT 应用于视频生成需要在架构上增加时间维度的建模能力。
1. 时间一致性建模
为了生成连贯的视频，模型需要捕捉帧间的时间依赖。通常的做法是引入因果时间注意力（Causal Time Attention）模块。当生成第 $i$ 帧时，该模块允许其关注历史帧 $k < i$ 的内容，从而维护时间一致性。同时，局部空间注意力（Local Spatial Attention）负责收集单帧内的空间信息，MLP 模块则负责时空特征的融合。
2. 多分辨率支持
为了支持不同长宽比和分辨率的视频，可以采用类似 NaViT 的序列打包技术。通过设置 0/1 注意力掩码矩阵，确保同一帧内的 Patch 可以相互注意，而不同帧之间的 Patch 互不可见。这样既避免了跨帧污染，又支持了变长序列的高效处理。
相关工作对比
1. U-ViT
在 DiT 提出之前，清华朱军团队已发布 U-ViT 工作，同样尝试用 Transformer 替代 U-Net。U-ViT 引入了长跳跃连接以保留低级特征，有助于噪声预测网络的训练。两者在 patch size 选择（2x2）和 Scale 特性上结论一致，但 U-ViT 在小数据集上的实验更为丰富。
2. Simple Diffusion
Google Research 提出的 Simple Diffusion 也采用了类似的 Transformer 骨干。其特点是在低分辨率处使用 Transformer，高分辨率处结合卷积残差块，并通过 Dropout 提升性能。该方法实现了端到端的高分辨率图像生成。
3. U-DiT
2024 年，北大与华为联合提出 U-DiT，重新审视 U-Net 架构的价值。研究发现 U-Net 主干中的全局自注意力存在冗余，因此提出了下采样自注意力机制，在降低算力消耗的同时保持了生成质量。
机器人动作预测：PAD 框架
DiT 在具身智能领域的应用同样引人注目。清华大学等机构提出的 PAD（Prediction with Action）框架，利用 DiT 架构实现了预测与动作的联合去噪。
1. 核心思想
传统的扩散策略仅根据观测预测动作，而 PAD 认为图像预测能力可以增强策略学习。PAD 框架在统一的去噪过程中，同时预测未来图像帧和机器人动作。这种联合训练使得模型能够利用大规模互联网视频数据预训练视觉先验，再通过少量机器人演示数据进行微调。
