文章：(van den Oord et al., 2017)；苏神博客：(苏剑林, 2019)。
先把 \(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\)，但是需要解决两个问题：
- 取最近邻这个步骤不可导；
- 如何更新 \(e\)。
对于第一个问题，文章提出了一种叫 Straight-Through Estimator 的方法，参考的是 (Bengio et al., 2013)。本质上其实是在求导的时候，将离散操作看成某个平滑的连续函数。具体操作上，我们令 \(\mathrm{sg}[x]\) 表示 stop gradient，相当于 python 里的 detach。那么我们其实算的是：
而写代码的时候，我们只需要在 forward 的时候写
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>\gamma\)，论文中是 \(\gamma=\frac{1}{4}\beta\)。
所以最终 loss 是