摘要
在计算机视觉领域，生成式Transformer在合成高保真度和高分辨率图像方面获得了迅速普及。 然而，到目前为止，最好的生成式Transformer模型仍然将图像简单地视为一系列符元，并根据光栅扫描顺序（即逐行）对图像进行顺序解码。 我们发现这种策略既不理想也不高效。 本文提出了一种使用双向Transformer解码器的新型图像合成范式，我们将其称为MaskGIT。 在训练期间，MaskGIT学习通过关注所有方向的符元来预测随机掩码的符元。 在推理时，模型从同时生成图像的所有符元开始，然后根据之前的生成迭代地细化图像。 我们的实验表明，MaskGIT在ImageNet数据集上显著优于最先进的Transformer模型，并将自回归解码速度提高了64倍。 此外，我们说明MaskGIT可以轻松地扩展到各种图像编辑任务，例如修复、外推和图像操作。
图1
：
MaskGIT在图像合成和操作任务上的示例生成
. 我们展示了MaskGIT是一个灵活的模型，它可以在以下方面生成高质量的样本：（a）类别条件合成，（b）类别条件图像操作，
例如
用给定类别中的对象替换边界框中的选定对象，以及（c）图像外推。 这里显示的示例在三列中的分辨率分别为512
×
512、512
×
512和512
×
2560。
放大查看细节。
1
引言
图 2
:
顺序解码和 MaskGIT 的计划并行解码之间的比较。
第 1 行和第 3 行是每次迭代的输入潜在掩码，第 2 行和第 4 行是每个模型在该迭代时生成的样本。 我们的解码从所有未知代码（用浅灰色标记）开始，并逐渐用越来越分散的预测（用深灰色标记）填充潜在表示，其中预测符元的数量随着迭代次数的增加而急剧增加。MaskGIT 在 8 次迭代中完成了解码，而顺序方法需要 256 轮。
近年来，深度图像合成领域取得了很大进展。 目前，最先进的结果是生成对抗网络 (GAN)，它们能够以极快的速度合成高保真图像。 但是，它们也存在一些众所周知的问题，包括训练不稳定和模式崩溃，导致样本多样性不足。 解决这些问题仍然是开放的研究问题。
受 Transformer
[48]
和 GPT
[5]
在 NLP 中取得成功的启发，生成式 Transformer 模型在图像合成中越来越受到关注
[7, 15, 37]
。 通常，这些方法旨在将图像建模为一个序列，并利用现有的自回归模型来生成图像。 图像生成分为两个阶段：第一阶段是将图像量化为一系列离散符元（或视觉词）； 第二阶段是学习一个自回归模型（例如，Transformer）来基于先前生成的结果按顺序生成图像符元（
即
自回归解码）。 与 GAN 中使用的微妙的极小极大优化不同，这些模型是通过最大似然估计学习的。 由于设计上的差异，现有工作已经证明了它们在提供稳定训练和改进分布覆盖率或多样性方面优于 GAN 的优势。
现有的生成式 Transformer 工作主要集中在第一阶段，
即
如何量化图像以最大限度地减少信息损失，并共享从 NLP 中借鉴的相同第二阶段。 因此，即使是最先进的生成式 Transformer
[15, 35]
仍然天真地将图像视为一个序列，其中图像被展平为一个 1D 的符元序列，遵循光栅扫描顺序，
即
从左到右逐行 (
参见
图
2
）。 我们发现这种表示对于图像既非最佳也非高效。 与文本不同，图像不是按顺序排列的。 想象一下一件艺术作品是如何创作的。 画家从素描开始，然后通过填充或调整细节来逐步完善它，这与之前工作
[7, 15]
中使用的逐行打印形成鲜明对比。 此外，将图像视为一个平面序列意味着自回归序列长度呈二次方增长，很容易形成一个极长的序列——比任何自然语言句子都要长。 这不仅给建模长期相关性带来了挑战，而且还使解码变得难以处理。 例如，使用 32x32 个符元自回归地在 GPU 上生成一张图像需要相当可观的 30 秒时间。
本文介绍了一种新的用于图像合成的双向 Transformer，称为掩码生成图像 Transformer (MaskGIT)。 在训练过程中，MaskGIT 在与 BERT
[11]
中的掩码预测类似的代理任务上进行训练。 在推理时，MaskGIT 采用一种新颖的非自回归解码方法，以恒定步数合成图像。 具体来说，在每次迭代中，模型并行地预测所有符元，但只保留最自信的符元。 其余符元被屏蔽，并在下一轮迭代中重新预测。 掩码比率逐渐减小，直到所有符元在经过几次迭代的细化后生成。 如图
2
所示，MaskGIT 的解码速度比自回归解码快一个数量级，因为它只需要 8 步而不是 256 步来生成图像，并且每一步的预测都是可并行的。 此外，双向自注意力机制允许模型从所有方向的生成令牌中生成新的令牌，而不是仅根据光栅扫描顺序的先前令牌进行条件化。 我们发现掩码调度（
i.e
.
每次迭代被掩盖的图像比例）会显著影响生成质量。 我们建议使用余弦调度，并在消融研究中证实了其有效性。
在 ImageNet 基准测试中，我们通过实验证明，MaskGIT 比最先进的自回归 Transformer（
i.e
.
VQGAN）快得多（最多快 64 倍），并且能够生成更高质量的样本，用于 256
×
256 和 512
×
512 分辨率的类条件生成。 即使与领先的 GAN 模型（
i.e
.
BigGAN）和扩散模型（
i.e
.
ADM
[12]
）相比，MaskGIT 也提供了相当的样本质量，同时产生更有利的多样性。 值得注意的是，我们的模型在分类准确性得分 (CAS)
[36]
和 FID
[23]
上为合成 512
×
512 图像建立了新的最先进水平。 据我们所知，本文首次提供了证明掩码建模在常见 ImageNet 基准测试上进行图像生成的有效性的证据。
此外，MaskGIT 的多方向性质使其很容易扩展到图像操作任务，而这些任务对于自回归模型来说是困难的。 图
1
显示了类条件图像编辑的新应用，其中 MaskGIT 根据给定类别重新生成边界框内的内容，同时保持上下文（框外）不变。 这个任务对于自回归模型来说要么不可行，要么对于 GAN 模型来说很困难，但对于我们的模型来说却很简单。 在数量上，我们通过将 MaskGIT 应用于图像修复和任意方向上的图像外推来证明这种灵活性。 即使我们的模型不是为这些任务而设计的，它在每个任务上也获得了与专用模型相当的性能。
2
相关工作
2.1
图像合成
深度生成模型
[29, 45, 17, 53, 41, 12, 46, 34]
在图像合成任务中取得了许多成功。 基于 GAN 的方法在生成高保真样本方面展现出惊人的能力
[17, 4, 27, 53, 44]
. 相比之下，基于似然的模型，例如变分自动编码器 (VAE)
[29, 45]
、扩散模型
[41, 12, 24]
和自回归模型
[46, 34]
，提供了分布覆盖范围，因此可以生成更多样化的样本
[41, 45, 46]
.
但是，直接在像素空间中最大化似然可能很具有挑战性。 因此， VQVAE
[47, 37]
提议分两个阶段在潜在空间中生成图像。 首先是
标记化
，它试图将图像压缩成离散的潜在空间，主要包含三个部分：
•
一个编码器
E
，它学习将图像
x
∈
ℝ
H
×
W
×
3
标记化成潜在嵌入
E
(
x
)
，
•
一个码本
𝐞
k
∈
ℝ
D
,
k
∈
1
,
2
,
⋯
,
K
，用于最近邻查找，将嵌入量化为视觉标记，以及
•
一个解码器
G
，它根据视觉标记
𝐞
预测重建后的图像
x
^
。
在第二阶段，它首先使用深度自回归模型预测视觉标记的潜在先验，然后使用第一阶段的解码器将标记序列映射到图像像素。 由于这种两阶段方法的有效性，许多方法都遵循这种范式。DALL-E
[35]
使用 Transformer
[48]
来改进第二阶段中的标记预测。VQGAN
[15]
在第一阶段添加了对抗性损失和感知损失
[26, 54]
，以提高图像保真度。 我们的同期研究 VIM
[51]
提议使用 VIT 主干
[13]
来进一步改进标记化阶段。 由于这些方法仍然使用自回归模型，因此第二阶段的解码时间与标记序列长度成正比。
2.2
使用双向 Transformer 进行掩码建模
Transformer 架构
[48]
最初是在 NLP 中提出的，最近已扩展到计算机视觉
[13, 6]
。Transformer 由多个自注意力层组成，允许捕获序列中所有元素对之间的交互。 特别是，BERT
[11]
引入了掩码语言建模 (MLM) 任务用于语言表示学习。BERT
[11]
中使用的双向自注意力允许利用来自两个方向的上下文预测 MLM 中的掩码符元。 在视觉方面，BERT
[11]
中的掩码建模已被扩展到图像表示学习
[21, 2]
，其中图像被量化为离散符元。 然而，由于使用双向注意力执行自回归解码的难度，很少有工作成功地将相同的掩码建模应用于图像生成
[56]
。 据我们所知，本文提供了第一个证据，证明了掩码建模在常见 ImageNet 基准上用于图像生成的有效性。 我们的工作灵感来自 NLP 中的双向机器翻译
[16, 20, 19]
，我们的新颖之处在于提出的新的掩码策略和解码算法，正如我们的实验所证实的那样，这些策略和算法对于图像生成至关重要。
图 3
:
管道概述。
MaskGIT 遵循一个两阶段设计，包括 1) 将图像标记化为视觉符元的标记器，以及 2) 执行 MVTM 的双向 Transformer 模型，即学习预测随机掩码的视觉符元。
3
方法
我们的目标是设计一个新的图像合成范式，利用并行解码和双向生成。
我们遵循
2.1
中讨论的两阶段方案，如图
3
所示。 由于我们的目标是改进第二阶段，因此我们对第一阶段采用与 VQGAN 模型
[15]
相同的设置，并将标记步骤的潜在改进留待将来工作。
对于第二阶段，我们建议通过
掩码视觉符元建模
(MVTM) 来学习双向 Transformer。 我们在
3.1
中介绍了 MVTM 训练，在
3.2
中介绍了采样过程。 然后，我们将在
3.3
中讨论掩码设计这一关键技术。
3.1
训练中的 MVTM
令
𝐘
=
[
y
i
]
i
=
1
N
表示将图像输入到 VQ 编码器后获得的潜在符元，其中
N
是重塑后的符元矩阵的长度，
𝐌
=
[
m
i
]
i
=
1
N
是相应的二进制掩码。 在训练期间，我们对符元的子集进行采样，并将它们替换为特殊的
[MASK]
符元。 如果
m
i
=
1
，则用
[MASK]
替换符元
y
i
；否则，当
m
i
=
0
时，
y
i
将保持不变。
采样过程由掩码调度函数
γ
(
r
)
∈
(
0
,
1
]
参数化，并按如下方式执行：我们首先从
0
到
1
采样一个比率，然后在
𝐘
中均匀地选择
⌈
γ
(
r
)
⋅
N
⌉
个符元来放置掩码，其中
N
是长度。 掩码调度会显著影响图像生成的质量，这将在
3.3
中讨论。
将
Y
𝐌
¯
表示为将掩码
𝐌
应用于
𝐘
后得到的结果。 训练目标是最小化掩码符元的负对数似然：
具体来说，我们将掩码后的
Y
𝐌
¯
输入到多层双向 Transformer 中，以预测每个掩码符元的概率
P
(
y
i
|
Y
𝐌
¯
)
，其中负对数似然被计算为真实 one-hot 符元和预测符元之间的交叉熵。 注意与自回归建模的关键区别：MVTM 中的条件依赖关系有两个方向，这允许图像生成通过关注图像中的所有符元来利用更丰富的上下文。
3.2
迭代解码
在自回归解码中，符元是根据先前生成的输出按顺序生成的。 这个过程不可并行化，因此对于图像来说非常慢，因为图像标记长度，例如
。
256 或 1024，通常比语言的长度大得多。
我们引入了一种新颖的解码方法，其中图像中的所有符元都以并行方式同时生成。 由于 MTVM 的双向自注意力，这是可行的。
从理论上讲，我们的模型能够在一遍推断中推断所有符元并生成整个图像。 我们发现这具有挑战性，因为与训练任务不一致。 下面介绍了所提出的迭代解码。 为了在推断时生成图像，我们从一个空白画布开始，所有符元都被屏蔽，
即
.
Y
𝐌
(
0
)
. 对于迭代
t
，我们的算法按如下方式运行：
1.
预测。
给定当前迭代中被屏蔽的符元
Y
𝐌
(
t
)
，我们的模型预测所有被屏蔽位置的概率，表示为
p
(
t
)
∈
ℝ
N
×
K
，并行。
2.
采样。
在每个被屏蔽位置
i
，我们根据其预测概率
p
i
(
t
)
∈
ℝ
K
在代码本中所有可能的符元上采样一个符元
y
i
(
t
)
。 采样一个符元
y
i
(
t
)
后，其相应的预测分数用作“置信度”分数，指示模型对此预测的置信度。 对于
Y
𝐌
(
t
)
中未被屏蔽的位置，我们简单地将其置信度分数设置为
1.0
。
3.
屏蔽计划。
我们根据屏蔽计划函数
γ
通过
n
=
⌈
γ
(
t
T
)
N
⌉
计算要屏蔽的符元数量，其中
N
是输入长度，
T
是迭代总数。
4.
屏蔽。
我们通过在
Y
𝐌
(
t
)
中屏蔽
n
个符元来获得
Y
𝐌
(
t
+
1
)
。 迭代
t
+
1
的屏蔽
𝐌
(
t
+
1
)
是根据以下公式计算的：
