#U-ViT
别只盯着DiT，国产模型也在暗自发力！U-ViT：ViT 架构实现 Diffusion 的开山之作！

本文的提出比 DiT 还要早一点，而且本文是用 ViT 架构实现 Diffusion 模型的开山之作，使用 ViT 来替换 Diffusion 模型中基于卷积的 U-Net 架构，同时维持模型宏观的 U 形架构不变。

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 很重要，但是上下采样就没那么重要。

51c视觉~合集25_视觉

图1：U-ViT 模型架构。其特点将所有输入 (包括 time、condition 和 noised image patches) 视为 tokens，并在浅层和深层之间使用 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 扩散模型简介

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

51c视觉~合集25_视觉_02

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

51c视觉~合集25_视觉_03

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

51c视觉~合集25_视觉_04

其中,  是条件或其连续嵌入。在以往的图像建模工作中, 扩散模型主要依赖于基于 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 的方法:

将它们 Concat 起来, 然后过个线性投影, 如图1所示, 即 Linear  。

直接把它们加起来, 即  。

对  执行线性投影，然后把它们加起来，即  。

加起来, 然后执行线性投影, 即 Linear  。

也对比了不使用 long skip connections 的情况。

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

51c视觉~合集25_视觉_05

图2：long skip 分支消融实验结果

2) AdaLN 的使用

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

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

51c视觉~合集25_视觉_06

图3：time 融入网络的方式消融实验

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

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

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

51c视觉~合集25_视觉_07

图4：添加额外卷积的消融实验

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

4) Patch Embedding 的变体

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

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

51c视觉~合集25_视觉_08

图5：Patch Embedding 的变体消融实验

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

5) 位置编码的变体

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

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

51c视觉~合集25_视觉_09

图6：位置编码的变体消融实验

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

1.4 深度、宽度和 Patch Size 的影响

作者通过在 CIFAR10 研究深度 (即层数)、宽度 (即隐藏大小  ) 和 Patch Size 的影响来展示 U-ViT 的缩放性能。如图 7 所示，随着深度 (即层数) 从 9 增加到 13，性能会提高。但是，更大的深度就不会再带来性能提升了。把宽度 (隐藏大小) 从 256 提升到 512 性能会提升，但是进一步增加到 768 就没有增益。将 Patch Size 从 8 减少到 2 可以提高性能，进一步降低到 1 不会带来增益。

51c视觉~合集25_视觉_10

图7：深度、宽度和 Patch Size 对 U-ViT 的影响

作者也发现 Patch Size 为 2 是获得比较好性能的前提，这可能是由于扩散模型的噪声预测任务是 low-level 的，需要小的 Patch Size。这不同于 high-level 比如分类任务。

由于使用较小的 Patch Size 会带来比较昂贵的计算代价，因此 U-ViT 首先把特征转换到低维的 latent 空间中，并使用 U-ViT 对 latent representation 进行建模。

1.5 实验结果

作者在 3 种任务上做了实验：unconditional image generation，class-conditional image generation，text-to-image generation。

实验设置

unconditional image generation：CIFAR10, CelebA 64×64。

class-conditional image generation：ImageNet 64×64, 256×256, 512×512。

text-to-image generation：MS-COCO 256×256。

对于高分辨率的图像生成，作者使用了 Stable Diffusion[3]提供的预训练图像自动编码器，分别将它们转换为 32×32 和 64×64 分辨率下的 latent representation。然后使用所提出的 U-ViT 对这些 latent representation 进行建模。

在 MS-COCO 上，作者使用 Stable Diffusion 的 CLIP text Encoder 将离散文本转换为一系列 Embedding。然后将这些 Embedding 输入 U-ViT 作为 token 序列。

下图 8 所示为不同 U-ViT 的大小。

51c视觉~合集25_视觉_11

图8：不同 U-ViT 的配置

U-ViT 在训练时使用 AdamW 优化器，在 CIFAR10 和 CelebA 64×64 数据集上以 bs=128 训练了 500K iterations，在 ImageNet 64×64 和 ImageNet 256×256 数据集上以 bs=1024 训练了 300K iterations，ImageNet 512×512 数据集上以 bs=1024 训练了 500K iterations，MS-COCO 数据集上以 bs=256 训练了 1M iterations。

unconditional image generation 和 class-conditional image generation 实验结果

如图 9 所示，U-ViT 在 unconditional image generation 任务中，在 CIFAR10 和 CelebA 64×64 数据集上与 U-Net 相当，同时性能优于 GenViT。

在 class-conditional ImageNet 64×64 上，作者最初尝试使用 131M 参数的 U-ViT-M 配置。如图 9 所示，FID 为 5.85，优于使用 100M 参数的 U-Net 的 IDDPM 的 6.92。为了进一步提高性能，采用了具有 287M 参数的 U-ViT-L，FID 从 5.85 提高到了 4.26。

51c视觉~合集25_视觉_12

图9：CIFAR10 和 CelebA 64×64 上 unconditional 图像生成的 FID 结果，以及 ImageNet 64×64、256×256 和 512×512 上的 class conditional 图像生成实验结果。SR：超分模块，AE：图像自编码器，Cls：分类器

同时，U-ViT 在 latent space 中表现得特别好，在应用扩散模型之前，首先将图像转换为 latent representation。在 class conditional ImageNet 256×256 上，U-ViT 获得了 2.29 的 FID，优于所有先前的扩散模型。

Text-to-Image Generation 实验结果

与之前的文献一致，作者从 MS-COCO 验证集中随机抽取 30K 个 prompt，并在这些提示上生成样本以计算 FID。

实验结果如图 10 所示，U-ViT-S 在不需要访问大型外部数据集的情况下，已经在所有的方法中实现了最好的 FID。通过进一步将层数从 13 增加到 17，U-ViT-S (Deep) 甚至可以实现更好的 FID 5.48。图 11 显示了使用相同随机种子的 U-Net 和 U-ViT 生成的样本进行比较。可以发现 U-ViT 能够生成更高质量的样本，而且可以更好匹配文本。例如，给定文本 "a baseball player swinging a bat at a ball"，U-Net 既不生成蝙蝠也不生成球。相比之下，U-ViT-S 生成 "ball"，U-ViT-S (Deep) 进一步生成 "bat"。作者认为这是因为文本和图像在 U-ViT 的每一层交互，这比仅在交叉注意力层交互的 U-Net 更频繁。

51c视觉~合集25_视觉_13

图10：MS-COCO 验证集 (256×256) 上不同模型的 FID 结果

51c视觉~合集25_视觉_14

图11：MS-COCO 上的文生图结果。除了 Backbone 之外，所有其他设置都是相同的。U-Net 和 U-ViT 使用相同的随机种子生成样本进行公平比较

参考

  1. ^https://openai.com/index/sora/