VQVAE：Neural Discrete Representation Learning

论文名称：Neural Discrete Representation Learning
开源地址
发表时间：NIPS2017
作者及组织：Aaron van den Oord,Oriol Vinyals和Koray Kavukcuoglu, 来自DeepMind。

1、VAE

  简单回顾下VAE的损失函数，ELBO的下界为：
$$\begin{array}{c}LowerBound=E_{q_{\phi }(z|x)}[log{p}_{\theta }(x|z)]-D_{KL}(q_{\phi }(z|x)||p(z))\end{array}\text{\hspace{1em}(0)}$$
 其中第一项为解码器的重构损失(regression loss) ;第二项为正则项，用KL散度来使Encoder----后验概率 $q_{\phi }(z|x)$ 和 先验 $p(z)$ 分布近似，通常 $p(z)$ 假设为多元标准正太分布，该项主要防止VAE坍塌到一个点，毕竟是生成模型。
 而VQVAE和VAE主要不同：Encoder输出是离散的，而不是连续的隐变量z。

1、方法

1.1.模型结构

 为了实现离散化编码，VQVAE引入了一个可学习的codebook，即上图中的EmbeddingSpace。大概说下流程：输入一张图像，经过CNN得到 $z_e(x)\in {\mathbb{R}}^{H\ast W\ast D}$ ，然后计算 $z_e$ 中每条特征向量跟codebook的最接近的向量的索引，得到 $q(z|x)\in {\mathbb{R}}^{H\ast W}$ , 然后用codebook中向量 $e_i$ 来替换 $z_e(x)$ 得到 $z_q(x)$ 。最后经过Decoder得到 $x$ 。

1.2.训练

 先说下总体损失函数，其实跟VAE的损失函数类似：
$$\begin{array}{c}L=logp(x|z_q(x))+||sg[z_e(x)]-e|{|}_2^2+\beta ||z_e(x)-sg[e]|{|}_2^2\end{array}\text{\hspace{1em}(1)}$$

 其中第一项就是VAE中的重构损失，但有个问题：在用L2 Loss计算重构损失后，反向传播时，由于在codebook中argmin这个操作是不可导的，这样就优化不了Encoder，于是本文直接将 $z_q(x)$ 节点的梯度拷贝给了 $z_e(x)$ ，使得反向传播得以继续。具体的表达式如下：
$$\begin{array}{c}logp(x|z_q(x))=||x-decoder(z_e(x)+sg(z_q(x)-z_e(x)))|{|}_2^2\end{array}\text{\hspace{1em}(2)}$$
 式中的 $sg$ 表示 .detach() 操作，由于VQVAE多了一个可学习的codebook，而重构损失并没有梯度传过去。因此损失第二项就是让 $e$ 逼近 $z_e(x)$ ，这项仅更新codebook。

  由于训练过程中，Encoder相较于codebook，肯定易于优化，也就是Encoder收敛快，而codebook收敛慢 ，为了让Encoder别距离codebook太远，于是增加了第三项损失，让 $z_e(x)$ 逼近 $e$ 。

 在回过头来跟VAE的式子比较下 ，发现缺少了KL散度项：这是因为在VQVAE中，在根据 $x$ 取得 $e$ 的概率非0即1: $q(z=e|x)=1,q(z=other|x=0)$ ，相当于二项分布，同时假设 $p(z)$ 是均匀分布，两个均匀分布的KL散度是常数，在损失中可忽略。

1.3.生成

 在训练集上训练完VQVAE后，VQVAE学习到的是一个有效的低维度的离散表示。然后将VQVAE置为推理阶段，用自回归模型PixCNN来拟合 $q(z|x)$ ，训练完成后，PixCNN就能生成有意义的索引矩阵，然后去codebook中拿到对应的张量，送去VQVAE的Decoder中解码生成图像。

2、实验

  生成的小图还是可以的。

思考

  替换更强的自回归模型Transformer也就是后来VQGAN的工作了。