Variational Autoencoder (VAE)
最开始的文章是 (Kingma & Welling, 2014)。
给定样本 \(x_i\)，后验分布 \(p_{\phi}(z\mid x_i)\sim \mathcal{N}(\mu_i, \sigma^2_i I)\)。其中我们希望 \(p_\phi(z)\sim\mathcal{N}(0, I)\)。
构建网络 \(\left<\mu_k, \log \sigma_k^2\right>=\phi(x_i)\)。
如果酱紫直接训的话，因为取正态分布的时候会带来噪声，所以模型肯定是希望 \(\sigma^2_k\) 越小越好，逐渐的会退化为普通的 autoencoder。
所以我们希望 \(p_\phi(z)\sim\mathcal{N}(0, I)\)，这样子加上一个
就行。
然后写代码的时候有个 reparameterization trick，其实就是如果我们算出 \(\mu_i\) 和 \(\sigma_i^2\) 然后采样，那么是不可导的。方法其实就是 \(\mathcal{N}(\mu, \sigma^2 I)\sim \sigma\times\mathcal{N}(0, I)+\mu\) 就行。
Vector Quantized-Variational Autoencoder (VQ-VAE)
文章：(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\)，论文中是 \(\beta=4\gamma\)。
所以最终 loss 是