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

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

TL; DR：本文提出了 Imagen，并介绍了多个训练 pixel-based 级联图像生成模型的技巧。最关键的两个点是：一：本文指出，在文生图模型中，scaling 文本编码器比 scaling 扩散生图模型带来的提升要显著得多。用纯文本预训练的超大语言模型，（相比于 CLIP 文本编码器）更适合作为文生图模型的文本编码器。二：介绍了动态阈值和噪声条件增强等方法，提升生图的质量。分别对应标题中的 “deep language understanding” 和 “photorealistic”。

方法

文生图模型中的文本编码器

现有的文生图模型一般使用 CLIP text encoder 作为文本编码器，是因为 CLIP 通过对比学习的方式在图文对数据上进行训练，多模态语义对齐能力更好，适合于在文生图模型编码文本 prompt。本文提出，大语言模型可能是更适合在文生图模型中作为文本编码器。

根据实验结果（如下图所示），(a)、(b) 分别是 scaling 文本编码器和 scaling 文生图 UNet 的性能提升，横轴是 clip score 表征语义一致性，越高越好；纵轴是 FID，越低越好，图中越靠近右下，整体性能越好。可以看到，scaling 文本编码器比 scaling 文生图 UNet 模型对生图结果带来的提升，无论是在生图质量还是语义一致性上。都更显著。

作者认为，纯文本数据比图文对数据更易获得，使用纯文本数据预训练一个大语言模型，能够 scaling 到更大的规模，而更大规模的文本编码器，对文生图模型生图质量和语义一致性的提升更显著。在 Imagen 中，作者选用了 T5-XXL 模型作为文本编码器。

更大的CFG scale值与采样thresholding

Classifier-free Guidance（CFG）是目前大部分主流文生图模型中添加文本条件的方式，可表示为：
$$\widetilde{ϵ}(z_t,c)=w{ϵ}_{\theta }(z_t,c)+(1-w){ϵ}_{\theta }(z_t,\varnothing )$$
其中 ${ϵ}_{\theta }$ 是条件去噪模型，$c$ 是文本条件，$\varnothing$ 表示空文本，$w$ 是 CFG scale 值。当 $w=1$ 时，等价于未使用 CFG，简单地直接引入文本条件 $c$，当 $w>1$ 时，其值越大，文本条件的影响越强，生图结果与文本提示词的一致性也越好。然而，更大的 $w$ 通常会导致生图的质量有所下降，作者认为，这是由于 训练测试不匹配(train-test mismatch) 导致的。在每个采样步时，模型预测的归一化像素值会与训练时对齐，分布在一个范围之内，比如 $[-1,1]$ 内，然而，作者通过实验发现，过大的 CFG scale $w$ 会导致采样生图时模型的预测值超出这个范围，导致训练测试不匹配。而且由于采样生成过程时一步一步迭代进行的，每一步的输入是上一步的输出，这会导致这种数值问题越来越严重，最终降低生图质量，甚至可能不收敛。

为了解决这一问题，作者提出了两种方法：静态阈值 (static thresholding) 和动态阈值 (dynamic thresholding)。静态阈值是指将模型的预测值强行截断在 [-1, 1] 范围之内。动态阈值则是在每一个采样步，引入一个绝对像素值的百分比值 $s$，如果 $s>1$ ，则将像素值归一化到 $[-s,s]$ 范围内，并除以 $s$​ 。

级联模型结构与超分噪声条件增强

整体上，Imagen 是一种级联的文生图模型，在 T5-XXL 模型对输入 prompt 进行编码之后，分别通过交叉注意力层注入到各级和超分模型中。首先由超分模型生成一张 64x64 的低分辨率图片，然后经过两级超分扩散模型，分别将图像超分至 256x256 和 1024x1024。

基础生成模型接收经过池化的文本 embedding 作为输入，并加入扩散模型的时间步 embedding，另外，该模型也在不同分辨率下通过交叉注意力注入文本 embedding。在超分模型的具体结构上，Imagen 还是沿用扩散模型中常用的 UNet 模型，但是做出了一些改进，如移除自注意力层等，降低模型的显存占用，最终构建了一种新的、高效的 UNet 模型，称为 Efficient UNet。

两个超分模型同样以文本 embedding 作为条件。此外，在超分阶段，作者还使用了噪声条件增强。即训练时，对基础生成模型生成的 64x64 的低分辨率图片，先随机进行数据增强（如加高斯噪声、高斯模糊等），并记录下所加破坏的强度，作为条件同时输入给超分模型；在推理生图时，遍历各种破坏强度，选择一张质量最好的图片作为结果。这里是否和后来 SDXL 中的 micro conditioning 有点像？

总结

技术上，Imagen 的贡献主要有以下几点：

1. Imagen 是一种级联的 pixel-based 的图像生成模型，先生成 64x64 的小图，再经过两级超分模型，生成 256x256、1024x1024 的大图；
2. scaling 文本编码器比 scaling 扩散生图模型带来的提升要显著得多。用纯文本预训练的超大语言模型，（相比于 CLIP 文本编码器）更适合作为文生图模型的文本编码器；
3. 动态阈值有助于解决 pixel-based 条件生成模型在生图时 CFG scale 太大导致的 训练推理不匹配 的问题；
4. 噪声条件增强有助于提升级联超分模型的稳健性；

最核心的贡献，还是在于指出了大规模、纯文本预训练的语言模型对文生图模型在生图质量和语义一致性上的提升。但是不知为何目前的 SD 等模型似乎还是执着于使用 CLIP 作为文本编码器。最近腾讯也提出了使用大规模的 T5 模型代替 SD 中的 CLIP 作为文本编码器，在语义理解上提升显著。