VQ-VAE

​ 本次我们来介绍的是自动变分编码器家族中的一员——VQ-VAE。之前我们介绍过最基本的变分自动编码器的原理及代码，需要回顾的读者不妨看这里：

《链接》

由于笔者近来现实事务繁忙，原定于本周发布的GAN对抗生成网络代码解析延期，预计将于下周释出，请各位读者见谅。

介绍

​ VQ-VAE由Google团队在论文《Neural Discrete Representation Learning》中提出，在19年DeepMind 团队使用 VQ-VAE-2 算法生成了以假乱真的高清大图，生成效果比肩BigGAN，本期让我们一探究竟。

​ 传统的VAE如下图所示，编码器模块将图像编码为隐层表示，解码器部分通常由解卷积组成，从隐层表示恢复出二维图像。

​ VQ-VAE在生成模型部分使用的是PixelCNN的思路，就像自然语言处理中的生成任务一样，逐个像素地生成一张图片。以cifar10的图像为例，它是32×32大小的3通道图像，换言之它是一个32×32×3的矩阵，矩阵的每个元素是0～255的任意一个整数，这样一来，我们可以将它看成是一个长度为32×32×3=3072的像素序列，而每个像素值限定在0-255之间的整数。因此可以像循环神经网络一样，来逐像素地、递归地生成一张图片（传入前面的所有像素，来预测下一个像素），这就是所谓的自回归方法：

意思是整张图片的概率分布是从第1个像素x0开始，按照已生成的像素作为条件，计算新像素的概率分布，依次抽取剩余的像素的值进行生成。

​ 自回归的方法很稳妥，也能有效地做概率估计，但它有一个最致命的缺点：慢。因为它是逐像素地生成的，所以要每个像素地进行随机采样，上面举例的cifar10已经算是小图像的，目前做图像生成好歹也要做到128×128×3的才有说服力了吧，这总像素接近5万个（想想看要生成一个长度为5万的句子），真要逐像素生成会非常耗时。而且这么长的序列，不管是RNN还是CNN模型都无法很好地捕捉这么长的依赖。

​ 原始的自回归还有一个问题，就是割裂了类别之间的联系。虽然说因为每个像素是离散的，所以看成256分类问题也无妨，但事实上连续像素之间的差别是很小的，纯粹的分类问题捕捉到这种联系。更数学化地说，就是我们的目标函数交叉熵是−logpt，假如目标像素是100，如果我预测成99，因为类别不同了，那么pt就接近于0，−logpt就很大，从而带来一个很大的损失。但从视觉上来看，像素值是100还是99差别不大，不应该有这么大的损失。

​ 针对以上自回归模型常见的毛病，VQ-VAE给出的解决方案是，将照片压缩到低维空间，在低维空间训练自回归神经网络，再解码到高维。

​

图2.整体流程

​ 如上图所示，如果图像编码器使用卷积神经网络，那么潜在表示z是 d （卷积核数量）个 feature map 构成的立方体, 而神经网络深层产生的 feature map 与输入图像的横平面空间位置存在对应关系。因此 feature map 上横平面每个坐标对应的纵向 d个通道构成潜在表示的一个分量。这样潜在表示就大体保留了位置信息。

​ 在图2中，上方部分是一个Embedding空间，也可以说是一个K维的编码表，其中每一个基 ei 都是大小为d的向量。在图像编码器获得隐层向量z之后，通过最近邻搜索，将z映射为K个向量之一：

$\begin{equation}z\to e_k,\quad k = \mathop{\arg\min}_j \Vert z - e_j\Vert_2\end{equation}$

图二中下方的小矩阵里的数字i，代表该位置的隐层向量被替换为基ei。

其中 x 表示是输入图像，D 是解码器，E是编码器，sg 表示stop gradient，训练时不对 sg[] 中的变量计算梯度，误差反向传播时也不会向其中的变量更新梯度。

拟合编码分布

​ 读者可能会发现损失函数中并没有KL散度损失项。在之前对标准VAE的介绍中，我们说过隐层表示和标准高斯分布之间的KL散度损失是为了在没有图像编码器的情况下，从高斯分布中采样生成隐层表示z。那么VQ-VAE的隐层向量应该如何生成呢？

​ 由于这个m×m的矩阵一定程度上也保留了原来输入图片的位置信息，所以可以用自回归模型来对编码矩阵进行拟合（即建模先验分布）。通过PixelCNN得到编码分布后，就可以随机生成一个新的编码矩阵，然后通过编码表E映射为浮点数矩阵zq，最后经过解码器得到一张图片。

以上。