生成模型的核心任务是学习数据分布 p(x)p(x),然后从中采样生成新样本。对于观测变量 xx,一个自然的框架是引入隐变量 zz,通过联合分布 p(x,z)=p(xz)p(z)p(x, z) = p(x|z)p(z) 建模数据生成过程。然而,一旦模型建好,真正困难的步骤浮现出来:如何从观测 xx 推断隐变量 zz 的后验分布 p(zx)p(z|x)

这里需要澄清一个常见的混淆:先验假设与后验假设是两件不同的事。VAE 假设先验 p(z)=N(0,I)p(z) = \mathcal{N}(0, I),这只是对边际分布的约束——它规定了隐空间的整体结构,但并未约束给定 xxzz 的条件分布。真正核心的建模假设是:编码器输出的 qϕ(zx)=N(μϕ(x),σϕ2(x)I)q_\phi(z|x) = \mathcal{N}(\mu_\phi(x), \sigma_\phi^2(x) I) 近似真实后验 p(zx)p(z|x)。换言之,VAE 假设的核心不在于边际分布层面的"zz 是正态的",而在于条件分布层面的"给定 xx 后,zz 的后验是正态的"——并且用参数化的编码器来逼近这个后验。

这个区分至关重要。如果只假设先验为正态,就无法解释为什么 VAE 的采样是有效的——正是因为后验被约束为正态,才能使用重参数化技巧将随机性从参数中剥离,才能计算 KL 散度的解析形式 KL(N(μ,σ2)N(0,I))=12(μ2+σ2logσ21)\text{KL}(\mathcal{N}(\mu, \sigma^2) \| \mathcal{N}(0, I)) = \frac{1}{2}(\mu^2 + \sigma^2 - \log \sigma^2 - 1)。苏剑林在《变分自编码器(二):从贝叶斯观点出发》中反复强调:VAE 的本质是用一个参数化的正态后验去逼近不可解的真实后验,而非仅仅给先验指定一个正态形式。

这个看似简单的问题,恰恰是概率推断的根本难题。变分自编码器(Variational Autoencoder, VAE)提供了一种优雅的近似方案——用一个可参数化的分布 qϕ(zx)q_\phi(z|x) 去逼近不可解的真实后验 p(zx)p(z|x),而优化的目标,就是证据下界(Evidence Lower Bound, ELBO)。

这篇文章将从生成模型的推断困境出发,推导 ELBO 的两种等价形式,解释重参数化技巧为何不可或缺,并分析 VAE 的信息瓶颈视角与后验坍塌问题。

生成模型的推断难题

考虑一个隐变量生成模型 p(x)=p(xz)p(z)dzp(x) = \int p(x|z)p(z)dz。给定观测 xx,我们关心后验分布

p(zx)=p(xz)p(z)p(x)p(z|x) = \frac{p(x|z)p(z)}{p(x)}

问题出在分母 p(x)p(x) 上。它是所有可能 zz 值上联合概率的边际化:

p(x)=p(xz)p(z)dzp(x) = \int p(x|z)p(z)dz

对于高维连续隐变量,这个积分没有解析解。更具体地说,当 p(xz)p(x|z) 由神经网络参数化时,p(x)p(x) 的计算需要对所有可能的隐编码求积分,这在高维空间中是不可行的。

这就是变分推断的出发点:既然精确后验不可得,我们就寻找一个分布族 {qϕ(zx)}\{q_\phi(z|x)\},在其中找到与 p(zx)p(z|x) 最接近的那个。度量"接近"的工具是 KL 散度,而优化 KL 散度等价于最大化 ELBO。

ELBO 推导的两种路径

路径一:Jensen 不等式

对数边际似然 logp(x)\log p(x) 可以写成:

logp(x)=logp(x,z)dz=logq(zx)p(x,z)q(zx)dz\log p(x) = \log \int p(x, z)dz = \log \int q(z|x) \frac{p(x, z)}{q(z|x)} dz

这里我们引入了任意分布 q(zx)q(z|x)(只要 q(zx)=0p(x,z)=0q(z|x) = 0 \Rightarrow p(x, z) = 0),它不改变积分值,只是对被积函数做了一个重要性采样的变形。

由于对数函数是凹函数,Jensen 不等式给出:

logp(x)=logEq(zx)[p(x,z)q(zx)]Eq(zx)[logp(x,z)q(zx)]\log p(x) = \log \mathbb{E}_{q(z|x)}\left[\frac{p(x, z)}{q(z|x)}\right] \geq \mathbb{E}_{q(z|x)}\left[\log \frac{p(x, z)}{q(z|x)}\right]

右边就是 ELBO:

ELBO=Eq(zx)[logp(x,z)q(zx)]=Eq(zx)[logp(xz)]KL(q(zx)p(z))\text{ELBO} = \mathbb{E}_{q(z|x)}\left[\log \frac{p(x, z)}{q(z|x)}\right] = \mathbb{E}_{q(z|x)}[\log p(x|z)] - \text{KL}(q(z|x) \| p(z))

Jensen 不等式取等号的条件是 p(x,z)q(zx)\frac{p(x,z)}{q(z|x)} 为常数,即 q(zx)=p(zx)q(z|x) = p(z|x)。这意味着当变分分布完全匹配真实后验时,ELBO 退化为 logp(x)\log p(x),下界变紧。

路径二:KL 散度分解

第二条路径从 KL 散度的定义出发,更加直观地揭示了 ELBO 与推断质量的关系:

KL(q(zx)p(zx))=Eq(zx)[logq(zx)p(zx)]\text{KL}(q(z|x) \| p(z|x)) = \mathbb{E}_{q(z|x)}\left[\log \frac{q(z|x)}{p(z|x)}\right]

p(zx)=p(x,z)p(x)p(z|x) = \frac{p(x,z)}{p(x)} 代入:

KL(q(zx)p(zx))=Eq(zx)[logq(zx)p(x)p(x,z)]=logp(x)Eq(zx)[logp(x,z)q(zx)]\text{KL}(q(z|x) \| p(z|x)) = \mathbb{E}_{q(z|x)}\left[\log \frac{q(z|x) \cdot p(x)}{p(x, z)}\right] = \log p(x) - \mathbb{E}_{q(z|x)}\left[\log \frac{p(x, z)}{q(z|x)}\right]

整理得:

logp(x)=Eq(zx)[logp(x,z)q(zx)]ELBO+KL(q(zx)p(zx))0\log p(x) = \underbrace{\mathbb{E}_{q(z|x)}\left[\log \frac{p(x, z)}{q(z|x)}\right]}_{\text{ELBO}} + \underbrace{\text{KL}(q(z|x) \| p(z|x))}_{\geq 0}

这个分解的含义非常清晰:logp(x)\log p(x) 被分解为 ELBO 加上一个非负的 KL 散度项。由于 logp(x)\log p(x) 对于固定的 xx 是常数,最大化 ELBO 等价于最小化 KL(q(zx)p(zx))\text{KL}(q(z|x) \| p(z|x))。两条路径殊途同归。

ELBO 的物理意义

将 ELBO 展开为更标准的形式:

ELBO=Eq(zx)[logp(xz)]重构质量KL(q(zx)p(z))正则约束\text{ELBO} = \underbrace{\mathbb{E}_{q(z|x)}[\log p(x|z)]}_{\text{重构质量}} - \underbrace{\text{KL}(q(z|x) \| p(z))}_{\text{正则约束}}

第一项是重构对数似然的期望:从 q(zx)q(z|x) 采样隐编码 zz,解码器根据 zz 重建 xx 的质量越高,这项越大。第二项是变分后验与先验的 KL 散度:它惩罚 q(zx)q(z|x) 偏离先验 p(z)p(z) 的程度,鼓励隐编码服从先验分布。

这两项之间存在张力。如果 q(zx)q(z|x) 完全忽略 xx 而退化为先验 p(z)p(z),正则项为零,但重构项会很大(因为隐编码不包含关于 xx 的信息)。反之,如果 q(zx)q(z|x) 为每个 xx 编码极多信息,重构项变小,但正则项变大。VAE 训练的本质就是在重构质量与正则约束之间寻找最优平衡点。

ELBO 的重构-正则权衡
Loading visualization...
横轴为 KL 散度(正则项),纵轴为重构误差。不同 β 值下的 Pareto 前沿展示了信息瓶颈的权衡。拖动滑块调整 β,观察 latent space 的形态变化。

重参数化技巧

ELBO 的梯度计算涉及 Eqϕ(zx)[logp(xz)]\mathbb{E}_{q_\phi(z|x)}[\log p(x|z)],即从 qϕ(zx)q_\phi(z|x) 采样后对解码器输出求值。问题在于采样操作不可微zqϕ(zx)z \sim q_\phi(z|x) 是一个随机过程,梯度无法穿过采样节点回传到 ϕ\phi

这是 VAE 训练的核心计算障碍。重参数化技巧(Reparameterization Trick)通过将随机性从参数中"剥离"出来解决了这个问题。

假设变分后验为对角高斯:qϕ(zx)=N(z;μϕ(x),diag(σϕ2(x)))q_\phi(z|x) = \mathcal{N}(z; \mu_\phi(x), \text{diag}(\sigma_\phi^2(x)))。直接采样 zqϕ(zx)z \sim q_\phi(z|x) 的梯度 ϕEqϕ[f(z)]\nabla_\phi \mathbb{E}_{q_\phi}[f(z)] 无法用标准反向传播计算。重参数化将采样过程改写为:

z=μϕ(x)+σϕ(x)ϵ,ϵN(0,I)z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)

现在随机性完全来自 ϵ\epsilon,它不依赖于 ϕ\phiμ\muσ\sigma 是确定性函数(由编码器网络输出),梯度可以沿着加法和乘法操作顺利回传:

ϕEqϕ[f(z)]=EϵN(0,I)[ϕf(μϕ(x)+σϕ(x)ϵ)]\nabla_\phi \mathbb{E}_{q_\phi}[f(z)] = \mathbb{E}_{\epsilon \sim \mathcal{N}(0,I)}\left[\nabla_\phi f(\mu_\phi(x) + \sigma_\phi(x) \odot \epsilon)\right]

这正是 Kingma 和 Welling (2014) 的核心贡献。它将一个不可微的蒙特卡罗估计转化为一个标准的可微计算图,使得 ELBO 可以用标准的随机梯度下降优化。

一个采样点就够了

实践中,VAE 训练时对每个 xx 只从 qϕ(zx)q_\phi(z|x) 采样一个 zz——为什么一次采样就足够了?初看之下,蒙特卡罗估计的方差应随样本数减少而增大,只用一个样本似乎过于粗糙。

苏剑林在《变分自编码器(三):这样做为什么能成?》中给出了关键论证。实际使用时,解码器 p(xz)p(x|z) 的输出取均值,是一个近乎确定性的映射——给定 zz,输出几乎唯一确定。这意味着 xxzz 几乎一一对应,真实后验 p(zx)p(z|x) 的方差很小。低方差的后验意味着一次采样就足够代表分布:当 qϕ(zx)q_\phi(z|x) 本身就很集中时,从它采一个样本与采多个样本再取平均几乎没有区别。

一个更直观的类比是"资源争夺"。如果从先验 p(z)p(z) 采样,所有数据点 xx 共享同一批 zz,每个 xx 都在争夺有限的 zz 资源,导致训练失败——这正是普通自编码器的问题。而 VAE 从后验 qϕ(zx)q_\phi(z|x) 采样,每个 xx 有自己专属的 zz 区域,不存在竞争。这种"按需分配"的机制使得每个数据点只需一个专属采样点就能获得有效的梯度信号。

重参数化技巧的几何
Loading visualization...
左侧:直接采样 z ~ q(z|x),梯度路径断开(红色叉号);右侧:重参数化 z = μ + σ·ε,梯度可穿过(绿色箭头)。点击任意位置观察采样过程。

β\beta-VAE 与信息瓶颈

标准 VAE 的训练目标是最大化 ELBO,即最小化

L(θ,ϕ;x)=Eqϕ(zx)[logpθ(xz)]+KL(qϕ(zx)p(z))\mathcal{L}(\theta, \phi; x) = -\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] + \text{KL}(q_\phi(z|x) \| p(z))

Higgins 等人 (2017) 提出的 β\beta-VAE 在 KL 项前引入系数 β\beta

Lβ(θ,ϕ;x)=Eqϕ(zx)[logpθ(xz)]+βKL(qϕ(zx)p(z))\mathcal{L}_\beta(\theta, \phi; x) = -\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] + \beta \cdot \text{KL}(q_\phi(z|x) \| p(z))

β>1\beta > 1 时,模型受到更强的正则约束,q(zx)q(z|x) 被推向先验 p(z)p(z),隐编码被迫更独立、更解耦(disentangled)。当 β<1\beta < 1 时,重构质量更高但隐空间可能缺乏结构。

这个框架与信息瓶颈理论(Information Bottleneck)有深层联系。信息瓶颈的目标是找到表示 ZZ,在压缩关于输入 XX 的信息的同时最大化关于目标 YY 的信息:

minq(zx)I(Z;X)s.t.I(Z;Y)I0\min_{q(z|x)} I(Z; X) \quad \text{s.t.} \quad I(Z; Y) \geq I_0

其中 I(Z;X)KL(q(zx)p(z))I(Z; X) \geq \text{KL}(q(z|x) \| p(z)) 给出了 ELBO 中 KL 项的信息论下界。β\beta-VAE 的 Lagrangian 形式恰好是信息瓶颈问题的一个松弛:β\beta 扮演了 Lagrange 乘子的角色,控制压缩与保留信息的权衡。

从 rate-distortion 理论的角度看,β\beta-VAE 的 Pareto 前沿定义了率失真曲线:横轴(rate)是隐编码携带的信息量,纵轴(distortion)是重构误差。不同 β\beta 值对应前沿上的不同操作点。

动态噪声调节:KL 项的微观机制

前面的讨论将 KL 散度视为一个整体的正则项,但若深入到每个潜变量维度的微观层面,会发现 KL 项是模型自身通过方差 σ2(x)\sigma^2(x) 动态调节的内在平衡机制,而非外部强加的约束。这一视角是理解 VAE 为何不会退化为普通自编码器的关键。

VAE 编码器输出的是每个维度的均值 μ(x)\mu(x) 和方差 σ2(x)\sigma^2(x),其中方差是模型动态学习的,而非固定常数。这正是 VAE 与普通自编码器的核心区别:普通自编码器的编码是确定性的,没有噪声;VAE 的编码则叠加了由 σ2(x)\sigma^2(x) 控制的高斯噪声,而噪声的大小由模型自主决定。

考虑 KL 散度的解析形式。对于第 jj 个潜变量维度:

KLj=12(μj2+σj2logσj21)\text{KL}_j = \frac{1}{2}\left(\mu_j^2 + \sigma_j^2 - \log \sigma_j^2 - 1\right)

当某个维度对重构不重要时,模型会增大该维度的 σj2\sigma_j^2。此时后验 q(zjx)=N(μj,σj2)q(z_j|x) = \mathcal{N}(\mu_j, \sigma_j^2) 在该维度上变得宽泛,逐渐退化为先验 p(zj)=N(0,1)p(z_j) = \mathcal{N}(0, 1)——大量噪声淹没均值信号,KL 项趋近零,该维度"自由通过",不被约束。极端情况下 σj21\sigma_j^2 \to 1μj0\mu_j \to 0,该维度完全放弃编码信息。

当某个维度对重构重要时,模型会减小该维度的 σj2\sigma_j^2。后验变得尖锐地集中在均值 μj\mu_j 附近,编码信息更精确,重构质量提升,但 KL 项增大——因为一个远离先验的窄分布必然偏离 N(0,1)\mathcal{N}(0, 1)

这形成了一种内在的对抗平衡

  • 解码器训练不足时,重构误差主导损失函数,梯度信号驱使模型降低噪声(减小 σj2\sigma_j^2),使后验更尖锐、编码更精确,从而更容易重构。解码器逐步提升。
  • 解码器训练充分后,KL 项的相对权重上升,梯度信号驱使模型增大噪声(增大 σj2\sigma_j^2),逼迫解码器从带噪声的编码中提取信息,而非过度依赖精确编码。这正是生成能力提升的来源——解码器学会了对噪声鲁棒,从先验采样时也能生成有意义的样本。

换言之,σ2(x)\sigma^2(x) 充当了一个自适应的"噪声阀门":模型根据重构需求与正则压力的实时对比,自主决定每个维度注入多少噪声。苏剑林在《变分自编码器(一):原来是这么一回事》中用一张精妙的表格展示了这个动态平衡——VAE 就是在常规自编码器的基础上,对编码器结果加上高斯噪声,使得解码器对噪声有鲁棒性;KL 项是对编码器的正则。这张表格的核心洞见是:KL 项是噪声水平的校准器,而非惩罚。σ2\sigma^2 的梯度 KLjσj2=12(11σj2)\frac{\partial \text{KL}_j}{\partial \sigma_j^2} = \frac{1}{2}\left(1 - \frac{1}{\sigma_j^2}\right) 清楚地表明,当 σj2>1\sigma_j^2 > 1 时梯度为正(抑制进一步增大),当 σj2<1\sigma_j^2 < 1 时梯度为负(抑制进一步减小),σj2=1\sigma_j^2 = 1 是先验对应的自然平衡点。

这个视角解释了为什么 VAE 不会退化为普通自编码器:如果 σ20\sigma^2 \to 0,KL 项趋于无穷大,梯度会强力将 σ2\sigma^2 拉回;如果 σ2\sigma^2 \to \infty,重构误差会失控。模型被锁定在这两个极端之间的动态平衡中,而平衡点的位置由重构质量与正则约束的相对强度决定——这正是 ELBO 中两项张力的微观体现。

后验坍塌与模糊生成

VAE 在实践中面临两个显著问题。

后验坍塌(Posterior Collapse):当解码器 p(xz)p(x|z) 足够强大时(例如自回归解码器),它可能忽略 zz 而直接建模 p(x)p(x),使得 q(zx)p(z)q(z|x) \approx p(z)。此时隐编码不携带关于 xx 的信息,VAE 退化为一个没有意义的自编码器。β\beta 越大,后验坍塌的风险越高。

缓解策略包括:

  • KL 退火(KL Annealing):训练初期 β\beta 从 0 线性增长到目标值
  • Free bits:设置 KL 项的最低阈值,防止某些维度完全坍塌
  • 弱化解码器:使用不那么强大的解码器结构

模糊生成:VAE 生成的图像通常比 GAN 和扩散模型更模糊。原因在于 VAE 的损失函数是最小化负对数似然的下界,而 Eq(zx)[logp(xz)]\mathbb{E}_{q(z|x)}[\log p(x|z)]p(xz)p(x|z) 为高斯分布时等价于最小化 MSE。MSE 损失对"平均"所有可能性的偏好,导致生成结果趋向于训练样本的均值,表现为模糊。

更深层的原因是,VAE 使用单步从 zzxx 的映射,缺乏迭代精炼的能力。扩散模型通过多步去噪逐步细化生成结果,每一步只需做小幅修正,天然更适合产生清晰的细节。这是扩散模型在图像生成领域"胜出"的关键原因之一。

应用:视频与图像生成中的 VAE

尽管 VAE 作为独立生成模型已被扩散模型超越,它在现代生成系统中仍然扮演着关键角色——作为潜空间编码器。

Stable Diffusion 使用一个预训练的 VAE 将 512×512×3512 \times 512 \times 3 的图像压缩到 64×64×464 \times 64 \times 4 的潜空间,扩散过程在潜空间中进行。VAE 编码器将高维像素空间映射到低维潜空间,大幅降低了扩散模型的计算开销。解码器则将潜表示恢复为像素空间图像。这种设计的前提是 VAE 的潜空间具有足够的重建质量,同时压缩比足够大。

Wan2.2 的 3D VAE(参考 Wan2.2 视频世界模型)进一步将 VAE 从 2D 空间扩展到时空联合空间。3D VAE 在时间维度上也进行压缩,将视频帧序列编码为更紧凑的潜表示。这使得视频扩散模型能在合理的计算预算下处理长视频序列。3D VAE 的训练需要同时优化空间重建质量和时间一致性,其 KL 正则项的权重 β\beta 需要更精细的调整,以避免空间细节丢失或时间闪烁。

这些应用表明,VAE 的价值不在于"生成最好的图像",而在于"学习最好的潜空间"。一个良好的潜空间应当紧凑、连续、有结构,使得下游生成模型(如扩散模型)能够高效运作。这正是 ELBO 中重构-正则权衡的工程体现。

参考文献

  1. 苏剑林. 变分自编码器(一):原来是这么一回事. https://kexue.fm/archives/5253
  2. 苏剑林. 变分自编码器(二):从贝叶斯观点出发. https://kexue.fm/archives/5343
  3. 苏剑林. 变分自编码器(三):这样做为什么能成?. https://kexue.fm/archives/5383
  4. Kingma, D. P., & Welling, M. (2014). Auto-Encoding Variational Bayes. ICLR.
  5. Higgins, I., et al. (2017). β\beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR.
  6. Alemi, A. A., et al. (2017). Deep Variational Information Bottleneck. ICLR.
  7. Bowman, S. R., et al. (2016). Generating Sentences from a Continuous Space. CoNLL.