Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding (Paper reading)

Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding

Chitwan Saharia, Google Research, Brain Team, Neurips2022, Cited:619, Code, Paper

1. 前言

我们提出了Imagen，一个文本到图像的扩散模型，具有前所未有的写实主义程度和深度的语言理解。Imagen建立在理解文本的大型变压器语言模型的基础上，并依赖于在高保真图像生成中扩散模型的强度。我们的关键发现是，在纯文本语料库上预训练的通用大型语言模型(例如T5)在编码用于图像合成的文本方面惊人地有效:在Imagen中增加语言模型的大小比增加图像扩散模型的大小更能提高样本保真度和图像文本对齐。Imagen在没有经过COCO训练的情况下，在COCO数据集上获得了7.27的最新FID分数，并且人类评分者发现Imagen样本在图像-文本对齐方面与COCO数据本身不相上下。为了更深入地评估文本到图像模型，我们引入了DrawBench，这是一个全面且具有挑战性的文本到图像模型基准测试。使用DrawBench，我们将Imagen与最近的方法进行了比较，包括VQ-GAN+CLIP, Latent Diffusion Models, GLIDE和DALL-E 2，并发现人类评分者在并排比较中更喜欢Imagen，无论是在样本质量还是图像-文本对齐方面。

2. 介绍

近年来，多模态学习逐渐兴起，其中以文本-图像合成和图像-文本对比学习最为突出。这些模型改变了研究界，并通过创造性的图像生成和编辑应用获得了广泛的公众关注。为了进一步探索这一研究方向，我们引入了Imagen，这是一种文本到图像的扩散模型，它结合了transformer语言模型(LMs)和高保真扩散模型的力量，在文本到图像合成中提供了前所未有的摄影真实感和更深层次的语言理解。与之前仅使用图像-文本数据进行模型训练的工作相比，Imagen背后的关键发现是，来自大型Lm的文本嵌入，在纯文本语料库上预训练，对于文本-图像合成非常有效。

虽然Imagen概念简单且易于训练，但它产生了令人惊讶的强大结果。Imagen在COCO上的zero-shot的FID-30K为7.27，优于其他方法，显著优于先前的工作，如GLIDE(12.4)和DALL-E 2(10.4)的并行工作。我们的zero-shot的FID评分也优于COCO训练的最先进的模型，例如，Make-A-Scene(7.6)。此外，人工评分者指出，从Imagen生成的样本在图像-文本对齐方面与COCO标题上的参考图像一致。

我们将介绍DrawBench，这是一种用于文本到图像评估的新的结构化文本提示suite。DrawBench通过对文本到图像模型的多维评估，提供了更深入的见解，文本提示旨在探测模型的不同语义属性。这些包括组合性、基数性、空间关系、处理复杂文本提示或带有罕见单词的提示的能力，还包括创造性提示，这些提示推动了模型生成高度不可信场景的能力的极限，远远超出了训练数据的范围。使用DrawBench，广泛的人类评估表明Imagen优于其他近期的方法。我们进一步展示了使用大型预训练语言模型作为Imagen文本编码器的一些明显优势，而不是使用CLIP等多模态嵌入。

主要贡献有：

1. 我们发现，仅在文本数据上训练的大型冻结语言模型对于文本到图像的生成是非常有效的文本编码器，并且缩放冻结文本编码器的大小比缩放图像扩散模型的大小更能显著提高样本质量。
2. 我们引入了动态阈值，一种新的扩散采样技术，以利用高引导权重，生成比以前更逼真和详细的图像。
3. 我们强调了几个重要的扩散架构设计选择，并提出了Efficient U-Net，这是一个新的架构变体，它更简单，收敛更快，内存效率更高。
4. 我们实现了最新的COCO FID 7.27。人类评分者发现Imagen在图像-文本对齐方面与参考图像不相上下。
5. 我们引入了DrawBench，这是一种新的全面且具有挑战性的文本到图像任务评估基准。在DrawBench人工评估中，我们发现Imagen优于所有其他工作，包括DALL-E 2[54]的并发工作。

3. Imagen

Imagen由一个将文本映射到一系列嵌入的文本编码器和一个将这些嵌入映射到分辨率不断增加的图像的条件扩散模型级联组成(见图)。在下面的小节中，我们将详细描述每个组件:

3.1 预训练文本编码器

文本到图像模型需要强大的语义文本编码器来捕获任意自然语言文本输入的复杂性和组成性。在当前的文本到图像模型中，训练成对图像-文本数据的文本编码器是标准流程;它们可以从头开始训练，也可以对图像-文本数据进行预训练(例如，CLIP)。图像-文本训练目标表明，这些文本编码器可以编码视觉语义和有意义的表示，特别是与文本-图像生成任务相关。对文本进行编码的另一种选择是大型语言模型。大型语言模型(例如BERT ， GPT ， T5)的最新进展显示了文本理解和生成能力的飞跃。语言模型是在纯文本语料库上训练的，语料库比成对的图像-文本数据大得多，因此暴露在非常丰富和广泛的文本分布中。这些模型通常也比当前图像-文本模型中的文本编码器大得多。

因此，探索文本到图像任务的两种文本编码器家族就变得很自然了。Imagen探索了预训练的文本编码器:BERT， T5和CLIP。为了简单起见，我们冻结了这些文本编码器的权重。冻结具有嵌入的离线计算等优点，在文本-图像模型训练过程中计算量或内存占用可以忽略不计。在我们的工作中，我们发现一个明确的conviction，缩放文本编码器的大小可以提高文本到图像的生成质量

3.2 扩散模型和无分类器引导classifier-free guidance

在这里，我们简要介绍了扩散模型。扩散模型是一类生成模型，通过迭代去噪过程将高斯噪声转换为来自学习数据分布的样本。这些模型可以是有条件的，例如基于类标签、文本或低分辨率图像。损失函数的形式：
$${\mathbb{E}}_{\mathbf{x},\mathbf{c},\mathbf{ϵ},t}\left[w_t{\Vert {\hat{\mathbf{x}}}_{\theta }\left({\alpha }_t\mathbf{x}+{\sigma }_t\mathbf{ϵ},\mathbf{c}\right)-\mathbf{x}\Vert }_2^2\right]$$
这里和平时看到的不太一样，我说一下，这里是预测的$x_0=x$，$c$是$x_0$对应的条件。权重是跟$t$相关的，意思是说对于不同的$t$值赋予不同权重，训练时对不同权重的$t$的损失就不一样，简而言之就是给$t$分了难易程度。

分类器引导是一种在采样期间使用来自预训练模型的梯度修正来提高样本质量，同时减少条件扩散模型中的多样性的技术，详细见Guided Diffusion/Diffusion Models Beat GANs on Image Synthesis (Paper reading)。无分类器引导是一种替代技术，它通过在训练过程中随机丢弃条件$c$（例如，10%的概率），在有条件和无条件的目标上联合训练单个扩散模型，从而避免了这种预训练模型。注意，本文使用调整后的$x$预测，也就是说，同样预测的是噪声${ϵ}_{\theta }$，但是先通过公式把噪声转变成$\hat{x}$然后作为损失函数：
$${\widetilde{ϵ}}_{\theta }(x_t,c)=w{ϵ}_{\theta }(x_t,c)+(1-w){ϵ}_{\theta }(x_t)$$
其实就是有条件和无条件的均衡，$w=1$就是有条件的，如果$w>1$，会对条件进行增强。

3.3 Large guidance weight samplers

我们证实了最近文本引导扩散工作的结果，并发现增加无分类器引导权重可以改善图像-文本对齐，但会损害图像保真度，从而产生高度饱和和不自然的图像。我们发现，这是由于高引导权重引起的训练测试不匹配。在每个采样步骤$t$，$x$的预测必须在与训练数据$x$相同的边界内，即在[-1，1]内，但我们根据经验发现，高引导权重导致$x$的预测超过这些边界。这是一种训练-测试失配，由于扩散模型在整个采样过程中反复应用于其自身的输出，因此采样过程会产生不自然的图像，有时甚至会出现发散。为了解决这个问题，我们研究了静态阈值和动态阈值。效果可视化参见如下：

“宇航员骑马的照片”的256×256样本阈值技术。从上到下，引导权重从1增加到5。没有阈值的处理会导致具有高引导权重的较差图像。静态阈值是一种改进，但仍会导致样本过饱和。我们的动态阈值处理可获得最高质量的图像。

3.4 鲁棒级联扩散模型

Imagen利用基础64×64模型和两个文本条件超分辨率扩散模型的pipeline，将64×64生成的图像上采样为256×256图像，然后再上采样为1024×1024图像。具有噪声调节增强的级联扩散模型在逐步生成高保真图像方面非常有效（Cascaded diffusion models for high fidelity image generation）。此外，通过噪声水平调节，使超分辨率模型意识到添加的噪声量，显著提高了样本质量，并有助于提高超分辨率模型的鲁棒性，以处理低分辨率模型生成的伪影。Imagen对这两个超分辨率模型都使用了噪声调节增强。我们发现这对于生成高保真图像至关重要。

给定条件化的低分辨率图像和增强级别（也称为$aug_level）（例如，高斯噪声或模糊的强度），我们用增强（对应于aug_level）破坏低分辨率图像，并在aug_level上条件化扩散模型。在训练过程中，aug_level是随机选择的，而在推理过程中，我们对其不同的值进行扫描，以找到最佳的样本质量。在我们的案例中，我们使用高斯噪声作为一种增强形式，并应用类似于扩散模型中使用的前向过程的保方差高斯噪声增强。

4. 评估和实验

我们为64×64的文本到图像合成训练20亿参数模型，分别训练600M和400M参数模型用于64到256和256到1024。我们对所有模型使用2048和2.5M的批量训练步骤。我们的基本64×64模型使用256个TPU-v4芯片，两种超分辨率型号都使用128个TPU-v4芯片。我算了下，一天大概消费6百万人民币，就这GPU。

5. Efficient U-Net

我们为我们的超分辨率模型引入了一种新的架构变体，我们称之为高效U-Net。我们发现，与之前的一些实现相比，我们的高效U-Net更简单、收敛更快、内存效率更高，尤其是对于高分辨率。我们对U-Net架构进行了几项关键修改，例如将模型参数从高分辨率块转移到低分辨率块，将跳跃连接缩放$1/\sqrt{2}$，并颠倒下采样/上采样操作的顺序，以提高前向通过的速度。高效U-Net对典型U-Net模型进行了几个关键修改：

1. 我们通过为较低分辨率添加更多残差块，将模型参数从高分辨率块转移到低分辨率块。由于分辨率较低的块通常有更多的通道，这使我们能够通过更多的模型参数来增加模型容量，而不会产生惊人的内存和计算成本。
2. 当以较低分辨率使用大量残差块时（我们以较低的分辨率使用8个残差块），我们发现将跳过连接缩放$1/\sqrt{2}$显著提高了收敛速度。
3. 在典型的U-Net的下采样块中，下采样操作发生在卷积之后，而在上采样块中的上采样操作则发生在卷积之前。我们颠倒了下采样和上采样块的顺序，以显著提高U-Net的前向通过速度，并且没有发现性能下降。

有了这些关键的简单修改，与之前的一些U-Net实现相比，高效U-Net更简单、收敛更快、内存效率更高。下图显示了高效U-Net的完整架构，而图A.28和A.29分别显示了高效U-Net的下采样和上采样块的详细描述。