VQ-VAE（2018-05：Neural Discrete Representation Learning）

本篇参考：

轻松理解 VQ-VAE：首个提出 codebook 机制的生成模型（周弈帆教授）

近两年，有许多图像生成类任务的前沿工作都使用了一种叫做"codebook"的机制。追溯起来，codebook机制最早是在VQ-VAE论文中提出的。相比于普通的VAE，VQ-VAE能利用codebook机制把图像编码成离散向量，为图像生成类任务提供了一种新的思路。VQ-VAE的这种建模方法启发了无数的后续工作，包括声名远扬的Stable Diffusion。

在这篇文章中，我将先以易懂的逻辑带领大家一步一步领悟VQ-VAE的核心思想，再介绍VQ-VAE中关键算法的具体形式，最后把VQ-VAE的贡献及其对其他工作的影响做一个总结。通过阅读这篇文章，你不仅能理解VQ-VAE本身的原理，更能知道如何将VQ-VAE中的核心机制活学活用。

1.从 AE 到 VQ-VAE

让我们来整理一下VQ-VAE的工作过程。

训练VQ-VAE的编码器和解码器，使得VQ-VAE能把图像变成「小图像」，也能把「小图像」变回图像。
训练PixelCNN，让它学习怎么生成「小图像」。
随机采样时，先用PixelCNN采样出「小图像」，再用VQ-VAE把「小图像」翻译成最终的生成图像。

到这里，我们已经学完了VQ-VAE的核心思想。让我们来总结一下。VQ-VAE不是一个VAE，而是一个AE。它的目的是把图像压缩成离散向量。或者换个角度说，它提供了把大图像翻译成「小图像」的方法，也提供了把「小图像」翻译成大图像的方法。这样，一个随机生成大图像的问题，就被转换成了一个等价的随机生成一个较小的「图像」的问题。有一些图像生成模型，比如PixelCNN，更适合拟合离散分布。可以用它们来完成生成「小图像」的问题，填补上VQ-VAE生成图片的最后一片空缺。

2.VQ-VAE 设计细节

在上一节中，我们虽然认识了VQ-VAE的核心思想，但略过了不少实现细节，比如：

VQ-VAE的编码器怎么输出离散向量。
VQ-VAE怎么优化编码器和解码器。
VQ-VAE怎么优化嵌入空间。

在这一节里，我们来详细探究这些细节。

2.1输出离散编码

想让神经网络输出一个整数，最简单的方法是和多分类模型一样，输出一个Softmax过的概率分布。之后，从概率分布里随机采样一个类别，这个类别的序号就是我们想要的整数。比如在下图中，我们想得到一个由3个整数构成的离散编码，就应该让编码器输出3组logit，再经过Softmax与采样，得到3个整数。

但是，这么做不是最高效的。得到离散编码后，下一步我们又要根据嵌入空间把离散编码转回一个向量。可见，获取离散编码这一步有一点多余。能不能把编码器的输出张量（它之前的名字叫logit）、解码器的输入张量embedding、嵌入空间直接关联起来呢？

2.2优化编码器和解码器

为了优化编码器和解码器，我们先来制订一下VQ-VAE的整体优化目标。由于VQ-VAE其实是一个AE，误差函数里应该只有原图像和目标图像的重建误差。

L = x - decoder(z_e + (z_q - z_e).detach())

通过这一技巧，我们完成了梯度的传递，可以正常地训练编码器和解码器了。

2.3优化嵌入空间

3.总结

VQ-VAE是一个把图像编码成离散向量的图像压缩模型。为了让神经网络理解离散编码，VQ-VAE借鉴了NLP的思想，让每个离散编码值对应一个嵌入，所有的嵌入都存储在一个嵌入空间（又称"codebook"）里。这样，VQ-VAE编码器的输出是若干个「假嵌入」，「假嵌入」会被替换成嵌入空间里最近的真嵌入，输入进解码器里。

VQ-VAE的优化目标由两部分组成：重建误差和嵌入空间误差。重建误差为输入图片和重建图片的均方误差。为了让梯度从解码器传到编码器，作者使用了一种巧妙的停止梯度算子，让正向传播和反向传播按照不同的方式计算。嵌入空间误差为嵌入和其对应的编码器输出的均方误差。为了让嵌入和编码器以不同的速度优化，作者再次使用了停止梯度算子，把嵌入的更新和编码器的更新分开计算。

训练完成后，为了实现随机图像生成，需要对VQ-VAE的离散分布采样，再把采样出来的离散向量对应的嵌入输入进解码器。VQ-VAE论文使用了PixelCNN来采样离散分布。实际上，PixelCNN不是唯一一种可用的拟合离散分布的模型。我们可以把它换成Transformer，甚至是diffusion模型。如果你当年看完VQ-VAE后立刻把PixelCNN换成了diffusion模型，那么恭喜你，你差不多提前设计出了Stable Diffusion。

可见，VQ-VAE最大的贡献是提供了一种图像压缩思路，把生成大图像的问题转换成了一个更简单的生成「小图像」的问题。图像压缩成离散向量时主要借助了嵌入空间，或者说"codebook"这一工具。这种解决问题的思路可以应用到所有图像生成类任务上，比如超分辨率、图像修复、图像去模糊等。所以近两年我们能看到很多使用了codebook的图像生成类工作。