综述由AI生成Diffusion Transformer (DiT) 将扩散模型中的 U-Net 替换为 Vision Transformer (ViT),显著提升了可扩展性。文章详细解析了 DiT 的三种条件策略,并探讨了其在视频生成领域的改造方案,如引入时间注意力机制。此外,重点介绍了 DiT 在机器人动作预测中的应用,特别是清华大学提出的 PAD 框架,该框架通过联合去噪过程同时预测未来图像和机器人动作,利用大规模视频数据增强视觉策略学习,相比传统两阶段方法具有更高的一致性和性能。
DataScient18 浏览
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。
但它把 DPPM 中的卷积架构 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 上做两项改造:
首先,把噪音 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)。
意思是说,εθ(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.
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)。
之前的工作 (涉及 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) 或产生直接的低层次,比如电机控制信号 (30,33,42–44)。许多研究利用预训练的大型语言模型(LLMs)和视觉语言模型(VLMs)的推理能力来创建高层次的计划,然后是运动原语。
此外,一些方法通过添加动作头来调整预训练模型以发出低层次的电机控制信号。
除了直接作为策略使用之外,预训练模型还可以从多个方面指导策略学习,例如提供良好的表示 (涉及 R3m 等)、提供奖励信号 (涉及 Eureka: Human-level reward design via coding large language models)、合成目标图像 (Zero-shot robotic manipulation with pretrained image-editing diffusion models, 51),以及预测未来序列 (Learning universal policies via text-guided video generation)。
如果可用,深度图像将被下采样并分词为 t_E。
机器人姿态 x_A 使用多层感知器 (MLP) (Martin Riedmiller and A Lernen. Multi layer perceptron. Machine Learning Lab Special Lecture, University of Freiburg, 24, 2014) 编码为 ε_A(x_A),并线性转换为具有相同嵌入大小 h 的 token t_A。
对于 Diffusion Transformer (DiT) 骨干网络,作者采用了 Diffusion Transformer (DiT)[20] 作为 PAD 的模型骨干,这相较于以往工作中常用的 U-net 骨干网络 (Zero-shot robotic manipulation with pretrained image-editing diffusion models, 17-Learning universal policies via text-guided video generation) 提供了若干优势。
首先,对于初始化。
按照 [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 次以编码堆叠的潜在变量,并且机器人状态的编码器和解码器也进行零初始化。
其次,对于策略执行细节。
作者的策略以当前图像 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 框架作为图像编辑模型。