Diffusion Transformer(DiT)——将扩散过程中的U-Net换成ViT：近频繁用于视频生成与机器人动作预测(含清华PAD详解)

前言

本文最开始属于此文《视频生成Sora的全面解析：从AI绘画、ViT到ViViT、TECO、DiT、VDT、NaViT等》

但考虑到DiT除了广泛应用于视频生成领域中，在机器人动作预测也被运用的越来越多，加之DiT确实是一个比较大的创新，影响力大，故独立成本文

第一部分 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，Peebles在𝕏上用昵称Bill，在Linkedin上及论文署名时用大名William)、Saining Xie(当时在纽约大学)的两人通过论文《Scalable Diffusion Models with Transformers》提出了一种叫 DiT 的神经网络结构

其结合了视觉 transformer 和 diffusion 模型的优点，即DiT = DDPM + ViT
但它把DPPM中的卷积架构U-Net换成了Transformer架构(We train latent diffusion models of images, replacing the commonly-used U-Net backbone with a transformer that operates on latent patche，至于什么是U-Net，请参见此文的2.1节)

1.1.2 Diffusion Transformer(DiT)的3个不同的条件策略

如下图所示，便是扩散transformer(DiT)的架构

左侧：训练条件潜在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
右侧：DiT blocks的细节，作者试验了标准transformer块的变体，这些变体通过自适应层归一化、交叉注意和额外输入token来加入条件(incorporate conditioning via adaptive layer norm, cross-attention and extra input tokens，这个coditioning相当于就是带条件的去噪)，其中自适应层归一化效果最好

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

自适应层归一化，即Adaptive layer norm (adaLN) block
鉴于自适应归一化层在GANs和具有U-Net骨干的扩散模型中的广泛使用，故用自适应层归一化(adaLN)替换transformer块中的标准层归一化层
不是直接学习维度方向的缩放和偏移参数 $\gamma$ 和 $\beta$ ，而是从 $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
adaLN-Zero block
关于ResNets的先前工作发现，将每个残差块初始化为恒等函数是有益的。例如，Goyal等人发现，在每个块中将最终批量归一化尺度因子 γ零初始化可以加速大规模训练在监督学习设置中[13]
扩散U-Net模型使用类似的初始化策略，在任何残差连接之前将每个块的最终卷积层零初始化。作者对adaLN DiT块的修改，它做了同样的事情。除了回归 γ和 β，还回归在DiT块内的任何残差连接之前立即应用的维度方向的缩放参数 α
交叉注意力块
将t和c的嵌入连接成一个长度为二的序列，与图像token序列分开，transformer块被修改为：在多头自注意块之后，包含一个额外的多头交叉注意层，类似于LDM用于根据类标签进行条件处理的设计。交叉注意力使模型增加了最多的Gflops，大约增加了15%的开销
上下文条件化(In-context conditionin)
将 $t$ 和 $c$ 的向量嵌入作为两个额外的token追加到输入序列中，与图像token无异地对待它们，这类似于ViTs中的 cls token，它允许我们使用标准ViT块而无需修改。在最后一个块之后，从序列中移除条件token。这种方法没有新增模型多少Gflops，可忽略

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

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

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

首先，需要设计一定的模型结构，用来支持不同长宽比和分辨率的视频
以保证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导致的每帧不同分辨率和长宽比的问题
第二，需要把图片生成模型改成视频生成模型，本质就是加上一个时间的维度
具体而言，需要加入一个Casual Time Attention子模块
Causal Time Attention模块的作用是在生成第 $i$ 帧的时候，收集历史Time信息，也就是通过Attention让第 $i$ 帧看到之前的比如 $k$ 帧内容，这是用来维护生成视频的时间一致性的
最终，可以得到类似下文即将介绍的VDT的一个结构，如下图所示(图源：张俊林)
首先，把噪音Patch线性化后，并入Text Prompt Condition和Time Step Condition，一起作为Transformer的输入
其次，Transformer内部由三个子模块构成：
$\rightarrow$ Local Spatial Attention模块负责收集视频帧空间信息
$\rightarrow$ Causal Time Attention模块负责收集历史时间信息
$\rightarrow$ MLP模块负责对时间和空间信息通过非线性进行融合

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

// 待更

1.3 相关工作

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

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

其实早在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参数化了噪声预测网络 $\boldsymbol{\epsilon}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t, c\right)$ ，它接受时间 $t$ 、条件 $c$ 、噪声图像 $x_t$ 作为输入，然后预测注入到 $x_t$ 中的噪声。按照ViT的设计方法，图像被分割成patch，并且U-ViT将所有输入(包括时间、条件和图像patch)都视为token
受到基于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)
意思是说， $\boldsymbol{\epsilon}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t, c\right)$ 的目标是像素级预测任务，对低级特征敏感。长跳跃连接为低级特征提供了捷径，因此有助于噪声预测网络的训练
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在输出之前可选地添加一个3×3的卷积块。这旨在防止transformer生成的图像中出现潜在的伪影(potential artifacts in images)

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

均是提出了将 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改造而成)：先预训练后微调，支持语言、图像、动作多种输入》所述，RDT所用的扩大大模型架构便是基于DiT改造
且再如此文《视觉语言动作模型VLA的持续升级：从π0之参考基线Octo、OpenVLA到CogACT、TinyVLA》的第四部分所述，CogACT是VLM作为认知基础，DiT作为动作模块

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

其对应的论文为：Prediction with Action: Visual Policy Learning via Joint Denoising Process
其对应的作者为：郭延江12∗, 胡宇成13∗, 张建科1, 王彦仁14, 陈晓宇12, 陆超超3†, 陈剑宇12†

具体而言，如下图

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

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

2.1.2 提出背景与相关工作

最近，扩散模型[4–6，涉及DDPM、DDIM、inproved DDPM，关于这几个的模型的介绍，详见此文《图像生成发展起源：从VAE、扩散模型DDPM、DDIM到DETR、ViT、Swin transformer》]通过在大规模网络图像和视频数据集上进行训练，在视觉生成任务中取得了令人瞩目的成功[7–9，涉及Imagenet、Bridgedata v2等]

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

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

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

故第一方面相关的工作是，针对用于具身控制的预训练，具身控制的视觉-语言预训练模型，编码了物理知识，可以从多个方面增强具身控制。首先，预训练模型可以直接作为策略，通过生成高层次计划 [37–41，涉及Do as i can, not
as i say: Grounding language in robotic affordances、Code as policies: Language model programs for embodied control] 或产生直接的低层次

比如电机控制信号[30,33,42–44]。许多研究利用预训练的大型语言模型（LLMs）和视觉语言模型（VLMs）的推理能力来创建高层次的计划，然后是运动原语

此外，一些方法通过添加动作头来调整预训练模型以发出低层次的电机控制信号

除了直接作为策略使用之外，预训练模型还可以从多个方面指导策略学习，例如提供良好的表示[45–47，涉及R3m等]、提供奖励信号[48–50，涉及 Eureka: Human-level reward design
via coding large language models]、合成目标图像[18- Zero-shot robotic manipulation with pretrained image-editing diffusion models, 51]，以及预测未来序列[17-Learning universal policies via text-guided video generation]

第二方面相关的工作是：用于具身控制的扩散模型。最近，扩散模型已被采用来解决具身控制中的挑战。一部分研究集中于训练条件扩散模型，以在低维状态输入设置下，根据所需的奖励和约束引导行为合成[52–55]

Diffusion Policy[16]训练了一种基于RGB观测的视觉-运动策略，并且可以更好地表达多模态动作分布。然而，这些方法从头开始开发任务特定的策略，错过了利用互联网数据进行预训练的好处

另一部分研究利用大规模预训练的扩散模型对训练数据进行数据增强，例如GenAug[56]、ROSIE[57]和CACTI[58]

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

比如也有一些工作利用未来图像预测来辅助Policy学习。GR-1[36] 使用自回归transformer顺序预测未来图像和动作
相比之下，本节的PAD采用联合扩散架构，预测更准确的未来图像，可能导致改进的政策学习性能
UniPi[17]和SuSIE [18] 采用两阶段政策学习过程，最初使用扩散生成模型预测未来图像或视频序列，随后基于这些目标图像训练逆动力学模型或低级扩散政策

与这些两阶段方法相比，PAD的方法具有明显优势。首先，虽然以前的方法利用基于CNN 的 U-net 主干的扩散模型 [23]，主要设计用于图像生成并限于视觉预测，PAD的方法采用Diffusion Transformer(DiT) 架构 [20]

首先，该架构通过简单的token连接和注意力掩码机制巧妙地同时处理多种模态，使得能够同时联合预测未来和动作
其次，使用图像作为预测和动作之间的接口可能无法充分利用预训练扩散模型内部编码的特征。这些两阶段方法的有效性在很大程度上依赖于生成图像的质量。相比之下，PAD的模型在统一去噪过程中整合图像生成和机器人动作

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

2.2.1 预备知识与PAD概述

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

首先，表示一个机器人数据集 $D_{\text {robot }}=\left\{\zeta_{1}, \zeta_{2}, \ldots \zeta_{n}\right\}$ ，包含 $n$ 个演示
第 $i$ 个演示 $\zeta_{i}=\left(I_{i}, l_{i}, \tau_{i}\right)$ 包含一个像素输入序列 $I_{i}$ 、一个自然语言指令 $l_{i}$ 、机器人位姿序列 $p_{i}^{1: T}$ 组成的机器人轨迹 $\tau_{i}$
然而，由于收集机器人数据既有风险又昂贵， $D_{\text {robot }}$ 的规模将受到限制。因此，作者也考虑容易在互联网上获取的RGB 视频数据集 $D_{\text {video }}$ ， $D_{\text {video }}$ 中的一个实例可以表示为 $\zeta_{j}=\left(I_{j}\right)$

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

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

令 $z_{0}=\varepsilon\left(x_{0}\right)$ 表示从真实数据编码的潜在样本
加噪过程逐渐在 $T$ 步中向 $z_{0}$ 添加标准高斯噪声 $(\mathcal{N})$ ，从而得到一组噪声样本 $Z=\left\{z_{t} \mid t \in[1, T]\right\}$ ，这相当于从以下分布中采样： $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)$
其中 $\left\{\alpha_{t} \mid t \in[1, T]\right\}$ 是预定义的超参数，用于控制噪声的幅度

令 $\bar{\alpha}_{t}=\prod_{i=1}^{t} \alpha_{i}$ ，根据DDIM [5]， $z_{t}$ 可以通过向 $z_{0}$ 添加高斯噪声 $\epsilon_{t} \mathrm{t}$ 直接获得： $z_{t}=\sqrt{\bar{\alpha}_{t}} z_{0}+\sqrt{1-\bar{\alpha}_{t}} \epsilon_{t}$
进一步地，去噪过程从最噪声的潜在样本 $z_{T}$ 开始，并逐步减少噪声以恢复带有条件 $c$ 的真实样本 $z_{0}$
这基于给定条件 $q\left(z_{t-1} \mid z_{t}, c\right)$ 概率的变分近似：
$\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}$
噪声估计器 $\epsilon_{\theta}\left(z_{t}, t, c\right)$ 被实现为一个神经网络，并被训练以逼近噪声数据分布的对数密度的梯度[22-Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022]，即

$\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 框架，该框架在联合潜在去噪过程中同时预测未来帧和动作

作者主要关注RGB 图像模态 $M_{I}$ 和机器人动作模态 $M_{A}$ ，每个机器人动作可以通过一个机器人姿态来表征，包括末端执行器的位置和旋转，以及夹爪状态
值得注意的是，该框架可以轻松扩展到额外的模态 $M_{E}$ ——例如，作者在实验部分中额外加入了深度图像模态，这提供了更准确的距离测量
且在提出的PAD 框架中，预测和动作是基于多模态当前观测进行的，包括RGB图像 $c_{I}$ 、机器人姿态 $c_{A}$ 、额外的深度图 $c_{E}$ (在现实世界任务中)，以及自然语言指令文本 $l$

该框架同时输出相应的未来预测 $x_{I}$ , $x_{E}$ 和机器人动作 $x_{A}$ 「为方便大家一目了然，这三种输出，我在上面分别用了不同的字体颜色表示——即紫色的图像模态 $x_{I}$ 、绿色的深度图像模态 $x_{E}$ 、黄色的机器人动作模态 $x_{A}$ 」
且PAD 不仅预测单一未来步骤，还可以预测 $k$ 个未来步骤 $x_{I}^{1: k}, x_{A}^{1: k}, x_{E}^{1: k}$ ，这可以看作是机器人的k 步规划
比如机器人只执行第一个预测动作 $x_{A}^{1}$ ，然后触发一个新的预测周期。这个迭代预测和执行过程使机器人能够以闭环方式持续规划和行动

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

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

类似地

机器人姿态 $x_{A}$ 使用多层感知器(MLP)[24-Martin Riedmiller and A Lernen. Multi layer perceptron. Machine Learning Lab Special Lecture, University of Freiburg, 24, 2014]编码为 $\varepsilon_{A}\left(x_{A}\right)$ ，并线性转换为具有相同嵌入大小 $h$ 的token $t_{A}$
如果可用，深度图像将被下采样并分词为 $t_{E}$
自然语言指令通过冻结的CLIP 编码器[25-Learning transferable visual models from natural language supervision，CLIP的介绍详见《AI绘画原理解析：从CLIP、BLIP到DALLE、DALLE 2、DALLE 3、Stable Diffusion(含ControlNet详解)》的第一部分] 处理，以生成文本嵌入 $c_{l}$

对于Diffusion Transformer(DiT)骨干网络，作者采用了Diffusion Transformer(DiT)[20] 作为PAD的模型骨干，这相较于以往工作中常用的U-net 骨干网络[18-Zero-shot robotic manipulation with pretrained image-editing diffusion models, 17-Learning universal policies via text-guided video generation] 提供了若干优势

特别是，DiT 架构通过自注意力机制有效地整合了各种模态。输入如RGB 图像、机器人姿态和附加数据——被转化为长度为 $T_{I}, T_{A}, T_{E}$ 的token序列 $t_{I}, t_{A}, t_{E}$
这些来自不同模态的token序列被连接在一起，并经历一个联合的潜在去噪过程
此外，DiT 架构能够适应缺失的模态。例如，在缺少机器人动作的视频数据集中，DiT 的输入仅包含图像token $t_{I}$ ，作者只需将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] 中类似的策略，作者在通道维度上连接条件潜变量和噪声潜变量

具体来说，在获得编码潜变量 $\varepsilon_{I}\left(c_{I}\right), \varepsilon_{A}\left(c_{A}\right), \varepsilon_{E}\left(c_{E}\right)$ 后，作者将这些潜变量与噪声连接以获得条件噪声潜变量 $L_{I}=\left[\varepsilon_{I}\left(c_{I}\right), z_{t}^{I}\right],L_{A}=\left[\varepsilon_{A}\left(c_{A}\right), z_{t}^{A}\right], L_{E}=\left[\varepsilon_{E}\left(c_{E}\right), z_{t}^{E}\right]$
例如，如果编码潜变量 $\varepsilon_{I}\left(c_{I}\right)$ 的形状为 $c \times d \times d$ ，那么 $z_{t}^{I^{\prime}}$ 的形状将为 $k c \times d \times d$ 以表示 $k$ 个未来帧，最终的潜变量 $L_{I}$ 的形状为 $(k+1) c \times d \times d$ ，其他模态经历类似的过程
将潜在变量连接后，这些来自不同模态的条件化噪声潜在变量被tokenized为相同嵌入大小的token序列 $t_{I}, t_{A}, t_{E}$
图像潜在变量 $L_{I}$ 的tokenization遵循与[20-Diffusion Transformer(DiT)] 相同的切块过程，而机器人姿态的tokenization采用简单的线性投影
最后，这些token被输入到多层DiT 中，以预测未来帧的潜在表示