变分自编码器 VAE 超详解，从简单公式推导到模型结构到模型理解

参考文献：

[1] Kingma D P, Welling M. Auto-encoding variational bayes[J]. arXiv preprint arXiv:1312.6114, 2013.

[2] Doersch C. Tutorial on variational autoencoders[J]. arXiv preprint arXiv:1606.05908, 2016.

[3] 变分自编码器（一）：原来是这么一回事 - 科学空间|Scientific Spaces

[4] 一个例子搞清楚（先验分布/后验分布/似然估计）_一个例子搞清楚(先验分布/后验分布/似然估计)-CSDN博客

[5] 直观解读KL散度的数学概念 - 简书 (jianshu.com)

[6] Kullback-Leibler Divergence Explained — Count Bayesie

[7] VAE变分自编码机详解——原理篇 - 知乎 (zhihu.com)

[8] 贝叶斯估计浅析 - xueliangliu - 博客园 (cnblogs.com)

基本概念介绍

从 AutoEncoder 到 VAE

自编码器

变分自编码器做出的改变

VAE 原理思路入门

我们想做什么

变分下界引入与推导

变分下界的理解

VAE 模型架构介绍

根据潜在空间确定模型架构

对模型架构进行调整

重参数技巧

VAE 模型总结

基本概念介绍

首先我们需要明白或者回顾一些基本的概念，下面是部分重要概念以及随附文章：
- 先验/后验概率，证据，似然估计，MAP/ML/贝叶斯估计等：
  - 一个例子搞清楚（先验分布/后验分布/似然估计）_一个例子搞清楚(先验分布/后验分布/似然估计)-CSDN博客
  - 贝叶斯估计浅析 - xueliangliu - 博客园 (cnblogs.com)
- KL 散度（包括二项分布）：
  - 直观解读KL散度的数学概念 - 简书 (jianshu.com)
  - Kullback-Leibler Divergence Explained — Count Bayesie

上述文章有先后顺序，如果对相关概念不了解，推荐从上到下简单看一看。

从 AutoEncoder 到 VAE

让我们默认了您之前已经简单了解了编码器/解码器架构，这将更好理解下面的文章。当然，不了解也没关系

自编码器

书面上，自编码器是一种无监督学习的神经网络结构，它试图学习输入数据的紧凑表示（编码），然后通过解码器将该表示还原为输入数据。自编码器包括一个编码器网络和一个解码器网络。
- 编码器接收输入的数据后对其编码，我们可以理解为它在提取输入数据的关键信息。其生成的编码，也就是中间变量，我们称之为潜在表示，用 h 来表示。h 之后会进入解码器，得到最终的输出。即：
  - 编码器（Encoder）：将输入数据映射到低维潜在空间，捕捉输入数据的重要特征。
  - 解码器（Decoder）： 将编码的潜在表示映射回原始输入空间，重构原始数据。

变分自编码器做出的改变

那么变分自编码器（VAE）和自编码器差别在哪里呢？我认为最大的改变就是潜在表示。在自编码器中，潜在表示是一个固定值，而 VAE 中，潜在表示则是一个不确定变量，或者换一个说法，是一个概率分布。
在 VAE 中，编码器不再只学习提取输入数据的编码信息，而是去学习获取输入数据的概率分布。在原始论文中，我们设定这一概率分布为正态分布，模型也就变成了这样：
- 与此同时，中间量我们不再叫潜在变量，而是称为潜在空间（也称潜在分布，隐变量等），并使用 Z 来表示，用来凸显其是一个变化的量。
- 解码器要做的工作就是如何从这样的概率分布中采样并还原成最终的输出。

VAE 原理思路入门

我们从熟知的模型延伸到 VAE，从而对它的思路以及独特之处做了简单的了解。下面我们将从头介绍一下 VAE 的大体思路。

我们想做什么

首先我们看生成式模型的“梦想”。首先我们有一批数据样本 {X1, …, Xn}，其整体用 X 来描述，我们本想根据 {X1,…,Xn} 得到 X 的分布 p(X) ，如果能得到的话，那我直接根据 p(X) 来采样，就可以得到所有可能的 X 了（包括 {X1, …, Xn} 以外的），这就算是一个终极理想的生成模型了。
当然，这个“梦想”很难实现，于是我们通过一个迂回的策略，做一个概率分布映射，从一个简单的分布（如高斯分布）映射到 p(X) ，这样貌似也不错。于是我们将分布改一改：

$\begin{aligned} p(X)& = \sum_Z{p(x|z)p(z)} \\ &=\int{p(x|z)p(z)}\space\mathrm{d}z \end{aligned}$
然而，Z 所在的潜在空间通常是高维且包含复杂相互作用的，直接去积分来计算 p(X) 几乎是不可能的，于是乎，我们就需要将这个问题简化、近似化。于是，变分下界（Evidence Lower BOund，ELBO）的概念呼之欲出。

变分下界引入与推导

那么 VAE 是怎么解决这个问题的呢？事实上，我们再看看原来的 p(X) 式子：

$p(X)=\int{p(x|z)p(z)}\space\mathrm{d}z$
我们可以说，对于 z~P(z) 的大部分采样，也就是大部分 z，对于 P(X) 的计算都没有贡献，也就是它们的 P(z|X) 都近乎为0，根本就跟 X 没什么关系。因此如果需要近似，我们可以从这个角度切入，剔除那些没有什么贡献的 z，只需要计算使得 P(z|X) 更大的那部分 z 即可。
怎么找到这部分贡献大的 z 呢？在给定 X 的情况下 z 的后验分布 P(z|X) 就约等于使得 X 的生成概率最大的 z 的分布，所以问题就变成了如何计算后验分布 P(z|X)。而直接去计算 z 的后验分布是很困难的。于是我们引入一个新的近似后验分布 q，使其近似代替真实后验分布 P(z|X)，式子如下：

$q(z|x)$
好，就此打住千万别再绕远了。我们还是得求 p(X) 使其最大化啊，如何将这两个式子联系起来呢？我们在这里使用 KL 散度来将这个近似后验分布同我们要求的联系起来。由 KL 散度公式，我们可以有下式：

$D[q(z|x)||p(z|x)] = \mathbb{E}_{z\sim q}[\log{q(z|x)} - \log{p(z|x)}]$
使用贝叶斯公式转化右式的 p(z|x)，式子变形如下：

$\begin{aligned} D[q(z|x)||p(z|x)] &= \mathbb{E}_{z\sim q}[\log{q(z|x)} - \log{p(z|x)}] \\ &=\mathbb{E}_{z\sim q}[\log{q(z|x)} - \log{\frac{p(x|z)*p(z)}{p(x)}}] \\ &=\mathbb{E}_{z\sim q}[\log{q(z|x)} - \log{p(x|z)-\log{p(z)}+\log{p(x)}}] \\ \end{aligned}$
p(X) 是确定量，我们将其抽出，并进行移项：

$\begin{aligned} D[q(z|x)||p(z|x)] &=\mathbb{E}_{z\sim q}[\log{q(z|x)} - \log{p(x|z)}-\log{p(z)}+\log{p(x)}] \\ &=\mathbb{E}_{z\sim q}[\log{q(z|x)} - \log{p(x|z)}-\log{p(z)}] + \log{p(x)} \\ \end{aligned}$

$\log{p(x)} - D[q(z|x)||p(z|x)] = \mathbb{E}_{z\sim q}[-\log{q(z|x)} + \log{p(x|z)} +\log{p(z)}] \\ \log{p(x)} - D[q(z|x)||p(z|x)] = \mathbb{E}_{z\sim q}[\log{p(x|z)}] - \mathbb{E}_{z\sim q}[\log{q(z|x)} - \log{p(z)}] \\$
再使用 KL 散度公式进行重写，我们就得到了这个式子：

$\log{p(x)} - D[q(z|x)||p(z|x)] = \mathbb{E}_{z\sim q}[\log{p(x|z)}] - D[q(z|x)||p(z)]$
由于 KL 散度非负，因此我们可以说：

$\log{p(x)} \ge \mathbb{E}_{z\sim q}[\log{p(x|z)}] - D[q(z|x)||p(z)]$
上面式子的 RHS（右式）很明显可以算是 log p(x) 的下界，我们一般称其为变分下界（Variational Lower Bound，ELBO）。实际上，我们的目标和 ELBO 的差距就是近似后验分布 q(z|x) 和真实后验分布 p(z|x) 的差距。至此，我们将 “使 p(x) 最大化” 这一问题转化为了将变分下界最大化。

变分下界的理解

既然目标是让变分下界最大化，那么我们就需要仔细研究一下这个变分下界。我们不难看出，变分下界主要由两部分组成：
- z 在 q(z|x) 的分布下，x given z 的概率分布的期望值
- z given x 的近似后验分布与 z 的真实分布（先验分布）的 KL 散度
首先是第一项，要想最大化 ELBO，那我们自然是想让第一项尽可能的大，也就是 x given z 的概率分布期望值更大。这很明显就是由 z 到 x 重组的过程，也就是 AutoEncoder 中的 Decoder，从潜在空间 Z 中重组 x。模型想做的是尽可能准确地重组
其次是第二项，要想最大化 ELBO，我们自然需要让这项 KL 散度尽可能小，也就是 潜在空间 z 的近似后验分布尽可能接近于 z 的先验分布！这一项我们可以理解为，模型想让 z 尽可能避免过拟合，而让 z 的近似后验分布尽可能接近其先验分布。那如何做到这一效果呢？只能在生成 z 的时候，也就是从 Encoder 上下手了。

VAE 模型架构介绍

根据潜在空间确定模型架构

理解了我们的核心公式和要做的事情后，我们就可以着手搭建模型了。首先我们需要做的是确定 z。VAE 是如何去描述潜在空间 z 的呢？它认为，z 根本没有一种合适的阐述方法，而是直接假定 z 的样本可以从简单的分布中抽取，即标准正态分布 N(0, I)，其中 I 是单位矩阵。因为，任何 d 维度的分布都可以用一组 d 个服从正态分布的变量，通过足够复杂的函数进行映射从而生成。

其中的原理可以参考[1606.05908] Tutorial on Variational Autoencoders (arxiv.org)这篇论文。此外，也有采用正态分布可以更方便 KL 散度计算的说法。
那么怎么找到映射函数，并获取其输出，也就是正态分布的平均值和方差？这种复杂的任务自然是交给神经网络学习啦，我们可以使用两个神经网络分别输出平均值和标准差，如下：

$\mu = f_1(X) \\ \log \sigma^2 = f_2(X)$
- 其中 f_1 和 f_2 分别代表两个独立的神经网络
- 选择拟合 log σ^2 是因为原本 σ^2 为非负，需要加激活函数，如果去拟合其 log 值就不用加激活函数了。
这下我们就可以将我们的近似后验分布 q 描述为下面的公式：

$q(z|x) = \mathcal{N}(z;\mu, \sigma ^2\mathbf{I})$
- 其中，μ 和 σ^2 就是通过两个不同的神经网络，根据样本学习到的平均值和方差。
- I 为单位矩阵。
好的，那模型貌似就很简单了。我们来看看整个模型：
- 首先，样本进入 Encoder，Encoder 学习并输出潜在空间的平均值和标准差，得到潜在空间 z，
- 然后再从潜在空间中采样，并进入 Decoder，Decoder 根据采样进行重构，重新生成样本。
- 根据我们的目标函数，只需要保证重构尽可能准确 + 潜在空间接近标准正态分布即可，即 ELBO 最大化。那么我们可以直接对 ELBO 进行取反作为我们的损失函数：
  
  $\mathcal{L} = - \mathbb{E}_{z\sim q}[\log{p(x|z)}] + D[q(z|x)||p(z)]$
根据上面的理论，我们可以画出我们的模型架构图：
好像看似很完美？那就开始训练吧！损失函数就是 ELBO 取反。根据我们的损失函数，我们需要对比 X 和 X’，即原样本和生成样本。但是随即问题就出现在了我们的眼前：如何对比？我们来看看一些常规的想法：
- KL 散度？KL 散度作为一种常见的比较两种分布相近程度的方法，在这里却失去了作用。其原因在于 KL 散度是根据两个概率分布的表达式来算它们的相似度的，然而我们只有一批从构造的分布 Z 采样重构而来的数据 {X_1‘, X_2’ ,…, X_n‘}，还有一批从真实的分布采样而来的数据（也就是训练集）{X_1, X_2, …, X_n}，我们只知道样本本身，没有分布表达式，当然也就没有方法算KL散度。
- 直接作比较？那既然我们有了生成的数据集和训练集，为什么不能直接算 D（某种距离函数）这样的值直接进行比较呢？这样也有困难！其原因在于我们并不知道哪一个生成的样本对应着哪一个真实样本，也就是我们并不知道生成样本集和训练集的对应关系，X_1' 不一定对应着 X_1，自然就不能马马虎虎直接算 D(X_k', X_k)。
坏了，这些常规的想法都不行！那我们应该怎么办？

对模型架构进行调整

遇到问题我们就去解决问题。我们来看看同为生成模型，GAN（生成式对抗网络）是如何解决的：
- GAN 主要由两部分：生成器和判别器两部分组成。通过一个生成器网络生成样本，同时通过一个判别器网络对生成的样本进行判别。训练时，我们将真样本也给到判别器网络，让其进行训练判别能力。而生成器的训练目标就是能尽可能欺骗过判别器，让它识别不出来自己生成的样本。
- 也就是说，GAN的生成过程是一个对抗性的过程，生成器和判别器相互竞争，最终生成器试图生成逼真的样本，判别器试图准确地区分真实和生成的样本。
GAN 的思路很简单。既然比较有困难，那我直接把比较的方法也一并训练出来不就行了？这里 GAN 训练出来的比较的方法就是判别器。
那么 VAE 呢？VAE 则不是这样粗暴，而是采用了一种精妙的迂回方式。既然直接比较是找不到对应关系的，那我干脆就一对一进行训练，直接一个真实样本训练一个潜在空间，这样重构产生的样本自然就对应着它的老爸啦。修改后的模型示意图如下：
如此改动后，我们就有了专属于每一个 Xk 的正态分布，从中采样并进行重构的自然就是对应的 Zk，如此我们就可以直接进行比较。于是乎，我们的损失函数就可以进一步完善成下式：

$\mathcal{L}(\theta, \phi; x^{(i)}) = - \mathbb{E}_{z\sim q_\phi(z|x^{(i)})}[\log{p_\theta(x^{(i)}|z)}] + D_{KL}[q_\phi(z|x^{(i)})||p(z)]$
- 其中，θ 和 φ 分别是 Decoder 和 Encoder 模型参数
- 我们针对每一个样本都单独计算其独有的均值和标准差，构造属于其专有的潜在空间。
我们也重新对我们的近似后验分布进行描述，由于在损失函数中，我们需要其对数形式，所以我们将对数形式添加上去，最终的公式如下：

$\log q(z|x^{(i)}) = \log \mathcal{N}(z;\mu^{(i)}, \sigma^{2(i)}\mathbf{I})$
值得一提的是，有一个很有意思的地方在于损失函数的第二项。我们在之前的变分下界解析中曾经提到，这一项主要目的是为了防止潜在空间过拟合。当知晓了模型架构后，我们重新来看损失函数。第一项很自然的，就是重构项和对应的样本之间的差距，越小越好。然而，在整个模型中，潜在表示，也就是重构材料 Z_k 是采样得到的，并不是像常规 AutoEncoder 那样由 Encoder 直接计算得到，因此这部分也就相当于噪声，它的随机性在干扰重构的过程。因此在训练过程中，为了更好重构，模型会尽可能让潜在空间的方差变为0，进而退化成普通的 AutoEncoder，所谓的“生成式模型”可就名存实亡了。这可万万不能！此时，损失函数的第二项就起到了作用，它让潜在空间的后验分布逼近标准正态分布 N(0, I)，从而避免了随机性消失，也就是方差变成0的情况，相当于对训练过程的正则化。因此，我们通常称损失函数中：

$- \mathbb{E}_{z\sim q_\phi(z|x^{(i)})}[\log{p_\theta(x^{(i)}|z)}]$

这一项为重构项（Reconstruction Term），而另外一项：

$D_{KL}[q_\phi(z|x^{(i)})||p(z)]$

这一项为正则化项（Regularization Term）。

重参数技巧

貌似万事大吉了？等等！这能训练吗？整个过程中，你可是从一个概率分布（潜在空间）中采样才得到潜在表示 Z_k 的，而恰恰“采样”这个操作是不可导的，这也就导致在训练过程中无法进行反向传播，就无法训练了，怎么办呢？
为了解决这个问题，VAE 使用了重参数技巧（reparameterization trick）。我们重新观察我们的潜在空间 Z，也就是后验分布 q(z|x)，利用正态分布标准化，我们可以得到下面的结论：

$Z \sim \mathcal{N}(\mu, \sigma^2) \Rightarrow \frac{Z-\mu}{\sigma} \sim \mathcal{N}(0, \mathbf{I})$
即，(z − μ) / σ 服从均值为 0，方差为 1 的标准正态分布。因此我们引入噪声项 ε，令 ε = (z − μ) / σ，从而将原本 从 N(μ, σ^2) 中采样得到 z 的操作转换成了 从 N(0, I) 中采样一个 ε，令 z = μ + ε*σ。如此一来，梯度计算可以通过 μ 和 σ 直接传递，而不再涉及对随机采样的梯度。

VAE 模型总结

现在我们整体看一下整个 VAE 的模型结构，实际上它就是从普通的自编码器发展过来的变种。普通的 AutoEncoder，其 Encoder 生成的是有关样本的“信息”，也就是潜在表示，这样算出来的值是确定的。而 VAE 的 Encoder 却生成的是一个潜在空间，更具体地说是生成专属于当前样本的平均值和方差，这样算出来的“值”是不确定的。而从这样的潜在空间采样后再进行重构，这一过程本身就相当于“给潜在表示添加噪声”，相当于对 Encoder 的正则化，这样就使得 Decoder 能够对噪声具有一定的鲁棒性。
从这个角度来看，生成方差的那个神经网络又何尝不可以理解为对噪声大小的调节器呢？方差越大，采样的结果就会越分散，这就变相增加了重构的难度；方差越小，采样的结果就越集中，就会降低重构的难度。引用变分自编码器（一）：原来是这么一回事 - 科学空间|Scientific Spaces这篇文章中的内容来说，就是：

当decoder还没有训练好时（重构误差远大于KL loss），就会适当降低噪声（KL loss增加），使得拟合起来容易一些（重构误差开始下降）；反之，如果decoder训练得还不错时（重构误差小于KL loss），这时候噪声就会增加（KL loss减少），使得拟合更加困难了（重构误差又开始增加），这时候decoder就要想办法提高它的生成能力了。