变分自编码器 (Variational Autoencoder, VAE)
2013 年,Kingma 与 Welling 在论文 "Auto-Encoding Variational Bayes" 中提出了变分自编码器 (VAE)。VAE 将深度学习与贝叶斯推断相结合,在自编码器框架上引入概率建模,使隐空间具备连续、可采样的特性,从而成为第一个真正意义上兼具推断能力与生成能力的深度生成模型。
1. 背景与动机
1.1 生成模型的目标
生成模型的核心目标是学习数据的真实分布 \(p(x)\),并从中采样生成新的、合理的数据。给定一组训练数据 \(\{x^{(1)}, x^{(2)}, \dots, x^{(N)}\}\),我们希望找到一个参数化模型 \(p_\theta(x)\),使其尽可能接近真实数据分布。
1.2 Autoencoder 回顾
传统自编码器 (AE) 由 Encoder 和 Decoder 两部分组成:
- Encoder \(f_\phi\):将输入 \(x\) 映射到低维隐表示 \(z = f_\phi(x)\)
- Decoder \(g_\theta\):将隐表示 \(z\) 重建为 \(\hat{x} = g_\theta(z)\)
训练目标是最小化重建误差,如 \(\|x - \hat{x}\|^2\)。
传统 AE 的局限
传统 AE 的隐空间是不规则的:不同训练样本被映射到隐空间中零散的点上,点与点之间的区域没有语义含义。因此,我们无法在隐空间中随机采样一个点并期望 Decoder 输出有意义的数据——AE 不是一个合格的生成模型。
1.3 VAE 的核心创新
VAE 的关键思想是:让 Encoder 输出的不是一个确定的点,而是一个概率分布的参数。具体来说,Encoder 对每个输入 \(x\) 输出一个高斯分布 \(q_\phi(z|x) = \mathcal{N}(\mu, \sigma^2 I)\),然后从这个分布中采样得到 \(z\)。同时,通过 KL 散度正则项,强制这些分布接近标准正态分布 \(\mathcal{N}(0, I)\)。
这样做的结果是:隐空间变得连续、平滑、可采样。生成新数据时,只需从 \(\mathcal{N}(0, I)\) 中采样 \(z\),再通过 Decoder 即可。
2. VAE 架构详解
2.1 整体架构
Encoder Decoder
┌───────────────┐ ┌───────────────┐
│ │──→ μ ──┐ │ │
Input x ──────→ │ Neural Net │ ├─→ z ────→ │ Neural Net │ ──→ x̂
│ │──→ σ ──┘ │ │
└───────────────┘ ↑ └───────────────┘
│
ε ~ N(0, I)
(Reparameterization)
z = μ + σ ⊙ ε
整个流程可以分为三个阶段:
- 编码 (Encoding):Encoder 网络接收输入 \(x\),输出隐变量分布的参数 \(\mu\) 和 \(\sigma\)(或 \(\log \sigma^2\))
- 重参数化采样 (Reparameterization):从 \(\mathcal{N}(\mu, \sigma^2 I)\) 中采样得到 \(z\)
- 解码 (Decoding):Decoder 网络接收 \(z\),输出重建数据 \(\hat{x}\)
2.2 Encoder:从确定性到概率性
传统 AE 的 Encoder 输出一个确定性向量 \(z\)。VAE 的 Encoder 则输出两个向量:
- \(\mu = f_\mu(x)\):隐变量分布的均值
- \(\log \sigma^2 = f_\sigma(x)\):隐变量分布的对数方差(实践中输出 \(\log \sigma^2\) 而非 \(\sigma\),因为前者取值范围为整个实数轴,数值更稳定)
这定义了一个条件分布:
2.3 Reparameterization Trick
问题:我们需要从 \(q_\phi(z|x) = \mathcal{N}(\mu, \sigma^2 I)\) 中采样 \(z\),但"采样"这一操作是随机的,不可微分,无法通过反向传播计算梯度。
解决方案:将随机性从计算图中剥离出来。具体做法是:
- 先从标准正态分布中采样噪声:\(\epsilon \sim \mathcal{N}(0, I)\)
- 然后通过确定性变换得到 \(z\):
其中 \(\odot\) 表示逐元素乘法。
重参数化的本质
这个技巧的关键在于:\(z\) 现在是 \(\mu\) 和 \(\sigma\) 的确定性函数(给定 \(\epsilon\)),因此梯度可以经由 \(z\) 流回 \(\mu\) 和 \(\sigma\),再流回 Encoder 的参数 \(\phi\)。随机性被"外包"给了 \(\epsilon\),它不依赖于任何可学习参数。
2.4 Decoder
Decoder 接收隐变量 \(z\),输出重建数据 \(\hat{x}\)。从概率角度看,Decoder 定义了条件分布 \(p_\theta(x|z)\):
- 连续数据(如灰度图像像素值归一化到 \([0,1]\)):可以假设 \(p_\theta(x|z) = \mathcal{N}(x; \hat{x}, I)\),此时最大化对数似然等价于最小化 MSE
- 二值数据(如二值化 MNIST):可以假设 \(p_\theta(x|z) = \text{Bernoulli}(x; \hat{x})\),此时最大化对数似然等价于最小化二元交叉熵 (BCE)
3. 数学推导
这是 VAE 理论的核心部分。
3.1 目标:最大化边际似然
我们的目标是学习参数 \(\theta\),使得观测数据的边际对数似然最大化:
其中 \(p(z) = \mathcal{N}(0, I)\) 是隐变量的先验分布。
积分不可解
上述积分涉及对所有可能的 \(z\) 求积分。在高维空间中,这个积分没有解析解,也无法通过蒙特卡洛采样高效估计(因为大部分 \(z\) 对 \(p_\theta(x|z)\) 的贡献极小)。
3.2 引入变分推断
既然 \(p_\theta(z|x)\)(真实后验)不可解,我们引入一个参数化的近似后验 \(q_\phi(z|x)\) 来逼近它。接下来推导 ELBO(Evidence Lower Bound,证据下界)。
第一步:写出边际对数似然
第二步:引入 \(q_\phi(z|x)\),利用重要性采样的思想
第三步:应用 Jensen 不等式(\(\log\) 是凹函数,\(\log \mathbb{E}[\cdot] \geq \mathbb{E}[\log(\cdot)]\))
这个下界就是 ELBO:
3.3 分解 ELBO
将联合分布拆分为 \(p_\theta(x, z) = p_\theta(x|z) \, p(z)\):
ELBO 的两项直觉解释
- 重建项:鼓励 Decoder 从 \(z\) 尽可能准确地重建原始输入 \(x\)。这项越大,重建越好。
- KL 散度项:惩罚近似后验 \(q_\phi(z|x)\) 与先验 \(p(z) = \mathcal{N}(0, I)\) 之间的差异。这项迫使隐空间保持规整,使得从先验采样有意义。
3.4 ELBO 与边际似然的关系
也可以从另一个角度理解。对边际对数似然做如下分解:
由于 KL 散度恒非负,所以 \(\text{ELBO} \leq \log p_\theta(x)\),即 ELBO 确实是边际对数似然的下界。最大化 ELBO 同时在做两件事:
- 最大化边际似然 \(\log p_\theta(x)\)(让模型更好地拟合数据)
- 最小化 \(D_{KL}(q_\phi(z|x) \| p_\theta(z|x))\)(让近似后验更接近真实后验)
3.5 损失函数的具体形式
VAE 的损失函数为负 ELBO:
重建损失
实践中,重建项的期望通过单次采样(\(L=1\))的蒙特卡洛估计:
- MSE 损失(对应高斯 Decoder):\(\mathcal{L}_{\text{recon}} = \|x - \hat{x}\|^2\)
- BCE 损失(对应伯努利 Decoder):\(\mathcal{L}_{\text{recon}} = -\sum_i [x_i \log \hat{x}_i + (1 - x_i) \log(1 - \hat{x}_i)]\)
KL 散度的解析解
当 \(q_\phi(z|x) = \mathcal{N}(\mu, \text{diag}(\sigma^2))\),\(p(z) = \mathcal{N}(0, I)\) 时,KL 散度有如下解析解(假设隐空间维度为 \(J\)):
推导过程:
两个高斯分布之间的 KL 散度一般公式为:
代入 \(\mu_1 = \mu\),\(\Sigma_1 = \text{diag}(\sigma^2)\),\(\mu_2 = 0\),\(\Sigma_2 = I\):
4. 前向传播数值示例
为了加深理解,我们用一个极简的例子走一遍 VAE 的前向传播。
设定:输入维度 \(d = 4\),隐空间维度 \(J = 2\)。
第一步:编码
假设输入 \(x = [0.8, 0.3, 0.9, 0.5]\),经过 Encoder 网络后得到:
由此可得:
第二步:重参数化采样
从标准正态分布采样 \(\epsilon = [0.42, -0.15]\),则:
第三步:解码
将 \(z = [0.755, -0.417]\) 送入 Decoder 网络,假设输出:
第四步:计算损失
重建损失(MSE):
KL 损失:
总损失:
观察
在训练初期,KL 损失通常远大于重建损失,因为隐分布还未被正则化到接近标准正态。随着训练推进,两项会逐渐趋于平衡。
5. 训练过程与技巧
5.1 基本训练流程
for each epoch:
for each batch x:
1. μ, log σ² = Encoder(x)
2. ε ~ N(0, I)
3. z = μ + σ ⊙ ε
4. x̂ = Decoder(z)
5. L_recon = reconstruction_loss(x, x̂)
6. L_KL = -0.5 * sum(1 + log σ² - μ² - σ²)
7. L = L_recon + L_KL
8. 反向传播,更新 φ 和 θ
5.2 Posterior Collapse
现象:在训练过程中,Decoder 学会忽略隐变量 \(z\),仅依靠自身能力(例如强大的自回归结构)生成数据。此时 \(q_\phi(z|x)\) 退化为先验 \(p(z) = \mathcal{N}(0, I)\),KL 散度降为零,隐变量变成了"死代码"。
原因:KL 正则项在训练初期过于强势,迫使 Encoder 在还没学会有用编码之前就向先验靠拢,Decoder 只好独自完成重建任务。
5.3 KL Annealing
策略:在训练初期给 KL 项一个较小的权重 \(\beta\),随训练逐步增加到 1:
- 线性退火:\(\beta\) 从 0 线性增长到 1(例如在前 10 个 epoch 内完成)
- 周期退火:\(\beta\) 周期性地从 0 增长到 1,每个周期重新开始
这使得 Encoder 在早期有充分的自由度学习有意义的编码,之后再逐渐正则化隐空间。
5.4 Beta-VAE
Higgins 等人 (2017) 提出 Beta-VAE,将 KL 权重固定为 \(\beta > 1\):
更大的 \(\beta\) 迫使隐空间更加紧凑、更符合先验,有助于学习解纠缠表示 (disentangled representation)——即隐空间的每个维度独立地编码数据的一个语义因子(如人脸的表情、姿态、发型等)。
Beta-VAE 的权衡
\(\beta\) 越大,解纠缠程度越高,但重建质量会下降(因为重建项的相对权重降低了)。实践中需要根据任务需求调节 \(\beta\)。
5.5 其他训练技巧
- Free Bits:为 KL 散度的每个维度设置一个最低值,防止某些维度完全坍缩
- Warm-up with Learning Rate:训练初期使用较小的学习率,稳定隐空间结构
- 输出 \(\log \sigma^2\) 而非 \(\sigma\):避免对 \(\sigma\) 进行 exp 操作时的数值溢出
6. VAE 的重要变体
6.1 Conditional VAE (CVAE)
CVAE 在编码和解码过程中引入条件变量 \(c\)(如类别标签):
ELBO 变为:
通过在生成时指定不同的 \(c\),可以控制生成数据的类别或属性。例如,在 MNIST 上训练 CVAE 后,可以指定 \(c = 7\) 来专门生成数字"7"。
6.2 VQ-VAE (Vector Quantized VAE)
van den Oord 等人 (2017) 提出的 VQ-VAE 用离散隐空间替代 VAE 的连续高斯隐空间:
- 维护一个可学习的码本 (codebook) \(\{e_1, e_2, \dots, e_K\}\),包含 \(K\) 个嵌入向量
- Encoder 输出连续向量 \(z_e\),通过最近邻查找将其量化到码本中最近的嵌入:
- 用直通估计器 (straight-through estimator) 解决量化操作不可微的问题
- 不使用 KL 散度,而是用码本损失和承诺损失
VQ-VAE 的重要性
VQ-VAE 是许多现代生成模型的基础组件。DALL-E (OpenAI, 2021) 使用 dVAE(一种 VQ-VAE 的变体)将图像编码为离散 token,再用 Transformer 建模。Stable Diffusion 则在 VQ-VAE / KL-VAE 的隐空间中运行扩散过程,大幅降低了计算成本。
6.3 VAE-GAN
VAE-GAN 结合了 VAE 和 GAN 的优势:
- VAE 部分提供稳定的训练和有意义的隐空间
- GAN 的判别器提供感知级别的损失信号,替代逐像素的重建损失
这种混合架构缓解了 VAE 生成图像模糊的问题,同时保留了隐空间的连续性和可解释性。
6.4 其他变体一览
| 变体 | 核心思想 | 年份 |
|---|---|---|
| Ladder VAE | 多层次隐变量,自底向上推断 | 2016 |
| IAF-VAE | 逆自回归流增强后验灵活性 | 2016 |
| WAE | 用 Wasserstein 距离替代 KL 散度 | 2018 |
| Optimus | 将 VAE 与预训练 Transformer 结合用于文本生成 | 2020 |
| NVAE | 深层层次化 VAE,达到极高图像生成质量 | 2020 |
7. VAE vs AE vs GAN 对比
| 特性 | AE | VAE | GAN |
|---|---|---|---|
| 目标 | 重建输入 | 最大化 ELBO | 对抗博弈 |
| 隐空间 | 不规则,不可采样 | 连续、正则化、可采样 | 无显式隐空间结构 |
| 生成能力 | 无(或很弱) | 有,但图像偏模糊 | 有,图像锐利 |
| 训练稳定性 | 稳定 | 稳定 | 不稳定,模式坍缩 |
| 似然估计 | 无 | 有(ELBO 下界) | 无 |
| 数学框架 | 确定性映射 | 变分推断 | 博弈论 |
| 插值能力 | 差 | 好,隐空间连续平滑 | 可行但不保证平滑 |
| 典型应用 | 降维、去噪 | 生成、表示学习 | 高质量图像生成 |
8. 思考与讨论
8.1 为什么 VAE 生成的图像偏模糊?
VAE 生成图像模糊的根本原因在于其重建损失的形式:
- 逐像素损失的局限:MSE/BCE 计算的是像素级误差的平均值。对于一张图片的某个区域,如果多个训练样本对应不同的细节(比如不同的头发纹理),VAE 倾向于输出这些可能性的平均值,这就导致了模糊
- KL 正则化的代价:KL 项迫使隐空间保持规整,这本质上在压缩信息容量。隐空间的"平滑性"意味着微小的 \(z\) 变化会导致输出的微小变化,这使得尖锐的高频细节难以编码
- 单次采样估计:训练时通常只采样一个 \(z\),这使得重建损失的梯度估计有方差,进一步导致优化偏向保守(模糊)的解
8.2 VAE 在现代生成模型中的地位
虽然在端到端图像生成方面,VAE 已被扩散模型和自回归模型超越,但 VAE 的思想和组件在现代架构中仍然发挥着关键作用:
- Stable Diffusion 的隐空间:使用预训练的 KL-VAE(一种正则化的自编码器)将图像压缩到低维隐空间,然后在隐空间中运行扩散过程。这大幅降低了计算成本——从像素空间(如 \(512 \times 512 \times 3\))压缩到隐空间(如 \(64 \times 64 \times 4\))
- DALL-E 的离散编码:使用 dVAE 将图像编码为离散 token 序列,然后用 Transformer 对图像-文本的联合 token 序列建模
- 表示学习:VAE 的隐空间在药物发现(分子生成)、音乐生成、异常检测等领域仍有广泛应用
8.3 从 VAE 到 VQ-VAE 再到 Latent Diffusion
现代高质量生成模型的发展路线中,VAE 扮演了关键的桥梁角色:
VAE (2013)
│ 连续隐空间 + 变分推断
│
├──→ VQ-VAE (2017)
│ 离散隐空间 + 码本量化
│ │
│ ├──→ VQ-VAE-2 (2019): 层次化离散隐空间
│ │
│ └──→ DALL-E (2021): dVAE + Transformer
│
└──→ KL-VAE / 正则化 AE
│
└──→ Latent Diffusion / Stable Diffusion (2022)
在 VAE 隐空间中运行扩散
核心洞见
VAE 最持久的贡献不在于它作为一个独立生成模型的生成质量,而在于它提出的"学习一个结构化的隐空间"这一思想。这个思想被后续的 VQ-VAE、Latent Diffusion 等模型继承和发展,成为了现代生成式 AI 的核心组件之一。
参考文献
- Kingma, D. P., & Welling, M. (2013). Auto-Encoding Variational Bayes. arXiv:1312.6114
- Rezende, D. J., Mohamed, S., & Wierstra, D. (2014). Stochastic Backpropagation and Approximate Inference in Deep Generative Models. ICML 2014
- Higgins, I., et al. (2017). beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework. ICLR 2017
- van den Oord, A., Vinyals, O., & Kavukcuoglu, K. (2017). Neural Discrete Representation Learning. NeurIPS 2017
- Rombach, R., et al. (2022). High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022