133

文章目录

  1. 背景
  2. 模型介绍
    1. 模型训练
      1. Gumbel-softmax trick
      2. dVAE训练
      3. 自回归模型的训练
    2. 训练细节
      1. 训练数据
      2. 混合精度训练
      3. 分布式优化
  3. 引用

2021年1月份,google 发布了一个新的模型DALL-E,该模型拥有120亿参数,可以从一段文本描述生成生动的图片。

背景

在文本生成图片的任务中,人们一直致力于在确定的小数据集下寻找更好的网络结构,然而一些假设通常会引入复杂的模型结构或者繁琐的损失函数。本文提出一种简单的自回归策略,利用transformer将文本和图片同等建模,即使在zero shot情况下也可以达到SOTA效果。GPT3证明了在文本上使用大规模参数和大规模数据对模型带来的提高,本文使用120亿参数的模型在2.5亿数据上进行了训练,证明了大规模模型在文本到图像生成任务上也有很好的效果。

模型介绍

本文目标是利用transformer结构同时编码文本和图像,但是图像是无法像文本一样编码的,若将像素作为单位则会导致序列长度太长,极大的浪费了计算资源,并且过长的序列也不利于transformer的学习。为了使用较短的序列表达图像,我们使用了一个两阶段训练的策略:

  1. 第一阶段我们利用离散VAE(dVAE)将256 256大小的图片编码为32 32个图像字符,每个字符有8192个可能的值
  2. 第二阶段我们将文本的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)$,这样我们实现将采样过程放进带有参数的函数内部,从而避免了梯度无法传播的情况。

Gumbel概率密度函数

但是,当我们$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/