引言
上一篇文章中，我们介绍了变分推断及其数学原理。我们已经知道，变分推断就是在用一个可解的带参分布 $q(z;\lambda)$ 来近似难解的后验分布 $p(z\mid x)$。通过最大化对数似然函数下界估计（即ELBO），我们可以找到最优的参数 $\lambda^$，并用 $q(z;\lambda^)$ 来作为后验分布的近似。
变分自编码器 (Variational Autoencoders, VAEs) 就是基于变分推断的最简单的生成模型。VAE使用一个可学习的编码器 (encoder) 来作为变分分布 $q(z;\lambda)$，将观测数据 $x$ 映射到隐变量 $z$ 上，然后使用一个可学习的解码器 (decoder) 从隐变量 $z$ 还原出原始数据 $x$。
基于这种思想，研究者们还提出了层级VAEs (Hierarchical VAEs, HVAEs) 捕捉不同尺度的数据特征。
一、VAE
自编码器 (Autoencoders, AEs) 是深度生成模型的最简单的版本，其中包含两个模块：确定性 (deterministic) 的编码器用于从原始数据映射到隐变量，确定性的解码器用于从隐变量还原回原始数据，并通过最小化重建误差 (reconstruction error) 来训练编码器和解码器。
虽然这种方法能够保证较好的重建效果，但我们完全不知道隐空间的形式，也没有对其加以约束。因此在生成阶段，我们很难从隐空间中进行采样并利用解码器来生成一个新样本。
为了解决这个问题，VAE为隐分布引入了概率结构，即使用概率化 (probalistic) 的编码器和解码器来分别学习隐变量的分布 $q(z\mid x)$ 以及原始数据的分布 $p(x\mid z)$。这使得VAE具有优秀的生成能力，而不仅仅只是对数据进行编码和重建。
1.1. Probalistic Encoder-Decoder
解码器
VAE假设任意一个数据 $x$ 都可以由某个隐变量 $z$ 生成，且这个隐变量服从某个简单的先验分布 (prior distribution)，比如说标准高斯分布：$z\sim p_{prior}:=\mathcal{N}(0,1)$。
VAE使用一个解码器 $p_\phi(x\mid z)$ 将 $z$ 映射到 $x$ 中。在实际应用中，解码器的分布一般是一个带参数的简单分布，最常用的是高斯分布。在生成阶段，我们首先从先验分布中采样一个隐变量 $z\sim p_{prior}$，然后使用解码器来生成新样本 $x\sim p_\phi(x\mid z)$。
从上面的过程可以看到，VAE实际上是一个隐变量生成模型，通过解码器分布和先验分布来建模原始数据的似然：
\[\begin{equation} p_\phi(x)=\int p_\phi(x\mid z)p_{prior}(z)\mathrm{d}z \end{equation}\]然而，由于这个积分是难解的，因此我们一般通过最大化ELBO来最大化似然函数 $p_\phi(x)$，正如我们在前面介绍的那样。
编码器
为了将隐分布与真实数据联系起来，我们需要通过贝叶斯定理来计算后验分布：
\[\begin{equation} p_\phi(z\mid x) = \frac{p_\phi(x\mid z)p_{prior}(z)}{p_\phi(x)} \end{equation}\]正如上面提到的，由于似然函数 $p_\phi(x)$ 是难解的，因此VAE中引入了一个编码器 $q_\theta(z\mid x)$ 作为变分分布来近似这个后验分布：
\[\begin{equation} q_\theta(z\mid x) \approx p_\phi(z\mid x) \end{equation}\]其中，$\theta$ 是编码器的参数。我们可以通过训练参数 $\theta$ 来是的变分分布尽可能接近真实后验分布。
1.2. 使用ELBO进行训练
正如前一篇文章提到的，VAE通过最大化ELBO来训练变分参数 $\theta$。具体来说，ELBO的形式如下：
\[\begin{equation} \begin{aligned} \mathcal{L}_{ELBO} &:=\int q_\theta(z\mid x)\log \frac{p_\phi(x, z)}{q_\theta(z\mid x)}\mathrm{d}z\\ &=\mathbb{E}_{z\sim q_\theta(z\mid x)}[\log p_\phi(x,z)-\log q_\theta(z\mid x)]\\ &=\mathbb{E}_{z\sim q_\theta(z\mid x)}[\log p_\phi(x\mid z)]-\mathcal{D}_{KL}(q_\theta(z\mid x)\|p_{prior}(z)) \end{aligned} \end{equation}\]通过Jensen不等式我们可以知道，ELBO是对数似然函数的下界估计，即：
\[\begin{equation} \log p_\phi(x)\ge\mathcal{L}_{ELBO} \end{equation}\]当且仅当 $q_\theta(z\mid x) = p_\phi(z\mid x)$ 时等号成立。具体的证明见上一篇文章。
从公式 $(4)$ 中我们看到，ELBO可以拆分为两项：
- 第一项是重构误差项，用于鼓励模型从隐变量 $z$ 中准确重构得到原始数据 $x$。后面我们将看到，这一项就是自编码器的损失函数。然而，只优化重构误差会导致模型只记住了训练数据，无法提升生成能力。
- 第二项是隐分布KL散度项，用于鼓励编码器的分布 $q_\theta(z\mid x)$ 尽可能接近隐变量的先验分布 $p_{prior}(z)$。这个正则化项将隐空间变得更加平滑和连续，使得在生成时我们可以从 $p_{prior}(z)$ 中采样到有效的隐变量 $z$ 来生成逼真的样本。
1.3. Gaussian VAE
在这一节中我们介绍VAE在实际应用中的做法。最简单和最标准的实践是称为高斯VAE的方法，它在编码器和解码器中都使用高斯分布。
编码器
编码器 $q_\theta(z\mid x)$ 建模了一个参数可学习的高斯分布：
\[\begin{equation} q_\theta(z\mid x) :=\mathcal{N}\left( z;\mu_\theta(x),\text{diag}\left( \sigma_\theta^2(x) \right) \right) \end{equation}\]其中，$\mu_\theta(x)$ 和 $\sigma_\theta^2(x)$ 是神经网络的两个输出。
解码器
同样的，解码器 $p_\phi(x\mid z)$ 则是建模了一个均值可学习的高斯分布，这里的方差被固定为一个小常量 $\sigma^2$，用于保证稳定生成：
\[\begin{equation} p_\phi(x\mid z) :=\mathcal{N}\left( x;\mu_\phi(z),\sigma^2 \right) \end{equation}\]ELBO的形式
在这种设定下，ELBO中的重构误差项可以被改写为如下形式：
\[\begin{equation} \begin{aligned} \mathbb{E}_{q_\theta(z\mid x)}[\log p_\phi(x\mid z)] &=\mathbb{E}_{q_\theta(z\mid x)}\left[ \log\left( \frac{1}{(2\pi\sigma^2)^{D/2}}\exp\left( -\frac{\|x-\mu_\phi(z)\|^2}{2\sigma^2} \right) \right) \right]\\ &=\mathbb{E}_{q_\theta(z\mid x)}\left[ -\frac{D}{2}\log(2\pi\sigma^2)-\frac{\|x-\mu_\phi(z)\|^2}{2\sigma^2} \right]\\ &=-\frac{D}{2}\log(2\pi\sigma^2)-\frac{1}{2\sigma^2}\mathbb{E}_{q_\theta(z\mid x)}\left[\|x-\mu_\phi(z)\|^2\right]\\ &=-\frac{1}{2\sigma^2}\mathbb{E}_{q_\theta(z\mid x)}\left[\|x-\mu_\phi(z)\|^2\right]+C \end{aligned} \end{equation}\]其中，$C=-\frac{D}{2}\log(2\pi\sigma^2)$ 是独立于 $\phi$ 和 $\theta$ 的常量。
因此，Gaussian VAE的目标函数可以写为：
\[\begin{equation} \begin{aligned} \min_{\theta,\phi}\mathbb{E}_{q_\theta(z\mid x)}\left[\frac{1}{2\sigma^2}\|x-\mu_\phi(z)\|^2\right]+\mathcal{D}_{KL}(q_\theta(z\mid x)\|p_{prior}(z)) \end{aligned} \end{equation}\]1.4. VAE的缺点
即使VAE理论上很完善，但实际应用中仍有一些缺陷。一个最明显的问题就是VAE会生成比较模糊的数据。
为了理解这一现象，我们将编码器 $q_{enc}(z\mid x)$ 固定，解码器的形式为：
\[\begin{equation} p_{dec}(x\mid z) :=\mathcal{N}\left( x;\mu(z),\sigma^2 \right) \end{equation}\]此时目标函数可以简化为如下的最小二乘优化问题：
\[\begin{equation} \begin{aligned} \min_{\mu}\mathbb{E}_{p_{data}(x),q_{enc}(z\mid x)}\left[\|x-\mu(z)\|^2\right]\\ \end{aligned} \end{equation}\]这个优化问题的最优解为：
\[\begin{equation} \begin{aligned} \mu^*(z) &=\mathbb{E}_{q_{enc}(x\mid z)}[x] \end{aligned} \end{equation}\]其中，$q_{enc}(x\mid z)$ 是编码器的后验分布，其形式由贝叶斯定理给出：
\[\begin{equation} q_{enc}(x\mid z)=\frac{q_{enc}(z\mid x)p_{data}(x)}{p_{prior}(z)} \end{equation}\]最优解有另外一种等价形式：
\[\begin{equation} \begin{aligned} \mu^*(z) &=\frac{\mathbb{E}_{p_{data}(x)}[q_{enc}(z\mid x)\cdot x]}{\mathbb{E}_{p_{data}(x)}[q_{enc}(z\mid x)]} \end{aligned} \end{equation}\]现在我们考虑两个不同的输入 $x\neq x’$，且这两个输入被映射到隐空间中的某块重叠区域，即 $q_{enc}(\cdot\mid x)$ 和 $q_{enc}(\cdot\mid x’)$ 的支撑集有交集。
对于重叠区域中的某个 $z$，公式 $(12)$ 的最优解中的期望项包含来自两个不同输入的贡献。因此，VAE有可能会生成模糊的结果。其根本原因就在于最优解的形式中包含了一个条件期望项 $q_{enc}(x\mid z)$。
二、Hierarchical VAEs
2.1. HVAE的数学建模
为了增强VAE对于更加复杂和高维分布的拟合能力，研究者们提出了层级VAEs (HVAEs)，通过引入多个隐变量 $z_1,z_2,\dots,z_L$ 将VAE堆叠为一个自上而下的层级结构。
HVAE建模了如下的似然函数：
\[\begin{equation} \begin{aligned} p_{HVAE}(x) &:=\int p_{\phi}(x,z_{1:L})\mathrm{d}z_{1:L}\\ &=\int p_{\phi}(x\mid z_1)\prod_{i=2}^{L}p_{\phi}(z_{i-1}\mid z_i)p(z_L) \mathrm{d}z_{1:L} \end{aligned} \end{equation}\]在生成阶段，我们先采样 $z_L\sim p(z_L)$，并逐层解码得到新的生成数据 $x$。
HVAE同样引入了一个可学习的编码器 $q_\theta(z_{1:L}\mid x)$ 作为变分分布，最常见的形式如下：
\[\begin{equation} q_\theta(z_{1:L}\mid x)=q_\theta(z_1\mid x)\prod_{i=2}^L q_\theta(z_{i}\mid z_{i-1}) \end{equation}\]2.2. HVAE的ELBO形式
和VAE中类似，我们可以从对数似然函数出发，定义HVAE的ELBO项：
\[\begin{equation} \begin{aligned} \log p_{HVAE}(x) &=\log \int p_{\phi}(x,z_{1:L})\mathrm{d}z_{1:L}\\ &=\log \int \frac{p_{\phi}(x,z_{1:L})}{q_{\theta}(z_{1:L}\mid x)}q_{\theta}(z_{1:L}\mid x)\mathrm{d}z_{1:L}\\ &=\log \mathbb{E}_{q_{\theta}(z_{1:L}\mid x)}\left[\frac{p_{\phi}(x,z_{1:L})}{q_{\theta}(z_{1:L}\mid x)}\right]\\ &\ge \mathbb{E}_{q_{\theta}(z_{1:L}\mid x)}\left[\log \frac{p_{\phi}(x,z_{1:L})}{q_{\theta}(z_{1:L}\mid x)}\right]\\ &:=\mathcal{L}_{ELBO}(\phi) \end{aligned} \end{equation}\]将 $p_{\phi}(x,z_{1:L})$ 和 $q_{\theta}(z_{1:L}\mid x)$ 进行拆分，可以得到ELBO的另一个等价形式：
\[\begin{equation} \begin{aligned} \mathcal{L}_{ELBO}(\phi) &:=\mathbb{E}_{q_{\theta}(z_{1:L}\mid x)}\left[\log \frac{p_{\phi}(x,z_{1:L})}{q_{\theta}(z_{1:L}\mid x)}\right]\\ &=\mathbb{E}_{q_{\theta}(z_{1:L}\mid x)}\left[\log \frac{p_{\phi}(x\mid z_1)\prod_{i=2}^{L}p_{\phi}(z_{i-1}\mid z_i)p(z_L)}{q_\theta(z_1\mid x)\prod_{i=2}^L q_\theta(z_{i}\mid z_{i-1})} \right]\\ \end{aligned} \end{equation}\]2.3. HVAE v.s. VAE
正如深度神经网络依靠更深的结构、更大的参数量在众多场景下超越了传统机器学习方法一样，HVAE的有效性也展示了更深的生成模型的有效性。这种思想是现代生成式建模的基石所在，贯穿了后续的众多方法（如Score-based Models、DDPM、Normalizing Flows等）。
这种思想的核心在于：堆叠更深的层数可以让模型逐步生成数据，首先生成一些粗粒度的轮廓，再逐步添加细粒度的细节。这个过程让模型更能捕捉到高维数据中的复杂结构。
这里必须指出，使用HVAE与单纯地加深VAE中的编码器和解码器是有区别的，这是因为VAE本身有一些局限性，无法通过加深模型层数来改进。
局限1: 变分分布的形式无法改变
标准VAE中，变分分布的形式为：
\[\begin{equation} q_\theta(z\mid x) :=\mathcal{N}\left( z;\mu_\theta(x),\text{diag}\left( \sigma_\theta^2(x) \right) \right) \end{equation}\]更深的编码器只会让 $\mu_\theta(x)$ 和 $\sigma_\theta^2(x)$ 更加准确，但并不能改变变分族的形式，变分分布仍然是一个高斯分布。
当真实的后验分布具有多个峰值时，变分分布的形式就决定了无法很好的拟合，这使得ELBO项和真正的似然函数之间存在比较大的距离，也就削弱了生成能力。
局限2: 后验崩溃
当解码器变得太强时，我们往往会遭遇后验崩溃 (posterior collapse) 的现象。
为了解释这个现象，我们对VAE对训练目标进行一定的等价变形：
\[\begin{equation} \begin{aligned} \mathbb{E}_{p_{data}(x)} [\mathcal{L}_{ELBO}(x)] &=\mathbb{E}_{p_{data}(x), q_\theta(z\mid x)}[\log p_\phi(x\mid z)]-\mathbb{E}_{p_{data}(x)}[\mathcal{D}_{KL}(q_\theta(z\mid x)\|p(z))]\\ &=\mathbb{E}_{p_{data}(x), q_\theta(z\mid x)}[\log p_\phi(x\mid z)]-\mathcal{I}_q(x;z)-\mathcal{D}_{KL}(q_\theta(z\mid x)\|p(z)) \end{aligned} \end{equation}\]其中，$q_\theta(z)=\int p_{data}(x)q_\theta(z\mid x)\mathrm{d}x$ 是聚合后验概率。第二项是随机变量 $x$ 和 $z$ 的互信息 (mutual information)：
\[\begin{equation} \begin{aligned} \mathcal{I}_q(x;z) &:=\mathbb{E}_{q_\theta(x,z)}\left[ \log\frac{q_\theta(z\mid x)}{q_\theta(z)} \right]\\ &=\mathbb{E}_{p_{data}(x)}[\mathcal{D}_{KL}(q_\theta(z\mid x)\|q_\theta(z))] \end{aligned} \end{equation}\]当解码器变得太强，不需要隐变量 $z$ 就能很好地建模数据分布，即解码器学到了一个与 $z$ 无关的分布 $r(x)\approx p_{data}(x)$ 时，模型为了最大化ELBO，可以让变分分布 $q_\theta(z\mid x)=p(z)$，此时 $\mathcal{I}_q(x;z)=0$。
因此ELBO的优化本质是 重建精度 与 正则化强度 的权衡：如果解码器能忽略 $z$ 仍重建出高质量样本，则ELBO会偏好这种解，因为此时正则化项为零。这导致变分分布 $q_\theta(z\mid x)$ 被推向先验分布 $p(z)$，隐变量变得冗余。
