综述由AI生成DiT 架构将扩散模型中的 U-Net 替换为 Transformer，提升了可扩展性。文章介绍了 DiT 的基本原理、条件策略（如 adaLN-Zero），以及其在视频生成领域的改造方法（加入时间注意力）。重点阐述了 DiT 在机器人动作预测中的应用，特别是清华提出的 PAD 框架，通过联合去噪同时预测未来图像和动作，利用大规模视频数据增强视觉策略学习，并在 Metaworld 等任务上验证了有效性。
念念不忘6 浏览
Diffusion Transformer(DiT)：将扩散过程中的 U-Net 换成 ViT
1.1 什么是 DiT
1.1.1 DiT：在 VAE 框架之下扩散去噪中的卷积架构换成 Transformer 架构
在 ViT 之前，图像领域基本是 CNN 的天下，包括扩散过程中的噪声估计器所用的 U-Net 也是卷积架构。但随着 ViT 的横空出世，人们开始考虑这个噪声估计器可否用 Transformer 架构来代替。
2022 年 12 月，William Peebles、Saining Xie 通过论文《Scalable Diffusion Models with Transformers》提出了一种叫 DiT 的神经网络结构：
其结合了视觉 transformer 和 diffusion 模型的优点，即 DiT = DDPM + ViT。
但它把 DDPM 中的卷积架构 U-Net 换成了 Transformer 架构。
1.1.2 Diffusion Transformer(DiT) 的 3 个不同的条件策略
扩散 transformer(DiT) 的架构如下所示：
右侧：DiT blocks 的细节，作者试验了标准 transformer 块的变体，这些变体通过自适应层归一化、交叉注意和额外输入 token 来加入条件（incorporate conditioning via adaptive layer norm, cross-attention and extra input tokens），其中自适应层归一化效果最好。
左侧：训练条件潜在 DiT 模型 (conditional latent DiT models)，潜在输入被分解成 patch 并通过几个 DiT blocks 处理。
经过 N 个 Dit Block（基于 transformer）通过 mlp 进行输出，从而得到噪声'Noise 预测'以及对应的协方差矩阵。
After the final DiT block, we need to decode our sequence of image tokens into an output noise prediction and an output diagonal covariance prediction，最后经过 T 个 step 采样，得到 32x32x4 的降噪后的 latent。
Rather than directly learn dimension-wise scale and shift parameters γ and β, we regress them from the sum of the embedding vectors of t and c.
1.2 DiT 在视频生成领域中的应用
1.2.1 如何改造 DiT，以使其可以做视频生成
当然，DiT 只是用于做图像生成，如果基于其做视频生成，则需要在 DiTs 上做两项改造：
最终，可以得到类似下文即将介绍的 VDT 的一个结构，如下图所示（图源：张俊林）：
首先，把噪音 Patch 线性化后，并入 Text Prompt Condition 和 Time Step Condition，一起作为 Transformer 的输入。
其次，Transformer 内部由三个子模块构成：
Local Spatial Attention 模块负责收集视频帧空间信息。
Causal Time Attention 模块负责收集历史时间信息。
MLP 模块负责对时间和空间信息通过非线性进行融合。
第二，需要把图片生成模型改成视频生成模型，本质就是加上一个时间的维度。具体而言，需要加入一个 Casual Time Attention 子模块。
Causal Time Attention 模块的作用是在生成第 i 帧的时候，收集历史 Time 信息，也就是通过 Attention 让第 i 帧看到之前的比如 k 帧内容，这是用来维护生成视频的时间一致性的。
首先，需要设计一定的模型结构，用来支持不同长宽比和分辨率的视频。以保证 Transformer 在做 Local Spatial Attention 的时候，属于某帧图片的 Patch 只能相互之间看到自己这帧图片内的其它 Patch，但不能看到其它帧图片的内容。
好在根据此篇论文《Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance》可知，可采用'0/1 Attention Mask 矩阵'来达成目标，即假设 Batch 内序列最大长度是 8，就可以设置一个 8*8 的 0/1 Attention Mask，只有对角线正方形子 Block 位置全是 1，其它地方都设置成 0。
其实早在 2022 年 9 月，清华朱军团队（背后关联的公司为生数科技，后于 24 年 3 月上旬拿到一笔数亿元的融资）就发布了一篇名为《All are Worth Words: A ViT Backbone for Diffusion Models》的论文（其一作为 Fan Bao，比 DiT 早提出两个多月，后被 CVPR 2023 收录），这篇论文提出了用「基于 Transformer 的架构 U-ViT」替代基于卷积架构的 U-Net。
U-ViT 在输出之前可选地添加一个 3×3 的卷积块。这旨在防止 transformer 生成的图像中出现潜在的伪影 (potential artifacts in images)。
受到基于 CNN 的 U-Net 在扩散模型中的成功启发，U-ViT 采用了类似的浅层和深层之间的长跳跃连接 (Inspired by the success of the CNN-based U-Net in diffusion models [65], U-ViT also employs similar long skip connections between shallow and deep layers)。
Intuitively,the objective in Eq. (1) is a pixel-level prediction task and is sensitive to low-level features. The long skip connections provide shortcuts for the low-level features and therefore ease the training of the noise prediction network。
1.3.2 比 DiT 更晚之又一个 U-ViT 的工作：Google Research 提出 Simple Diffusion
Google Research 下的 Brain Team，即 Emiel Hoogeboom 等人在 2023 年 1 月份发布的这篇论文《Simple diffusion: End-to-end diffusion for high resolution images》中，也提出了类似 U-ViT 的架构。
如下图所示，本质上，在较低的层次，卷积层被具有自注意力的 MLP 块 (即 transformer) 所取代，且通过残差连接连接起来，只有较高层次的 ResBlocks 使用跳跃连接 (the convolutional layers are replaced by MLP blocks on levels with self-attention. These now form transformer blocks which are connected via residual connections, only the ResBlocks on higher levels use skip connections)。
且在提出的 PAD 框架中，预测和动作是基于多模态当前观测进行的，包括 RGB 图像 $c_{I}$、机器人姿态 $c_{A}$、额外的深度图 $c_{E}$ (在现实世界任务中)，以及自然语言指令文本 $l$。
注意，该框架同时输出相应的未来预测 $x_{I}$、$x_{E}$ 和机器人动作 $x_{A}$。
且 PAD 不仅预测单一未来步骤，还可以预测 k 个未来步骤 $x_{I}^{1: k}, x_{A}^{1: k}, x_{E}^{1: k}$，这可以看作是机器人的 k 步规划。比如机器人只执行第一个预测动作 $x_{A}^{1}$，然后触发一个新的预测周期。这个迭代预测和执行过程使机器人能够以闭环方式持续规划和行动。
图像潜在变量 $L_{I}$ 的 tokenization 遵循与 DiT 相同的切块过程，而机器人姿态的 tokenization 采用简单的线性投影。最后，这些 token 被输入到多层 DiT 中，以预测未来帧的潜在表示。
例如，如果编码潜变量 $\varepsilon_{I}\left(c_{I}\right)$ 的形状为 $c \times d \times d$，那么 $z_{t}^{I^{'}}$ 的形状将为 $k c \times d \times d$ 以表示 k 个未来帧，最终的潜变量 $L_{I}$ 的形状为 $(k+1) c \times d \times d$，其他模态经历类似的过程。
