【论文精读】DALLE: Zero-Shot Text-to-Image Generation零样本文本到图像生成

---

一、前言

https://openai.com/research/dall-e

DALL·E 是GPT-3的 120 亿参数版本， 经过训练，可以使用文本-图像对的数据集从文本描述生成图像。我们发现它具有多种功能，包括创建动物和物体的拟人化版本、以合理的方式组合不相关的概念、渲染文本以​​及对现有图像应用转换。

DALL·E 2可以生成更逼真、更准确的图像，分辨率提高 4 倍。

二、摘要

研究现状的缺点： 文本到图像生成传统上专注于寻找更好的建模假设以在固定数据集上进行训练，而这些假设会涉及到复杂的架构、辅助损失或辅助信息。

作者的研究：

1. 基于transformer
2. 将文本和图像标记自回归建模为单个数据流
3. 凭借足够的数据和规模，当以零样本方式评估时，我们的方法与以前的特定领域模型具有竞争力。

三、方法

（一）主要目标

1. 训练一个 Transformer 将文本和图像标记作为单个数据流 进行自回归建模

（图像是连续的，可以采用VQVAE离散化）

2. 图片的分辨率很大，如果把单个pixel当成一个token处理，会导致计算量过于庞大，于是DALL·E引入了一个dVAE模型来降低图片的分辨率。
3. 似然目标倾向于优先考虑对像素之间的短程依赖性进行建模，因此大部分建模能力将用于捕获高频细节，而不是使我们在视觉上可识别对象的低频结构。

（二）stage 1：训练离散变分自动编码器（dVAE）

• 将每个 256×256 RGB 图像压缩为 32×32 图像令牌网格
• 每个图像令牌可以取到8192个可能的离散值之一
• 这将transformer的上下文大小减少了 192 倍，而视觉质量没有大幅下降
  $$192=(256÷32)\times (256÷32)\times 3$$

（三）stage 2：训练自回归转换器

• 文本token：256个（BPE编码）
• 图像token：32 × 32 = 1024 个
• 连接文本token和图像token
• 训练自回归转换器来对文本和图像标记上的联合分布进行建模

BPE（Byte Pair Encoding）编码是一种常用的无损数据压缩算法，也常被用于自然语言处理中的词汇表示和分词任务。它基于统计的方法，通过不断合并数据中出现频率最高的字节对来构建编码表。

（四）公式表达

整个过程可以被视为最大化 x、y 和z 上模型分布的联合似然的证据下界 (ELB)

• x：图像images
• y：图像的文字描述captions
• z：编码的RGB 图像

$$p_{\theta ,\psi }(x,y,z)=p_{\theta }(x|y,z)p_{\psi }(y,z)$$

• $q_{\phi }$：在给定 RGB 图像的情况下，dVAE 编码器 生成的 32 × 32 图像令牌上的分布；
• $p_{\theta }$：给定图像标记的 dVAE 解码器 生成的 RGB 图像上的分布；
• $p_{\psi }$：由transformer 建模的文本和图像标记上的联合分布。

四、模型细节

（一）Stage One: Learning the Visual Codebook

1、最大化证据下界ELB

证据下界是用于评估概率模型的变分推断方法中的一个重要概念。

• 在变分推断中，我们通常试图近似一个复杂的后验分布，以使其更容易计算和处理。而证据下界是一个下界，用于评估我们近似的后验分布和真实后验分布之间的接近程度。如果我们的近似分布可以很好地拟合真实后验分布，那么证据下界将会比较高。
• 在机器学习中，通常使用 “Evidence Lower Bound” 作为损失函数，用于训练概率生成模型，如变分自编码器（VAE）和概率图模型。通过最大化证据下界，我们可以使生成模型更好地拟合数据，并学习到数据背后的潜在结构。

普通的变分自编码器VAE的证据下界（Evidence Lower Bound，ELB）的推导过程：（利用率琴生不等式）

DALL-E的Stage1中的证据下界可以写为：

2、模型优化

问题1： $q_{\psi }$是离散分布（one-hot编码），我们无法使用重新参数化梯度来最大化它。

重新参数化梯度是一种常用于训练变分自编码器（VAE）等生成模型的技术。

• 在训练VAE时，我们希望从一个分布中采样出一些隐变量，以生成模型的输出。然而，由于采样操作是不可导的，因此通常不能直接对采样操作求梯度。为了解决这个问题，我们可以使用重新参数化技术。
• 重新参数化技术的基本思想是，将采样过程拆分为两步：首先从一个固定的分布中采样一些固定的随机变量，然后通过一个确定的函数将这些随机变量转换为我们所需的随机变量。这样，我们就可以对这个确定的函数求导，从而能够计算出采样操作对于损失函数的梯度。

解决方案： 使用gumbel-softmax 松弛

Gumbel-Softmax 松弛的思想是，从 Gumbel 分布中采样一些随机噪声，并与离散分布的 logits 相结合，然后通过 softmax 操作将结果归一化为一个概率分布。这样，我们就可以用一个连续、可导的操作来近似离散分布的采样过程，并且可以通过反向传播来训练模型。

DALL-E中的gumbel-softmax：向sotfmax中引入超参数${\rm Y}$使得argmax可导。（设置${\rm Y}=\frac{1}{16}$）

超参数${\rm Y}$在深度学习中有一个专业术语叫温度，可以通过调整softmax曲线的平滑程度来实现不同的功能。

• 当${\rm Y}$的值大于1时，我们可以得到的更加平滑的softmax曲线，这种方式可以得到更加平滑的置信度分布；
• 当${\rm Y}$的值小于1时，得到的softmax曲线更加陡峭；
• 当${\rm Y}$的值趋近0时，可以得到近似argmax的效果，但是这时softmax还是可导的。

问题2： 在构建生成图像时，图像的像素是有值域范围的，而VAE中通过拉普拉斯分布或者高斯分布得到的值域是整个实数集，这就造成了模型目标和实际生成内容的不匹配问题。

解决方案： DALL-E提出了拉普拉斯分布的变体：log-拉普拉斯分布。它的核心思想是将sigmoid作用到拉普拉斯分布的随机变量上，从而得到一个值域是(0,1)的随机变量。

3、模型结构

DALL-E的离散VAE的编码器和解码器都是基于残差网络构建的，DALL-E保持了残差网络的基础结构，但也有其针对性的调整，核心修改如下：

• 编码器的输入层的卷积核的大小是7x7；
• 编码器的最后一个卷积的卷积核的大小是1x1，用于产生大小是32x32x8192的Feature Map；
• 使用最大池化而非原来的平均池化进行降采样；
• 解码器的的第一个卷积核最后一个卷积的卷积核的大小均为1x1；
• 解码器使用了最近邻的方式进行上采样；
• 为了提升训练速度，MAE使用了混合精度训练。

4、输入归一化和输出还原

输入归一化：图像像素范围是0到255，DALL-E中将输入像素的范围映射到$(ϵ,1-ϵ)$的范围。

输出还原：在解码器中通过log-拉普拉斯分布预测得到的结果的范围是$(0,1)$，通过$\phi$ 的逆运算即可得到解码器的生成内容。

（二）Stage Two: Learning the Prior

1、模型输入

阶段2的输入是拼接的文本特征和图像特征以及各自的位置编码等信息。

如图一个模型输入的例子，在这个示例中，文本嵌入的长度是6，图像嵌入的长度是2。文本的输入是文本的嵌入和文本的位置编码，图像的输入是图像的标志嵌入，行位置编码以及列位置编码。这些编码的长度均是3968。

2、模型结构（稀疏Transformer）

DALL-E使用的Transformer是稀疏Transformer，它的特点是只关注Top-k个贡献最大的特征的状态，因此比普通的Transformer更能关注重要的特征。DALL-E的Transformer有64个自注意力层，每个层的头数是62，每个注意力头的维度是64。

DALLE共使用了3个不同形式的稀疏自注意力编码：

• (a)是行注意力，每个标志只关注top-5的相关标志；
• (b)是列注意力，(c)是列注意力的转置，它能够更好的利用GPU；
• (d)是卷积自注意力。
  

（三）图像生成

图像生成过程：

1. 将输入文本编码成特征向量，然将特征向量送入到自回归的Transformer中生成图像的token
2. 将图像的token送入到dVAE的解码器中得到生成图像
3. 通过CLIP对生成样本进行评估，得到最终的生成结果

（四）混合精度训练

问题： 为了提升GPU的计算效率，DALL-E的大量参数以及激活都使用了16位的低精度存储。这种低精度模型的最大挑战是梯度下溢（underflow）的问题，也就是计算的梯度值超出了16位浮点数能表示的最低值。

传统的低精度训练通过将梯度值限制在一个模型能表示的范围内来避免梯度下溢。但是这种粗暴的限制每一个梯度的范围的方法并不适合DALL-E这种文本到图像更复杂的任务，它需要更过的精度表示。

解决方案： DALL-E使用了大量的技术来解决这一问题，其中最重要的一个点：每个残差块的梯度缩放（per-resblock gradient scaling）。

• 对每个残差块使用单独的梯度缩放比例，因此命名为混合精度训练
• 对于每个残差块进行缩放替代传统的对损失函数进行缩放；
• 只在有必要使用低精度浮点数来提速的地方才使用16位精度表示；
• 除以梯度的时候避免下溢。

DALL-E的一个残差块的混合精度训练：（实线表示前向传播的计算顺序，虚线表示反向传播的计算流程）

• 在前向计算时，对于单位映射，我们先将其缩放到16位，当进行完卷积运算时，我们再将其放到到32位。
• 在进行反向运算时，我们先对梯度进行过滤和缩放，其中过滤操作会将所有NaN和Inf的值置0，经过卷积操作权值的更新后再将其放大到32位。
  

（五）分布式运算

DALL-E的模型即使使用16位的精度来存储，也要占用大约24G的显存，这超过了他们训练环境的单卡（NVIDIA V100 16G）的硬件显存，这里他们使用了参数分片（Parameter Sharding）来解决显存不足的问题。在进行模型的参数分片训练时，一个问题是不同机器的通信问题，它们之间的带宽是远小于同一台机器的不同显卡之间的带宽的，这成为了多机多卡训练。

五、实验结果

（一）对比效果

1. 在不同程度的可靠性下，我们的模型似乎能够以合理的方式组合不同的概念，创建动物的拟人化版本，呈现文本，并执行某些类型的图像到图像的翻译：
   

2. 将模型中的样本与MS-COCO中先前方法的样本进行比较。每个模型样本都是由对比模型排名的512个样本中最好的：
   

3. 人类对模型(在没有温度降低的情况下评估零射击)与先前对MS-COCO标题的工作(DF-GAN)的评估。在五选一的投票中，模型的样本在90.0%的情况下被选为最真实的，在93.3%的情况下被选为最匹配共享标题的图像。
   

4. MS-COCO和CUB的定量结果。实线表示针对原始验证集计算的FID，虚线表示针对删除重叠图像的验证集计算的FID。对于MS-COCO，我们在从验证集中采样的30,000个标题的子集上评估所有模型。对于CUB，我们在测试集中的所有唯一标题上评估所有模型。
   

（二）消融研究

1. 我们展示了模型大小和梯度的最小压缩等级(最高128的倍数)之间的关系，这是避免在训练的前10%的训练损失中出现差距所必需的。这些结果表明，在我们的设置中，我们可以实现约85%的压缩率，与模型大小无关。
   
2. 增加对比重排序过程中图像数量对MS-COCO标题的影响。
   

参考：
【一起读论文】OpenAI文本驱动的图像生成DALL-E (DALLE)
DALL-E:Zero-Shot Text-to-Image Generation
如何评价DALL-E模型的实现？
[PMLR 2021] Zero-Shot Text-to-Image Generation：零样本文本到图像生成
非官方实现：DALLE-pytorch