Vector Quantized Diffusion Model for Text-to-Image Synthesis

Shuyang Gu, University of Science and Technology of China, Microsoft, CVPR2022, Cited: 340, Code, Paper

1. 前言

我们提出了用于文本到图像生成的矢量量化扩散(Vector Quantized Diffusion Model；VQ-Diffusion)模型。该方法基于矢量量化变分自编码器(VQ-VAE)，其潜在空间由最近开发的去噪扩散概率模型(DDPM)的条件变体建模。我们发现这种潜在空间方法非常适合文本到图像生成任务，因为它不仅消除了现有方法的单向偏差，还允许我们合并掩码和替换扩散策略以避免错误的积累，这是现有方法的严重问题。我们的实验表明，与传统的具有相似参数数量的自回归 (AR) 模型相比，VQ-Diffusion 产生了明显更好的文本到图像生成结果。与以往的基于gan的文本到图像方法相比，我们的VQ-Diffusion可以处理更复杂的场景，大大提高了合成的图像质量。最后，我们表明我们的方法中的图像生成计算可以通过重新参数化来高效。使用传统的AR方法，文本到图像的生成时间与输出图像分辨率呈线性增加，因此即使对于正常大小的图像也非常耗时。VQDiffusion 允许我们在质量和速度之间实现更好的权衡。我们的实验表明，重新参数化的VQDiffusion模型比传统的AR方法快15倍，同时获得了更好的图像质量。

2. 整体思想

把VQVAE的token用作扩散模型训练生成，很早的一篇文章，不太好评价。

3. 方法

给定文本图像对，我们使用预训练的 VQ-VAE 获得离散图像标记 $x\in Z^N$，其中$N=hw$ 表示标记的序列长度。假设VQ-VAE码本的大小为$K$，位置$i$处的图像标记$x_i$取指定码本中条目的索引，即 x i ∈ { 1 , 2 , . . . , K } x_i∈\{1, 2, ..., K\} xi​∈{1,2,...,K}。扩撒模型逐步破坏$x_0$，离散的情况下不是加入高斯噪声而是随机替换一些$x_t$的tokens，最后变成纯噪声。采样阶段预测后验分布$q(x_{t-1}|x_t,x_0)$。然而$x_0$在离散情况下是无法估计的，这里训练了一个Transformer来估计转移分布$p_{\theta }(x_{t-1}|x_t,y)$，这里可以以$y$为条件。

具体来说，考虑位置$i$处$x_0$的单个图像标记$x_0^i$，它采用指定码本中条目的索引，即 $x_0^i\in 1,2,...,K$。在不引入混淆的情况下，我们在下面的描述中省略了上标$i$。我们使用矩阵$[Q_t{]}_{mn}=q(x_t=m|x_{t-1}=n)\in R^{K\times K}$ 定义$x_{t-1}$过渡到$x_t$的概率。那么整个token序列的前向马尔可夫扩散过程可以写成:
$$q(x_t|x_{t-1})=v^T(x_t)Q_{t}v(x_{t-1})$$
这里$v(x)$是一个长度为$K$的one-hot列向量，只有条目$x$是 1。重要的是，由于马尔可夫链的性质，可以边缘化中间步骤，直接从$x_0$推导出任意时间步$x_t$的概率为:
$$q(x_t|x_0)=v^T(x_t){\bar{Q}}_{t}v(x_0),with{\bar{Q}}_t=Q_t\cdot \cdot \cdot Q_1$$
值得注意的是，在基于$x_0$时，后验扩散过程是tractable：

这里可以看到，转移矩阵$Q$对于建模是十分重要的，作者提出，具体过程可以看原文：

因此，这里可以表示：
$${\bar{Q}}_{t}v(x_0)={\bar{\alpha }}_{t}v(x_0)+({\bar{\gamma }}_t-{\bar{\beta }}_t)v(K+1)+{\bar{\beta }}_t$$
算法流程图如下：

作者在这里是预测$x_0$，然后重参数得到$x_{t-1}$。网络结构如下，包含了文本编码器和扩散解码器，扩散图像解码器采用图像令牌$x_t$和时间步长$t$，并输出无噪声令牌分布$p_{\theta }(x_0|x_t,y)$。解码器包含几个transformer块和一个softmax层。每个transformer块包含一个full注意力，一个cross注意，以结合文本信息和前馈网络块。使用自适应层归一化(AdaLN)算子将当前时间步$t$注入到网络中，即$AdaLN(h,t)=a_{t}LayerNorm(h)+b_t$，其中$h$为中间激活，$a_t$和$b_t$由时间步嵌入的线性投影得到:

4. 实验

我们的VQ-VAE编码器和解码器遵循VQGAN的设置，它利用GAN的损失来获得更真实的图像。我们直接采用OpenImages数据集上训练的公开可用的VQGAN模型进行所有文本到图像的合成实验。它将256×256图像转换为32×32令牌。去掉无用代码后的码本大小K = 2886。我们采用CLIP模型的一个公开可用的标记器作为文本编码器，产生长度为77的条件序列。

---