Diffusion Transformer(DiT)：用 ViT 替代 U-Net 用于视频生成与机器人动作预测

一、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（当时在 UC Berkeley）、Saining Xie（当时在纽约大学）通过论文《Scalable Diffusion Models with Transformers》提出了一种叫 DiT 的神经网络结构。

其结合了视觉 transformer 和 diffusion 模型的优点，即 DiT = DDPM + ViT。
但它把 DDPM 中的卷积架构 U-Net 换成了 Transformer 架构 (We train latent diffusion models of images, replacing the commonly-used U-Net backbone with a transformer that operates on latent patches)。

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 处理 (The input latent is decomposed into patches and processed by several DiT blocks)。

本质就是噪声图片减掉预测的噪声以实现逐步复原。

$\rightarrow$

比如当输入是一张 256x256x3 的图片，对图片做切 patch 后经过投影得到每个 patch 的 token，得到 32x32x4 的 Noised Latent（即加噪的图片，在推理时输入直接是 32x32x4 的噪声），结合当前的 Timestep t、Label y 作为输入。

$\rightarrow$

经过 N 个 Dit Block（基于 transformer）通过 mlp 进行输出，从而得到噪声'Noise 预测'以及对应的协方差矩阵。

$\Sigma$

(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。

接下来，仔细介绍下上图右侧 4 个不同的条件策略 (说白了，就是怎么加入 conditioning)。

adaLN-Zero block 关于 ResNets 的先前工作发现，将每个残差块初始化为恒等函数是有益的。例如，Goyal 等人发现，在每个块中将最终批量归一化尺度因子 γ零初始化可以加速大规模训练在监督学习设置中 [13]。扩散 U-Net 模型使用类似的初始化策略，在任何残差连接之前将每个块的最终卷积层零初始化。作者对 adaLN DiT 块的修改，它做了同样的事情。除了回归 γ和 β，还回归在 DiT 块内的任何残差连接之前立即应用的维度方向的缩放参数 α。
交叉注意力块 将 t 和 c 的嵌入连接成一个长度为二的序列，与图像 token 序列分开，transformer 块被修改为：在多头自注意块之后，包含一个额外的多头交叉注意层，类似于 LDM 用于根据类标签进行条件处理的设计。交叉注意力使模型增加了最多的 Gflops，大约增加了 15% 的开销。

上下文条件化 (In-context condition) 将 t 和 c 的向量嵌入作为两个额外的 token 追加到输入序列中，与图像 token 无异地对待它们，这类似于 ViTs 中的 cls token，它允许使用标准 ViT 块而无需修改。在最后一个块之后，从序列中移除条件 token。这种方法没有新增模型多少 Gflops，可忽略。

自适应层归一化，即 Adaptive layer norm (adaLN) block 鉴于自适应归一化层在 GANs 和具有 U-Net 骨干的扩散模型中的广泛使用，故用自适应层归一化 (adaLN) 替换 transformer 块中的标准层归一化层。不是直接学习维度方向的缩放和偏移参数 γ和 β，而是从 t 和类别标签 c 的嵌入向量之和中回归它们。 (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)

至于 DiT 的代码实现，请见相关开源项目。

1.2 DiT 在视频生成领域中的应用

1.2.1 如何改造 DiT，以使其可以做视频生成

当然，DiT 只是用于做图像生成，如果基于其做视频生成，则需要在 DiTs 上做两项改造：

最终，可以得到类似下文即将介绍的 VDT 的一个结构，如下图所示。

文章配图

首先，把噪音 Patch 线性化后，并入 Text Prompt Condition 和 Time Step Condition，一起作为 Transformer 的输入。其次，Transformer 内部由三个子模块构成：

$\rightarrow$ Local Spatial Attention 模块负责收集视频帧空间信息。

$\rightarrow$ Causal Time Attention 模块负责收集历史时间信息。

$\rightarrow$ MLP 模块负责对时间和空间信息通过非线性进行融合。

第二，需要把图片生成模型改成视频生成模型，本质就是加上一个时间的维度。具体而言，需要加入一个 Causal 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。上图左侧标为绿色的某帧三个 Patch，如果看矩阵前三行，易看出，针对其它帧的 Attention Mask 由于都是 0，所以加上 Mask 后就看不到其它图片 (比如橙色、紫色的帧)，而对于它 (三个绿色 patch) 对应的 3∗3 都是 1 的 Attention Mask，又可以保证三个 Patch 相互都能看到。如此，通过设置 Attention Mask，就可以很方便地支持 NaVIT 导致的每帧不同分辨率和长宽比的问题。

1.2.2 DiT 在视频生成中的典型应用：Cogvideo

// 待更

1.3 相关工作

1.3.1 比 DiT 更早之类似架构 U-ViT：虽也带文本条件融合，但和 DiT 都只局限在 2D 图像生成

对于扩散过程中噪声估计器的选择，国外的研究者很自然的想到了用 Transformer 代替 U-net，那国内就没人想到么？

其实早在 2022 年 9 月，清华大学朱军团队发布了一篇名为《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)。意思是说，εθ(xt, t, c) 的目标是像素级预测任务，对低级特征敏感。长跳跃连接为低级特征提供了捷径，因此有助于噪声预测网络的训练。 (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)

U-ViT 参数化了噪声预测网络 εθ(xt, t, c)，它接受时间 t、条件 c、噪声图像 xt 作为输入，然后预测注入到 xt 中的噪声。按照 ViT 的设计方法，图像被分割成 patch，并且 U-ViT 将所有输入 (包括时间、条件和图像 patch) 都视为 token。

对比来看，两项工作在架构路线上完全一致：

均是提出了将 Transformer 与扩散模型融合的思路。
并且在具体的实验路径上也一致，比如采用了相同的 patch embedding、patch size；都得出了同样的结论 ——patch size 为 2*2 是最理想的。
在模型参数量上，两者都在 50M-500M 左右的参数量上做了实验，最终都证实了 scale 特性。

不过 DiT 仅在 ImageNet 上做了实验，U-ViT 在小数据集 (CIFAR10、CelebA)、ImageNet、图文数据集 MSCOCO 上均做了实验。此外，相比传统的 Transformer，U-ViT 提出了一项「长连接」的技术，大大提升了训练收敛速度。

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)。

文章配图

他们指出：

对于加噪，应根据图像的大小进行调整，随着分辨率的增加，添加更多的噪声。
仅将 U-Net 架构缩放到 16×16 的分辨率即可改善性能。进一步地，也同样提出了 U-ViT 架构，它是具有 Transformer 骨干的 U-Net。本质上，这个 U-Vision Transformer（U-ViT）架构可以看作是一个小型卷积 U-Net，通过多个级别的下采样到达 16 × 16 的分辨率，然后经过一个大型 Transformer 之后，最后通过卷积 U-Net 进行上采样。
为了提高性能，应添加 Dropout，但不要在最高分辨率的特征图上添加。
对于更高的分辨率，可以进行下采样而不会降低性能。最重要的是，这些结果仅使用单个模型和端到端的训练设置获得。在使用现有的蒸馏技术后，这些技术现在只需应用于单个阶段，模型可以在 0.4 秒内生成一张图像。

1.3.3(选读) 24 年 11 月：北大&华为提出扩散架构 U-DiT(六分之一算力即可超越 DiT)

Sora 的发布让广大研究者及开发者深刻认识到基于 Transformer 架构扩散模型的巨大潜力。作为这一类的代表性工作，DiT 模型抛弃了传统的 U-Net 扩散架构，转而使用直筒型去噪模型。

鉴于直筒型 DiT 在隐空间生成任务上效果出众，后续的一些工作如 PixArt、SD3 等等也都不约而同地使用了直筒型架构。
然而令人感到不解的是，U-Net 结构是之前最常用的扩散架构，在图像空间和隐空间的生成效果均表现不俗；可以说 U-Net 的 inductive bias 在扩散任务上已被广泛证实是有效的。

因此，北大和华为的研究者们产生了一个疑问：能否重新拾起 U-Net，将 U-Net 架构和 Transformer 有机结合，使扩散模型效果更上一层楼？带着这个问题，他们提出了基于 U-Net 的 DiT 架构 U-DiT，对应的论文为《U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers》。

如下图所示，左图即为 DiT 的原始架构 (相当于 DDPM + ViT 的直筒型扩散架构)、中图相当于 U-net 版的 DiT：具有 Transformer 骨干的 U 型扩散架构、右图则是本次最新的工作 U-DiT。

文章配图

首先，研究者开展了一个小实验，在实验中尝试着将 U-Net 和 DiT 模块简单结合 (如上图中部)——具有 Transformer 骨干的 U 型扩散架构。然而，在相似的算力比较下，U-Net 的 DiT（DiT-UNet）仅仅比原始的 DiT 有略微的提升。
根据先前的工作，在扩散中 U-Net 的主干结构特征图主要为低频信号。由于全局自注意力运算机制需要消耗大量算力，在 U-Net 的主干自注意力架构中可能存在冗余。这时作者注意到，简单的下采样可以自然地滤除噪声较多的高频，强调信息充沛的低频。既然如此，是否可以通过下采样来消除对特征图自注意力中的冗余？

由此，作者提出了下采样自注意力机制，在自注意力之前。

$\rightarrow$ 首先，需将特征图进行 2 倍下采样。为避免重要信息的损失，生成了四个维度完全相同的下采样图，以确保下采样前后的特征总维度相同。

$\rightarrow$ 随后，在四个特征图上使用共用的 Q、K、V 映射，并分别独立进行自注意力运算。

$\rightarrow$ 最后，将四个 2 倍下采样的特征图重新融为一个完整特征图。和传统的全局自注意力相比，下采样自注意力可以使得自注意力所需算力降低 3/4。

// 待更

二、DiT 在机器人动作预测中的典型应用：清华 RDT、CogACT、PAD

2.1 预测与动作扩散器 PAD：通过联合去噪同时预测未来图像和动作

2.1.1 什么是 PAD

首先，DiT 作为机器人的动作预测模型已经不是新鲜事了。

相关研究指出，RDT 所用的扩大大模型架构便是基于 DiT 改造而成。且再相关研究所述，CogACT 是 VLM 作为认知基础，DiT 作为动作模块。

而除了 RDT、CogACT 之外，24 年 11 月，来自清华大学 IIIS、上海启智研究院、上海人工智能实验室、加州大学伯克利分校的研究者提出了：预测与动作扩散器 PAD——作为一个统一的策略学习框架，该框架在 Diffusion Transformer(DiT) 架构下整合了预测和动作。

其对应的论文为：Prediction with Action: Visual Policy Learning via Joint Denoising Process
其对应的作者为：郭延江，胡宇成，张建科，王彦仁，陈晓宇，陆超超，陈剑宇。

具体而言，如下图。

文章配图

DDPM 是根据加噪的图像做去噪生成清晰的图像，如上图图 2(a) 所示。
diffusion policy 是根据观察图像和机器人本体状态去去噪预测机器人未来状态，如上图图 2(b) 所示。
而 PAD 则是利用 Diffusion Transformer(DiT) 模型无缝合并所有模态输入，并通过联合去噪同时预测未来图像和动作，如上图图 2(c) 所示。

此外，灵活的 DiT 骨干网络还允许 PAD 在大规模视频数据上进行联合训练，并扩展到其他机器人模态，如深度图像。

2.1.2 提出背景与相关工作

最近，扩散模型 (涉及 DDPM、DDIM、improved DDPM) 通过在大规模网络图像和视频数据集上进行训练，在视觉生成任务中取得了令人瞩目的成功 (涉及 Imagenet、Bridgedata v2 等)。

例如，图像编辑模型可以根据用户指令预测结果 (涉及 DALLE 2、ControlNet)，而视频生成模型可以生成未来图像序列 (涉及 Imagen video、Cogvideo、Latte)，这代表了对物理世界的良好理解。
另一方面，扩散模型在机器人控制任务中也显示出有效性，通过在机器人观察条件下对动作去噪，这被称为扩散策略 (详见 Diffusion Policy——斯坦福 UMI 所用的动作预测算法：基于扩散模型的扩散策略)。

虽然扩散生成模型和扩散策略在两个领域中发挥着不同的功能，但作者相信图像预测的能力可以显著增强机器人策略学习，因为它们共享相同的基本物理法则。

之前的工作 (涉及 Learning universal policies via text-guided video generation、Zero-shot robotic manipulation with pretrained image-editing diffusion models、Compositional foundation models for hierarchical planning) 以现成的方式使用图像编辑模型，首先合成目标图像，然后学习目标条件策略。
然而，这种两阶段的方法将预测和动作学习过程分离，忽视了预测和动作之间更深层的联系。这样一来，动作并没有利用预测模型中预训练的表示，这些模型编码了关于物理世界的丰富知识。

故第一方面相关的工作是，针对用于具身控制的预训练，具身控制的视觉 - 语言预训练模型，编码了物理知识，可以从多个方面增强具身控制。首先，预训练模型可以直接作为策略，通过生成高层次计划 (涉及 Do as i can, not as i say: Grounding language in robotic affordances、Code as policies: Language model programs for embodied control) 或产生直接的低层次电机控制信号 (涉及 R3m 等)。许多研究利用预训练的大型语言模型（LLMs）和视觉语言模型（VLMs）的推理能力来创建高层次的计划，然后是运动原语。此外，一些方法通过添加动作头来调整预训练模型以发出低层次的电机控制信号。除了直接作为策略使用之外，预训练模型还可以从多个方面指导策略学习，例如提供良好的表示、提供奖励信号 (涉及 Eureka: Human-level reward design via coding large language models)、合成目标图像，以及预测未来序列。

第二方面相关的工作是：用于具身控制的扩散模型。最近，扩散模型已被采用来解决具身控制中的挑战。一部分研究集中于训练条件扩散模型，以在低维状态输入设置下，根据所需的奖励和约束引导行为合成。 Diffusion Policy 训练了一种基于 RGB 观测的视觉 - 运动策略，并且可以更好地表达多模态动作分布。然而，这些方法从头开始开发任务特定的策略，错过了利用互联网数据进行预训练的好处。另一部分研究利用大规模预训练的扩散模型对训练数据进行数据增强，例如 GenAug、ROSIE 和 CACTI。

第三方面相关的工作是：Policy 学习的未来预测。

比如也有一些工作利用未来图像预测来辅助 Policy 学习。GR-1 使用自回归 transformer 顺序预测未来图像和动作。相比之下，本节的 PAD 采用联合扩散架构，预测更准确的未来图像，可能导致改进的政策学习性能。
UniPi 和 SuSIE 采用两阶段政策学习过程，最初使用扩散生成模型预测未来图像或视频序列，随后基于这些目标图像训练逆动力学模型或低级扩散政策。与这些两阶段方法相比，PAD 的方法具有明显优势。首先，虽然以前的方法利用基于 CNN 的 U-net 主干的扩散模型，主要设计用于图像生成并限于视觉预测，PAD 的方法采用 Diffusion Transformer(DiT) 架构。首先，该架构通过简单的 token 连接和注意力掩码机制巧妙地同时处理多种模态，使得能够同时联合预测未来和动作。其次，使用图像作为预测和动作之间的接口可能无法充分利用预训练扩散模型内部编码的特征。这些两阶段方法的有效性在很大程度上依赖于生成图像的质量。相比之下，PAD 的模型在统一去噪过程中整合图像生成和机器人动作。

2.2 PAD：通过联合去噪过程进行预测与行动

2.2.1 预备知识与 PAD 概述

首先，在模仿学习的情境下，咱们来对基于像素输入且受语言条件制约的机器人控制进行考量 (We consider pixel-input language-conditioned robotic control under the imitation learning setting)。

然而，由于收集机器人数据既有风险又昂贵，D_robot 的规模将受到限制。因此，作者也考虑容易在互联网上获取的 RGB 视频数据集 D_video，D_video 中的一个实例可以表示为 ζj=(Ij)。

首先，表示一个机器人数据集 D_robot={ζ1,ζ2,…ζn}，包含 n 个演示。第 i 个演示 ζi=(Ii,li,τi) 包含一个像素输入序列 Ii、一个自然语言指令 li、机器人位姿序列 p_i^1:T 组成的机器人轨迹 τi。

虽然 D_video 缺乏机器人动作数据，但作者提出的 PAD 框架能够在机器人数据集 D_video 和视频数据集 D_video 上进行协同训练，利用大规模的 D_video 数据来增强视觉策略学习。

其次，扩散模型的核心思想是不断添加高斯噪声以使样本成为高斯分布，并利用去噪过程生成数据 [4-DDPM]。

噪声估计器 εθ(z_t, t, c) 被实现为一个神经网络，并被训练以逼近噪声数据分布的对数密度的梯度 (Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022)，即。

进一步地，去噪过程从最噪声的潜在样本 z_T 开始，并逐步减少噪声以恢复带有条件 c 的真实样本 z_0。这基于给定条件 q(z_{t-1}|z_t,c) 概率的变分近似：

$\begin{aligned} p\left(z_{t-1} \mid z_{t}, c\right) & =\mathcal{N}\left(z_{t-1} ; \sqrt{\bar{\alpha}{t-1}} \mu{\theta}\left(z_{t}, t, c\right),\left(1-\bar{\alpha}{t-1}\right) \mathbb{I}\right) \ \mu{\theta}\left(z_{t}, t, c\right) & =\left(z_{t}-\sqrt{1-\bar{\alpha}{t}} \epsilon{\theta}\left(z_{t}, t, c\right)\right) / \sqrt{\bar{\alpha}_{t}} . \end{aligned}$

令 z_0=ε(x_0) 表示从真实数据编码的潜在样本。加噪过程逐渐在 T 步中向 z_0 添加标准高斯噪声 (N)，从而得到一组噪声样本 Z={z_t|t∈[1,T]}，这相当于从以下分布中采样：

$q\left(z_{t} \mid z_{t-1}\right)=\mathcal{N}\left(z_{t} ; \sqrt{\alpha_{t}} z_{t-1},\left(1-\alpha_{t}\right) \mathbb{I}\right)$

其中 {α_t|t∈[1,T]} 是预定义的超参数，用于控制噪声的幅度。

令 \bar{α}t=∏{i=1}^{t} α_i，根据 DDIM [5]，z_t 可以通过向 z_0 添加高斯噪声 ε_t 直接获得：

$z_{t}=\sqrt{\bar{\alpha}{t}} z{0}+\sqrt{1-\bar{\alpha}{t}} \epsilon{t}$

$\epsilon_{\theta}\left(z_{t}, t, c\right) \approx-\sqrt{1-\bar{\alpha}{t}} \nabla{z_{t}} \log p\left(z_{t} \mid c\right)$

而对于本文的 PAD 框架，该框架在联合潜在去噪过程中同时预测未来帧和动作。

且在提出的 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，然后触发一个新的预测周期。这个迭代预测和执行过程使机器人能够以闭环方式持续规划和行动。

作者主要关注 RGB 图像模态 M_I 和机器人动作模态 M_A，每个机器人动作可以通过一个机器人姿态来表征，包括末端执行器的位置和旋转，以及夹爪状态。值得注意的是，该框架可以轻松扩展到额外的模态 M_E——例如，作者在实验部分中额外加入了深度图像模态，这提供了更准确的距离测量。

2.2.2 PAD 的模型架构：选择 DiT 作为模型骨干然后联合去噪预测未来图像和动作

鉴于原始数据可能具有多种格式和高维度，作者首先将所有模态映射到潜在空间，并进行潜在扩散过程。根据 Diffusion Transformer(DiT)[20] 中的过程，RGB 图像 x_I 最初通过预训练的、冻结的 VAE[23] 编码器 ε_I 处理，以获得潜在表示 ε_I(x_I)。然后，通过分词器将此潜在表示转换为嵌入大小为 h 的 token 序列 t_I。

类似地，自然语言指令通过冻结的 CLIP 编码器处理，以生成文本嵌入 c_l。如果可用，深度图像将被下采样并分词为 t_E。机器人姿态 x_A 使用多层感知器 (MLP) 编码为 ε_A(x_A)，并线性转换为具有相同嵌入大小 h 的 token t_A。

对于 Diffusion Transformer(DiT) 骨干网络，作者采用了 Diffusion Transformer(DiT)[20] 作为 PAD 的模型骨干，这相较于以往工作中常用的 U-net 骨干网络提供了若干优势。

此外，DiT 架构能够适应缺失的模态。例如，在缺少机器人动作的视频数据集中，DiT 的输入仅包含图像 token t_I，作者只需将 token 序列扩展到组合长度 T_I+T_A+T_E，并在自注意力块中引入注意力掩码以排除填充 token。输出中仅保留有效的预测，丢弃任何填充部分。

特别是，DiT 架构通过自注意力机制有效地整合了各种模态。输入如 RGB 图像、机器人姿态和附加数据——被转化为长度为 T_I, T_A, T_E 的 token 序列 t_I, t_A, t_E。这些来自不同模态的 token 序列被连接在一起，并经历一个联合的潜在去噪过程。整个过程的简要示意图如下图图 3 右侧所示。

文章配图

这一设计选择使得 PAD 能够同时在仅有 RGB 的视频数据集和机器人数据集上进行训练。

对于联合条件生成，作者将未来观测初始化为白噪声，并旨在重建未来观测帧和期望的机器人动作，基于当前观测 c_I, c_A, c_E 进行条件化，遵循与 [26-Instructpix2pix: Learning to follow image editing instructions] 中类似的策略，作者在通道维度上连接条件潜变量和噪声潜变量。将潜在变量连接后，这些来自不同模态的条件化噪声潜在变量被 tokenized 为相同嵌入大小的 token 序列 t_I, t_A, t_E。图像潜在变量 L_I 的 tokenization 遵循与 [20-Diffusion Transformer(DiT)] 相同的切块过程，而机器人姿态的 tokenization 采用简单的线性投影。最后，这些 token 被输入到多层 DiT 中，以预测未来帧的潜在表示。

例如，如果编码潜变量 ε_I(c_I) 的形状为 c×d×d，那么 z_t^{I'} 的形状将为 kc×d×d 以表示 k 个未来帧，最终的潜变量 L_I 的形状为 (k+1)c×d×d，其他模态经历类似的过程。

具体来说，在获得编码潜变量 ε_I(c_I), ε_A(c_A), ε_E(c_E) 后，作者将这些潜变量与噪声连接以获得条件噪声潜变量 L_I=[ε_I(c_I), z_t^I], L_A=[ε_A(c_A), z_t^A], L_E=[ε_E(c_E), z_t^E]。整体过程的示意图见下图图 3。

文章配图

2.2.3 训练过程

首先，对于初始化。按照 [15-Latte] 中的初始化过程，作者还从'在 ImageNet 上针对基于类别的图像生成任务进行预训练的 DiT 模型'中初始化了 PAD 权重 (We also initialize the PAD weights from the DiT model pre-trained on ImageNet for the image generation task conditioned on class [20].)。然而，由于有缺失或不兼容的模型参数，不能直接加载模型。故舍弃 DiT 中的标签嵌入层，并对新的文本嵌入层进行零初始化，且将图像潜在 tokenizer 的权重复制 k + 1 次以编码堆叠的潜在变量，并且机器人状态的编码器和解码器也进行零初始化。

其次，对于训练目标。扩散过程向目标编码的潜在变量 {ε_I(x_I), ε_A(x_A), ε_E(x_E)} 添加噪声，并产生噪声潜在变量 Z_{I,A,E}={z_t^I, z_t^A, z_t^E}。别忘了，上面所说的——PAD 框架同时输出相应的未来预测，和机器人动作。

作者训练 PAD 模型以同时预测添加到样本数据上的噪声 ε^I, ε^A, ε^E，这些噪声是基于当前观察 C_{I,A,E}={c_I, c_A, c_E} 和指令 l 的条件。这个去噪器使用 DDPM [5] 损失进行训练：

$\mathcal{L}{\text {diff }}^{\delta}(\theta)=\mathbb{E}{\epsilon^{\delta} \sim \mathcal{N}(0,1), t, C, l}\left[\left|\epsilon^{\delta}-\epsilon_{\theta}^{\delta}\left(z_{t}^{\delta}, t, C, l\right)\right|_{2}^{2}\right]$

其中 δ∈{I,A,E} 表示不同类型的输入模态。去噪损失 \mathcal{L}_{diff}^{δ} 旨在最大化证据下界 (ELBO)[5-DDIM]，同时近似条件分布 p(ε_δ(x_δ)|C,l)。且作者共同最小化以下潜在扩散目标，并使用超参数 λ_I, λ_A, λ_E 来平衡不同模态之间的预测损失。

从而最终的训练目标为：

$\mathcal{L}(\theta)=\lambda_{I} \mathcal{L}{\text {diff }}^{I}+\lambda{A} \mathcal{L}{\text {diff }}^{A}+\lambda{E} \mathcal{L}_{\text {diff }}^{E}$

2.3 实验

2.3.1 策略训练的细节与相关 baseline

首先，对于策略训练的细节。正如上文详细描述的那样，灵活的 PAD 框架可以在各种互联网 RGB 视频数据和机器人演示上进行联合训练。

为了节省计算资源并避免在每个机器人领域中从头开始联合训练模型的需要，作者首先在互联网数据上预训练模型，以建立更好的图像预测先验。
然后，将这个预训练的模型适应于各种机器人领域，包括模拟的 Metaworld 和现实世界的熊猫操作。

根据经验，首先在 BridgeData-v2 数据集 [9] 上预训练了 20 万步，该数据集包含 60,000 个轨迹。之后，将模型适应于每个领域，继续使用机器人演示进行额外的 10 万步训练。预训练和适应阶段大约需要 2 天和 1 天，使用 4 个 NVIDIA A100 GPU。

此外，作者发现，在早期适应阶段增加图像预测损失的权重可以加速收敛，因为图像先验已经在预训练模型中建立。具体来说，作者在整个训练过程中将图像预测损失系数 λ_I 保持在 1.0，并在 100k 训练步骤中将 λ_A 和 λ_E 从 0.0 线性增加到 2.0。

其次，对于策略执行细节。作者的策略以当前图像 c_I 和机器人姿态 c_A 为条件，预测未来的 k 帧图像和动作。具体而言，作者将预测范围配置为 k = 3，并将帧间间隔设置为 i = 4，适用于 Metaworld 和现实世界任务。

然后，机器人使用简单的线性插值运动规划器移动到第一个期望姿态，触发下一次预测循环。

在策略执行过程中，作者使用 75 步的 DDIM 采样 [5] 来对未来图像 x_I^{1:K} 和动作 x_A^{1:k} 的 k 步进行去噪。这些 k 步预测可以视为 k 步规划，只有第一个预测动作 x_A^1 由机器人执行。

在实验的过程中，作者选择与以下方法进行比较：

扩散策略 [16]——一种新颖的视觉控制策略，通过动作扩散器生成机器人动作。考虑到扩散策略在所有任务上使用单一文本条件视觉策略进行训练。故作者对原始扩散策略模型进行了增强，添加了指令条件以解决多任务设置。且使用 CLIP 编码器 [25-Learning transferable visual models from natural language supervision. In International conference on machine learning] 作为指令编码器，参考相关工作 [29-Scaling up and distilling down: Language-guided robot skill acquisition]。
SuSIE [18] 一种两阶段方法，利用预训练的图像编辑模型 [26] 为机器人任务生成图像目标，然后是目标条件的低级扩散策略。作者在相同的数据集上微调图像编辑扩散模型，并使用扩散策略进行目标条件行为克隆。为了确保公平比较，作者也使用更强大的 DiT 框架作为图像编辑模型。
RT-1 [30] 一种端到端的机器人控制策略，利用 FiLM 条件化 [31] 的 EfficientNet[32] 融合视觉输入和语言输入，然后通过 transformer 块输出动作。
RT-2*33 一个大规模的具身模型，直接微调视觉 - 语言模型（VLMs）以生成机器人动作。原始 RT-2 模型在 PaLM 模型 [34] 上进行了微调，该模型未公开。根据原始论文中概述的规格，作者使用 InstructBlip-7B[35] 骨干重新实现了 RT-2 模型。
GR-1 [36] 一种利用图像预测来辅助策略学习的方法。与 PAD 不同，他们通过自回归架构生成图像和动作。

在实际实验中，PAD 比 GR-1 方法预测出更精确的未来图像（下图图 6），这可能是因为扩散模型在图像生成任务中的卓越能力。

文章配图

这些精确的图像可能更有效地促进策略学习，尤其是在需要精确和准确操作的任务中，如在 Metaworld 中拾取小块、插入、篮球等，导致更高的成功率。

2.3.2 消融研究

由于可适应的 DiT 架构，他们的框架能够容纳额外的模态。在实际操控实验中，作者结合了额外的深度图像输入，并联合预测未来的 RGB 图像、深度图像和机器人动作，称为 PAD-depth。

且观察到在 PAD 的联合去噪框架下，不同模态之间的预测结果高度一致，部分结果如图 9 所示。

文章配图

深度输入的加入提升了操控任务的性能，如表 2 所示。这种改进可能源于深度信息的精确预测，这有助于代理识别距离变化，从而提升性能。

文章配图

此外，PAD 框架可以扩展到预测其他与机器人控制相关的模态，如触觉力或点云，这部分留待未来工作。

// 待更

Diffusion Transformer(DiT)：用 ViT 替代 U-Net 用于视频生成与机器人动作预测