（DALL-E）Zero-Shot Text-to-Image Generation

引用： Ramesh A, Pavlov M, Goh G, et al. Zero-shot text-to-image generation[C]//International conference on machine learning. Pmlr, 2021: 8821-8831.

论文链接： [2102.12092] Zero-Shot Text-to-Image Generation (arxiv.org)

代码链接： https://github.com/openai/DALL-E

简介

传统上，文本到图像生成侧重于为固定数据集上的训练找到更好的建模假设。这些假设可能涉及复杂的架构、辅助损耗或侧面信息，例如训练期间提供的对象的部分标签或分割掩码。论文提出了一种基于transformer的简单方法，transformer将文本和图像tokens自回归的建模为单个数据流。

Method

直接使用像素作为图像tokens对于高分辨率图像将需要过多的内存。可能性目标倾向于优先考虑像素之间的短程依赖关系建模，因此大部分建模能力将用于捕获高频细节，而不是在视觉上可识别对象的低频结构。因此，使用了一个两阶段的方法：

1. 训练一个离散变分自动编码器（dVAE），将每个256×256 RGB图像压缩成一个32×32的图像标记网格，其中每个元素可以假设8192个可能的值。这将转换器的上下文大小减小了 192 倍，而不会大幅降低视觉质量，如图1所示。
2. 将多达 256 个 BPE 编码的文本tokens与 32 × 32 = 1024 个图像tokens连接起来，并训练一个autoregressive transformer来模拟文本和图像标记上的联合分布。

整个过程可以被视为最大化关于编码 RGB 图像的图像 $x$、标题 $y$ 和toknes $z$ 上的模型分布的联合似然的置信下限（evidence lower bound，ELB）。使用因式分解 $p_{\theta ,\psi }(x,y,z)=p_{\theta }(x|y,z)p_{\psi }(y,z)$ 来模拟这个分布，从而产生下界：

$$\ln p_{\theta ,\psi }(x,y)\ge E_{z\sim q_{\varphi }(z|x)}(\ln p_{\theta }(x|y,z)-\beta D_{KL}(q_{\varphi }(y,z|x),p_{\psi }(y,z)))$$

• 假设 $y$ 在给定 $z$ 的情况下条件独立于 $x$，$q_{\varphi }$ 表示给定 RGB 图像 $x$ 的 dVAE 编码器生成的 $32\times 32$ 个图像tokens的分布;

• $p_{\theta }$ 表示给定图像tokens的 dVAE 解码器生成的 RGB 图像的分布;

• $p_{\psi }$ 表示 Transformer 建模的文本和图像tokens上的联合分布。

即第一项希望提升重建质量，第二项衡量图像与文本tokens的差异程度，希望KL散度越来越小。注意，边界仅适用于 $\beta =1$，而在实践中，使用更大的值会有所帮助。

Stage One: Learning the Visual Codebook

dVAE编码器和解码器是基于卷积的ResNets，具有bottleneck-style的残差块。这些模型主要使用 3 × 3 个卷积，其中 1 × 1 个卷积沿跳跃连接，其中特征图的数量在重新块的输入和输出之间变化。编码器的第一次卷积是 7 × 7，编码器的最后一个卷积（产生 32 × 32 × 8192 输出，用作图像标记分类分布的对数）为 1 × 1。解码器的第一个卷积和最后一个卷积都是 1 × 1。编码器使用最大池化（比平均池化产生更好的 ELB）来对特征图进行下采样，而解码器使用最近邻上采样。

第一阶段将初始先验 $p_{\theta }$ 设置为 $K=8192$ codebook向量上的均匀类别分布，并将 $q_{\varphi }$ 设置为编码器输出的 $32\times 32$ 网格中相同空间位置的 $8192$ logits 参数化的类别分布。

因为 $q_{\varphi }$ 是一个离散分布，所以不能使用重参数化梯度来最大化它。以前的工作使用online cluster assignment procedure和straight-through estimator解决了这个问题。这里改用gumbel-softmax松弛，用对 $q_{\varphi }^{\tau }$ 的期望代替对 $q_{\varphi }$ 的期望，其中当温度 $\tau \to 0$ 时，弛豫变得紧密。$p_{\theta }$ 的似然使用对数-拉普拉斯分布进行评估。使用Adam和指数加权迭代平均法最大化松弛的ELB，以下内容对于稳定训练尤为重要：

• 退火 $\tau$ 至 $1/16$ 足以缩小宽松验证ELB与 $q_{\varphi }$ 替换为 $q_{\varphi }^{\tau }$ 的真实验证ELB之间的差距（在前 150000 次更新中从 1 退火到 1/16）。

• 在编码器的末尾和解码器的开头使用 $1\times 1$ 卷积。减小弛豫周围卷积的感受野大小会导致它更好地泛化到真正的 ELB。

• 将编码器和解码器残差块的输出激活乘以一个小常数，以确保初始化时的稳定训练。

• 将 $KL$ 权重增加到 $\beta =6.6$ （在前 5000 次更新中，从 0 增加到 6.6）可以促进更好的codebook使用，并最终导致训练结束时重建误差较小。

Stage Two: Learning the Prior

在第二阶段，固定 $\varphi$ 和 $\theta$，并通过最大化相对于 $\psi$ 的 ELB 来学习文本和图像tokens的先验分布。这里，$p_{\psi }$ 由一个 120 亿参数的稀疏transformer表示。

给定文本-图像对，使用最多$256$个词汇大小为$16384$的标记对小写caption进行BPE编码，并使用3$2\times 32=1024$个词汇大小$8192$的tokens对图像进行编码。图像标记是使用dVAE编码器logits的最大化采样获得的，不添加任何gumbel噪声。最后，将文本和图像标记连接起来，并自回归建模为单个数据流。transformer是一个仅有解码器的模型，其中每个图像tokens都可以注意其 64 个自注意力层中的任何一个中的所有文本tokens，使用交叉熵损失。由于主要对图像建模感兴趣，因此将文本的交叉熵损失乘以 1/8，将图像的交叉熵损失乘以 7/8。

实验结果