#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 很重要,但是上下采样就没那么重要。
图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 扩散模型简介扩散模型逐渐向数据注入噪声,然后反转这个过程以从噪声中生成数据。噪声注入过程,也称为前向过程,形式化为马尔可夫链:
式中, 为数据, 和 表示 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 的方法:
将它们 Concat 起来, 然后过个线性投影, 如图1所示, 即 Linear 。
直接把它们加起来, 即 。
对 执行线性投影,然后把它们加起来,即 。
加起来, 然后执行线性投影, 即 Linear 。
也对比了不使用 long skip connections 的情况。
实验结果如图 2 所示, 第 1 种方案的表现最好。直接加起来的表现还不如不用 long skip connections。作者分析原因是 Transformer 本身就有 skip connection, 因此, 的唯一影响是以线性形式增加 的系数, 这不会改变网络的性质。相比之下, 结合 的所有其他方法在 上执行线性投影, 并且在没有 long skip connections 的情况下提高了性能。
图2:long skip 分支消融实验结果
2) AdaLN 的使用
作者考虑了 2 种方法将 输入网络:
- 将其视为 tokens。
- 在 Transformer Block 中的 LayerNorm 之后使用自适应组归一化 (adaptive group normalization) 。这个方法称为自适应层归一化 (AdaLN)。形式上表述为: , 其中 是 Transformer Block 的输入, 是从 time embedding 中获得的。这种方法虽然简单, 但是将时间视为 token 的简单方法比 AdaLN 表现更好。
图3:time 融入网络的方式消融实验
3) 在 Transformer 之后添加额外卷积
作者考虑了 3 种方法在 Transformer 之后添加额外卷积:
- 在线性投影后添加一个 3×3 卷积,将 token embedding 映射到 image patches。
- 在线性投影之前添加一个 卷积块, 这需要首先将 token embedding 的 1D 序列 重新排列为形状 的 2D 特征, 其中 是 Patch Size。
- 不使用 3×3 卷积块。
图4:添加额外卷积的消融实验
如图 4 所示,在线性投影后添加一个 3×3 卷积效果略微好一些。
4) Patch Embedding 的变体
作者考虑了 Patch Embedding 的 2 种变体。
- 原始的 Patch Embedding 采用线性投影,把 patches 映射为 token embedding。
- 使用 3×3 卷积的堆叠,然后使用 1×1 卷积把图像映射到 token embedding。
图5:Patch Embedding 的变体消融实验
如图 5 所示,原始的 Patch Embedding 表现更好。
5) 位置编码的变体
作者考虑位置编码的 2 种变体。
- 原始 ViT 中提出的一维可学习位置编码,这是本文中的默认设置。
- 二维正弦位置编码,它是通过连接位置 处的 Patch 和 的正弦 Embedding 获得的。
图6:位置编码的变体消融实验
如图 6 所示,一维可学习位置编码的表现更好。作者还尝试使用了任意位置编码,发现模型无法生成有意义的图像,这意味着位置信息在图像生成中至关重要。
1.4 深度、宽度和 Patch Size 的影响作者通过在 CIFAR10 研究深度 (即层数)、宽度 (即隐藏大小 ) 和 Patch Size 的影响来展示 U-ViT 的缩放性能。如图 7 所示,随着深度 (即层数) 从 9 增加到 13,性能会提高。但是,更大的深度就不会再带来性能提升了。把宽度 (隐藏大小) 从 256 提升到 512 性能会提升,但是进一步增加到 768 就没有增益。将 Patch Size 从 8 减少到 2 可以提高性能,进一步降低到 1 不会带来增益。
图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 的大小。
图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。
图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 更频繁。
图10:MS-COCO 验证集 (256×256) 上不同模型的 FID 结果
图11:MS-COCO 上的文生图结果。除了 Backbone 之外,所有其他设置都是相同的。U-Net 和 U-ViT 使用相同的随机种子生成样本进行公平比较
参考
- ^https://openai.com/index/sora/