Stable Diffusion——U-ViT用于扩散建模的 ViT 主干网

Stable Diffusion——U-ViT用于扩散建模的 ViT 主干网

引用

CSDN

1.

https://blog.csdn.net/matt45m/article/details/139204093

扩散模型是近年来兴起的深度生成模型，广泛应用于图像、视频、语音等多个领域。U-ViT作为基于视觉变换器（ViT）的UNet架构，在扩散建模中展现出优异的性能。本文将详细介绍U-ViT的架构设计、关键实现细节以及在多个生成任务中的实验结果。

1. 概述

扩散模型是一种强大的深度生成模型，能够生成高质量的图像。近年来，扩散模型在多个领域得到广泛应用，包括文本到图像生成、图像到图像生成、视频生成、语音合成和3D合成等。

除了算法的改进，骨干网络的改进在扩散建模中也发挥着重要作用。基于卷积神经网络（CNN）的U-Net是扩散模型中常用的骨干网络，已经在多个研究中得到应用。然而，视觉变换器（ViTs）在各种视觉任务中也取得了令人瞩目的成果，这促使研究者思考是否可以在扩散模型中使用基于ViT的U-Net。

本评论文章提出了一种基于ViT的U-Net（U-ViT）。所提出的方法在ImageNet和MS-COCO上生成的图像达到了最高的FID（衡量图像质量的标准）。

论文地址：https://arxiv.org/abs/2209.12152
源码地址：https://github.com/baofff/U-ViT.git

2. U-ViT架构

上图是U-ViT的概览。该网络获取扩散过程的时间（t）、条件（c）和噪声图像（x_t），并预测要注入到图像中的噪声（x_t）。按照ViT的设计方法，图像被划分为多个补丁，U-ViT处理所有输入，包括时间、条件和图像补丁，并将所有输入视为标记（词）。

与基于CNN的U-Net一样，U-ViT在浅层和深层之间使用长跳转连接。训练扩散模型是一项像素级的预测任务，对低层特征非常敏感。长跳转连接为低层特征提供了捷径，便于噪声预测网络的训练。

此外，U-ViT还可以选择在输出前添加一个3x3卷积块。这是为了防止变换器生成的图像中出现潜在的伪影。

以下各小节将提供有关U-ViT各部分的更多详细信息。

3.1 实施细节

本节将通过CIFAR10中生成图像的图像质量（FID）来优化U-ViT的结构。整体结果概览见下图。

如何结合长跳转连接

首先让h_m, h_s \in \mathbb{R}^{L \times D}成为主分支和长跳分支的嵌入。在将它们送入下一个变换块之前，考虑一些将它们组合起来的方法：

1. 将h_m, h_s连接起来，然后执行线性投影：Linear(Concat(h_m, h_s))。
2. 直接添加h_m、h_s：h_m + h_s。
3. 在h_s上进行线性投影，并将其添加到h_m：h_m + Linear(h_s)。
4. 添加h_m、h_s，然后进行线性投影：Linear(h_m + h_s)。
5. 删除长跳线连接。

如上图所示，在这些方法中，第一种使用连接Linear(Concat(h_m, h_s))的方法效果最好。特别是，与不使用长跳转连接的方法相比，生成图像的质量明显提高。

如何输入时间条件

将时间条件t输入网络的方法有两种。方法(1)是将它们视为标记，如图1所示。方法(2)是在转换器模块中将层归一化后的时间纳入，这与U-Net中使用的自适应组归一化类似；第二种方法称为自适应层归一化（AdaLN）。将时间视为标记的方法（1）比AdaLN性能更好。

如何在变换器后添加卷积块

在变换器后添加卷积块有两种方法。(1)在将标记嵌入映射到图像补丁的线性投影之后添加3×3卷积块。(2)在线性投影之前添加一个3×3卷积块。此外，还可以将其与去掉附加卷积块的情况进行比较，在线性投影后添加3×3卷积块的方法(1)的性能略优于其他两种方案。

补丁嵌入法

传统的补丁嵌入是将补丁映射到标记嵌入的线性投影。除了这种方法，我们还考虑了另一种将图像映射到标记嵌入的方法，即使用3×3卷积块堆叠，然后再使用1×1卷积块。但是，如图2(d)所示，传统的补丁嵌入法效果更好，因此最终模型采用了这种方法。

位置嵌入方法

本文使用的是原ViT中提出的一维可学习位置嵌入。也有一种替代方法，即二维正弦位置嵌入，一维可学习位置嵌入的效果更好。我们也尝试过不使用位置嵌入，但该模型无法生成清晰的图像，这表明位置信息对图像生成非常重要。

3.2 网络深度、宽度和补丁大小的影响

这里，我们在CIFAR10中研究了U-ViT的缩放特性，以考察层数、宽度和贴片尺寸的影响。如上图所示，将层数从9层增加到13层可提高性能，但对于深度超过17层的模型则没有影响。同样，增加宽度也能提高性能，但超过一定宽度就没有效果了。

减小贴片尺寸可以提高性能，但低于一定尺寸时性能就会下降。在扩散模型中，小补丁尺寸被认为适合低水平噪声预测任务。另一方面，对于高分辨率图像来说，使用小尺寸贴片的成本较高，因此必须先将图像转换为低维潜在表示，然后再用U-ViT进行建模。

4. 实验

4.1 数据集和设置

U-ViT的有效性在三个任务中进行了测试：无条件图像生成、类条件图像生成和文本到图像生成。

无条件图像生成实验在CIFAR10（50,000幅图像）和CelebA 64×64（162,770幅图像）上进行。对于类别条件图像生成，在64×64和256×256ImageNet数据集（包含来自1,000个不同类别的1,281,167幅训练图像）和512×512分辨率数据集上进行了实验。MS-COCO（82,783幅训练图像和40,504幅验证图像）用于文本到图像的训练。

在生成256×256和512×512分辨率的高分辨率图像时，使用由潜在扩散模型（Latent diffusion models，LDM）[Rombach et.al, 2022]提供的预训练图像自动编码器，分别将32×32和64×64分辨率的潜在表征分别生成32 x 32和64 x 64分辨率的潜表征。然后使用U-ViT对这些潜表征进行建模。

在MS-COCO中生成文本到图像时，使用CLIP文本编码器将离散文本转换为嵌入序列，然后将这些嵌入序列作为标记序列输入U-ViT。

4.2 无条件和类条件图像生成

表1.无条件图像生成和类别条件图像生成的结果

在这里，U-ViT与之前基于U-Net的扩散模型和GenViT进行了比较，GenViT是一种较小的ViT，它没有长跳接，并在归一化层之前加入了时间。FID分数用于衡量图像质量。

如表1所示，U-ViT在无条件的CIFAR10和CelebA 64×64中表现出与U-Net相当的性能，并且比GenViT性能更好。对于有类别条件的ImageNet 64×64，我们首先尝试了U-ViT-M配置，参数为131M。如表1所示，其FID为5.85，优于使用U-Net且参数为100M的IDDPM 6.92。为了进一步提高性能，我们采用了U-ViT-L配置（287M个参数），将FID从5.85提高到4.26。

在有类别条件的ImageNet 256×256中，U-ViT的最佳FID为2.29，优于之前的扩散模型。表2显示，在使用相同采样器的不同采样步骤中，U-ViT的表现优于LDM。U-ViT的表现也优于VQ-扩散模型，后者是一种以变换器为骨干的离散扩散模型。同样，在参数和计算成本相同的情况下，U-ViT也优于UNet。

对于带有类别条件的ImageNet 512×512，U-ViT的表现优于直接对图像像素建模的ADM-G。图4显示了ImageNet 256×256和512×512的部分样本以及其他数据集的随机样本，证实了图像的高质量和清晰度。

表2：ImageNet 256×256不同采样步数下的FID结果。

4.3 使用MS-COCO生成文本到图像

这里，我们使用MS-COCO数据集来评估U-ViT在文本到图像生成任务中的表现。我们还使用U-Net训练了另一个潜在扩散模型，模型大小与U-ViT相同，并与U-ViT进行了比较。

FID分数用于衡量图像质量：从MS-COCO验证集中随机选取30K个提示，并利用这些提示生成样本来计算FID。如表3所示，即使在生成模型的训练过程中不需要访问大型外部数据集，U-ViT也能获得最先进的FID。通过将层数从13层增加到17层，U-ViT-S（Deep）可以获得更好的FID。

图6显示了U-Net和U-ViT使用相同的随机种子生成的样本，以进行定性比较；U-ViT生成的样本质量更高，图像内容与文本的匹配度更高。

例如，给定文本“棒球运动员挥棒击球”，U-Net不会生成球棒或球，而U-ViT会生成球，U-ViT-S（Deep）会进一步生成球棒。这可能是由于与U-Net相比，U-ViT中文本和图像之间每一层的交互更为频繁。

表3.MS-COCO的实验结果

5. 结论

U-ViT将所有输入（时间、条件和噪声图像片段）视为标记，并在浅层和深层之间采用长跳转连接。U-ViT已经在无条件和有条件图像生成以及文本到图像生成等任务中进行了评估。

U-ViT的性能与类似规模的基于CNN的U-Nets不相上下，甚至更好。这些结果表明，长跳接对于基于扩散的图像建模非常重要，而基于CNN的U-Nets并不总是需要向下向上采样运算符。

U-ViT可以为未来的扩散建模骨干研究提供信息，并有利于在具有不同模式的大型数据集中进行生成建模。