本文主要内容是解读《Auto-encoding variational bayes》这篇论文。一直很喜欢这种有（数学或统计）理论基础的研究，比凑出来的深度学习模型不知道要高到哪里去。
问题背景
假设一组数据 $X = \lbrace\mathbf{x}^{(i)}\rbrace_{i=n}^{N}$ 是随机生成的，并且随机生成的过程与一个隐藏变量（Latent Variable）$\mathbf{z}$ 有关，我们想要得到数据的一个（近似的）概率密度函数 $p(\mathbf{x})$ 来代表这组数据。这个随机的生成过程有以下两步：(1) 从先验分布 $p_\mathbf{\theta}(\mathbf{z})$ 中生成 $\mathbf{z}^{(i)}$；(2) 从条件概率 $p_\mathbf{\mathbf{\theta}}(\mathbf{x}|\mathbf{z})$ 中生成 $\mathbf{x}^{(i)}$。由此，问题转换为求解 $\mathbf{\theta}$。这种假设在一定程度上是合理的，比如说在MNIST手写数字数据集中，我们想生成数字“1”的图片。首先我们要确定要生成数字是“1”，这对应第一步，然后我们在目标是数字“1”的条件下，通过某些方法生成相应的图片，这对应第二步。当然，实际的情况可能要更复杂，我们往往不能确定 $\mathbf{z}$ 究竟有什么实际含义，只知道它包含了一类数据的信息。我们对数据的概率分布 $p(\mathbf{x})$ 感兴趣，是因为有很多实际应用需要数据的 prior，比如 denoising[2], inpainting 还有 super-resolution。除此之外我们也对后验分布 $p_\mathbf{\mathbf{\theta}}(\mathbf{z}|\mathbf{x})$ 也感兴趣，因为 latent variable $\mathbf{z}$可以当作一个数据的 embedding。
一般情况下，我们可以用 Maximum Likelihood，即$\mathbf{\theta} = \text{argmax}_{\mathbf{\theta}} p_\mathbf{\theta}(\mathbf{x})$ 或者 Expectation Maximization 的方法求解$\mathbf{\theta}$。但是当模型很复杂时，$p_\mathbf{\theta}(\mathbf{x}) = \int p_\mathbf{\theta}(\mathbf{z})p_\mathbf{\mathbf{\theta}}(\mathbf{x}|\mathbf{z}) d\mathbf{z}$ 给不出具体的公式表达（intractable），所以也就没法用ML来求解，因为我们不知道 gradient 的表达式。后验分布 $p_\mathbf{\mathbf{\theta}}(\mathbf{z}|\mathbf{x}) = p_\mathbf{\mathbf{\theta}}(\mathbf{x}|\mathbf{z}) p_\mathbf{\theta}(\mathbf{z}) / p_\mathbf{\theta}(\mathbf{x})$ 也是如此，所以EM也不能用（不知道$q(\mathbf{z})$）。概率分布 intractable 的情况并不少见，比如说我们用神经网络+sigmoid/softmax 来表达概率函数。Variational inference 可以用来解决 intractabililty 的问题。
大体上来说，variational inference 是指用一个简单的，可计算的概率分布$q$来近似估计复杂分布$p$，最小化 KL divergence $D_{KL}(q || p) = \int q(\mathbf{z}) \log \frac{q(\mathbf{z})}{p(\mathbf{z})} d\mathbf{z}$。为了解决上面提到的问题，我们用一个函数 $q_{\mathbf{\phi}}(\mathbf{z} | \mathbf{x})$ 来估计 intractable 的后验分布 $p_\mathbf{\mathbf{\theta}}(\mathbf{z}|\mathbf{x})$，这样我们既可以近似的估计后验分布，在选择适当的条件概率 $p_{\mathbf{\theta}}(\mathbf{x}|\mathbf{z})$ 和隐藏变量的先验分布 $p_\mathbf{\theta}(\mathbf{z})$ 后，我们能够得到 $p(\mathbf{x})$ 的一个近似估计。在论文[1]中，作者提出了一个 variational bayes 的方法同时求解近似分布的参数 $\mathbf{\phi}$ 和条件概率参数 $\mathbf{\theta}$。算法叫 autoencoder 的原因是，可以把 $q_{\mathbf{\phi}}(\mathbf{z} | \mathbf{x})$ 当作一个 encoder ，$p_{\mathbf{\theta}}(\mathbf{x}|\mathbf{z})$ 当作 decoder，最后作者结合神经网络给出的例子也确实像一个 autoencoder。
Variational Bound
利用近似后验分布 $q$，我们可以构建一个数据概率(log) $p_\mathbf{\theta}(\mathbf{x})$ 的下界：
$$\begin{eqnarray}
log p_\mathbf{\theta}(\mathbf{x}^{(i)}) &=& D_{KL} (q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) || p_\mathbf{\theta}(\mathbf{z} | \mathbf{x}^{(i)})) + \mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)}) \\\
&\geq& \mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)})
\end{eqnarray} \tag{1}$$
不等式中的大于等于是因为 KL divergence 恒大于等于0，其中 $\mathcal{L}$ 具体表达式有两种形式，
$$\begin{eqnarray}
\mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)}) &=& \mathbb{E}_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} [-\log q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) + \log p_\mathbf{\theta}(\mathbf{x}^{(i)}, \mathbf{z})] \\\
&=& -D_{KL} (q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) || p_\mathbf{\theta}(\mathbf{z})) + \mathbb{E}_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} [\log p_\mathbf{\theta}(\mathbf{x}^{(i)} | \mathbf{z})]
\end{eqnarray} \tag{2}$$
其中第二种形式更加符合 autoencoder 的定义，等号右边第二项期望 reconstruction error，而第一项相当于一个 regularization，使 $q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})$ 尽量靠近先验分布 $p_\mathbf{\theta}(\mathbf{z})$。这两项优化的方向实际上是矛盾的，毕竟为了得到一个很好的 reconstruction，$\mathbf{z}$ 关于 $\mathbf{x}$ 的条件概率不可能与其先验概率一样。
这样，我们可以通过不断优化（提高）$\mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)})$ w.r.t. ${\bf \theta}$ 和 ${\bf \phi}$ 来得到我们想要的结果。优化之前有两个问题需要解决，一是如何选择 $q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})$和$p_{\mathbf{\theta}}(\mathbf{x}|\mathbf{z})$，二是怎么处理公式中的期望。一般来说我们会使用 Monte Carlo sampling 的方法近似估计期望值，不过这种方法会带来很大的方差，计算效率也很低。
Auto-Encoding Variational Bayes Algorithm
论文[1]中提出了一种全新的优化 $\mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)})$ w.r.t. ${\bf \theta}$和${\bf \phi}$ 的方法，简单来说就是 Monte Carlo 方法 + Reparameterization trick。
Monte Carlo 方法用来估计期望的值。具体到本问题中，$\mathbb{E}_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} [f(\mathbf{z})]$ 及其 gradient 是我们需要估计的，$f(\mathbf{z})$ 用来代指期望内的公式。估计方法如下
$$\mathbb{E}_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} [f(\mathbf{z})] \simeq \frac{1}{L} \sum_{l=1}^L f(\mathbf{z}^{(l)}) \tag{3}$$
其中 $\mathbf{z}^{(l)}$ 是从概率分布 $q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})$ 里抽样出来的。这是一个无偏估计。但是考虑到 gradient $\nabla_{\mathbf{\phi}} \mathbb{E}_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} [f(\mathbf{z})]$, $\nabla_{\mathbf{\phi}}$ 不能直接作用到期望里面，因为积分 $\int f(\mathbf{z}) \nabla_{\mathbf{\phi}} q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})$ 不再代表期望，进而不能使用 Monte Carlo 估计。一种解决方法是用一个数学变换，可以得到
$$\int f(\mathbf{z}) \nabla_{\mathbf{\phi}} q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) = \int f(\mathbf{z}) q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) \nabla_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} \log q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) = \mathbb{E}_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} [f(\mathbf{z}) \nabla_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} \log q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})] \tag{4}$$
但是这种方法会带来很大的 variance。 reparameterization 可以用来解决这个问题。
上面说到，$\nabla_{\mathbf{\phi}}$ 不能直接放入期望中是因为期望的概率分布 $q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})$和变量$\mathbf{\phi}$ 有关，reparameterization 就是用来把变量 $\mathbf{\phi}$ 移出概率分布。具体的，我们可以引入一个随机变量 $\epsilon$ 以及它的概率分布 $p(\epsilon)$，使得之前的随机变量 $\mathbf{z}$ 是 $\epsilon$ 的一个与 $\phi$ 有关的可导的变换（一一对应），$\mathbf{z} = g_{\phi}(\epsilon, \mathbf{x})$。这样，我们就有 $q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) = p(\epsilon, \mathbf{x}^{(i)})$, $\mathbb{E}_{q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})} [f(\mathbf{z})] = \mathbb{E}_{p(\epsilon)} [f(g_{\phi}(\epsilon, \mathbf{x}))]$，进而可以把$\nabla_{\mathbf{\phi}}$ 直接放到期望内，利用 Monte Carlo 估计求得 gradient。
运用到 $\mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)})$ 的第一种形式里去，我们得到
$$\mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)}) \simeq \frac{1}{L} \sum_{l=1}^{L} -\log q_\mathbf{\phi}(\mathbf{z}^{(i,l)} | \mathbf{x}^{(i)}) + \log p_\mathbf{\theta}(\mathbf{x}^{(i)}, \mathbf{z}^{(i,l)}) \tag{5}$$
$$\text{where } \mathbf{z}^{(i,l)} = g_{\phi}(\epsilon^{(l)}, \mathbf{x}^{(i)}) \text{ and } \epsilon^{(l)} \sim p(\epsilon) \tag{6}$$
并且我们可以直接对上式关于$\phi$和 $\theta$ 求导，利用SGD训练。
如果 $-D_{KL} (q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) || p_\mathbf{\theta}(\mathbf{z}))$ 是解析的（analytically），即有明确的数学表达，我们也可以利用 $\mathcal{L}$ 的第二种形式
$$\mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)}) \simeq -D_{KL} (q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)}) || p_\mathbf{\theta}(\mathbf{z})) + \frac{1}{L} \sum_{l=1}^{L} \log p_\mathbf{\theta}(\mathbf{x}^{(i)} | \mathbf{z}^{(i,l)}) \tag{7}$$
$$\text{where } \mathbf{z}^{(i,l)} = g_{\phi}(\epsilon^{(l)}, \mathbf{x}^{(i)}) \text{ and } \epsilon^{(l)} \sim p(\epsilon) \tag{8}$$
解决了期望计算求导的问题，我们就得到了如下的算法
Variational Autoencoder
关于 $q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})$，$p(\epsilon)$，$p_\mathbf{\theta}(\mathbf{z})$ 和 $p_{\mathbf{\theta}}(\mathbf{x}|\mathbf{z})$ 如何选择，文中给出了一个具体的例子，即本文的主角 variational autoencoder。
首先，$q_\mathbf{\phi}(\mathbf{z} | \mathbf{x}^{(i)})$ 是一个正态分布 $\mathcal{N}(\mathbf{z}; \mu^{(i)}, \sigma^{(i)2}\mathbf{I})$。其中均值 $\mu^{(i)}$ 和标准差 $\sigma^{(i)}$ 是一个以 $\mathbf{x}^{(i)}$ 为输入的 neural network 的输出。$\epsilon$ 服从一个标准正态分布 $\mathcal{N}(\epsilon;\mathbf{0},\mathbf{I})$，所以我们有
$$\mathbf{z}^{(i,l)} = g_{\phi}(\epsilon^{(l)}, \mathbf{x}^{(i)}) = \mu^{(i)} + \sigma^{(i)} \odot \epsilon^{(l)} \tag{9}$$
先验分布 $p_\mathbf{\theta}(\mathbf{z})$ 也服从标准正态分布 $\mathcal{N}(\mathbf{z};\mathbf{0},\mathbf{I})$，这么选择是因为这样可以极大的简化 KL divergence 的计算。
最后，$p_{\mathbf{\theta}}(\mathbf{x}|\mathbf{z})$ 是一个 neural network + non-linear layer 的输出，以 $\mathbf{z}^{(i,l)}$ 为输入。
综上，我们可以得到 variational autoencoder 的 objective function, 需要训练的参数 $\phi$ 和 $\theta$ 分别是作为 encoder 和 decoder 的 neural network 的参数，用 minibatch SGD 的方法训练模型最大化下面的函数
$$\mathcal{L}({\bf \theta},{\bf \phi}; {\bf x}^{(i)}) \simeq \frac{1}{2} \sum_{j=1}^{J} (1 + log((\sigma^{(i)}_j)^2) - (\mu^{(i)}_j)^2 - (\sigma^{(i)}_j)^2) + \frac{1}{L} \sum_{l=1}^{L} \log p_\mathbf{\theta}(\mathbf{x}^{(i)} | \mathbf{z}^{(i,l)}) \tag{10}$$
$$\text{where } \mathbf{z}^{(i,l)} = \mu^{(i)} + \sigma^{(i)} \odot \epsilon^{(l)} \text{ and } \epsilon^{(l)} \sim \mathcal{N}(\mathbf{\epsilon};\mathbf{0},\mathbf{I}) \tag{11}$$
pytorch 官方提供了一个简单明了的VAE实现https://github.com/pytorch/examples/tree/master/vae，对理解AEVB算法的整个流程很有帮助。
参考文献
[1] Kingma D P, Welling M. Auto-encoding variational bayes[J]. arXiv preprint arXiv:1312.6114, 2013.
[2] Im D J, Ahn S, Memisevic R, et al. Denoising Criterion for Variational Auto-Encoding Framework[C]//AAAI. 2017: 2059-2065.