本文目录

1 扩散模型解读 (六)：U-ViT：基于 ViT 的扩散模型的先驱之作
(来自清华大学，中国人民大学，北京智源)
1 U-ViT 论文解读
1.1 使用 ViT 实现扩散模型
1.2 扩散模型简介
1.3 U-ViT 的具体做法
1.4 深度、宽度和 Patch Size 的影响
1.5 实验结果

太长不看版

OpenAI Sora [1] 酷炫的效果及其与 Transformer 架构的关联使得领域开始关注基于 Transformer 的架构的扩散模型 Diffusion Transformer (DiT) [2] 。DiT 的特点是使用 vanilla 的 Transformer 架构，即相同的 Transformer Block 从头到尾，顺序排列，没有大 Shortcut 连接，极其简单的宏观架构。

本文的提出比 DiT 还要早一点，而且本文是用 ViT 架构实现 Diffusion 模型的开山之作，使用 ViT 来替换 Diffusion 模型中基于卷积的 U-Net 架构，同时维持模型宏观的 U 形架构不变。本文的模型因此命名 U-ViT，其特点是将 Diffusion 模型中的一些必要的信息，比如 time, condition 以及 noisy image patches 全部视为是 tokens，而且在浅层和深层之间使用 long skip connection 连接。

U-ViT 评测的任务包括：unconditional 或者 class-conditional 图像生成任务，文生图任务。U-ViT 的性能与相似尺寸的 U-Net 相当。U-ViT 的 latent 扩散模型在 ImageNet 256×256 上的 class-conditioned 图像生成中实现了 2.29 的 FID，在 MS-COCO 上的文生图任务中实现了 5.48 的 FID，同时没有使用大型外部数据集。

U-ViT 还给出了很多有价值的观察，比如对于基于扩散模型的图像建模，long skip connection 很重要，但是上下采样就没那么重要。

1 U-ViT：基于 ViT 的扩散模型的先驱之作

论文名称：All are Worth Words: A ViT Backbone for Diffusion Models (CVPR 2023)

论文地址：

http://arxiv.org/pdf/2209.12152

代码链接：

http://github.com/baofff/U-ViT

1.1 使用 ViT 实现扩散模型

Diffusion 模型在文生图任务 [3][4] ，视频生成任务 [5] 中的应用迅速增长。模型骨干架构的发展起到了至关重要的作用。基于 CNN 的 U-Net [6] 的特点是一组下采样块、一组上采样块以及两组之间的 long skip connections，这种架构主导了图像生成任务的扩散模型。另一方面，Vision Transformer (ViT [7] ) 在各种视觉任务中展示出了前景。很多基于 ViT 的视觉架构甚至优于基于 CNN 的方法。因此，出现了一个非常自然的问题： 基于 CNN 的 U-Net 的依赖是否在扩散模型中是必要的？

本文提出了一个简单而通用的基于 ViT 的架构，称为 U-ViT，如图 1 所示。按照 Transformer 的设计方法，U-ViT将所有输入 (包括时间、条件和噪声图像 Patches) 视为 tokens。U-ViT 在受 U-Net 启发的浅层和深层之间采用了 long skip connections。low-level 的特征对扩散模型中的像素级预测目标很重要，这种连接可以使得预测噪声的网络的训练更加容易。而且，U-ViT 在输出之前可以添加额外的 3×3 卷积以获得更好的视觉质量。

1.2 扩散模型简介

扩散模型逐渐向数据注入噪声，然后反转这个过程以从噪声中生成数据。噪声注入过程，也称为前向过程，形式化为马尔可夫链：

式中, 为数据, 和 表示 noise schedule,且满足 。为了反转这个过程, 使用一个高斯模型 来逼近 Ground Truth 的反向过程 。且最优的均值为:

式中, 是注入 的标准高斯噪声。这个时候, 神经网络的学习就可以等效为一个预测噪声的过程。形式上, 通过最小化噪声预测目标, 即 来学习 。其中 在 1 和 之间均匀分布。为了学习条件扩散模型, 例如类条件模型或文生图模型, 条件信息被进一步馈送到噪声预测目标：

其中, 是条件或其连续嵌入。在以往的图像建模工作中, 扩散模型主要依赖于基于 CNN 的 U , 这是一个卷积主干, 其特征是有一组下采样块、一组上采样块和两组之间的长跳跃连接。 通过自适应组归一化 和交叉注意 等机制输入给 U-Net 中。

1.3 U-ViT 的具体做法

U-ViT 的架构如图1 所示。U-ViT 代表的就是式 3 中的 。U-ViT 以时间 、条件 和噪声图像 作为输入并预测注入 的噪声。按照 ViT 的设计方法, 将图像分割成小块, U-ViT 将所有输入 (包括时间、条件和图像 patches) 全部视为 tokens。U-ViT 在浅层和深层之间也采用了类似的 long skip connections。

式 3 很像一个像素级预测任务，并且对低级特征很敏感。long skip connections 为低级特征提供了快捷的连接，从而简化了噪声预测网络的训练。

此外，U-ViT 在输出之前可以添加一个 3×3 卷积。目的是防止变压器产生的图像中潜在的伪影。根据本文实验，这个 3×3 卷积提高 U-ViT 生成样本的质量。

尽管 U-ViT 在概念上很简单，但作者还是精心设计了它的实现。为此，作者对 U-ViT 中的关键设计进行了比较系统的实证性研究。作者在 CIFAR10 上做消融实验，在 10K 生成的样本 (不是生成 50K 样本，为了提高效率) 上评估 FID 分数。

1) 结合 long skip 分支的方法

设 分别为主分支和 long skip 分支的嵌入。作者考虑了几种把二者结合馈送到下一个 Transformer Block 的方法:

1. 将它们 Concat 起来, 然后过个线性投影, 如图1所示, 即 Linear 。
2. 直接把它们加起来, 即 。
3. 对 执行线性投影，然后把它们加起来，即 。
4. 加起来, 然后执行线性投影, 即 Linear 。
5. 也对比了不使用 long skip connections 的情况。

实验结果如图 2 所示, 第 1 种方案的表现最好。直接加起来的表现还不如不用 long skip connections。作者分析原因是 Transformer 本身就有 skip connection, 因此, 的唯一影响是以线性形式增加 的系数, 这不会改变网络的性质。相比之下, 结合 的所有其他方法在 上执行线性投影, 并且在没有 long skip connections 的情况下提高了性能。

2) AdaLN 的使用

作者考虑了 2 种方法将 输入网络：

1. 将其视为 tokens。
2. 在 Transformer Block 中的 LayerNorm 之后使用自适应组归一化 (adaptive group normalization) 。这个方法称为自适应层归一化 (AdaLN)。形式上表述为: , 其中 是 Transformer Block 的输入, 是从 time embedding 中获得的。这种方法虽然简单, 但是将时间视为 token 的简单方法比 AdaLN 表现更好。

3) 在 Transformer 之后添加额外卷积

作者考虑了 3 种方法在 Transformer 之后添加额外卷积：

1. 在线性投影后添加一个 3×3 卷积，将 token embedding 映射到 image patches。
2. 在线性投影之前添加一个 卷积块, 这需要首先将 token embedding 的 1D 序列 重新排列为形状 的 2D 特征, 其中 是 Patch Size。
3. 不使用 3×3 卷积块。

如图 4 所示，在线性投影后添加一个 3×3 卷积效果略微好一些。

4) Patch Embedding 的变体

作者考虑了 Patch Embedding 的 2 种变体。

1. 原始的 Patch Embedding 采用线性投影，把 patches 映射为 token embedding。
2. 使用 3×3 卷积的堆叠，然后使用 1×1 卷积把图像映射到 token embedding。

如图 5 所示，原始的 Patch Embedding 表现更好。

5) 位置编码的变体

作者考虑位置编码的 2 种变体。

1. 原始 ViT 中提出的一维可学习位置编码，这是本文中的默认设置。
2. 二维正弦位置编码，它是通过连接位置 处的 Patch 和 的正弦 Embedding 获得的。

如图 6 所示，一维可学习位置编码的表现更好。作者还尝试使用了任意位置编码，发现模型无法生成有意义的图像，这意味着位置信息在图像生成中至关重要。