VAE 笔记

Posted on 2022-07-05 Views:

最近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)$ 。更为广义地来说，不一定要是一个确定性的过程，所以我们要求的实际上是 $P(\boldsymbol x|\boldsymbol z)$ 。如果我们假设 $P(\boldsymbol x|\boldsymbol z)$ 有固定的形式（比如说正态），那么也可以说我们寻找的是一个将 $\boldsymbol z$ 映射到 $P(\boldsymbol x|\boldsymbol z)$ 分布的参数的函数 $p(\boldsymbol z)$ 。在机器学习的背景下，的形式是固定的，我们需要寻找或者学习的其实是一组参数 $\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提供的解决方案之前，我在这一节先简要介绍一下相关的一些背景知识。

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

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

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

$\int -P(X)\log P(X)\,dX.$

这一个式子就是人们常说的信息熵——这个式子长成这样是有理由的！

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

假设一个人相信理论，那么事件对那个人来说的信息量为 $-\log Q(X)$ 。客观上来讲，事件是按照刻画的概率分布发生的，所以这个事件对这个人的平均信息量就是

$\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).$

上面这段话蕴含一个假设：所有基于

计算的信息熵一定不小于一个事件实际的信息熵——这样作差当距离才有意义。真的是这样的吗？换而言之，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距离中的和的地位就不一样。

回到VAE

为了更加精确地近似采样 $\boldsymbol z$ ，我们考虑从一个能给出“有代表性的 $\boldsymbol z$ ”的分布中采样。所谓“有代表性”，可以理解为应该近似于 $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)$ 足够小，即

足够逼近 $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。

现在还剩下一个问题，就是我们不知道这样的具体是什么，但是我们知道应该是和 $\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$ 。

接下来还有一个技术上的问题：我们在计算目标函数的过程中要从当中采样。但是采样是一个离散的过程，会打断梯度的反向传播，导致的参数不能被正确更新。解决这个的方法被称为reparameterization trick。我们可以从标准正态分布中采样，再把采样线性缩放到的分布。这样 $\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的推导，网络将一个样本映射到一个隐变量的分布，可以近似看作是一种“编码”的过程；网络将一个隐变量隐射到一个样本的分布，可以近似看作是一种“解码”的过程。两个网络是同步训练的，这点和传统意义上的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).$

除了最naive的RNN cell以外也可以是如LSTM、GRU之类的变体。如果是LSTM， $\boldsymbol h$ 实际上是hidden state和cell state的总称。

所谓的嵌入VAE，是指每一步除了以外，还多出了三个网络： $[\boldsymbol \mu_{\boldsymbol z, t}, \boldsymbol \sigma_{\boldsymbol z, t}]=q\bigl(\varphi_{\boldsymbol x}(\boldsymbol x_t),\boldsymbol h_{t-1}\bigr)$ 是上文的网络，即“encoder”； $[\boldsymbol \mu_{\boldsymbol x, t}, \boldsymbol \sigma_{\boldsymbol x, t}]=p\bigl(\varphi_{\boldsymbol z}(\boldsymbol z_t),\boldsymbol h_{t-1}\bigr)$ 是上文的网络，即“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].$