理解差分自动编码器 VAE：Variational AutoEncoder

引言

VAE 是一类 encoder-decoder 模型，给定一个高维数据 $\mathbf{x}\in {\mathbb{R}}^D$，encoder 将生成潜在空间对应的表达 $\mathbf{z}\in {\mathbb{R}}^d,d\ll D$，decoder 利用 $\mathbf{z}$ 恢复原始数据 ${\mathbf{x}}^{\prime }$，整个模型使用 $\mathbf{x}$ 和 ${\mathbf{x}}^{\prime }$ 之间的差异来训练。当训练好以后，随机从潜在空间中采样一个 $\mathbf{z}$，希望 decoder 能生成对应的数据。

---

动机

普通 encoder-decoder 用于生成任务时，容易造成严重的过拟合，使学习到的 潜在空间不规则。由于模型仅仅被要求在有限的训练数据上得到尽可能少的恢复损失，且缺少对潜在空间的约束，encoder 和 decoder 可能会学习到一种简单的一一对应关系，使得潜在空间中的一些点在被解码后的内容是无意义的。假设 $\mathbf{z}\in \mathbb{R}$，上述情况如下所示：

因此，为了使用 decoder 生成数据，必须保证潜在空间是规则的，一个可行的办法是在训练的时候对潜在空间添加正则化，这就是 VAE 的主要动机之一。

VAE 介绍

VAE 在训练的时候添加正则化以避免过拟合，同时使潜在空间更规则。普通的 encoder-decoder 将输入映射为一个数值，而 VAE 的 encoder-decoder 将输入映射为一个分布，在实作中，encoder 的输出分布被假设是高斯分布，encoder 实际预测的是均值和方差，流程如下：

目标函数为：$||\mathbf{x}-{\mathbf{x}}^{\prime }|{|}^2+KL[\mathcal{N}({\mu }_x,{\sigma }_x),\mathcal{N}(0,1)]$
第一项是恢复损失，第二项是正则项。

正则项的直观理解

规则的潜在空间应该具有两个属性：连续性，完备性，连续性指潜在空间中相邻点的解码内容应该相近，完备性指潜在空间中任意点的解码内容都有意义，如下图（引自[1]）：
如果仅要求 encoder 的输出是一个高斯分布，也难以满足连续性和完备性。例如，当学习到的均值差异很大，方差很小时，此时的模型和普通encoder-decoder没有区别，为了避免这种情况，必须同时对方差和均值进行约束。在论文中，作者要求 encoder 的输出服从 $\mathcal{N}(0,1)$，均值为0使各模式的均值接近，缓解均值差异大的问题，方差为1使各模式分布有一定的“重合”，增强空间连续性，如下图（引自[1]）：

数学理解

这一节从理论角度来分析 VAE。首先定义一些符号，其中 $\theta$ 代表生成模型参数：

• 先验 $p_{\theta }(\mathbf{z})$
• 似然 $p_{\theta }(\mathbf{x}|\mathbf{z})$
• 后验 $p_{\theta }(\mathbf{z}|\mathbf{x})$

当我们想要计算 $p_{\theta }(\mathbf{z}|\mathbf{x})$ 时，利用贝叶斯公式可知：

$p_{\theta }(\mathbf{z}|\mathbf{x})=\frac{p_{\theta }(\mathbf{x}|\mathbf{z})p_{\theta }(\mathbf{z})}{p_{\theta }(\mathbf{x})}$

$p_{\theta }(\mathbf{x})=\int p_{\theta }(\mathbf{x}|\mathbf{z})p_{\theta }(\mathbf{z})d\mathbf{z}$

在计算 $p_{\theta }(\mathbf{x})$ 时，由于无法遍历 $\mathbf{z}$，导致无法直接计算后验，故使用 $q_{\varphi }(\mathbf{z}|\mathbf{x})$ 来近似 $p_{\theta }(\mathbf{z}|\mathbf{x})$（这个涉及到概念 variational inference）。

目标函数

很自然的，可以使用 KL 散度来度量两个分布之间的距离：

$KL(q_{\varphi }(z|x)||p_{\theta }(z|x))$

$=-{\sum }_z{q}_{\varphi }(z|x)log(\frac{p_{\theta }(z|x)}{q_{\varphi }(z|x)})$

$=-{\sum }_z{q}_{\varphi }(z|x)log(\frac{p_{\theta }(x,z)}{q_{\varphi }(z|x)p_{\theta }(x)})$

$=-{\sum }_z{q}_{\varphi }(z|x)[log(\frac{p_{\theta }(x,z)}{q_{\varphi }(z|x)})-log(p_{\theta }(x))]$

$=log(p_{\theta }(x))-{\sum }_z{q}_{\varphi }(z|x)log(\frac{p_{\theta }(x,z)}{q_{\varphi }(z|x)})$

左右交换可以得到：

$log(p_{\theta }(x))=KL(q_{\varphi }(z|x)||p_{\theta }(z|x))+{\sum }_z{q}_{\varphi }(z|x)log(\frac{p_{\theta }(x,z)}{q_{\varphi }(z|x)})$

可以知道，$log(p_{\theta }(x))$ 是一个常数，且 KL 散度是非负的，于是最小化 KL 散度等价于最大化等式右边第二项，将这一项缩写为 $L(\theta ,\varphi ;x)$，论文中被称为 variational lower bound，继续对这一项进行变换：

$L(\theta ,\varphi ;x)$

$={\sum }_z{q}_{\varphi }(z|x)log(\frac{p_{\theta }(x|z)p_{\theta }(z)}{q_{\varphi }(z|x)})$ （1）

$={\sum }_z{q}_{\varphi }(z|x)[log(p_{\theta }(x|z))+log(\frac{p_{\theta }(z)}{q_{\varphi }(z|x)})]$ （2）

$=E_{q_{\varphi }(z|x)}[log(p_{\theta }(x|z))]-KL(q_{\varphi }(z|x)||p_{\theta }(z))$ （3）

其中第一项是重建损失，第二项是正则化项，整体是需要优化的目标函数，只需要求出其相对 $\varphi$ 的梯度，就可以用梯度下降法更新参数 $\varphi$ 了。

重参数化技巧

对于这种形式，一种典型的方法是 Monte Carlo 梯度估计法。记公式（2）中括号里的部分为 $f(z)$（虽然 $f(z)$ 与 $\varphi$ 有关，但这里先暂时忽略），这里仅展示结论性的推导结果：

$▽\varphi {\mathbb{E}}_{q_{\varphi }(z)}[f(z)]$

$={\mathbb{E}}_{q\varphi (z)}[f(z){▽}_{\varphi }log(q_{\varphi }(z))]$

$\simeq \frac{1}{L}{\sum }_{l=1}^{L}f(z){▽}_{\varphi }log(q_{\varphi }(z^l)),z^l\sim q_{\varphi }(z|x)$

如果使用该方法，会产生很大的 variance，使训练不稳定。对此，VAE 提出了 Generic Stochastic Gradient Variational Bayes (SGVB) 梯度估计，包含重参数化技巧（reparameterization trick），这里涉及一个很重要的讨论，为什么需要重参数化技巧？，在 Monte Carlo 梯度估计法中，我们假设 $f(z)$ 和参数 $\varphi$ 无关，但根据公式（1）明显是有关的，在有关情况下再次考虑梯度：

${▽}_{\varphi }{\mathbb{E}}_{q_{\varphi }(z|x)}[f_{\varphi }(z)]$

$={▽}_{\varphi }[{\sum }_z{q}_{\varphi }(z|x)f_{\varphi }(z)dz]$

$={\sum }_z{▽}_{\varphi }[q_{\varphi }(z|x)f_{\varphi }(z)]dz$

$={\sum }_z{f}_{\varphi }(z){▽}_{\varphi }{q}_{\varphi }(z|x)dz+{\sum }_z{q}_{\varphi }(z|x){▽}_{\varphi }{f}_{\varphi }(z)dz$

$={\sum }_z{f}_{\varphi }(z){▽}_{\varphi }{q}_{\varphi }(z|x)dz+{\mathbb{E}}_{q_{\varphi }(z|x)}[{▽}_{\varphi }{f}_{\varphi }(z)]$

由于第一项不能写成期望的形式，无法用求平均来近似，且无法遍历 z 进行积分，因此是不可解的，对此才提出重参数技巧。重参数化令 $z=g_{\varphi }(ϵ,x),ϵ\sim p(ϵ)$，从而使用 $g_{\varphi }(ϵ,x)$ 的分布来替代 $q_{\varphi }(z|x)$，该替换原理基于 variational inference：

${▽}_{\varphi }\mathbb{E}q\varphi (z|x)[f(z^i)]$

$={▽}_{\varphi }\mathbb{E}p(ϵ)[f(g\varphi (ϵ,x^i))]$

$=\mathbb{E}p(ϵ)[▽\varphi f(g_{\varphi }(ϵ,x^i))]$

$\approx \frac{1}{L}{\sum }_{l=1}^L{▽}_{\varphi }f(g_{\varphi }({ϵ}^l,x^i))$

此时就可以计算梯度了。对于为什么要使用重采样技巧，还有一种很受欢迎的解释是，从 $q_{\varphi }(z|x)$ 中采样的过程不可导，因此将抽样通过辅助随机变量 $ϵ$ 剥离出来，梯度就可以通过 $g_{\varphi }(ϵ,x)$ 进行反向传播，如下图（引自[3]）：

SGVB

根据公式（1），将 $z=g_{\varphi }(ϵ,x)$ 代入 $L(\theta ,\varphi ;x)$ 可以得到 SGVB 估计结果 ${\tilde{L}}^A(\theta ,\varphi ;x^i)\simeq L(\theta ,\varphi ;x^i)$：

$L(\theta ,\varphi ;x)$

$={\mathbb{E}}_{q_{\varphi }(z|x)}[-log{q}_{\varphi }(z|x)+log{p}_{\theta }(x,z)]$

${\tilde{L}}^A(\theta ,\varphi ;x^i)$

$=\frac{1}{L}{\sum }_{l=1}^{L}log{p}_{\theta }(x^i,z^{i,l})-log{q}_{\varphi }(z^{i,l}|x^i),$

$where{z}^{i,l}=g_{\varphi }({ϵ}^{i,l},x^i)and{ϵ}^l\sim p(ϵ)$

除此之外，论文还发现，由于我们假设 $q_{\varphi }(z|x)$ 和 $p_{\theta }(z)$ 都服从高斯分布，且 $q_{\varphi }(z|x)$ 的协方差矩阵为对角矩阵，在计算公式（3）中 KL 散度时可以利用高斯分布的性质来简化计算：

$\int q_{\varphi }(z|x)logp(z)dz$

$=\int \mathcal{N}(z;\mu ,{\sigma }^2)log\mathcal{N}(z;0,I)dz$

$=\int \mathcal{N}(z;\mu ,{\sigma }^2)(-\frac{1}{2}{z}^2-\frac{1}{2}log(2\pi ))dz$

$=-\frac{1}{2}\int \mathcal{N}(z;\mu ,{\sigma }^2)z^{2}dz-\frac{J}{2}log(2\pi )$

$=-\frac{J}{2}log(2\pi )-\frac{1}{2}{\mathbb{E}}_{z\sim \mathbb{E}(z;\mu ,{\sigma }^2)}(z^2)$

$=-\frac{J}{2}log(2\pi )-\frac{1}{2}({\mathbb{E}}_{z\sim \mathcal{N}(z;\mu ,{\sigma }^2)}^2(z)+Var(z))$

$=-\frac{J}{2}log(2\pi )-\frac{1}{2}{\sum }_{j=1}^J({\mu }_j^2+{\sigma }_j^2)$

其中 J 是 z 向量的维数，同理还可以得到：

$\int q_{\varphi }(z|x)log{q}_{\varphi }(z|x)dz$

$=-\frac{J}{2}log(2\pi )-\frac{1}{2}{\sum }_{j=1}^J(1+log({\sigma }_j^2))$

最终，KL 散度的计算可以简化为：

$-KL(q_{\varphi }(z|x)||p_{\theta }(z))$

$=\int q_{\varphi }(z|x)(log{p}_{\theta }(z)-log{q}_{\varphi }(z|x))dz$

$=\frac{1}{2}{\sum }_{j=1}^J(1+log{\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2)$

而为了优化公式（3）中的第一项，使用 MSE 损失即可，因此最终的计算过程可以总结为下图（引自[4]）：

参考

[1] Understanding Variational Autoencoders (VAEs)
[2] [论文简析]VAE: Auto-encoding Variational Bayes[1312.6114]
[3] From Autoencoder to Beta-VAE
[4] The Reparameterization Trick

---

补充自2022年10月18日

关于 evidence：p(x)
在很多任务里，需要计算数据的分布 $p(x)$，根据贝叶斯公式，可以知道：

$p(x)={\int }_z{p}_{\theta }(x|z)p(z)$

上面讨论过这个是 intractable 的，并引入了 $q_{\varphi }(z|x)$ 进行近似计算，因此可以进一步写成：

$p(x)=\int q_{\varphi }(z|x)\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}$

$logp(x)=log{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}]$

根据 Jenson 不等式，可以得到：

$logp(x)\ge {\mathbb{E}}_{z\sim q_{\varphi }(z|x)}log[\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}]$

不等式右边的项也就是 evidence lower bound（ELBO），VAE 的目标是尽可能最大化 ELBO，这样 $p(x)$ 也就会更大。这里就和前面联系起来了，为了计算 ELBO：

$ELBO$

$=E_{q_{\varphi }(z|x)}[log(p_{\theta }(x|z))]-KL(q_{\varphi }(z|x)||p_{\theta }(z))$

第一项用 MSE 处理，为了计算第二项 KL 散度，作者假设 $p_{\theta }(z)\sim \mathcal{N}(0,I),q_{\varphi }(z|x)\sim \mathcal{N}(\mu ,{\sigma }^2)$，为了在反向传播中使梯度能够通过不可导的 $q_{\varphi }(z|x)$ 采样这一步，引入了重参数化技巧，此时 $ϵ$ 仅仅作为网络的输入，不参与梯度传播，梯度通过相关系数 $\mu ,\sigma$ 进行传播。