【论文极速读】VQ-VAE：一种稀疏表征学习方法

【论文极速读】VQ-VAE：一种稀疏表征学习方法

FesianXu 20221208 at Baidu Search Team

前言

最近有需求对特征进行稀疏编码，看到一篇论文VQ-VAE，简单进行笔记下。如有谬误请联系指出，本文遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明并且联系笔者，谢谢 。

$\nabla$ 联系方式：

e-mail: [e-mail]

github: https://github.com/FesianXu

---

图片，视频等视觉模态有充足的冗余信息，可以通过稀疏编码进行编码，以减少储存消耗。Vector-Quantised Variational AutoEncoder (VQ-VAE) 就是进行图片稀疏编码的工作[1]。 如Fig 1. 所示，VQ-VAE有三大部分组成，Encoder，Decoder和储存稀疏编码的Embedding Space字典。其中的Embedding space字典的形状为$\mathcal{E}\in {\mathbb{R}}^{K\times D}$，其中的$K$为字典的大小，$D$为字典的特征维度，字典中每一个样本$e_i\in {\mathbb{R}}^D,i\in 1,\cdots ,K$表示了第$i$个稀疏编码的特征表达。

Fig 1. VQ-VAE的框架示意。

单从稀疏编码的角度看，如Fig 2.所示，整个工作中，将会考虑将中间特征图的$H\times W\times D$，通过用离散的稀疏编码表示，形状为$H\times W\times 1$，进行稀疏编码的方式可以通过简单的最近邻方法得到，如公式(1-1)所示
$$q(z=k|x)=\{\begin{array}{ll}1& fork=\arg {\min }_j||z_e(x)-e_j|{|}_2\\ 0& otherwise\end{array}\text{\hspace{1em}(1-1)}$$
其中的$x$为原始的图片输入，$z_e(x)$表示图片输入经过编码器后得到的feature map，而$q(z|x)$即是进行稀疏编码后的结果。通过式子(1-2)，可以将稀疏编码后的结果恢复为feature map（当然这个过程是有损的，只保留最为重要的特征信息）。整个过程可见Fig 2.示意图，应该比较容易理解。
$$z_q(x)=e_k,wherek=\arg \underset{j}{\min }||z_e(x)-e_j|{|}_2\text{\hspace{1em}(1-2)}$$

Fig 2. 通过最近邻方法在字典里面查找稀疏令牌，作为稀疏编码的结果，然后通过反查字典可以对feature map进行恢复。

整个框架中有若干参数需要学习，分别是encoder，decoder网络参数和Embedding space字典的参数。然而稀疏编码的过程由于出现了最近邻方法，这个过程显然是无法传递梯度的，为了实现编码器的更新，可以考虑将解码器的梯度直接拷贝到编码器中。假设对于编码后恢复的$z_q(x)$而言，其每个元素表示为$D_{i,j,k}$，那么对于其中某个元素的梯度表示为$\frac{\partial \mathcal{L}}{\partial D_{i,j,k}}$，同理，对于编码后的$z_e(x)$而言，同样有$\frac{\partial \mathcal{L}}{\partial E_{i,j,k}}$，令$\frac{\partial \mathcal{L}}{\partial E_{i,j,k}}=\frac{\partial \mathcal{L}}{\partial D_{i,j,k}}$。那么对于编码器的梯度就可以表示为$\frac{\partial \mathcal{L}}{\partial W_E}=\frac{\partial E_{i,j,k}}{\partial W_E}\frac{\partial \mathcal{L}}{\partial E_{i,j,k}}$。

Fig 3. 通过梯度拷贝，将decoder的梯度拷贝到encoder中。

最后的损失函数如(1-3)所示，其中的$sg(\cdot )$为停止梯度函数，表示该函数无梯度传导。decoder的参数通过第一项损失项进行更新（这部分损失可通过MSE损失$\mathcal{L}(\mathbf{x},\hat{\mathbf{x}})$建模），称之为重建损失。encoder参数通过第一项和第三项损失进行更新，其中第一项是重建损失，第三项是为了encoder编码产出和embedding space进行对齐而设计的，由于此时通过$sg(\cdot )$函数停止了梯度，因此此时$\mathcal{E}$的参数不会得到更新。Embedding space的参数通过第二项损失项进行更新，通过将encoder编码结果进行停止梯度，我们只对$\mathcal{E}$进行参数更新。

$$\mathcal{L}=\log (p(x|z_q(x)))+||sg[z_e(x)]-\mathcal{E}|{|}_2^2+\beta ||z_e(x)-sg[\mathcal{E}]|{|}_2^2\text{\hspace{1em}(1-3)}$$

作者在原论文中贴了不少图片稀疏编码的结果，如Fig 4.所示，将$128\times 128\times 3$的原始图片稀疏编码到$32\times 32\times 1$（K=512），信息压缩比为$128\times 128\times 3\times 8/(32\times 32\times 9)=42.6$。从效果上看，除了在高频细节，比如毛发等上有些模糊外，其他图片信息都得到了较好的保留。

Fig 4. 将原始图片稀疏编码后，再解码出来的图片与原始图片的对比，可以发现对于高频信号（毛发细节等）有些损失。

Reference

[1]. Van Den Oord, Aaron, and Oriol Vinyals. “Neural discrete representation learning.” Advances in neural information processing systems 30 (2017).