NeurIPS 2022 | MoVQ: 基于Modulating Quantized Vectors的高保真图像生成

原文标题：MoVQ: Modulating Quantized Vectors for High-Fidelity Image Generation

一、问题提出

虽然两级Vector Quantized (VQ)【指VQVAE-2】生成模型允许合成高保真和高分辨率图像，但它们的量化操作符将图像中的相似patch编码到相同的索引中，导致使用现有解码器架构的相似相邻区域重复产生伪影。

尽管现有模型的性能出色，VQ-VAE或VQGAN管道也有其缺点。例如，其第二阶段通常以自回归的方式建模为序列生成过程，即在不同的空间位置逐个生成每个离散的潜码，这是非常耗时的推理。MaskGIT通过在每个步骤中基于预测置信度预测多个令牌，很好地解决了这个问题，这大大减少了自回归令牌生成器所需的步骤数。

结合空间条件归一化来调制量化向量，以便在嵌入的索引映射中插入空间变化的信息，鼓励解码器生成更逼真的图像。在不增加模型和码本成本的前提下，利用多信道量化提高离散码的重组能力。

为了在第二阶段生成离散token，采用Masked Generative Image Transformer(MaskGIT)来学习压缩潜在空间中的潜在先验分布，这比传统的自回归模型快得多。在两个基准数据集上的实验表明，modulated VQGAN能够极大地提高重建图像的质量，并提供高保真的图像生成。

二、模型

首先，训练一个编码器-解码器架构的量化器，使用学习到的codebook中的离散token将图像embedding到紧凑序列中，然后学习一个先验网络来建模离散空间中的底层分布。与流行GAN模型中的极大极小博弈不同，基于vq的生成器通过优化训练集中所有示例的负对数似然来训练，从而导致稳定的训练并避免“模式崩溃”问题。

左：MoVQ的量化器架构。我们将空间条件归一化层合并到解码器中，其中两个卷积层以point-wise的方式预测调制参数γ和β，以调制学习到的离散结构信息。右：Masked image generation。一个双向transformer被用于估计具有多个通道的离散表示的基础先验分布。

1、Modulating Quantized Vector

模型基于VQVAE（关于VQVAE可以查看：详解VQVAE：Neural Discrete Representation Learning_羊飘的博客-CSDN博客）

Spatially conditional normalization（adaptive instance normalization）

输入是embedding的离散特征，其中包含学习到的紧凑内容。这是第一个在基于vq的图像生成框架上调制量化向量并使用多通道量化的工作。

量化算子为有损，相似的patch通常embedding为相同的codebook索引，导致通过现有解码器架构合成时会出现重复的伪影：

引入一个新的空间条件归一化层来将embedding的内容传播到activations.。动机是在离散映射中添加空间变量，这样模型就可以生成合理而多样的结果，即使是相邻区域的相同量化索引

通过常规Norm对F进行归一化，然后通过从embedding向量计算的学习尺度和偏差进行调制。具体地说，解码器Gθ中第i层的F为（可对照stage1的右图MoVQ结构）：

这种Spatially conditional normalization是由自适应实例归一化(AdaIN)派生而来的。由于这种调制算子注入了空间变异的信息，它鼓励相同的量化项在不同位置上产生可信且不同的结果。

Multichannel representation

与现有的量化器将图像映射为单通道索引映射不同，这里将其转换为具有共享码本的多通道索引映射，以进一步提高图像质量。其中1 × 1 × Nc随机向量能够生成结构合理的逼真图像，表明跨通道的值包含丰富的信息。在实践中，首先沿着通道维度将编码的连续特征z∈Rh×w×nz细分为多个块，即z = {z(1)，···，z(c)}。然后，这些块中的每一个都基于equation2到codebook中最接近的codevector进行量化。

对于256 × 256的图像，通过16到16 × 16个特征的固定因子对其进行降采样。使用c = 4个并行块，等效序列表示为16 × 16 × 4，比原始图像小192x倍。注意，这样codebook中的codevector数将降为nq = nz/c。由于降低了每个code的维数，总的计算成本与全维的单通道表示相似。事实上，codebook Z的尺寸更小，维数为nq = nz/c = 256/4 = 64。

虽然较高的分辨率表示，如32 × 32 × 1或64 × 64 × 1也可以提高重建质量，但由于序列s = h × w较长，第二阶段的计算成本将较高。相比之下，多通道表示保持了更小的序列长度，即16 × 16，尽管每个序列标记现在由四个片段重新组合，每个片段对应一个量化指标。对于每个空间网格特征zij，这样的指数组合容量Kc，使得多通道表示能力远大于原始码本中的K项。

2、Modeling Prior Distribution

直接用现有方法进行第二阶段，包括传统自回归模型和最新掩码生成图像transformer(MaskGIT)

Autoregressive token generation.

图像生成可以自然地表述为自回归下一符号预测问题。特别地，最大化似然函数:

Masked token generation.

为了高效地生成索引序列，遵循MaskGIT并行预测所有token。目的是根据固定长度内的可见指标来估计所有指标的分布:

将训练过程中带有随机掩码比的掩码图s¯m输入多层双向transformer，以估计每个掩码指标的可能分布，其中还在ground-truth one-hot指标与预测指标之间计算负对数似然(NLL)。在推理时，模型运行固定的步骤，并在每次迭代中基于先前产生的可见token同时并行地预测所有token。然后，选出最可信的token，剩余的masked token将在下一步条件下对先前预测的token进行预测。

Multichannel token generation.

对于每个位置，沿着通道将块反向拼接为输入的一个token，并预测输出的c个index。在实际操作中，为各个通道和位置独立地屏蔽每个索引。与MaskGIT相比，模型中的预测指标依赖于三个方向，即来自不同位置和通道的可见索引，由于组合能力的原因导致了更大的多样性。

三、实验Experiments

1、Experimental Details

Datasets. 在无条件和类条件图像生成任务上实例化MoVQ，分别使用FFHQ和ImageNet。训练和验证设置遵循基线模型VQGAN的默认设置。在256 × 256的图像上训练所有模型。

Evaluation metrics. 使用两个最常见的评估指标Fréchet初始距离(FID)和初始分数(IS)来评估生成图像和真实图像之间的质量和多样性。也评估重建图像与原始图像之间的质量，还报告了传统的patch-level图像质量指标，包括峰值信噪比(PSNR)、结构相似指数(SSIM)，以及最新学习到的配对重建图像和真实图像的特征级LPIPS

2、Image Quantization

Network structures and implementation details.

只在256 × 256图像上训练一个具有码本Z∈R1024×64的尺度量化器，即每个64维的1024个码向量。编码器-解码建立在VQGAN1上，唯一的区别是用spatially conditional normalization layer取代了Group Normalization。在解码器的前三个块中应用了新的归一化层。按照VQGAN中的默认设置，图像总是按16的固定因子进行下采样，即从256 × 256 × 3到一个大小为16 × 16 × 4的token grids。根据VQGAN设置超参数，并在4个V100 GPU上训练，per_gpu_batch_size为48，40个epoch。第一阶段的重建性能评估：

更少的参数量，更好的性能。学习的codebook只包含1024个64维的码向量，但有趣的是，它优于使用更大codebook的其他方法。这表明，一个更好的量化器可以提高码本的使用率，而不需要贪婪地增加码本的大小。

重建图像细节：

3、Image Generation

Network structures and implementation details.

第二阶段用transformer来估计底层先验分布。配置:  layers：24，heads：16，embedding_dim：1024和hidden_state：4096。网络架构建立在VQGAN基线之上，除了最后一层预测16 × 16 × 4 × 1024的张量，其中16 × 16 × 4是预测索引的数量，1024是codebook中的条目数量。在这里，将图像生成设置为两种场景:(1)“auto”场景，其中训练和推理与VQGAN基线完全相似;(2)“mask”场景，其中训练和推理受到MaskGIT的启发。默认使用MaskGIT中的token进行采样，以便更快地进行采样。超参数遵循VQGAN中的默认设置，在4个V100 gpu上训练200个epoch，batch_size为64。

无条件生成评估：

优于大多数基于VQ的方法，具有更小的参数和更快的推断时间（采样)。虽然模型比并行工作vit-vqgan表现得更差，但他们在第二阶段使用了更大的模型(比作者模型大5倍)，训练时间更长，这并不是一个好方法。（买点是快）

条件生成：

模型为每个类别生成50个样本，共50,000个样本，用于定量评价

与同样采用多通道表示的并发工作RQ-VAE相比，即使没有专门设计用于进一步预测多通道指标的RQ-transformer，性能也具有竞争力。由于MaskGIT和vit-vqgan使用更多的gpu进行更长时间的训练，因此与MoVQ进行比较并不直接公平。（这里的买点不是效果，而是效率）

条件生成实例：

4、Discussion

Architecture、Codebook size、Latent size：

最终：

在未来，希望探索学习codebook中每个条目的语义含义。

由于模型不仅能像最先进的GAN一样生成高保真图像，而且还能提供更好的重建图像，对基于vq的模型的未来感到兴奋，并计划将其应用于更多的图像反演、插值和翻译任务。

限制。

尽管MoVQ在相同压缩比下显著提高了图像表示质量，但模型有时会生成具有高频外观的图像，而结构信息缺失。这部分是由于在模型中采用的多通道表示。因此，需要进一步研究一种更好的多通道指标建模生成模型。

论文中最后提出的codebook中的语义是一个很好的想法，但是这个工作难度大，涉及可解释性，如果能做出来，那就很棒了。