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的条件序列。