Vector Quantized-Variational Autoencoder (VQ-VAE)

文章：(van den Oord et al., 2017)；苏神博客：(苏剑林, 2019)。

先把 $x \in R^{H \times W \times 3}$ 降采样成 $z \in R^{h \times w \times d}$ 。然后维护 embedding codebook $e \in R^{K \times d}$ 。找到每个 $z_{i j}$ 的最近邻 ${\hat{z}}_{i j}$ 进入 decoder。

我们希望的 loss 是 ${‖ x - Decoder (\hat{z}) ‖}_{2}^{2}$ ，但是需要解决两个问题：

取最近邻这个步骤不可导；
如何更新 $e$ 。

对于第一个问题，文章提出了一种叫 Straight-Through Estimator 的方法，参考的是 (Bengio et al., 2013)。本质上其实是在求导的时候，将离散操作看成某个平滑的连续函数。具体操作上，我们令 $sg [x]$ 表示 stop gradient，相当于 python 里的 detach。那么我们其实算的是：

{‖ x - Decoder (z + sg [\hat{z} - z]) ‖}_{2}^{2}

而写代码的时候，我们只需要在 forward 的时候写

python

z_q = z + (z_q - z).detach()

就可以了。其中 z_q 指的是 $\hat{z}$ 。

对于第二个问题，首先我们希望 $\hat{z}$ 和 $z$ 能够尽量接近，所以我们理论上讲可以加一个 $β ‖ z - \hat{z} ‖_{2}^{2}$ 的项。但是其实我们更希望是 $\hat{z}$ 去接近 $z$ 而不是 $z$ 去接近 $\hat{z}$ ，所以我们将这个项分解为

β ‖ sg [z] - \hat{z} ‖_{2}^{2} + γ ‖ z - sg [\hat{z}] ‖_{2}^{2}

一般情况下 $β > γ$ ，论文中是 $γ = \frac{1}{4} β$ 。

所以最终 loss 是

{‖ x - Decoder (z + sg [\hat{z} - z]) ‖}_{2}^{2} + β ‖ sg [z] - \hat{z} ‖_{2}^{2} + γ ‖ z - sg [\hat{z}] ‖_{2}^{2}

References

Bengio, Y., Léonard, N., & Courville, A. (2013). Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. arxiv.org

van den Oord, A., Vinyals, O., & Kavukcuoglu, K. (2017). Neural Discrete Representation Learning. Advances in Neural Information Processing Systems, 30, 6306–6315. proceedings.neurips.cc

苏剑林. (2019, June). VQ-VAE的简明介绍：量子化自编码器. spaces.ac.cn

References

Cite this post