模型综述 #
VQ-VAE（Vector Quantised - Variational AutoEncoder）首先出现在论《Neural Discrete Representation Learning》
VAE假设隐向量分布服从高斯分布
VQVAE假设隐向量服从类别分布
PixelCNN #
要追溯VQ-VAE的思想，就不得不谈到自回归模型。可以说，VQ-VAE做生成模型的思路，源于PixelRNN、PixelCNN之类的自回归模型，这类模型留意到我们要生成的图像，实际上是离散的而不是连续的。以cifar10的图像为例，它是32×32大小的3通道图像，换言之它是一个32×32×3的矩阵，矩阵的每个元素是0～255的任意一个整数，这样一来，我们可以将它看成是一个长度为32×32×3=3072的句子，而词表的大小是256，从而用语言模型的方法，来逐像素地、递归地生成一张图片（传入前面的所有像素，来预测下一个像素），这就是所谓的自回归方法：
p(x)=p(x1)p(x2|x1)…p(x3n2|x1,x2,…,x3n2−1)
其中p(x1),p(x2|x1),…,p(x3n2|x1,x2,…,x3n2−1)每一个都是256分类问题，只不过所依赖的条件有所不同。
自回归的方法很稳妥，也能有效地做概率估计，但它有一个最致命的缺点：慢。因为它是逐像素地生成的，所以要每个像素地进行随机采样，上面举例的cifar10已经算是小图像的，目前做图像生成好歹也要做到256×256×3的才有说服力了吧，这总像素接近20万个（想想看要生成一个长度为5万的句子），真要逐像素生成会非常耗时。而且这么长的序列，不管是RNN还是CNN模型都无法很好地捕捉这么长的依赖，transformer也会面临巨大计算复杂度（ O（n2） ）
原始的自回归还有一个问题，就是割裂了类别之间的联系。虽然说因为每个像素是离散的，所以看成256分类问题也无妨，但事实上连续像素之间的差别是很小的，纯粹的分类问题捕捉到这种联系。更数学化地说，就是我们的目标函数交叉熵是−logpt，假如目标像素是100，如果我预测成99，因为类别不同了，那么pt就接近于0，−logpt就很大，从而带来一个很大的损失。但从视觉上来看，像素值是100还是99差别不大，不应该有这么大的损失。
VQ-VAE #
针对自回归模型的固有毛病，VQ-VAE提出的解决方案是：先降维，然后再对编码向量用PixelCNN建模。
降维离散化 #
看上去这个方案很自然，似乎没什么特别的，但事实上一点都不自然。
因为PixelCNN生成的离散序列，你想用PixelCNN建模编码向量，那就意味着编码向量也是离散的才行。而我们常见的降维手段，比如自编码器，生成的编码向量都是连续性变量，无法直接生成离散变量。同时，生成离散型变量往往还意味着存在梯度消失的问题。还有，降维、重构这个过程，如何保证重构之后出现的图像不失真？如果失真得太严重，甚至还比不上普通的VAE的话，那么VQ-VAE也没什么存在价值了。
幸运的是，VQ-VAE确实提供了有效的训练策略解决了这两个问题。
最邻近重构 #
在VQ-VAE中，一张n×n×3的图片x先被传入一个encoder中，得到连续的编码向量z：
z=encoder(x)
这里的z是一个大小为d的向量。另外，VQ-VAE还维护一个Embedding层，我们也可以称为编码表(codebook)，记为
E=e1,e2,…,eK
这里每个ei都是一个大小为d的向量。接着，VQ-VAE通过最邻近搜索，将z映射为这K个向量之一：
z→ek,k=argminj∥z−ej∥2
我们可以将z对应的编码表向量记为zq，我们认为zq才是最后的编码结果。最后将zq传入一个decoder，希望重构原图x=decoder(zq)。
整个流程是：
x → encoder(z) → 最邻近zq → decoder(x) ->x’
这样一来，因为zq是编码表E中的向量之一，所以它实际上就等价于1,2,…,K这K个整数之一，因此这整个流程相当于将整张图片编码为了一个整数。
当然，上述过程是比较简化的，如果只编码为一个向量，重构时难免失真，而且泛化性难以得到保证。所以实际编码时直接用多层卷积将x编码为m×m个大小为d的向量
也就是说，z的总大小为m×m×d，它依然保留着位置结构，然后每个向量都用前述方法映射为编码表中的一个，就得到一个同样大小的zq，然后再用它来重构。这样一来，zq也等价于一个m×m的整数矩阵，这就实现了离散型编码。
自行设计梯度 #
我们知道，如果是普通的自编码器，直接用下述loss进行训练即可：
∥x−decoder(z)∥2
但是，在VQ-VAE中，我们用来重构的是zq而不是z，那么似乎应该用这个loss才对：
∥x−decoder(zq)∥2
但问题是zq的构建过程包含了argmin，这个操作是没梯度的，所以如果用第二个loss的话，我们没法更新encoder。
换言之，我们的目标其实是∥x−decoder(zq)∥2最小，但是却不好优化，而∥x−decoder(z)∥2容易优化，但却不是我们的优化目标。那怎么办呢？当然，一个很粗暴的方法是两个都用：
∥x−decoder(z)∥2+∥x−decoder(zq)∥2
但这样并不好，因为最小化∥x−decoder(z)∥2并不是我们的目标，会带来额外的约束。
VQ-VAE使用了一个很精巧也很直接的方法，称为Straight-Through Estimator，你也可以称之为“直通估计”，它最早源于Benjio的论文《Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation》
事实上Straight-Through的思想很简单，就是前向传播的时候可以用想要的变量（哪怕不可导），而反向传播的时候，用你自己为它所设计的梯度。根据这个思想，我们设计的目标函数是：
∥x−decoder(z+sg[zq−z])∥2
其中sg是stop gradient的意思，就是不要它的梯度。这样一来，前向传播计算（求loss）的时候，就直接等价于decoder(z+zq−z)=decoder(zq)，然后反向传播（求梯度）的时候，由于zq−z不提供梯度，所以它也等价于decoder(z)，这个就允许我们对encoder进行优化了。
顺便说一下，基于这个思想，我们可以为很多函数自己自定义梯度，比如x+sg[relu(x)−x]就是将relu(x)的梯度定义为恒为1，但是在误差计算是又跟relu(x)本身等价。当然，用同样的方法我们可以随便指定一个函数的梯度，至于有没有实用价值，则要具体任务具体分析了。
维护编码表 #
要注意，根据VQ-VAE的最邻近搜索的设计，我们应该期望zq和z是很接近的（事实上编码表E的每个向量类似各个z的聚类中心出现），但事实上未必如此，即使∥x−decoder(z)∥2和∥x−decoder(zq)∥2都很小，也不意味着zq和z差别很小（即f(z1)=f(z2)不意味着z1=z2）。
所以，为了让zq和z更接近，我们可以直接地将∥z−zq∥22加入到loss中：
∥x−decoder(z+sg[zq−z])∥22+β∥z−zq∥2
除此之外，还可以做得更仔细一些。由于编码表（zq）相对是比较自由的，而z要尽力保证重构效果，所以我们应当尽量“让zq去靠近z”而不是“让z去靠近zq”，而因为∥zq−z∥2的梯度等于对zq的梯度加上对z的梯度，所以我们将它等价地分解为
∥sg[z]−zq∥22+∥z−sg[zq]∥2
第一项相等于固定z，让zq靠近z，第二项则反过来固定zq，让z靠近zq。注意这个“等价”是对于反向传播（求梯度）来说的，对于前向传播（求loss）它是原来的两倍。根据我们刚才的讨论，我们希望“让zq去靠近z”多于“让z去靠近zq”，所以可以调一下最终的loss比例：
∥x−decoder(z+sg[zq−z])∥22+β∥sg[z]−zq∥22+γ∥z−sg[zq]∥2
其中γ<β，在原论文中使用的是γ=0.25β。
（注：还可以用滑动评论的方式更新编码表，详情请看原论文。）
拟合编码分布 #
经过上述一大通设计之后，我们终于将图片编码为了m×m的整数矩阵了，由于这个m×m的矩阵一定程度上也保留了原来输入图片的位置信息，所以我们可以用自回归模型比如PixelCNN，来对编码矩阵进行拟合（即建模先验分布）。通过PixelCNN得到编码分布后，就可以随机生成一个新的编码矩阵，然后通过编码表E映射为浮点数矩阵zq，最后经过deocder得到一张图片。
