（2021|CVPR，VQ-GAN，感知损失，对抗训练）驯服 Transformer 以进行高分辨率图像合成

Taming Transformers for High-Resolution Image Synthesis

公和众与号：EDPJ（添加 VX：CV_EDPJ 或直接进 Q 交流群：922230617 获取资料）

目录

0. 摘要

3. 方法

3.1 学习图像组成部分的有效码本，以在 transformer 中使用

3.2 使用 transformer 学习图像的构成

4. 实验

---

0. 摘要

设计用于在序列数据上学习长程交互的 Transformer 在各种任务中仍然展现出最先进的结果。与CNN 不同，它们不包含优先考虑局部交互的归纳偏差（inductive bias）。这使它们具有表达能力，但对于长序列（如高分辨率图像）来说，在计算上是不可行的。我们展示了如何将 CNN 的归纳偏差的有效性与 Transformer 的表达能力相结合，使它们能够对高分辨率图像进行建模和合成。我们展示了如何（i）使用 CNN 学习图像组成元素的上下文丰富词汇表，并且进而（ii）利用Transformer 在高分辨率图像中高效地建模它们的组合。我们的方法可以轻松应用于条件合成任务，在这些任务中，非空间信息（如对象类别）和空间信息（如分割）都可以控制生成的图像。特别是，我们展示了使用 Transformer 进行语义引导的百万像素图像的合成的初步结果，并在类别条件 ImageNet 上取得了自回归模型中的最先进成绩。代码和预训练模型可以在 https://git.io/JnyvK 找到。

3. 方法

我们的目标是利用 transformer模型 [74] 的高度有前途的学习能力，并将其引入到高分辨率图像合成，达到百万像素的范围。之前的工作 [55, 8] 将 transformer 应用于图像生成，对于 64 x 64 像素大小的图像显示了有前途的结果，但由于序列长度呈二次增加的成本，不能简单地扩展到更高的分辨率。

高分辨率图像合成需要一个理解图像全局构图的模型，使其能够生成局部逼真且全局一致的图案。因此，我们不再用像素表示图像，而是将其表示为从码本中感知丰富的图像组成部分的组合。通过学习一个有效的编码，如 3.1 节所述，我们可以显著减少组合的描述长度，从而能够有效地在图像中建模它们的全局相互关系，使用 transformer 体系结构，如 3.2 节所述。这种方法，如图 2 所总结的，能够在无条件和条件设置下生成逼真且一致的高分辨率图像。

3.1 学习图像组成部分的有效码本，以在 transformer 中使用

为了利用富有表现力的 transformer 架构进行图像合成，我们需要以序列的形式表达图像的组成部分。与建立在单个像素上的方法不同，复杂性要求使用一个离散的、学到的表示码本，使得任何图像 x ∈ R^(H * W * 3) 都可以由码本条目 z_q ∈ R^(h * w * n_z) 的空间集合表示，其中 n_z 是编码的维度。一个等效的表示是一系列 h * w 的索引，它们指定了学习码本中相应的条目。为有效地学习这样一个离散的空间码本，我们建议直接结合 CNN 的归纳偏差，并纳入神经离散表示学习 [72] 的思想。首先，我们学习一个由编码器 E 和解码器 G 组成的卷积模型，使它们一起学习用来表示图像的来自学到的离散码本

的编码。更准确地说，我们通过 ^x = G(z_q) 来近似给定的图像 x。我们通过使用编码 ^z = E(x) ∈ R^(h * w * n_z) 和后续的每个空间编码 ^z_ij ∈ R^(n_z) 的逐元素量化 q(·)，将其投射到最接近的码本条目 z_k 上，从而获得 z_q：

重建 ^x ≈ x 由以下公式给出：

不可微的量化操作在方程（3）中进行的反向传播是通过直通梯度估计器（straight-through gradient estimator）实现的，该估计器简单地将梯度从解码器复制到编码器 [3]，从而可以通过如下损失函数端到端地训练模型和码本。 

在这里，L_rec 是一个重构损失，sg[·] 表示 stop-gradient 操作，而第三项是所谓的 “承诺损失（commitment loss）” [72]。 

学习一个感知丰富的码本。使用 transformer 把图像表示为潜在图像组成部分的分布，需要我们推动压缩极限并学习一个丰富的码本。为此，我们提出了 VQGAN，这是原始 VQVAE 的一个变体，并使用鉴别器和感知损失 [40, 30, 39, 17, 47]，以在增加的压缩率下保持良好的感知质量。请注意，这与之前的作品形成对比，之前的作品仅在浅层量化模型的基础上应用基于像素的 [71, 61]和基于 transformer 的自回归模型 [8]。更具体地说，我们通过感知损失替换了 [72] 中 L_rec 使用的 L2 损失，并引入了一个基于 patch 的鉴别器 D [28] 的对抗训练过程，D 旨在区分真实和重建的图像：

找到最优压缩模型 Q* = {E*, G*, Z*} 的完整目标如下：

其中，我们根据下式计算自适应权重 𝜖： 

其中，L_rec 是感知重建损失 [81]，

表示相对于解码器的最后一层 L 的输入的梯度，而 δ = 10^(-6) 用于数值稳定性。为了从各处聚合上下文，我们在最低分辨率上应用一个单一的注意力层。这个训练过程显著减少了在展开潜在代码时的序列长度，从而使强大的 transformer 模型能够应用。

3.2 使用 transformer 学习图像的构成

潜在 Transformer：有了 E 和 G，我们现在可以用它们的编码的码本索引来表示图像。更准确地说，图像 x 的量化编码由 z_q = q(E(x)) ∈ R^(h * w * n_z) 给出，等效于从码本中得到的索引序列 s ∈ {0, ..., |Z|-1}^(h x w)，这是通过用码本 Z 中的索引替换每个编码获得的：

通过将序列 s 的索引映射回它们对应的代码簿条目，z_q = z_(s_ij) 可以轻松恢复，并解码为图像 ^x = G(z_q)。

因此，在选择序列 s 中索引的某种顺序后，图像生成可以被制定为自回归地预测下一个索引：给定索引 s_(< i)，transformer 学习预测可能的下一个索引的分布，即

以计算完整表示的似然性

这使我们能够直接最大化数据表示的对数似然： 

有条件的合成。在许多图像合成任务中，用户通过提供额外信息要求对生成过程进行控制，从而合成一个例子。这个我们将称之为 c 的信息可以是描述整体图像类别的单个标签，甚至是另一张图像。任务是学习给定此信息 c 时序列的似然性：

如果条件信息 c 具有空间范围，我们首先学习另一个 VQGAN，再次获得一个基于索引的表示

其中新获得的码本为 Z_c。由于 transformer 的自回归结构，我们可以简单地将 r 前置到 s，并将负对数似然的计算限制在条目

上。这种 “仅解码器” 策略也已成功用于文本摘要任务 [44]。

生成高分辨率图像。Transformer 的注意机制对其输入序列 s 的长度 h x w 有一定的限制。虽然我们可以调整 VQGAN 的下采样块数 m，将尺寸为 H x W 的图像缩小为 h = H/2^m x w = W/2^m，但我们观察到在 m 的临界值（取决于考虑的数据集）之后 ，重建质量出现了降级。因此，为了在百万像素的范围内生成图像，我们必须逐块进行操作，并在训练期间裁剪图像以将 s 的长度限制为最大可行的大小。为了采样图像，我们使用 transformer  以图 3 所示的滑动窗口方式。我们的 VQGAN 确保可用的上下文仍足够忠实地对图像进行建模，只要数据集的统计特性大致上是空间不变的，或者有空间的条件信息可用。实际上，这不是一个限制性的要求，因为当这被违反时，即在对齐数据上进行无条件图像合成时，我们可以简单地在图像坐标上进行条件设置，类似于[42]。

4. 实验

结果表明，在相同的训练时间内，transformer 在所有任务上始终优于 PixelSNAIL，并且当以相同的步数进行训练时，差距进一步增加。Transformer 的优势在我们提出的两阶段设置中得到了体现。 

我们报告了在 top-k / top-p 采样中的最佳 k/p 的 FID 和 Inception Scores。根据 [61]，我们还可以通过分类器拒绝进一步提高质量，该分类器仅保留在分类器得分方面排名前 m 个的 n 个样本，即接受率为 m/n。

我们观察到我们的模型在 FID 和 IS 方面优于其他自回归方法（VQVAE-2、DCTransformer），甚至在高接受率下也超过了 BigGAN 和 IDDPM，并在更低的接受率下得分接近最先进技术，参见图9。

正如预期的那样，更大版本的 VQGAN（无论是在更大的码本大小还是增加的编码长度方面）进一步提高了性能。在我们的模型中使用与 VQVAE-2 相同的分层码本设置可以提供最佳的重建 FID，尽管以非常长的、因此而不切实际的序列为代价。