Vector Quantized Diffusion Model for Text-to-Image Synthesis

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

具体来说，考虑位置 i i i处 x 0 x_0 x0的单个图像标记 x 0 i x^i_0 x0i，它采用指定码本中条目的索引，即 x 0 i ∈ 1 , 2 , . . . , K x^i_0 ∈ {1, 2,..., K} x0i∈1,2,...,K。在不引入混淆的情况下，我们在下面的描述中省略了上标 i i i。我们使用矩阵 [ Q t ] m n = q ( x t = m ∣ x t − 1 = n ) ∈ R K × K [Q_t]{mn} = q(x_t = m|x{t−1} = n) ∈R^{K×K} [Qt]mn=q(xt=m∣xt−1=n)∈RK×K 定义 x t − 1 x_{t−1} xt−1过渡到 x t x_t xt的概率。那么整个token序列的前向马尔可夫扩散过程可以写成:
q ( x t ∣ x t − 1 ) = v T ( x t ) Q t v ( x t − 1 ) q(x_t|x_{t−1}) = v^T(x_t)Q_tv(x_{t-1}) q(xt∣xt−1)=vT(xt)Qtv(xt−1)

这里 v ( x ) v(x) v(x)是一个长度为 K K K的one-hot列向量，只有条目 x x x是 1。重要的是，由于马尔可夫链的性质，可以边缘化中间步骤，直接从 x 0 x_0 x0推导出任意时间步 x t x_t xt的概率为:
q ( x t ∣ x 0 ) = v T ( x t ) Q ˉ t v ( x 0 ) , w i t h Q ˉ t = Q t ⋅ ⋅ ⋅ Q 1 q(x_t|x_0) = v^T(x_t)\bar Q_tv(x_{0}),~ with~ \bar Q_t = Q_t ···Q_1 q(xt∣x0)=vT(xt)Qˉtv(x0), with Qˉt=Qt⋅⋅⋅Q1

值得注意的是，在基于 x 0 x_0 x0时，后验扩散过程是tractable：

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

因此，这里可以表示：
Q ˉ t v ( x 0 ) = α ˉ t v ( x 0 ) + ( γ ˉ t − β ˉ t ) v ( K + 1 ) + β ˉ t \bar Q_tv(x_{0}) = \bar \alpha_tv(x_0)+(\bar \gamma_t - \bar \beta_t)v(K+1)+\bar \beta_t Qˉtv(x0)=αˉtv(x0)+(γˉt−βˉt)v(K+1)+βˉt

算法流程图如下：

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

4. 实验

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

---