最近 UROP 在做生成模型这方面的任务，自己作为一个机器学习的纯萌新自然是一窍不通，只是凭着自己的模糊记忆记得生成模型似乎有三大流派：VAE、GAN 和最近火起来的 Diffusion Model。其中 GAN 是直觉上最好理解的，如果只是需要实现的话对于数学的要求不太高，但是 VAE 和 DM 哪怕是要初步理解都需要一定的数学背景知识。我一个下学期才选概率相关的基础课的半瓶水读起来挺吃力的。硬啃论文一个礼拜总算是把这两个模型基本搞明白了。这篇文章先记录一下我对 VAE 的理解，主要参考了 Tutorial on Variational Autoencoders。

基本思路

VAE 是生成模型。什么是生成模型？生成模型的本质是从一个概率分布中采样。这种采样在我们明确知道目标分布的 PDF 或 CDF 的时候并不困难。然而，对于很多实际问题，需要采样的变量的维数非常高，而分布本身也无法显式求得，这些局限是我们需要一个个复杂而精妙的生成模型的根本原因。

VAE（以及 GAN）基于一个假设：如果我们的样本空间分布的复杂性来源于不同维之间复杂的相关性，那么样本空间真正有意义的维度应该少于样本空间形式上的维数。因此，复杂的、高维的样本可以看作是由一个相对简单的、低维的隐变量（latent variable）所决定的。我们不妨让这些隐变量 \(\boldsymbol z\) 服从标准正态分布。于是生成这些隐变量就很简单了，而难题也被我们转化为了求一个从隐变量到目标样本的映射 \(\boldsymbol x=f(\boldsymbol z)\)。更为广义地来说，\(f\) 不一定要是一个确定性的过程，所以我们要求的实际上是 \(P(\boldsymbol x|\boldsymbol z)\)。如果我们假设 \(P(\boldsymbol x|\boldsymbol z)\) 有固定的形式（比如说正态），那么也可以说我们寻找的是一个将 \(\boldsymbol z\) 映射到 \(P(\boldsymbol x|\boldsymbol z)\) 分布的参数的函数 \(p(\boldsymbol z)\)。在机器学习的背景下，\(p\) 的形式是固定的，我们需要寻找或者学习的其实是一组参数 \(\boldsymbol \theta\)。那么什么样的 \(p_{\boldsymbol \theta}\) 是好的？我们希望我们的模型能够生成类似于训练集的数据，所以从数学上讲，对于训练集中的数据 \(\boldsymbol x\)，我们希望最大化我们的模型的 likelihood，即 \(P(\boldsymbol x|\boldsymbol \theta)\)。

这里我觉得有必要提一下 likelihood 和 probability 的区别。我第一次读论文的时候看记号相似，就以为两个说的是一回事，犯了想当然的错误。Probability 是相对于一个固定的假设 / 分布的，关于结果的函数；Likelihood 是相对于一个固定结果的、关于假设 / 分布的函数。结果是互斥的，所以不同结果的 probability 一定和为 1；但是 likelihood 则不然。在生成模型的背景下，训练集的样本都是固定的而变化的是我们的模型本身，所以应该用 likelihood。

那么怎么计算 \(P(\boldsymbol x|\boldsymbol \theta)\) 呢？从定义上来讲它可以看作是一个 marginalization。不妨定义 \(p_{\boldsymbol \theta, \boldsymbol z}(\boldsymbol x)\) 为 \(p_{\boldsymbol \theta}(\boldsymbol z)\) 表示的分布的 PDF，则有 \[ P(\boldsymbol x | \boldsymbol \theta)=\int P(\boldsymbol x | \boldsymbol z,\boldsymbol \theta)P(\boldsymbol z)\,d\boldsymbol z = \mathbb{E}_{\boldsymbol z}\bigl[P(\boldsymbol x | \boldsymbol z,\boldsymbol \theta)\bigr] = \mathbb{E}_{\boldsymbol z}\bigl[p_{\boldsymbol \theta, \boldsymbol z}(\boldsymbol x)\bigr]. \] （这里我们假设 \(\boldsymbol z\) 独立于我们的参数，因为前文提过我们可以把前者看作是从标准正态分布中抽样而来的）

遗憾的是上式可以看作是正确的废话。注意 \(\boldsymbol x\) 和 \(\boldsymbol z\) 是强对应的，所以 \(P(\boldsymbol x | \boldsymbol z,\boldsymbol \theta)\) 是稀疏的。这导致如果我们在实践中用随机抽样的方法来估计 \(\mathbb{E}_{\boldsymbol z}\bigl[P(\boldsymbol x | \boldsymbol z,\boldsymbol \theta)\bigr]\)，我们随机到的 \(\boldsymbol z\) 很可能不会对我们的估计产生任何有意义的贡献，导致我们的估计值完全不可靠。

要解决这个问题，我们就需要确保我们抽样得到的都是有代表性的 \(\boldsymbol z\)。这样一来我们就不能盲目地从标准正态分布中采样了 —— 应该怎么办呢？

信息熵与 KL 距离

在介绍 VAE 提供的解决方案之前，我在这一节先简要介绍一下相关的一些背景知识。

思考一个问题：如何量化一个随机事件结果 \(X\) 的 “意外程度（surprise）”\(s(X)\)？这个意外程度应该满足如下三个要求

• 百分百发生的事情自然是毫不令人意外的，因此 \(P(X)=1\rightarrow s(X)=0\)；
• 一件事情发生的概率越低显然越让人意外，所以 \(s(X)\) 关于 \(P(X)\) 递减。
• 对于两个独立事件的结果。其总的意外程度应该是两个结果意外程度之和，即 \(s(X\cap Y)=s(X)+s(Y)\)。

香农论证了满足以上三个约束的 \(s(X)\) 只能取 \(s(X)=-\log P(X)\)（最多乘上一个正常数），这个结果是符合我们直觉的。同时注意到，“意外程度” 在某种意义上就是信息量：毫不意外的事情不会传达任何信息，而相反一件事越罕见其信息量越大。所以，香农接着定义一个随机事件的信息量为其所有结果信息量 / 意外程度的期望，即 \[ \int -P(X)\log P(X)\,dX. \] 这一个式子就是人们常说的信息熵 —— 这个式子长成这样是有理由的！

接下来思考一个问题：假设一个随机事件现实中的结果分布是 \(P\)，但是我们有一个理论说其结果分布是 \(Q\)。我们如何量化理论和现实的差距？

假设一个人相信理论 \(Q\)，那么事件 \(X\) 对那个人来说的信息量为 \(-\log Q(X)\)。客观上来讲，事件是按照 \(P\) 刻画的概率分布发生的，所以这个事件对这个人的平均信息量就是 \[ \mathbb E_{X\sim P}\bigl[-\log Q(x)\bigr]=\int -P(X)\log Q(X)\,dX. \] 上式和这个事件实际上的信息熵之差就可以看作是一种 “理论和现实的距离”，称之为相对熵（relative entropy）。更为常见的称呼是 Kullback-Leibler 距离： \[ \int -P(X)\log Q(X)\,dX-\int -P(X)\log P(X)\,dX = \int -P(X)\log \left(\frac{Q(X)}{P(X)}\right)\,dX = D_{\text{KL}}(P\|Q). \] 上面这段话蕴含一个假设：所有基于 \(Q\) 计算的信息熵一定不小于一个事件实际的信息熵 —— 这样作差当距离才有意义。真的是这样的吗？换而言之，KL 距离在任何情况下都是非负的吗？注意到 \(\log y\le y-1 \Rightarrow -\log y\ge 1-y\)： \[ \begin{aligned} D_{\text{KL}}(P\|Q) &= \int -P(X)\log \left(\frac{Q(X)}{P(X)}\right)\,dX\\ &\ge\int -P(X)\left(\frac{Q(X)}{P(X)}-1\right)\,dX\\ &=\int P(X)\,dX-\int Q(X)\,dX\\ &=1-1=0 \end{aligned} \] 所以 KL 距离作为一种 “距离”，在直觉上确实是合理的。

读 VAE 的文章之前我就不止一次地看到信息熵和 KL divergence 的式子，当时我也没有深究，觉得信就完事了，反正也不是特别难记（笑）。直到读 VAE 的文章，我终究觉得这种式子盲信起来总归有一种不踏实感，所以去找了一些相关的解释，发现如果不完全苛求严谨的话这两个式子还是比较好理解的，于是在这里记录一下。从这种解释中也能理解为什么 KL 距离是非对称的了，因为 KL 距离中的 \(P\) 和 \(Q\) 的地位就不一样。

回到 VAE

为了更加精确地近似采样 \(\boldsymbol z\)，我们考虑从一个能给出 “有代表性的 \(\boldsymbol z\)” 的分布 \(Q\) 中采样。所谓 “有代表性”，可以理解为 \(Q\) 应该近似于 \(P(\boldsymbol z|\boldsymbol x, \boldsymbol \theta)\)。换而言之，两者的 KL 距离应该尽可能小。于是考虑 \(D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\bigr)\)： \[ D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\bigr) = \mathbb E_{\boldsymbol z\sim Q}\bigl[\log Q(\boldsymbol z)-\log P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\bigr] \] 我们可以用贝叶斯公式展开 \(P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\)： \[ \begin{aligned} D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\bigr) &= \mathbb E_{\boldsymbol z\sim Q}\bigl[\log Q(\boldsymbol z)-\log P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\bigr] \\ &= \mathbb E_{\boldsymbol z\sim Q}\!\left[\log Q(\boldsymbol z)-\log \frac{P(\boldsymbol x|\boldsymbol z,\boldsymbol \theta)P(\boldsymbol z)}{P(\boldsymbol x|\boldsymbol \theta)}\right] \\ &= \mathbb E_{\boldsymbol z\sim Q}\bigl[\log Q(\boldsymbol z)-\log P(\boldsymbol x|\boldsymbol z,\boldsymbol \theta)-\log P(\boldsymbol z)+\log {P(\boldsymbol x|\boldsymbol \theta)}\bigr]. \end{aligned} \] 注意到 \(\log {P(\boldsymbol x|\boldsymbol \theta)}\) 一项和 \(\boldsymbol z\) 无关，所以可以提出来，然后再拆分一下期望里面的项， \[ \begin{aligned} D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\bigr) &= \mathbb E_{\boldsymbol z\sim Q}\bigl[\log Q(\boldsymbol z)-\log P(\boldsymbol x|\boldsymbol z,\boldsymbol \theta)-\log P(\boldsymbol z)+\log {P(\boldsymbol x|\boldsymbol \theta)}\bigr] \\ &= \log P(\boldsymbol x|\boldsymbol \theta) + \underbrace{\mathbb E_{\boldsymbol z\sim Q}\bigl[\log Q(\boldsymbol z)-\log P(\boldsymbol z)\bigr]}_{D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z)\bigr)} - \mathbb E_{\boldsymbol z\sim Q}\bigl[\log P(\boldsymbol x|\boldsymbol z,\boldsymbol \theta)\bigr]. \end{aligned} \] 移项，再利用 KL 距离非负的性质，我们就得到了 log likelihood 的一个下界： \[ \begin{aligned} \log P(\boldsymbol x|\boldsymbol \theta) &= \mathbb E_{\boldsymbol z\sim Q}\bigl[\log P(\boldsymbol x|\boldsymbol z,\boldsymbol \theta)\bigr] - D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z)\bigr) + D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\bigr) \\ &\ge \mathbb E_{\boldsymbol z\sim Q}\bigl[\log P(\boldsymbol x|\boldsymbol z,\boldsymbol \theta)\bigr] - D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z)\bigr). \end{aligned} \] 也就是说，只要 \(D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\bigr)\) 足够小，即 \(Q\) 足够逼近 \(P(\boldsymbol z|\boldsymbol x,\boldsymbol \theta)\)，优化 \(\mathbb E_{\boldsymbol z\sim Q}\bigl[\log P(\boldsymbol x|\boldsymbol z,\boldsymbol \theta)\bigr] - D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z)\bigr)\) 就近似等价于优化 \(\log P(\boldsymbol x|\boldsymbol z)\)，也就近似等价于优化 \(P(\boldsymbol x|\boldsymbol z)\)。这个下界被称为 variational lower bound，是 VAE 的 loss。所谓的 “variational”，是指我们之前对于贝叶斯公式的运用 —— 那种方法似乎被称为 variational Bayesian method。

现在还剩下一个问题，就是我们不知道这样的 \(Q\) 具体是什么，但是我们知道 \(Q\) 应该是和 \(\boldsymbol x\) 有关的。我们于是发挥 ML 的传统艺能：如果不知道就糊个 NN 上去。不妨假设 \(Q(\boldsymbol z)=\mathcal N\bigl(\boldsymbol\mu(\boldsymbol x),\operatorname{diag}\boldsymbol\sigma^2(\boldsymbol x)\bigr)\)，我们让神经网络给出 \(\boldsymbol \mu\) 和 \(\boldsymbol \sigma\) 就好了！注意到在 \(Q(\boldsymbol z)\) 和 \(P(\boldsymbol z)\) 都是正态分布（且我们假设 \(P(\boldsymbol z)=\mathcal N(0, I)\)）的情况下，其 KL 距离存在解析解： \[ D_{\text{KL}}\bigl(Q(\boldsymbol z)\|P(\boldsymbol z)\bigr)=\frac{1}{2}\left(\sum \boldsymbol\sigma^2(\boldsymbol x) + \bigl\|\boldsymbol \mu(\boldsymbol x)\bigr\|^2-\log \prod \boldsymbol \sigma^2(\boldsymbol x)-\dim \boldsymbol z\right) \] 这是一个关于 \(\boldsymbol\mu\) 和 \(\boldsymbol \sigma\) 的函数，于是我们在优化 \(\boldsymbol \theta\) 的同时可以顺便优化 \(\boldsymbol \mu\) 和 \(\boldsymbol \sigma\)。

接下来还有一个技术上的问题：我们在计算目标函数的过程中要从 \(Q\) 当中采样。但是采样是一个离散的过程，会打断梯度的反向传播，导致 \(Q\) 的参数不能被正确更新。解决这个的方法被称为 reparameterization trick。我们可以从标准正态分布中采样，再把采样线性缩放到 \(Q\) 的分布。这样 \(\boldsymbol \mu\) 和 \(\boldsymbol \sigma\) 就有梯度了。在数学上说，这就是把 loss 中的 \(\mathbb E_{\boldsymbol z\sim Q}\bigl[\log P(\boldsymbol x|\boldsymbol z,\boldsymbol \theta)\bigr]\) 改为 \(\mathbb E_{\boldsymbol \epsilon\sim \mathcal N(0,I)}\bigl[\log P(\boldsymbol x|\boldsymbol z=\boldsymbol \sigma (\boldsymbol x)\odot \boldsymbol\epsilon +\boldsymbol \mu(\boldsymbol x),\boldsymbol \theta)\bigr]\)。

最后还有一个问题：为什么 VAE 是一种 “auto-encoder”？回顾 VAE 的推导，\(Q\) 网络将一个样本映射到一个隐变量的分布，可以近似看作是一种 “编码” 的过程；\(P\) 网络将一个隐变量隐射到一个样本的分布，可以近似看作是一种 “解码” 的过程。两个网络是同步训练的，这点和传统意义上的 encoder-decoder 模型颇为相似。因为这几点，VAE 便得了一个 auto-encoder 的名字。但是说 “auto-encoder” 只是形似：encoder-decoder 的架构只能是事后的附会，而不是推导的灵感基础，何况 encoder-decoder 应该是确定性的，但是 VAE 两个网络输出的其实都是分布，这也是和 encoder-decoder 架构不一样的地方。

拓展：VRNN

VAE 的数学理论是优美的，但是有一个限制：只能够建模固定维度的样本。对于不定长序列的建模和生成，原版的 VAE 是无能为力的。为了解决这个问题，Chung et.al. 在 15 年发表了名为 Variational Recurrent Neural Network 的架构。论文很有意思，但是归根结底，VRNN 是给每一步的 RNN Cell 都嵌入一个 VAE，同时将我们之前的推导都 condition on 前一个 RNN cell 的 hidden state。

我们假设有两个比较简单的网络 \(\varphi_{\boldsymbol x}(\boldsymbol x)\) 和 \(\varphi_{\boldsymbol z}(\boldsymbol z)\) 可以从输入和 latent variable 中做一个初步的特征提取。令 RNN 上一个时刻的 hidden state 为 \(\boldsymbol h_{t-1}\)，则神经网络本身的 inference 可以表现为： \[ \boldsymbol h_t = f\bigl(\varphi_{\boldsymbol x}(\boldsymbol x_{t}),\varphi_{\boldsymbol z}(\boldsymbol z_{t}),\boldsymbol h_{t-1}\bigr). \] \(f\) 除了最 naive 的 RNN cell 以外也可以是如 LSTM、GRU 之类的变体。如果是 LSTM，\(\boldsymbol h\) 实际上是 hidden state 和 cell state 的总称。

所谓的嵌入 VAE，是指每一步除了 \(f\) 以外，还多出了三个网络：\([\boldsymbol \mu_{\boldsymbol z, t}, \boldsymbol \sigma_{\boldsymbol z, t}]=q\bigl(\varphi_{\boldsymbol x}(\boldsymbol x_t),\boldsymbol h_{t-1}\bigr)\) 是上文的 \(Q\) 网络，即 “encoder”；\([\boldsymbol \mu_{\boldsymbol x, t}, \boldsymbol \sigma_{\boldsymbol x, t}]=p\bigl(\varphi_{\boldsymbol z}(\boldsymbol z_t),\boldsymbol h_{t-1}\bigr)\) 是上文的 \(P\) 网络，即 “decoder”；\([\boldsymbol \mu_{\boldsymbol z, t}, \boldsymbol \sigma_{\boldsymbol z, t}]=\varphi_{\text{prior}}(\boldsymbol h_{t-1})\) 给出 \(\boldsymbol z\) 的先验分布 —— 我们不能再假设 \(\boldsymbol z\) 总是可以从标准正态分布中采样了。可以看到，和上文唯一的差异是三个网络的参数都加上了上一时刻的 hidden state。整个 VRNN 的 loss 就是每一步 VAE loss 之和： \[ \sum_{t=1}^T \biggl[\mathbb E_{\boldsymbol z \sim q\bigl(\varphi_{\boldsymbol x}(\boldsymbol x_t),\boldsymbol h_{t-1}\bigr)}\!\Bigl[\log P\Bigl(\boldsymbol x_t\Big |\boldsymbol x_{t}\sim p\bigl(\varphi_{\boldsymbol z}(\boldsymbol z_t),\boldsymbol h_{t-1}\bigr)\Bigr)\Bigr] - D_{\text{KL}}\Bigl(q\bigl(\varphi_{\boldsymbol x}(\boldsymbol x_t),\boldsymbol h_{t-1}\bigr)\|\varphi_{\text{prior}}(\boldsymbol h_{t-1})\Bigr)\biggr]. \]