在文本生成图片的任务中，人们一直致力于在确定的小数据集下寻找更好的网络结构，然而一些假设通常会引入复杂的模型结构或者繁琐的损失函数。本文提出一种简单的自回归策略，利用transformer将文本和图片同等建模，即使在zero shot情况下也可以达到SOTA效果。GPT3证明了在文本上使用大规模参数和大规模数据对模型带来的提高，本文使用120亿参数的模型在2.5亿数据上进行了训练，证明了大规模模型在文本到图像生成任务上也有很好的效果。
2021年1月份，google 发布了一个新的模型DALL-E，该模型拥有120亿参数，可以从一段文本描述生成生动的图片。
背景
在文本生成图片的任务中，人们一直致力于在确定的小数据集下寻找更好的网络结构，然而一些假设通常会引入复杂的模型结构或者繁琐的损失函数。本文提出一种简单的自回归策略，利用transformer将文本和图片同等建模，即使在zero shot情况下也可以达到SOTA效果。GPT3证明了在文本上使用大规模参数和大规模数据对模型带来的提高，本文使用120亿参数的模型在2.5亿数据上进行了训练，证明了大规模模型在文本到图像生成任务上也有很好的效果。
模型介绍
本文目标是利用transformer结构同时编码文本和图像，但是图像是无法像文本一样编码的，若将像素作为单位则会导致序列长度太长，极大的浪费了计算资源，并且过长的序列也不利于transformer的学习。为了使用较短的序列表达图像，我们使用了一个两阶段训练的策略：
- 第一阶段我们利用离散VAE(dVAE)将256 256大小的图片编码为32 32个图像字符，每个字符有8192个可能的值
- 第二阶段我们将文本的256个字符和32 * 32个图像字符拼接起来，使用transformer进行自回归建模。
形式化的，我们定义$x, y, z$分别为图像，文本和图像字符，定义$p_\theta(x|y, z)$为dVAE的解码器，$p_\psi(y, z)$为自回归的transformer模型，$q_\phi(y, z|x)$为dVAE的编码器部分，我们要进行建模的分布为$p_{\theta, \psi}(x, y)$，并且假定给定$x$和$y$在$z$下条件独立，那么我们有
$$ \begin{align} &lnp_{\theta, \psi}(x, y)\\\ =\ &E_{q_\phi(y, z| x)}(lnp_{\theta, \psi}(x, y)) \\\ =\ &E_{q_\phi(y, z| x)}(\frac{lnp_{\theta, \psi}(x, y, z)q_\phi(y, z| x)}{lnp_{\theta,\phi}(z|x, y)q_\phi(y,z|x)}) \\\ =\ &E_{q_\phi(y, z| x)}[ ln p_\theta(x|y, z)\\\ & + lnp_\psi(y, z) \\\ & + lnq_\phi(y, z|x) \\\ & - lnp_{\theta,\psi}(z|x,y)\\\ & - lnq_\phi(y,z|x) ] \\\ =\ &lnp_\theta(x|y, z) \\\ & - E_{q_\phi(y, z| x)}(lnq_\phi(y,z|x) - lnp_\psi(y,z))\\\ & + E_{q_\phi(y, z| x)}(lnq_\phi(y,z|x) - lnp_{\theta,\psi}(z|x, y)) \\\ =\ &lnp_\theta(x|y, z) \\\ & - KL(q_\phi(y, z| x)||p_\psi(y,z)) \\\ & + KL(q_\phi(y, z| x)||p_{\theta,\psi}(y,z|x) - E(lnp_{\theta,\psi}(y|x)) \\\ \geq\ & lnp_\theta(x|y, z)\\\ & - KL(q_\phi(y, z| x)||p_\psi(y,z)) \end{align} $$
这个下界限即作为我们优化的$loss$函数。
模型训练
本文使用一个两步的训练方式，首先训练dVAE，即参数$\theta$和$\phi$，然后训练transformer的参数$\psi$。
训练dVAE是一个很困难的过程，我们固定参数$\psi$为8192个图像字符的均匀分布，然后训练$\theta$和$\phi$。由于$p_\psi$是离散分布，重参数化（reparameterization trick）是无法使用的，所以这里需要使用Gumbel-softmax trick。
Gumbel-softmax trick
在正常的VAE中，我们训练过程中要使用重参数化（reparameterization trick）保证编码器梯度的正常传播，但当隐变量分布为离散分布时，重参数化是不能使用的，我们对此首先进行简单解释。给定随机变量$\epsilon \sim p(\epsilon)$和任意可测函数$g$，我们有：
$$ E_{\epsilon \sim p(\epsilon)}(g(\epsilon)) = \int g(\epsilon)d F(\epsilon) $$
然后假设随机变量$z\sim p_\phi(z)$并且$z$可以表达为$z = g(\epsilon, \phi)$。对于任意可测函数$f$，显然
$$ E_{z \sim p_\phi(z)}(f(z)) = E_{\epsilon \in p(\epsilon)}(f(g(\epsilon, \phi))) $$
在神经网络的反向传播下我们感兴趣的是$\nabla_\phi E_{z \sim p_\phi(z)}(f(z))$，注意到这个期望求解过程中包含对$z$采样过程，参数$\phi$无法用反向传播求解，所以我们交换积分和求导顺序
$$ \nabla_\phi E_{z \sim p_\phi(z)}(f(z)) = E_{\epsilon \sim p(\epsilon)}(\nabla f(g(\epsilon, \phi))) $$
通常我们取$p(\epsilon)$为一个已知的简单分布例如$N(0, I)$，这样我们实现将采样过程放进带有参数的函数内部，从而避免了梯度无法传播的情况。
但是，当我们$z$的值为离散时，重参数化显然无法使用，因为函数$g(\epsilon, \phi)$的值为离散，函数必定不连续也就不可能可导。这时我们需要借助Gumbel分布实现梯度传播。Gumbel分布的分布函数为$F(G) = e^{-e^{-G}}$，概率密度函数图像如上图所示。假设我们有均匀分布$u \sim U(0, 1)$，显然从Gumbel分布中采样可以转变为从均匀分布中采样然后进行变换$G=−log(−log(U))$。现在我们假设$X$为离散的随机变量且：
$$ P(X =k) = \alpha_k $$
若直接从这个分布中采样，同样存在无法传播梯度的情况，所以我们要将采样过程放入参数内部。假设$\{G_k\}_{k\leq K}$为从Gumbel分布中采样出的一组随机变量，那么我们有
$$ \hat{X} = \mathop{argmax}_k (log \alpha_k + G_k) $$
$\hat{X}$和$X$的分布是等价的，可以利用Gumbel分布的性质进行简单证明。首先假设
$$ z_k = log\alpha_k + G_k $$
显然$z_k$的概率分布函数为
$$ e^{-e^{log\alpha_k - z_k}} $$
概率密度函数为
$$ e^{log\alpha_k - z_k - e^{log\alpha_k - z_k}} $$
在已知所有$log\alpha_k$的情况下，我们可以计算对于某一$k$的条件概率分布函数为
$$ P(z_k \geq z_j;\forall j \neq k|z_k, \{\alpha_j\}_{j=1}^{K}) = \prod_{j \neq k}P(z_j \lt z_k) \\ = \prod_{j \neq k} e^{-e^{log\alpha_j - z_k}} $$
为了得到对于$z_k$的边缘分布，我们对$z_k$进行积分，然后我们有
$$ \begin{align} &\ P(z_k \geq z_j;\forall j \neq k) \\\ =&\ \int_{z_k}P(z_k \geq z_j;\forall j \neq k|z_k, \{\alpha_j\}_{j=1}^{K})e^{log\alpha_k - z_k - e^{log\alpha_k - z_k}}dz_k \\\ =&\ \int_{z_k}\prod_{j \neq k} e^{-e^{log\alpha_j - z_k}}e^{log\alpha_k - z_k - e^{log\alpha_k - z_k}}dz_k \\\ =&\ \int exp \left( -(\sum_{j=1}^Ke^{log\alpha_j})e^{-z_k} - z_k + log\alpha_k \right) dz_k \\\ =&\ \alpha_k \int_{z_k} exp\left(-\sum_{j=1}^{K}\alpha_j e^{-z_k} - z_k \right)dz_k \\\ =&\ \alpha_k \int_{z_k} exp\left(-\beta e^{-z_k} - z_k \right)dz_k \\\ =&\ \alpha_k \int_{z_k} exp\left(-e^{log\beta -z_k} + log \beta - z_k -log \beta \right)dz_k \\\ =& \frac{\alpha_k}{\beta}\int_{z_k}exp\left(-e^{log\beta -z_k} + log \beta - z_k \right) dz_k\\\ =&\ \frac{\alpha_k}{\beta} = \frac{\alpha_k}{\sum_{j=1}^{K}\alpha_j} \end{align} $$
所以说两个随机变量等价。
但$argmax$依旧不可导，所以我们再将$argmax$替换为$softmax$，那么我们定义
$$ X^\tau = \mathop{softmax}_k^\tau(log\alpha_k +G_k) \\ = \frac{e^{(log\alpha_k +G_k) / \tau}}{\sum_j e^{(log\alpha_j +G_j) / \tau}} $$
当温度$\tau\to \infty$，$X^{\tau}$为一个均匀分布；当温度$\tau\to 0$，$X^\tau$逐渐逼近于$X$，逐渐离散化，并且可以实现导数的传播。
dVAE训练
在了解Gumbel-softmax trick后，我们简单介绍dVAE的训练过程。我们固定自回归参数$\psi$，只考虑dVAE的编码器和解码器参数$\theta$和$\phi$。对于传入的图像$x\in R^{256 256 3}$，经过若干的卷积层将其变为形状$enc(x) \in R^{32 32 8192}$，最后一个维度代表该位置在8192个图像字符上的分布。
dVAE与标准VAE一样有两部分损失函数，一个是encoder编码的分布与我们假定隐变量分布之间的KL散度，dVAE中假定隐变量为8192个token上的均匀分布，我们将$enc(x)$与均匀分布计算KL散度既可以得到这个损失函数。
另一部分为重建损失函数，我们将从$enc(x)$给出的分布中采样，这里使用Gumbel-softmax trick进行采样从而保证梯度可以连接回编码器中的参数。采样结果输入到解码器中得到重构图像$\hat{x}$，然后定义重构损失函数即可。
自回归模型的训练
本文中*dVAE模型的作用是将图片进行离散化，并且减少序列长度。实际120亿参数量是指自回归模型。在dVAE训练结束后对于图片$x \in R^{256 256 3}$编码为整数矩阵$Z_{8192}^{32 32}$。然后将256个文本的token和这1024个图像token拼接后直接训练自回归模型。
训练细节
训练数据
模型训练使用MS-COCO数据集的一个扩展数据集 Conceptual Captions，该数据集包含330万个文本-图片对。
混合精度训练
为了节省GPU内存和提高吞吐量，大多数参数都被存储为16位精度的。使用16位精度数据进行训练，并且还要保持收敛性，是模型训练过程中最具挑战性的工作。
下图展示了resblock如何在32位精度和16位精度之间进行调节的。
图中实线是前向传播时数据调节过程，虚线是反向传播数据调节过程。
分布式优化
在训练拥有120亿参数的模型时，使用16位精度训练时，需要占用24GB左右的显存，这已经超过了 NVIDIA V100 GPU的16GB内存。文中使用了参数分割解决这个问题。使用参数分割时，在进行计算密集型操作时，同一机器上的各个GPU之间的延迟几乎可以忽略。下图展示了图片生成效果。
以下图片皆来源于openai.com[2]
标题: an illustration of a baby daikon radish in a tutu walking a dog
图片:
标题: an armchair in the shape of an avocado
图片:
标题: the exact same cat on the top as a sketch on the bottom
图片:
可以看到图片生成的效果还是非常不错的，特别是椅子那个，还有点艺术家设计的味道。大家可以到官网去实验一下模型效果[2]。
引用
[1] Ramesh, Aditya, et al. "Zero-shot text-to-image generation." arXiv preprint arXiv:2102.12092 (2021).
[2] https://openai.com/blog/dall-e/
更多推荐