为了高分辨率图像的自回归（AR）建模，矢量量化（VQ）将图像表示为离散代码序列。对于 AR模型来说，较短的序列长度对于减少计算成本和考虑编码的长程交互非常重要。然而，我们假设以前的 VQ 不能缩短代码序列并在速率-失真权衡方面生成高保真图像。在这项研究中，我们提出了两阶段框架，由残差量化 VAE（Residual-Quantized VAE，RQ-VAE）和 RQ-Transformer 组成，以有效生成高分辨率图像。在固定的码书大小下，RQ-VAE 可以精确地逼近图像的特征图，并将图像表示为离散编码的堆叠图。然后，通过预测下一个编码堆栈，RQ-Transformer 学习预测下一个位置的量化特征向量。由于 RQ-VAE 的精确逼近，我们可以将 256×256 的图像表示为 8×8 的特征图分辨率，而 RQ-Transformer 可以有效降低计算成本。因此，我们的框架在各种无条件和有条件图像生成的基准上优于现有的 AR 模型。我们的方法还具有比先前的 AR 模型更快的采样速度，以生成高质量图像。

3. 方法

我们提出了一个包含 RQ-VAE 和 RQ-Transformer 的两阶段框架，用于图像的 AR 建模（见图 2）。RQ-VAE 使用码书将图像表示为 D 个离散编码的堆叠图。然后，我们的 RQ-Transformer 自回归地预测下一个空间位置的下 D 个代码。我们还介绍了我们的 RQ-Transformer 如何解决 AR 模型训练中的曝光偏差（exposure bias） [38]。

3.1 阶段1：残差量化 VAE

在这一部分，我们首先介绍 VQ 和 VQVAE 的公式。然后，我们提出了 RQ-VAE，它可以在不增加码书大小的情况下精确逼近特征图，并解释了 RQ-VAE 如何将图像表示为离散编码的堆叠图。

3.1.1 VQ 和 VQ-VAE 的公式

其中，码本 C 由代码 k 及相应代码嵌入 e(k) 成对构成。z 是给定向量，Q(z; C) 是其量化的向量。

给定一个图像 X，VQ-VAE 提取特征图 Z = E(X)。通过对每个位置的每个特征向量应用 VQ，VQ-VAE 对 Z 进行量化，并返回其代码图 M 和其量化特征图 ^Z，如下所示：

其中 Z_hw 是在 (h, w) 处的特征向量，而 M_hw 是其代码。最后，输入被重构为 ^X = G(^Z)。

3.1.2 残差量化

与增加码书大小不同，我们采用残差量化（residual quantization，RQ）来离散化向量 z。在给定量化深度 D 的情况下，RQ 将 z 表示为一个有序的 D 个编码：

其中 C 是大小为 |C| = K 的码书，k_d 是深度 d 处 z 的编码。从第 0 个残差项 r_0 = z 开始，RQ递归计算 k_d，它是残差 r_(d-1) 的编码，下一个残差 r_d 如下所示：

其中，d = 1, ⋯ , D。此外，我们定义 ^z^(d) =

作为前 d 个代码嵌入的部分和，并且 ^z := ^z^(D) 是 z 的量化向量。

RQ 的递归量化以一种由粗到细的方式逼近向量 z。注意，^z^(1) 是码书中最接近 z 的代码嵌入 e(k_1)。然后，随后选择剩余的代码以减小每个深度的量化误差。因此，到 d 的部分和 ^z(d) 随着 d 的增加提供了更精细的逼近。

虽然我们可以为每个深度 d 单独构建一个码书，但是对于每个量化深度都使用一个共享的码书 C。共享码书对于使用 RQ 逼近向量 z 有两个优势。首先，使用单独的码书需要进行广泛的超参数搜索，以确定每个深度的码书大小，而共享码书只需要确定总的码书大小 K。其次，共享码书使得所有的代码嵌入对于每个量化深度都是可用的。因此，代码可以在每个深度上使用，以最大化其效用。

我们注意到当它们的码书大小相同时，RQ 可以比 VQ 更精确地逼近一个向量。虽然 VQ 将整个向量空间 R^(n_z) 划分为 K 个簇，但深度 D 的 RQ 将向量空间最多划分为 K^D 个簇。也就是说，深度 D 的 RQ 与具有 K^D 个代码的 VQ 有相同的划分容量。因此，我们可以增加 RQ 的 D 来替换具有指数增长的码书 VQ。

3.1.3 RQ-VAE

在图 2 中，我们提出了 RQ-VAE 来精确量化图像的特征图。RQ-VAE 也由 VQ-VAE 的编码器-解码器架构组成，但是 VQ 模块被替换为上面的 RQ 模块。具体而言，深度为 D 的 RQ-VAE 将特征图 Z 表示为代码的堆叠图 M ∈ [K]^(H×W×D) 并提取 ^Z(d) ∈ ℝ^(H×W×n_z)，它是深度为 d 的量化特征图，其中 d ∈ [D]，因此，

为简洁起见，在深度 D 处的量化特征图 ^Z^(D) 也用 ^Z 表示。最后，解码器 G 从 ^Z 重构输入图像为 ^X = G(^Z)。

我们的 RQ-VAE 可以使 AR 模型有效地生成具有低计算成本的高分辨率图像。对于固定的降采样因子 f，RQ-VAE 可以产生比 VQ-VAE 更真实的重建，因为 RQ-VAE 可以使用给定的码书大小精确逼近特征图。请注意，重建图像的保真度对于生成图像的最大质量至关重要。此外，RQ-VAE 的精确逼近允许 f 的增加和空间分辨率 (H, W) 的减小，同时保持重建质量。因此，RQ-VAE 使得 AR模型能够减少计算成本、提高图像生成速度，并更好地学习代码的长程交互。

RQ-VAE 的训练：为了训练 RQ-VAE 的编码器 E 和解码器 G，我们使用相对于损失 L = L_recon + β·L_commit 的梯度下降。重建损失 L_recon 和承诺损失 L_commit 定义为：

其中 sg[⋅] 是停梯度（stop-gradient）运算符，而直通估计器（straight-through estimator）[45] 用于通过 RQ 模块进行反向传播。请注意，L_commit 是来自每个 d 的量化误差的总和，而不是单一项。它的目标是使 ^Z^(d) 在 d 增加时顺序减少 Z 的量化误差。因此，RQ-VAE 以由粗到细的方式逼近特征图，并保持训练稳定。码书 C 通过聚类特征的指数移动平均值进行更新 [45]。

RQ-VAE 的对抗训练。RQ-VAE 还通过对抗学习进行训练，以提高重建图像的感知质量。采用基于补丁的对抗损失 [22] 和感知损失 [23]，如先前的研究 [14] 中所述。详细信息请参阅补充材料。

3.2 阶段2：RQ-Transformer

在这一部分，我们在图 2 中提出 RQ-Transformer，用于自回归地预测 RQ-VAE 的代码堆栈。在我们制定由 RQ-VAE 提取的代码的自回归建模之后，我们介绍了我们的 RQ-Transformer 如何有效学习离散代码的堆叠图。我们还提出了 RQ-Transformer 的训练技术，以防止在 AR 模型训练中的曝光偏差 [38]。

3.2.1 深度 D 的代码的自回归建模

在 RQ-VAE 提取了一个代码图 M ∈ [K]^(H×W×D) 之后，光栅扫描顺序 [34] 将 M 的空间索引重新排列为代码的 2D 数组 S ∈ [K]^(T×D)，其中 T = HW。也就是说，S_t，它是 S 的第 t 行，包含 D个代码，如下：

将 S 视为图像的离散潜在变量，AR 模型学习 p(S)，该模型按如下方式自回归分解：

3.2.2 RQ-Transformer 架构

一种朴素的方法是使用光栅扫描顺序将 S 展开为长度为 T·D 的序列，并将其馈送到传统的Transformer [46] 中。然而，这既没有利用 RQ-VAE 减小的 T 长度，也没有减少计算成本。因此，我们提出了 RQ-Transformer，以有效学习由深度 D 的 RQ-VAE 提取的代码。如图 2 所示，RQ-Transformer 由空间 Transformer 和深度 Transformer 组成。

空间 Transformer。空间 Transformer 是一堆掩码自注意块，用于提取一个总结先前位置信息的上下文向量。对于空间 Transformer 的输入，我们重复使用深度为 D 的 RQ-VAE 学习到的码书。具体而言，我们将空间 Transformer 的输入 u_t 定义为：

其中 PE_T(t) 是光栅扫描顺序中空间位置 t 的位置嵌入。请注意，第二项等于方程 5 中图像的量化特征向量。对于第一个位置的输入，我们将 u_1 定义为一个可学习的嵌入，它被视为序列的开始。在空间变换器处理序列完所有的 T 个 u_t 之后，上下文向量 h_t 对 S_(< t) 的所有信息进行编码，如下所示：

深度 Transformer：给定上下文向量 h_t，深度 Transformer 自回归地预测在位置 t 处的 D 个代码(S_t1; ⋯ ; S_tD)。在位置 t 和深度 d 处，深度变换器的输入 v_td 被定义为深度 d 之前的代码嵌入的和，如下所示：

其中 PE_D(d) 是深度 d 的位置嵌入，在每个位置 t 上共享。我们在 v_td 中不使用 PE_T(t)，因为位置信息已经编码在 u_t 中。对于 d = 1，我们使用 v_t1 = PE_D(1) + h_t。请注意，方程 12 中的第二项对应于方程 5 中深度为 d-1的量化特征向量 ^Z^(d-1)_(hw)。因此，深度 Transformer 基于先前的到 d - 1 的估计，预测了下一个代码，用于对 ^Z_t 进行更精细的估计。最后，深度 Transformer 预测条件分布 p_td(k) = p(S_td = k | S_(<t, d), S_(t, <d)) 如下：

RQ-Transformer 被训练以最小化 L_AR，即负对数似然（NLL）损失：

3.2.3 软标签和随机抽样

曝光偏差 [38] 已知会由于训练和推断中预测的差异而导致 AR 模型性能下降。在 RQ-Transformer的推断过程中，由于深度 D 的增加，预测误差也可能累积，因为随着 d 的增加，对特征向量的更精细估计变得更加困难。

因此，我们提出了从 RQ-VAE 进行软标签和随机抽样来解决曝光偏差。调度抽样 [2] 是减小差异的一种方法。然而，对于大规模的 AR 模型来说，这是不合适的，因为每个训练步骤需要多次推断，增加了训练成本。相反，我们利用了 RQ-VAE 中代码嵌入的几何关系。我们定义了在 R_(n_z) 中由向量 z 条件化的 [K] 上的分类分布，记作 Q_T(k|z)，其中 T > 0 是温度。

当 T 接近零时，Q_T 被锐化并收敛到一个 one-hot 分布 Q_0(k|z) = 1。

目标代码的软标签。基于代码嵌入之间的距离，软标签用于通过对 RQ-VAE 中代码之间的几何关系进行明确监督来改善 RQ-Transformer 的训练。对于位置 t 和深度 d，设 Z_t 是图像的特征向量，而 r_(t, d-1)是方程 4 中深度 d-1 的残差向量。然后，方程 14 中的 NLL 损失使用 one-hot 标签 Q_0(· | r_(t, d-1)) 作为 S_td 的监督。我们使用软化的分布 Q_T(· | r_(t, d-1)) 作为监督，而不是 one-hot 标签。

对 RQ-VAE 的代码进行随机抽样：除了上述的软标签外，我们提出从 RQ-VAE 中随机抽样代码图，以减小训练和推断中的差异。我们不使用方程 4 中 RQ 的确定性代码选择，而是通过从 Q_T(· | r_(t, d-1)) 中抽样来选择代码 S_td。请注意，我们的随机抽样在 T 趋近于零时等同于 RQ 的原始代码选择。随机抽样为给定图像特征映射提供了代码 S 的不同组合。

4. 实验

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/盐析白兔/article/detail/603407