VAE（Variational Autoencoder，变分自编码器）解决的是一个生成建模问题：怎样把容易采样的简单分布，变成与真实数据分布相似的复杂样本？

它的关键不只是“encoder + decoder”结构，而是三步数学转换：用近似后验绕过不可解积分，用 ELBO 把最大似然改写为可优化目标，再用重参数化让随机采样能够参与反向传播。

本文沿着这条推导主线展开。读完后应该能够回答：

  • 为什么通常选择 \(p(z)=\mathcal{N}(0,I)\)；
  • 为什么直接计算 \(p_\theta(x)\) 很困难；
  • \(q_\phi(z|x)\)、KL 散度和 ELBO 分别在做什么；
  • reconstruction loss 与 KL loss 为什么同时出现；
  • 重参数化技巧如何让 encoder 获得梯度；
  • 单层 VAE 如何扩展成层次 VAE。

从潜变量生成模型开始

现在我们已经有数据集\(\mathcal{D}\)。假设在高维空间有可以很容易从其概率密度函数采样的潜变量(latent variable（潜变量）) \(z \in \mathcal{Z}\)，一个通过\(\theta\)进行参数化的函数\(f(z;\theta)\)进行如下映射

\[ f: \mathcal{Z}\times \Theta \rightarrow \mathcal{D} \]

函数是确定的，\(z\)是随机变量，所以函数\(f\)的结果也是随机变量。

我们期望的是优化参数\(\theta\)使得\(f(z;\theta)\)尽量和现有的的数据集相似。

\[ P(X) = \int P(X|z;\theta)P(z)dz \]

为了解决上面的问题，需要解决两个问题：

  • 如何定义\(z\);
  • 如何处理积分。

对于一个具体问题，我们很难去定一个latent variable让其每个维度代表数据本质特征的一个方面。实际上，我们甚至都不知道任意给定的一组数据本征特征到底是需要多少维才能表示。尽管，对于直线\(y=x\)上的点我们可以知道本质上一个标量就可以表示；对于\(x^2+y^2=1\)上的点也只需要一个维度就可以表示(一个到x轴正方向的旋转角就可以表示)。对于任意的数据我们没有这样的先验。退一步说，哪怕我们知道几个维度，我们也不知道哪个维度表示什么才是最好的分解。例如MNIST手写数字，笔画的宽度是和书写速度相关的，而书写速度又会导致拐角的角度不一样。

事实上，一个更加简单的选择是直接假设latent variable \(z\)是从标准高斯分布采样得来，也就是说\(z \sim \mathcal{N}(0, I)\)。维度也是我们自己根据先验定一个大小。

为什么可以这么选择呢？这么定义的\(z\)能够捕捉到数据最本质的信息吗？

其实不能。

但是，除了latent variable之外，我们还有基于它的函数映射。

如上图，左边的点\(z\)是从高斯分布采样得来，而右边显示的是对应\(g(z) = z/10+a/||z||\)的图像。可以看到，高斯分布的采样点经过变换之后得到了接近圆环的分布。

这就说明，虽然假设的latent variable分布简单，但是只要有复杂的变换函数，一样可以得到任意复杂的分布。而事实上，神经网络可以拟合任意复杂的函数，那也就是说，一个简单分布加上一个神经网络可以拟合任意复杂的分布。

所以，即便数据的本征特征很复杂，我们也总是可以有若干层网络来使得一个简单的高斯分布变换成数据的本征分布。这就使得我们在网络的capacity足够的情况下可以使用简单的分布来当成数据的真实分布。

有了\(z\)分布，我们继续考虑\(P(X) = \int P(X|z;\theta)P(z)dz\)里面的积分。我们可以采样许多 \(z\)得到\(\{z_1, z_2, \cdots, z_n\}\)，然后计算\(P(X)\approx \frac{1}{n}\sum_iP(X|z_i)\)。(这里我们假设概率密度函数是已经知道的，实际上是一个训练出来的网络)。虽然理论上可行，但是我们会发现采样的大小\(n\)需要特别大才能实现。

为什么呢？

考虑我们经常使用的MNIST数据集合，里面是0到9的数字图像。

上面图像里面a是训练的数据集，b和c是取不同的\(z\)生成的。b左下角缺角了，c看起来虽然相似，但是图像却整体往右移动了。而我们一般度量两幅图像是否相似就是考虑它们之间的欧式距离，那么b和c都不算好，也就是\(P(X|z_i)\)的概率比较小。只有我们真的采样到非常接近能够生成图像a的\(z'\)，才有可能得到比较大的\(P(X|z')\)。由于\(z\)是一个标准的高斯分布，\(P(X|z)\)也是一个高斯分布，\(P(X|z) = \mathcal{N}(X|\mu, \sigma^2*I)\)。

在一个连续高斯分布空间要找到非常接近真实值的值，我们有两个方法：

  1. 采样非常非常多的点，点越多瞎猫能碰到死耗子的概率就越大；
  2. 让高斯分布的方差\(\sigma\)特别小，整个分布变得特别陡峭，这样采样的值都接近真实的值。

对于方法1，我们需要注意，真实的值不一定是均值，或者说大概率不是。它可能落在分布里面任何一个可能的地方。而对于大多数的z，\(P(X|z)\)的值都接近0，只有接近真实的值的\(z\)才会使得这个概率很大。假设P是通过一个网络学习到的(事实上也是，就是VAE的decoder)，哪怕这个网络已经训练得很好了，也需要找到真正的z让\(P(X|z)\)很大，进而让\(P(X)\)很大才能判别这个网络训练好了。

对于方法2，听起来是一个比较好的方法，但是只适用于理论而不适用于网络。因为这种情况下梯度太大，整个网络训练的时候不断震荡，非常难训练。

其实，除了上面两个方法，我们还有一个方法。方法2不可能是因为它没法让网络很好训练，方法1现在不可行是因为采样的数量太大了。为什么大？因为即便我们的概率密度函数就是真实的概率密度函数\(P\)，我们也只知道\(z\)取所有值的情况下的均值和方差，在一个连续空间里面要找到一点，就需要无穷多采样才可能。

如果我们知道给定的训练数据\(x_i\)对应的\(z\)在一个高斯分布里面的大致取值范围，我们就不用采样非常多的 \(z\)来试图得到跟\(x_i\)真实对应的\(z_i\)，而只需要在得到的\(z\)的大致范围内采样更少的数据就可以得到真实的\(z_i\)(或者非常接近的值)。也就是说，我们需要\(P(z|X)\)，给定训练数据的情况下\(z\)的条件分布。

通过贝叶斯公式

\[ P(z|X) = \frac{P(X|z)P(z)}{P(X)} \]

我们很容易看出，\(P(X)\)不可解的情况下\(P(z|X)\)也不可解。我们现在就是在试图通过一些方法让\(P(X)\)可解。这形成了循环依赖：为了得到后验需要边缘似然，但边缘似然正是我们尚未解决的积分。

既然无法直接得到，我们就想办法绕过。我们需要一个新的函数\(Q(z|X)\)，它的概率分布能够很好地拟合\(P(z|X)\)。

怎么度量两个概率分布相似？

KL散度。

如果\(q_\phi(z|x_i)\)是一个真实的条件概率\(p_\theta(z|x_i)\)的估计，那么我们有

\[ \begin{aligned} KL(q_\phi(z|x_i)|| p_\theta(z|x_i)) &= \int q_\phi(z|x_i) \log \frac{q_\phi(z|x_i)}{p_\theta(z|x_i)}dz \\ &= \int q_\phi(z|x_i)[\log q_\phi(z|x_i) - \log p_\theta(z|x_i)]dz \end{aligned} \]

使用贝叶斯公式我们可以有

\[ \begin{aligned} \log p_\theta(z|x_i) &= \log \frac{p_\theta(x_i|z)p(z)}{p_\theta(x_i)} \\ &= \log p_\theta(x_i|z) + \log p_\theta(z) - \log p_\theta(x_i) \end{aligned} \]

这样，KL散度就变成

\[ \begin{aligned} KL(q_\phi(z|x_i)|| p_\theta(z|x_i)) &= \int q_\phi(z|x_i)[\log q_\phi(z|x_i) - \log p_\theta(z|x_i)]dz \\ &= \int q_\phi(z|x_i)[\log q_\phi(z|x_i) - \log p_\theta(x_i|z) - \log p_\theta(z) + \log p_\theta(x_i)] dz \\ &= \log p_\theta(x_i) + \int q_\phi(z|x_i)[\log q_\phi(z|x_i) - \log p_\theta(x_i|z) - \log p_\theta(z)] dz \\ &= \log p_\theta(x_i) + \int q_\phi(z|x_i) \log \frac{q_\phi(z|x_i)}{p_\theta(x_i|z)p_\theta(z)}dz \\ &= \log p_\theta(x_i) + \int q_\phi(z|x_i) \log \frac{q_\phi(z|x_i)}{p_\theta(z)}dz - \int q_\phi(z|x_i)\log p_\theta(x_i|z) dz \\ &= \log p_\theta(x_i) + KL(q_\phi(z|x_i)||p_\theta(z)) - E_{z\sim q_{\phi}}[\log p_\theta(x_i|z)] \end{aligned} \]

整理可以得到

\[ \log p_\theta(x_i) -KL(q_\phi(z|x_i)|| p_\theta(z|x_i)) = E_{z\sim q_{\phi}}[\log p_\theta(x_i|z)] - KL(q_\phi(z|x_i)||p_\theta(z)) \]

公式右边的第一项是拟合函数\(q_\phi\)采样的z生成的样本的期望。我们对后验概率拟合越好，重建出来的图像和输入的样本越相似，概率也更大。第二项是测量学到的后验分布和我们先验分布的相似性。优化它能够让我们的拟合函数去学习到一个分布而不是一个其它的函数。因为在第一项可以优化的情况下，拟合函数可以是在真实的点上为1，其它地方为0的阶跃函数，而这显然是我们不想要的。

\(P(X) = \prod_i p_\theta(x_i)\)，也就是说\(\log P(X) = \sum_i \log p_\theta(x_i)\)。要最大化\(P(X)\)，就要最大化\(\log p_\theta(x_i)\)。上面公式中左边的第二项\(KL(q_\phi(z|x_i)|| p_\theta(z|x_i)) \)代表的是我们的拟合函数的分布和真实的分布的相似度，我们当然期望这个相似度越大越好，而相似度越大，它们之间的KL散度就会越小，完全一致的时候KL散度就等于0。而拟合的函数越接近真实的条件分布，从里面采样的z会让\(P(X)\)越大。也就是说，上面公式的左边就是我们想最大化的目标函数，它会让我们同时优化\(\theta, \phi\)。其中\(\theta\)是拟合真实的条件概率\(p(X|z)\)的参数，\(\phi\)是拟合真实的条件概率\(P(z|X)\)的参数。

我们记

\[ \zeta(\phi, \theta) = E_{z\sim q_{\phi}}[\log p_\theta(x_i|z)] - KL(q_\phi(z|x_i)||p_\theta(z)) \]

它有一个名称叫做Evidence Lower Bound (ELBO)。

事实上

\[ \begin{aligned} \zeta(\phi, \theta) &= -\int q_\phi(z|x_i) \log \frac{q_\phi(z|x_i)}{p_\theta(x_i|z)p_\theta(z)}dz \\ &= - \int q_\phi(z|x_i) \log \frac{q_\phi(z|x_i)}{p_\theta(x_i, z)}dz \\ &= \int q_\phi(z|x_i) \log \frac{p_\theta(x_i, z)}{q_\phi(z|x_i)}dz \\ &= E_{z\sim q_{\phi}}[\log \frac{p_\theta(x_i, z)}{q_\phi(z|x_i)} ] \end{aligned} \]

这是我们经常见到的ELBO的另外一种写法，实际上，这两种写法是等价的。

为什么叫做ELBO？

\[ \begin{aligned} \log p_\theta(x_i) &= E_{q_\phi(z|x_i)}[\log \frac{p_\theta(x_i, z)}{q_\phi(z|x_i)} ] + KL(q_\phi(z|x_i)|| p_\theta(z|x_i)) \\ &\ge E_{q_\phi(z|x_i)}[\log \frac{p_\theta(x_i, z)}{q_\phi(z|x_i)} ] \end{aligned} \]

所以ELBO是\(\log p_\theta(x_i)\)的下界。

当然了，ELBO还有另外两种推导方法。

\[ \begin{aligned} \log p(x) &= \log \int p(x, z)dz \\ &= \log \int \frac{p(x, z)q_\phi(z|x)}{q_\phi(z|x)}dz \\ &= \log E_{q_\phi(z|x)}[\frac{p(x, z)}{q_{\phi}(z|x)}] \\ &\ge E_{q_\phi(z|x)} [\log \frac{p(x, z)}{q_{\phi}(z|x)}] \end{aligned} \]

或者

\[ \begin{aligned} \log p(x) &= \log p(x) \int q_{\phi}(z|x)dz \\ &= \int q_{\phi}(z|x) \log p(x) dz \\ &= E_{ q_{\phi}(z|x)}[\log p(x)] \\ &= E_{ q_{\phi}(z|x)}[\log \frac{p(x,z)}{p(z|x)}] \\ &= E_{ q_{\phi}(z|x)}[\log \frac{p(x,z)q_{\phi}(z|x)}{p(z|x)q_{\phi}(z|x)}] \\ &= E_{ q_{\phi}(z|x)}[\log \frac{p(x,z)}{q_{\phi}(z|x)}] + E_{ q_{\phi}(z|x)}[\log \frac{q_{\phi}(z|x)}{p(z|x)} \\ &= E_{ q_{\phi}(z|x)}[\log \frac{p(x,z)}{q_{\phi}(z|x)}] + KL(q_{\phi}(z|x)|| p(z|x)) \\ &\ge E_{ q_{\phi}(z|x)}[\log \frac{p(x,z)}{q_{\phi}(z|x)}] \end{aligned} \]

对于整个数据集\(\mathcal{D}\)，我们的优化目标函数为：

\[ \begin{aligned} &E_{X\sim \mathcal{D}}[\log P(X) - KL(q(z|X)|| p(z|X))]=\\ &E_{X\sim \mathcal{D}}[E_{z\sim q}[\log p(X|z)] - KL(q(z|X)||p(z))] \end{aligned} \]

如何选择使用\(\phi\)参数化的拟合函数\(q\)？

记住我们要这个函数的终极目标是什么，是对我们训练集的数据限制\(z\)的选择范围，让我们能够采样更少的点就可以得到数据\(X\)对应的真实\(z\)。

那既然是用来限制范围的，直接选择一个高斯函数不就好了，高斯函数简单呀，从高斯函数采样也简单。

\[ q_{\phi} = \mathcal{N}(z|\mu(X), \Sigma(X)) \]

并且，我们让\(\Sigma\)​是一个对角矩阵，也就是选择的是一个各个维度不相关的高斯。这么做可以简化计算。

同样，这里有人会问，那如果真实的后验分布不是高斯怎么办？用一个高斯去拟合非高斯肯定是不合适的。这里跟前面的选择考虑是一样的，假设后验分布是一个很复杂的分布，同时网络容量足够(也就是有足够的参数)，那么我们一定可以用若干层参数来把一个复杂的分布变换成一个高斯分布。

这种选择下，

\[ \begin{aligned} KL(q_\phi(z|x_i)||p_\theta(z)) &= KL(\mathcal{N}(\mu_0, \Sigma_0)||\mathcal{N}(\mu_1, \Sigma_1)) \\ &= \frac{1}{2}\left[\operatorname{tr}(\Sigma_1^{-1}\Sigma_0)+(\mu_1-\mu_0)^T\Sigma_1^{-1}(\mu_1-\mu_0)-k + \log(\frac{\det \Sigma_1}{\det \Sigma_0})\right] \end{aligned} \]

其中\(k\)是分布的维度。

我们选择了\(z \approx \mathcal{N}(0, I)\)，带入上面

\[ \begin{aligned} KL(q_\phi(z|x_i)||p_\theta(z)) &= KL(\mathcal{N}(\mu_0, \Sigma_0)||\mathcal{N}(0, I)) \\ &= \frac{1}{2}\left[\operatorname{tr}(\Sigma_0)+\mu_0^T\mu_0-k - \log(\det \Sigma_0)\right] \end{aligned} \]

回顾一下我们的目标函数\(E_{z\sim q_{\phi}}[\log p_\theta(x_i|z)] - KL(q_\phi(z|x_i)||p_\theta(z))\)。到这里，第二项我们已经有了显式表达式。对于第一项，我们既然已经得到了\(q_{\phi}\)，自然就可以采样得到\(z\)，然后计算\(\log p_\theta (X|z)\)。期望的计算需要一系列的值，但是既然我们已经已经大致得到了给定一个x的情况下z的分布，实际做法是在这个分布里面任意采样一个点然后计算\(\log p_\theta(x_i|z)\)，把这个值直接作为期望。

到此，我们就可以画出VAE的结构图了

这时候还剩下一个关键问题：如果把随机采样直接放在计算图中，梯度无法按照普通确定性节点的方式传回 encoder。为了解决这个问题，可以采取"reparameterization trick"，也就是把重采样操作放入到输入，输入节点采样一个\(\epsilon \sim \mathcal{N}(0, I)\)，然后计算

\[ z = \mu(X) + \sigma(X) \odot \epsilon \]

训练完成后，生成阶段不再需要 encoder：直接从先验 \(p(z)\) 采样，再交给 decoder 生成样本。

层次 VAE：把潜变量继续分层

单级VAE可以归结为下图

主要元素是可见的图像(训练数据集或者生成的图像)\(x\)以及潜空间变量\(z\)。encoder主要用来估计给定的图像\(x_i\)对应的潜变量\(z_i\)在潜空间的分布；decoder主要是给定潜空间的一个变量，生成对应的图像。

我们可以把单级VAE进行扩展，潜变量是由更高层次的潜变量生成(Hierarchical VAE, HVAE)

我们encoder里面让\(z_n\)只依赖\(z_{n-1}\)，decoder里面\(z_n\)只依赖\(z_{n+1}\)，形成一个马尔科夫链。

\[ p(x, z_{1:T}) = p(z_T)p_\theta(x|z_1)\prod_{t=2}^T p_\theta(z_{t-1}|z_t) \]
\[ q_\phi(z_{1:T}|x) = q_\phi(z_1|x)\prod_{t=2}^Tq_\phi(z_t | z_{t-1}) \]

ELBO可以扩展成

\[ \begin{aligned} \log p(x) &= \log \int p(x, z_{1:T})d z_{1:T} \\ &= \log \int \frac{p(x, z_{1:T}) q_\phi(z_{1:T}|x)}{q_\phi(z_{1:T}|x)}d z_{1:T} \\ &= \log E_{q_\phi(z_{1:T}|x)}\left[ \frac{ p(x, z_{1:T})}{q_\phi(z_{1:T}|x)} \right] \\ &\ge E_{q_\phi(z_{1:T}|x)}\left[ \log \frac{ p(x, z_{1:T})}{q_\phi(z_{1:T}|x)} \right] \end{aligned} \]

把前面的\(p, q\)带入有

\[ E_{q_\phi(z_{1:T}|x)}\left[ \log \frac{ p(x, z_{1:T})}{q_\phi(z_{1:T}|x)} \right] = E_{q_\phi(z_{1:T}|x)}\left[ \log \frac{ p(z_T)p_\theta(x|z_1)\prod_{t=2}^T p_\theta(z_{t-1}|z_t)}{q_\phi(z_1|x)\prod_{t=2}^Tq_\phi(z_t | z_{t-1})} \right] \]

从 VAE 到今天的 AIGC

VAE 的意义并没有随着 Diffusion 和 Transformer 出现而消失。许多图像和视频生成系统仍然先用 VAE 把像素压缩到潜空间，再在更低维的 latent 上训练扩散模型或流模型。Stable Diffusion 中常说的 latent diffusion，正建立在这套思想之上。

从工程角度看，VAE 同时承担三件事：

  1. 压缩：降低后续生成模型处理的空间尺寸和计算量；
  2. 结构化潜空间：让相邻 latent 倾向于解码成语义相近的样本；
  3. 生成与重建接口：encoder 把真实图像映射到 latent，decoder 把生成模型输出还原为像素。

因此，理解 ELBO、KL 正则和重参数化，不只是为了读懂一个经典模型，也是理解现代 latent image/video generation 训练目标、重建误差和 latent scaling 的基础。