OpenAI如何训练一个拥有120亿个参数的文本到图像生成器：DALL-E

深度学习 30/03/2021

3个要点
✔️ 一个120亿参数的图像-文本生成模型和2.5亿图像标题数据集
✔️ 训练大型模型的一些技术。
✔️ 90%的真实度和准确度的MS-COCO字幕，小样本。

Zero-Shot Text-to-Image Generation
Written by Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, Ilya Sutskever
(Submitted 24 Feb 2021 (this version), latest version 26 Feb 2021 (v2))
Comments: Accepted to arXiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV); Machine Learning (cs.LG)

code:.

首先

自DRAW模型首次将文字引入图像生成以来，在这方面有了巨大的改进。从那时起，GAN集成、自注意力和辅助损失等方法使这些模型能够产生具有高视觉保真度和小样本泛化能力的图像。尽管如此，这些模型还是容易产生混淆，比如物体位置不正确、混合不自然、物体变形等。

当模型大小、数据和计算量的比例适当时，自回归变换器已经显示出优异的结果。以往的文本图像生成模型都是在相对较小的数据集上进行训练和评估的，如MS-COCO和CUB。因此，值得看看扩大模型大小和数据集是否会提高性能。

为了了解缩放是如何工作的，在这个案例中，OpenAI在2.5亿个取自互联网的文本-图像对上训练了一个具有120亿个参数的模型。如上图所示，这个模型非常稳健：在MS-COCO数据集上，它可以在没有任何训练标签的情况下产生优秀的小样本图像，与之前的自定义训练模型竞争。

大数据集

为了训练120亿参数模型，我们使用了2.5亿个图像/文本对。有的取自互联网，有的取自概念标题和YFCC100M数据集。该数据集不包括整个MS-COCO数据集，其中还包括无字幕的验证图像。

两步走的学习过程

在这里，我们训练一个自回归变换器，其中文本和图像作为单一数据流通过模型。然而，即使是分辨率为256x256的图像，也会形成一个长度为256x256x3的序列，这需要很高的计算量和内存。此外，与CNN不同的是，自观察不能有效地捕捉图像的局部特征，因此该模型不能形成视觉识别图像。因此，我们使用可变自动编码器（dVAE）将图像压缩成一个32x32的网格令牌，有8192个可能的值。

学习过程可以看作是最大化图像x、标题y和编码RGB图像token z的变异下限（ELBO）。这些变量的联合对数似然可以建立如下模型。

这里我们假设给定一个标记(z)，标题(y)是有条件地独立于图像(x)的。因此，下限在数学上由以下公式给出

p_θ是给定图像标记和标题，dVAE解码器产生的RGB图像的概率分布，_qφ是dVAE编码器产生的图像标记的分布。接下来，让ψ : Prior成为ELBO的模型。其中_pψ为文字和图像标记的联合分布。上述约束只在β=1的情况下成立，但我们发现，较大的数值是有益的。(DKL为KL分歧)

1）学习Visual Codebook

在这个阶段，我们只使用RGB图像来训练变量自动编码器（dVAE）。如前所述，dVAE编码器将RGB图像编码成32x32网格标记，K=8192Codebook向量，一个ADAM优化器用于最大化ELBO。_φ是一个离散分布，所以重参数化技术对于最大化是没有用的，所以为了使其可微分，我们使用了龈-软最大放松法是用来使其可微分的。_φ用^1/qφT代替，参数T控制松弛量。

训练时将T退火到1/16，使放松的验证ELBO更接近真实的ELBO。我们还发现，在编码器的末尾和解码器的开始（接近松弛操作）执行1x1卷积，可以提高对真实ELBO的逼近。此外，将编码器和解码器的复位块的输出激活量放大一个小常数，可以使学习更加稳定。

(2)对前期资料的研究

先验分布_pψ使用具有120亿个参数的稀疏变换模型进行建模，argmax用于从dVAE中抽取32×32=1024个图像标记。每个文本-图像对都采用BPE编码，BPE编码使用256个标记，因此词汇量为16384个。此外，图像令牌的编码也有8192个词条大小。文字令牌和图像令牌是连在一起处理的。

所用的换能器是一个仅有64个自注意层的解码器模型，每个自注意层有62个注意头，隐藏维度为64，使用了三种类型的自注意掩码。我们对标题使用标准的文间注意力掩码，对图像使用行、列和卷积注意力掩码。

一个标题的最大长度为256个令牌。文字和图像标记由两个填充标记隔开：256个位置中的每个位置都有一个特殊的填充标记，这个标记基本上是在训练中学习的。我们使用ADAM优化器进行训练，由于我们的目标主要是图像建模，所以我们取文本和图像的交叉熵损失的加权和，权重分别为1/8和7/8。

培训方法

1）使用混合精度训练的难点。

为了节省GPU内存，最好以16位精度存储亚当力矩和激活等参数。但是，要训练超过10亿个参数，精度参差不齐，难度很大。随着模型越来越深、越来越宽，激活梯度的指数就会落在16位格式的范围之外。指数部分在一般的GPU上大约占5位，模型不够用。因此，激活规范在以后的层数中会变小，并舍去零。这就是所谓的下流，是造成不稳定的最大原因之一。

这个问题可以通过对模型中的每个resblock使用不同的梯度尺度来解决。上面的图片可以用来说明这个问题。实线代表前进路径，虚线代表后退路径。每个输入梯度都被缩放/过滤，转换为16位，然后在离开resblock之前取消缩放/过滤。在过滤过程中，NaN或Inf值被设置为0。如果没有这个功能，如果发生非无限事件，所有前面的区块的梯度都会降低，会发生下溢。

2）分布式优化

我们使用参数共享来训练模型，混合精度训练需要24GB的内存。每个参数数组在一个节点中被分片（分区）成8个GPU。在每个GPU上，一旦计算出当前的激活，就会使用All-gather预取下一个参数块的参数碎片。同样，在后传中，计算当前区块的激活和梯度，同时使用all-gather预取前一个区块的参数碎片。一旦节点中的所有八个GPU都完成了对所有收集到的参数的梯度计算，reduce-scatter就会对所有八个GPU的梯度进行平均，让每个GPU只剩下其参数块的梯度。

同一节点上GPU之间的通信速度足够快，与这些繁重的计算重叠。主要的挑战是节点间的通信，在不同节点上计算出的梯度有延迟。我们通过使用梯度压缩技术来解决这个问题。每个GPU独立于邻居计算自己参数的低阶系数。然后，误差缓冲器会存储从低阶系数计算出的梯度与同一节点上GPU计算出的梯度平均值之间的差异。这样就把大的未压缩参数的通信开销降低到了低阶因子的小通信操作。

实验

该模型与其他几个SOTA模型如AttnGAN、DM-GAN和DF-GAN进行了比较。上图为这些模型生成的样本对比。为了进一步验证模型的图像的真实性，我们请一位人类评审员检查我们的模型DALL-E和其他模型生成的图像。

在MS-COCO数据集中，该模型在真实性和准确性方面都显著优于DF-GAN模型的标题。准确度是衡量标题与图片的匹配程度，匹配概率为93.3%。

然而，该模型在CUB数据集上表现不佳（见上文）。模型与性能最好的模型之间有很大的差距（40FID）。小样本法似乎并不能正确地适应CUB数据集的特殊分布，微调可能是解决这个问题的一个选择。

摘要

当OpenAI公布DALL-E模型并发布生成图片时，人们对DALL-E进行了大量的炒作，这也是理所当然的（还记得牛油果形状的扶手椅吗？与之前的模型相比，文本到图像转换器产生的结果非常惊人，即使是小样本任务，与GPT模型一样，我们表明，在适当的训练下，扩大模型、训练实例和计算量，可以显著提高深度学习模型的性能。通过扩大模型、训练实例和计算的规模，并通过适当的训练。在某些情况下，模特们制作了与标题相关的优秀图片，给我们带来了所谓"创意"的暗示。用不了多久，这样的文字图像模型就可以用来生成真实电影中的场景、广告海报、物体（如扶手椅）的设计等等。