本文目录

1 VIM：使用 ViT 改进的 VQGAN 进行矢量量化图像生成
(来自谷歌)
1 VIM 论文解读
1.1 VIM 的诞生背景
1.2 ViT-VQGAN 总体框架
1.3 VQGAN 编码器使用 Vision Transformer
1.4 ViT-VQGAN 的训练损失函数
1.5 矢量量化图像建模
1.6 图像矢量量化实验结果
1.7 图像生成实验结果
1.8 无监督学习实验结果

太长不看版

VQGAN 使用矢量量化的办法，把图片经过分词器 (tokenizer) 之后变成 image tokens，再通过自回归的方式去建模预测下一个 token。这种 next-token prediction 的范式是受到语言模型的启发，这样训练出的模型能够在一众任务上展示出现象级别的生成能力。

本文探索了在 VQGAN 里面，把图像的 Encoder 换成 Vision Transformer (之前是 ConvNet)，得到的模型称为 ViT-VQGAN。作者首先对原始的 VQGAN 进行了多种改进，包括从 architecture 到 codebook，并获得了更好的效率和重建的质量。本文改进之后的 ViT-VQGAN 进一步改进了矢量量化图像建模的任务，包括无条件图像生成、class-conditioned 图像生成和无监督的表征学习。比如对于 256×256 分辨率的 ImageNet 生成任务，ViT-VQGAN 获得了 175.1 的 IS 和 4.17 的 FID，相比于 VQGAN 的 70.6 的 IS 和 17.04 的 FID 有了显著的提升。

基于 ViT-VQGAN 和无监督预训练的范式，ViT-VQGAN 也可以做类似于 Image GPT (iGPT) 的训练。ImageNet 预训练的 VIM-L 在 Linear Probing 精度方面显著优于 iGPT-L，对于相似的模型大小，Linear Probing 精度度从 60.3% 提高到 73.2%。VIM-L 也优于 iGPT-XL。

本文做了哪些具体的工作

1. 指出提升图像生成和理解性能的关键因素是拥有一个更好的图像量化器 (image quantizer)，其计算效率和重建质量都很好。高效的图像量化器可以加快第2阶段的训练，首先对图像使用随机增强策略，然后使用图像量化器的 Encoder 来获得输入 image tokens。而且，一个重建质量较好的图像量化器可以减少信息丢失，这个性质对于图像理解任务至关重要。
2. 从3个角度评估了所提出的 ViT-VQGAN 和 VIM。1) 根据重建质量指标评估图像量化器，包括重建图像的 距离、 距离、log-laplace 距离以及 Inception Score (IS) 和 FID。2) 基于 FID 和 IS 评估了图像量化器的无条件和 class-conditioned 图像生成能力。3) 基于 linear-probe 精度评估了学习到的图像表征。

1 VIM：使用 ViT 改进的 VQGAN 进行矢量量化图像生成

论文名称：Vector-quantized Image Modeling with Improved VQGAN (ICLR 2022)

论文地址：

https://arxiv.org/pdf/2110.04627

1.1 VIM 的诞生背景

受到自然语言处理领域的启发，视觉领域也探索了使用 GPT 风格的模型对图片进行自回归预训练，这种训练的范式没有结合图像这种 2D 结构的先验。比如 iGPT[1]就是通过自回归预测像素进行图片的预训练。但是这种编码方式非常不适合长序列的高分辨率图像 (例如 224×224 分辨率的每张图像就有 50,176 个 token)。与语言模型相比，就需要更多的内存和计算进行训练。因此，iGPT 的最高分辨率为 64×64，用于大规模图像识别，严重限制了其表征能力。

一类很有效的图像生成模型，比如 VQVAE[2]，DALL-E[3]，VQGAN[4]等通过将图片量化为离散的 latent tokens，并对其进行自回归的建模，可以获得很好的图像生成结果。在这些方法中，第1阶段先通过一个卷积神经网络 (ConvNet) 对图片进行编码；第2阶段再通过 ConvNet 或者 Transformer 编码 latent 的表征。这个范式已被证明对于图像生成有效，但对判别任务研究有限。

1.2 ViT-VQGAN 总体框架

本文提出的 ViT-VQGAN 也属于一种矢量量化的方法，全称为 Vector-quantized Image Modeling (VIM)：将图片量化为离散的 tokens，它可以应用于图像生成和图像理解任务。VIM 可以分为两个阶段：

• 阶段1：图像量化。 给定一个分辨率为 256×256 的图像，基于 Vision Transformer 的 VQGAN 将其编码为32×32 的离散编码。codebook 的大小为 8192。作者对 VQGAN 提出了从 architecture 到 codebook 的多个改进。ViT-VQGAN 在图像重建指标 IS 和 FID 方面更加有效，并提高了重建的保真度。
• 阶段2：矢量量化图像建模。 作者训练了一个 Transformer 模型来自回归地预测 32×32 = 1024 个 image tokens，其中 image tokens由阶段 1 的 ViT-VQGAN 完成编码。对于无条件图像合成或无监督学习，作者预训练了一个 Decoder-Only 的 Transformer 模型来预测下一个 token。对于 class-conditioned 的图像合成任务，在 image token 之前添加了一个 class-id token。为了评估无监督学习的质量，作者对中间 Transformer 特征进行平均并使用 Linear Probing 的方式评测。

VQVAE 是一种基于 ConvNet 的自动编码器，其 latent space 是离散可学习矩阵，通过 straight-through estimation 进行端到端的训练。VQGAN 是一种通过引入判别器产生的对抗性损失来改进 VQVAE 的模型。本文对 VQGAN 进行了进一步的改进，以提高效率和提高重建质量。

1.3 VQGAN 编码器使用 Vision Transformer

VQVAE 和 VQGAN 使用的核心网络架构来编码和重建图像是 ConvNet。在 ViT-VQGAN 中，作者提出直接使用 Vision Transformer (ViT) 来替换 ConvNet 作为图像的编码器和解码器，如上图1所示。在给定足够数据的前提下，作者发现 ViT-VQGAN 受到卷积的归纳偏置的约束较少。而且，ViT-VQGAN 在加速器上产生更好的计算效率，并产生更高质量的重建。

ViT-VQGAN 的 Encoder 首先将 8×8 不重叠的图像 Patches 映射到 image tokens，然后是 Transformer Block，这一系列过程将 256×256 分辨率的图像编码为 32×32=1024 的 token 序列。Decoder 执行逆过程，将每个 image token 从 latent 变量映射回 8×8 的图像 Patch，并将它们重新分组重排为 256×256 的图像。在 Transformer 的输出端使用了一个2层的 Feed-Forward Network，并在中间使用 tanh 激活函数。使用 logit-laplace 损失，因此 sigmoid 激活函数应用于 Decoder。

由于码本 (codebook) 的初始化不佳，原始的 VQVAE 的 codebook 的使用率通常较低。因此在预训练期间一大部分的 code 很少使用。而且，有效码本大小的减小会导致阶段1的量化器训练和阶段2的多样性较差。本文的 codebook 尺寸为 8192。矢量量化的训练目标定义如下：

式中, 是 stop-gradient 算子, 是实验中设置为 0.25 的 commitment loss 的超参数, 是码本向量。那么到底是选择码本中的哪个向量? 是通过欧氏距离判断的 (取欧氏距离最小的): 。

在低维度进行 code 查找。 作者引入了一个从编码器的输出到更低维度的 latent 空间的线性投影 (Linear Projection)，用于 code 的索引和查找。比如说从 768 维度的向量降低到 32 或者 8 维的向量。作者发现这个做法可以促进码本的使用。这个做法相当于在低维空间上查找与输入编码最接近的变量，然后将匹配的 latent code 再投影到高维空间中。作者的实验表明将查找空间的维度从 256 减少到 32 可以不断地提高重建的质量。

归一化的码本。 作者也对 latent 的编码 和码本 latent 变量 使用了 范数操作。通过映射所有的 latent 变量到球上, 归一化使 和 之间欧几里得距离演化为余弦相似度 , 进一步提高了训练稳定性和重建质量。

1.4 ViT-VQGAN 的训练损失函数

最终的训练损失函数为：

式中, 的定义是式1, 为矢量量化损失。 是 logit-laplace 损失, Logit-Laplace 损失可以看作是 归一化损失, 它假设像素级的噪声是拉普拉斯分布的, 损失假设噪声是高发现 logit-laplace 损失有助于码本使用, 而 损失和感知损失显着有助于 FID。

使用基于 VGG 的感知损失的一个需要注意的地方是 VGG 的训练过程使用了分类损失，可能对于分类任务而言，泄露了监督信号。因此，对于所有的无监督学习结果，作者排除了 ViT-VQGAN 训练期间感知损失。对于所有无条件和 class-conditioned 的图像合成，作者使用用感知损失训练 ViTVQGAN 量化器，因为会得到高保真的重建结果。

1.5 矢量量化图像建模

第1阶段训练好 ViT-VQGAN 之后，图像被编码为一堆离散的 latent code ids。在第2阶段中，使用 Decoder-Only 的 Transformer 模型来对图像数据 进行自回归范式建模：

其中, 是可学习的权重。训练目标是最小化数据的负对数似然 。

图 2 总结了 Transformer 的架构配置。作者首先将离散的 image token id 嵌入到可学习嵌入空间中，并使用加性可学习的 2D 位置编码。将一堆 Transformer Block 应用于整个序列。在所有残差、激活和注意力输出中都使用了 0.1 的 dropout rate。在所有 Transformer Block 的最后一层，使用了一个额外的 Layer Norm。

图像生成

通过预训练的生成 Transformer 模型，通过从输出的 Softmax 分布中进行采样 token 来实现无条件图像生成。将采样的标记输入 ViT-VQGAN 的解码器以解码输出图像。阶段 1 的 ViT-VQGAN 将分辨率为 256×256 的输入图像编码为 32×32 的 latent code，码本大小为 8192；阶段 2 的 Transformer 输入是序列化的 image tokens，序列长度是 1024。

class-conditioned ImageNet 生成也是评判图像合成模型性能的一个广泛使用的基准。通过在 image tokens 之前添加一个 class id token，可以做 class-conditioned ImageNet 生成。在采样期间，在第1个位置提供了一个 class id token，以自回归的范式解码剩余的 image tokens。

无监督学习

对于图像理解任务，作者将输入的所有 image tokens 输入到预训练的 Transformer 中，并获得 1024 个 tokens 的序列。一共 LL 层，作者在特定的 ll 层处获取了输出，对特征序列取平均值，输入 Softmax 层把平均特征投影到 logits。这个做法和 iGPT 类似。但是作者只使用了一层的 Transformer Block 输出，而不是像 iGPT 那样连接不同 Blocks 的输出。作者发现 Linear Probing 的判别特征通常接近所有 Transformer 的中间层的特征。

1.6 图像矢量量化实验结果

作者分别在 CelebA-HQ、FFHQ 和 ImageNet 这 3 个数据集上训练了 ViT-VQGAN。对于 CelebA-HQ 和 FFHQ，作者遵循 VQGAN 默认的训练和验证集的分割方式。对于阶段 1 图像量化，Encoder 或者 Decoder 尝试了3种不同的架构大小，Small、Base 和 Large。最小的 ViT-VQGAN-SS 具有 Small-size Encoder 和 Small-size Decoder，而 ViT-VQGAN-BB 具有 Base-size Encoder 和 Base-size Decoder。最大的 ViT-VQGAN-SL 具有非对称的 Small-size Encoder 和 Large-size Decoder，其动机是第 2 阶段训练只需要 ViT-VQGAN Encoder 前向传播。

图 3 显示了 ImageNet、CelebA-HQ 和 FFHQ 数据集上验证集中重建图像和原始图像之间的 FID。与普通 VQGAN 相比，ViT-VQGAN 能够以 8192 的大码本大小实现更好的 FID。

下图4是 ImageNet 上的消融实验结果。在不同的模型大小中，ViT-VQGAN-SS 的性能比 ViTVQGAN-BB 和 ViT-VQGAN-SL 差，但实现了更好的吞吐量。与基于 ViT 的 VQGAN 相比，基于 CNN 的 VQGAN 架构在质量和吞吐量上都更差。对于码本学习，当 latent 维度从 256 减少到 16 或 8 时，始终获得更好的重建质量。

1.7 图像生成实验结果

在预训练的 ViT-VQGAN 图像量化器的基础上，作者训练阶段 2 的 Transformer 进行无条件和 class-conditioned 图像合成。作者对阶段 1 使用 ViT-VQGAN-SS 的默认模型大小，在第 2 阶段使用 VIM-Large，因为发现在保持阶段 1 转换器轻量级的同时，将更多的计算放在阶段 2 中是有益的。

在 CelebA-HQ 和 FFHQ 上的无条件图像生成结果如图 5 所示。ViT-VQGAN 在 CelebAHQ 上实现了 7.0 的 FID，在 FFHQ 上实现了 5.3 的 FID，明显优于 VQGAN。

图 6 为 ImageNet 上的 class-conditioned 图像生成模型。ViT-VQGAN-SS 实现了 175.1 的 IS 和 4.17 的 FID，比普通 VQGAN 提高了 70.6 的 IS 和 17.04 的 FID。当使用 classifier-based rejection sampling 时，基于 ViT-VQGAN 的模型进一步实现了 3.04 的最佳 FID 和 321.7 的最佳 IS。

下图7是定性采样的结果。

1.8 无监督学习实验结果

作者评估了学习到的图像表征，评估方法是基于 ImageNet 的 Linear Probing 的精度。作者遵循与 ImageNet 上的无条件图像合成模型相同的训练超参数，并使用 ViT-VQGAN-SS 图像量化器来更好地训练吞吐量。作者对一共 个 Transformer Block 中的第 个的特征做平均池化。与 iGPT 中的发现类似，中间 Block 的表征具有更好的 Linear Probing 精度。对于 VIM-Base，使用第 15 个 Transformer Block 的特征 (总共 36 个)。对于 VIM-Large，使用第 10 个 Transformer Block 的特征 (总共 24 个)。

实验结果如图 8 所示。本文提出的 VIM 与 ViT-VQGAN 的 Linear Probing 的精度优于其他生成式预训练方法，如 iGPT，并与 BYOL 和 DINO 等判别预训练方法具有竞争力。具体来说，ImageNet 预训练的 VIM-L 的性能显著超过了 iGPT-L，对于相似的模型大小，Linear Probing 的精度从 60.3% 增加到 73.2%。VIM-L 也优于 iGPT-XL，其更大，并使用了额外的网络图像数据进行训练。作者还比较了阶段 1 中不同图像量化器，比如基于 CNN 的 VQGAN，和预训练的 DALL-E dVAE 图像量化器。这些结果都比 ViT-VQGAN 量化器差。

参考

1. Generative Pretraining From Pixels
2. Neural discrete representation learning
3. Zero-shot text-to-image generation
4. Taming transformers for high-resolution image synthesis

公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT