正常来说图像的像素值是离散的（在0到255之间），所以我们其实可以把图像看成是一个长为\(H\times W\times C\)的序列，在预测当前像素值时，把之前的像素值作为已知信息参与到预测中，这样递归地逐像素预测下去，即自回归模型（autoregressive model）：
\[ p(x) = p(x_1)p(x_2\vert x_1)\dots p(x_n\vert x_1,x_2,\dots,x_{n-1}) \]
但上面这一过程是逐像素生成的，所以运行速度慢。VQ-VAE借鉴了自回归模型的思想，一个最明显的特征是通过encoder得到的隐变量是一个离散变量，将图像大小压缩到了一个小的特征空间。整体的流程大致如下图所示：
- 输入大小为\(H\times W\times 3\)的数据\(\mathbf{x}\)，通过encoder得到初步的中间变量\(\mathbf{z}=\text{encoder}(\mathbf{x})\)，大小缩小为\(h\times w\times d\)
- 模型会维护一个embedding层\(\mathbf{E}=[\mathbf{e}_1, \mathbf{e}_2, \dots, \mathbf{e}_K]\)，即codebook，其中每个\(\mathbf{e}_i\)的长度都是\(d\)。对于\(\mathbf{z}\)中的每个\(d\)，通过最近邻搜索得到相应的\(\mathbf{e}_i\)值，即有映射关系\(\mathbf{z}_q(\mathbf{x}) =\mathbf{e}_{argmin_j\vert\vert \mathbf{z}(\mathbf{x})-\mathbf{e}_j\vert\vert_2}\)。得到的编码向量用\(\mathbf{z}_q\)来表示，这样的编码向量其实就是一个离散向量
- 将\(\mathbf{z}_q\)输入到decoder中得到结果\(\hat{\mathbf{x}} = \text{decoder}(\mathbf{z}_q)\)
论文中给出了VQ-VAE的损失函数如下：
\[ L = -\log p(\mathbf{x}\vert\mathbf{z}_q(\mathbf{x}))+\vert\vert\text{sg}[\mathbf{z}(\mathbf{x})]-\mathbf{z}_q(\mathbf{x})\vert\vert_2^2+\beta\vert\vert\mathbf{z}(\mathbf{x})-\text{sg}[\mathbf{z}_q(\mathbf{x})]\vert\vert_2^2 \]
其中第一项是用来优化encoder和decoder的，即重建损失，从图中也可以看出，损失函数对\(\mathbf{z}\)的梯度是会推给encoder的。这一项通常是重建图像与输入图像的MSE损失\(\vert\vert \mathbf{x}-\text{decoder}(\mathbf{z}_q)\vert\vert_2^2\)。
由于我们在前向计算中有用到\(\argmin\)这个操作，这个是没有梯度的。对此，作者采用了Straight-Through Estimator方法人为地设计一个梯度，即\(\vert\vert \mathbf{x}-\text{decoder}(\mathbf{z} + \text{sg}[\mathbf{z}_q-\mathbf{z}])\vert\vert_2^2\)。式子中的\(\text{sg}\)代表的是stop gradient，即在反向传播时为0，前向计算为原值，数学描述如下：
\[ \text{sg}(x) = \begin{cases} x \quad \text{in forward process} \\ 0 \quad \text{backward process} \end{cases} \]
损失函数的后面两项与codebook有关。通常我们会希望\(\mathbf{z}_q\)与\(\mathbf{z}\)尽可能地靠近，所以优化目标是\(\vert\vert\mathbf{z}_q - \mathbf{z}\vert\vert_2^2\)。由于这个目标的梯度等于对\(\mathbf{z}\)的梯度加上对\(\mathbf{z}_q\)的梯度，所以我们可以将它等价为如下形式：
\[ \vert\vert\text{sg}[\mathbf{z}]-\mathbf{z}_q\vert\vert_2^2 + \vert\vert\mathbf{z}-\text{sg}[\mathbf{z}_q]\vert\vert_2^2 \]
等价在这里是指梯度等价，实际前向计算出来的loss值会是原来的两倍。上面式子中，第一项是\(\mathbf{z}_q\)的梯度，相当于固定\(\mathbf{z}\)，让\(\mathbf{z}_q\)靠近\(\mathbf{z}\)；而第二项是\(\mathbf{z}\)的梯度，相当于固定\(\mathbf{z}_q\)，让\(\mathbf{z}\)靠近\(\mathbf{z}_q\)。因为\(\mathbf{z}_q\)相对于\(\mathbf{z}\)来说更自由一些，\(\mathbf{z}\)是encoder的输出，为了尽力保证重建后的效果，应该尽量让\(\mathbf{z}_q\)去靠近\(\mathbf{z}\)。这一步可以用权重去调整损失各项的重要性。
综上，VQ-VAE最终损失函数如下：
\[ L = \vert\vert \mathbf{x}-\text{decoder}(\mathbf{z} + \text{sg}[\mathbf{z}_q-\mathbf{z}])\vert\vert_2^2 + \gamma\vert\vert\text{sg}[\mathbf{z}]-\mathbf{z}_q\vert\vert_2^2 +\beta \vert\vert\mathbf{z}-\text{sg}[\mathbf{z}_q]\vert\vert_2^2 \]
其中\(\beta \lt \gamma\)，论文中选取\(\gamma = 1\)，\(\beta = 0.25\)。
从上面的描述中，我们可以看出VQ-VAE其实跟VAE并没有太大关系，它只是一个AE模型，目标在于图像重建。
在VQ-VAE论文中，作者利用PixelCNN去预测\(\mathbf{z}_q\)，实现了随机图像的生成。训练时模型是共同训练VQ-VAE和PixelCNN的，输入数据通过encoder得到\(\mathbf{z}\)，查询codebook得到\(\mathbf{z}_q\)，再经过PixelCNN得到与\(\mathbf{z}_q\)同大小的变量\(\mathbf{z}_p\)作为decoder的输入。而在推理阶段则是直接将空图像输入到PixelCNN中得到\(\mathbf{z}_p^\prime\)，之后直接输入到decoder中生成图像。
关于PixelCNN模型，可以参考这篇文章。