# Vector Quantized-Variational Autoencoder (VQ-VAE)

Author: Fanyi Pu

Published: 2026-08-30

Canonical: <https://pufanyi.com/blog/ml/ml-revisit/ae/ml-revisit-vq-vae>

Notes for Vector Quantized-Variational Autoencoder (VQ-VAE)

文章：([van den Oord et al., 2017](https://pufanyi.com/blog/ml/ml-revisit/ae/ml-revisit-vq-vae#bib-vandenoord2017neural))；苏神博客：([苏剑林, 2019](https://pufanyi.com/blog/ml/ml-revisit/ae/ml-revisit-vq-vae#bib-kexuefm-6760))。

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

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

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

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

$$
\left\|x - \mathrm{Decoder}(z+\mathrm{sg}[\hat{z}-z])\right\|_2^2
$$

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

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

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

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

$$
\beta\|\mathrm{sg}[z]-\hat{z}\|_2^2+\gamma\|z-\mathrm{sg}[\hat{z}]\|_2^2
$$

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

所以最终 loss 是

$$
\left\|x - \mathrm{Decoder}(z+\mathrm{sg}[\hat{z}-z])\right\|_2^2+\beta\|\mathrm{sg}[z]-\hat{z}\|_2^2+\gamma\|z-\mathrm{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](https://arxiv.org/abs/1308.3432 "https://arxiv.org/abs/1308.3432")

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](https://proceedings.neurips.cc/paper/2017/hash/7a98af17e63a0ac09ce2e96d03992fbc-Abstract.html "https://proceedings.neurips.cc/paper/2017/hash/7a98af17e63a0ac09ce2e96d03992fbc-Abstract.html")

苏剑林. (2019, June). *VQ-VAE的简明介绍：量子化自编码器*. [spaces.ac.cn](https://spaces.ac.cn/archives/6760 "https://spaces.ac.cn/archives/6760")