ELBO 与变分推断

Variational Inference | Evidence Lower Bound | VAE Paper

---

概述

ELBO（Evidence Lower Bound，证据下界）是变分推断中的核心概念，用于近似难以计算的边缘似然 $\log p(x)$。ELBO 提供了一个可优化的目标函数，使得复杂概率模型的训练变得可行，是现代深度生成模型（如 VAE）的理论基础。

---

核心概念

什么是 ELBO

ELBO 是对数边缘似然 $\log p(x)$ 的下界：

$$\text{ELBO} = \mathbb{E}{q(z|x)}[\log p(x,z)] - \mathbb{E}{q(z|x)}[\log q(z|x)]$$

或等价形式：

$$\text{ELBO} = \mathbb{E}{q(z|x)}[\log p(x|z)] - D{KL}(q(z|x)||p(z))$$

其中：

• $x$ 是观测数据
• $z$ 是隐变量
• $q(z|x)$ 是近似后验分布（变分分布）
• $p(z)$ 是先验分布
• $p(x|z)$ 是似然函数

ELBO 的推导

目标：最大化边缘似然 $\log p(x)$

$$\log p(x) = \log \int p(x,z)dz$$

这个积分通常难以计算。引入变分分布 $q(z|x)$：

$$\begin{align}
\log p(x) &= \log \int p(x,z)dz \
&= \log \int \frac{q(z|x)}{q(z|x)} p(x,z)dz \
&= \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] \quad \text{[Jensen 不等式]} \
&= \text{ELBO}
\end{align}$$

关键关系：

$$\log p(x) = \text{ELBO} + D_{KL}(q(z|x)||p(z|x))$$

由于 $D_{KL} \geq 0$，所以 ELBO 是 $\log p(x)$ 的下界。

ELBO 的两种解释

解释 1：重构 + 正则化

$$\text{ELBO} = \underbrace{\mathbb{E}{q(z|x)}[\log p(x|z)]}{\text{重构项}} - \underbrace{D_{KL}(q(z|x)||p(z))}_{\text{正则化项}}$$

• 重构项：衡量模型拟合数据的能力
• 正则化项：约束后验分布接近先验分布

解释 2：似然 - 复杂度

$$\text{ELBO} = \underbrace{\mathbb{E}{q(z|x)}[\log p(x|z)]}{\text{数据似然}} - \underbrace{D_{KL}(q(z|x)||p(z))}_{\text{模型复杂度惩罚}}$$

这体现了奥卡姆剃刀原则：在解释数据的同时保持模型简单。

---

为什么要优化 ELBO

1. 直接计算困难

真实后验 $p(z|x) = \frac{p(x|z)p(z)}{p(x)}$ 难以计算，因为：

$$p(x) = \int p(x|z)p(z)dz$$

这个积分在高维空间中通常没有闭式解。

2. 等价优化

最大化 ELBO 等价于最小化 $D_{KL}(q(z|x)||p(z|x))$：

$$\begin{align}
\log p(x) &= \text{ELBO} + D_{KL}(q(z|x)||p(z|x)) \
\max_q \text{ELBO} &\Leftrightarrow \min_q D_{KL}(q(z|x)||p(z|x))
\end{align}$$

因为 $\log p(x)$ 不依赖于 $q$，最大化 ELBO 就是让 $q(z|x)$ 尽可能接近真实后验 $p(z|x)$。

3. 可计算性

ELBO 可以通过蒙特卡洛采样估计：

$$\text{ELBO} \approx \frac{1}{L}\sum_{i=1}^{L} \log p(x|z_i) - D_{KL}(q(z|x)||p(z))$$

其中 $z_i \sim q(z|x)$。

---

ELBO 在 VAE 中的应用

VAE 架构

输入 x → Encoder → μ(x), σ(x) → 采样 z ~ N(μ,σ²) → Decoder → 重构 x̂

VAE 的假设

• 先验：$p(z) = \mathcal{N}(0, I)$（标准正态分布）
• 后验：$q(z|x) = \mathcal{N}(\mu(x), \sigma^2(x)I)$（编码器输出）
• 似然：$p(x|z) = \mathcal{N}(\text{Decoder}(z), I)$（解码器输出）

ELBO 的具体形式

$$\text{ELBO} = \mathbb{E}{q(z|x)}[\log p(x|z)] - D{KL}(q(z|x)||p(z))$$

重构项（需要采样估计）：

$$\mathbb{E}{q(z|x)}[\log p(x|z)] \approx \frac{1}{L}\sum{i=1}^{L} \log p(x|z_i), \quad z_i \sim q(z|x)$$

通常 $L=1$（单样本估计）。对于伯努利分布的似然，这等价于二元交叉熵。

KL 项（有闭式解）：

$$D_{KL}(q(z|x)||p(z)) = -\frac{1}{2}\sum_{j=1}^{J}\left(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2\right)$$

其中 $J$ 是隐变量维度。

损失函数

VAE 的损失函数是 负 ELBO（因为要最小化）：

$$\mathcal{L} = -\text{ELBO} = \underbrace{-\mathbb{E}{q(z|x)}[\log p(x|z)]}{\text{重构损失}} + \underbrace{D_{KL}(q(z|x)||p(z))}_{\text{KL 损失}}$$

---

重参数化技巧

问题

直接从 $z \sim \mathcal{N}(\mu(x), \sigma^2(x))$ 采样不可微，无法反向传播。

解决方案：重参数化

核心思想：将随机性外部化

原始：z ~ N(μ, σ²)  ❌ 不可微

重参数化：
ε ~ N(0, I)         ✓ 随机性独立于参数
z = μ + σ ⊙ ε       ✓ 可微变换

梯度流

损失 → Decoder → z = μ + σε → μ, σ → Encoder
                  ↑ 可微！

• $\varepsilon$ 是常数（从参数角度），梯度不经过它
• $\mu$ 和 $\sigma$ 的梯度可以正常计算

---

代码实现

PyTorch 实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__init__()

        # Encoder
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)

        # Decoder
        self.fc3 = nn.Linear(latent_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, input_dim)

    def encode(self, x):
        """编码器：x → μ, log(σ²)"""
        h = torch.relu(self.fc1(x))
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)  # log(σ²) 数值稳定
        return mu, logvar

    def reparameterize(self, mu, logvar):
        """重参数化技巧：z = μ + σε"""
        std = torch.exp(0.5 * logvar)  # σ = exp(0.5*log(σ²))
        eps = torch.randn_like(std)     # ε ~ N(0,I)
        z = mu + eps * std              # z = μ + σε
        return z

    def decode(self, z):
        """解码器：z → x��"""
        h = torch.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h))

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon_x = self.decode(z)
        return recon_x, mu, logvar

def vae_loss(recon_x, x, mu, logvar):
    """
    VAE 损失函数 = -ELBO

    Args:
        recon_x: 重构的 x
        x: 原始 x
        mu: 编码器输出的均值
        logvar: 编码器输出的 log(σ²)

    Returns:
        总损失 = 重构损失 + KL 损失
    """
    # 重构损失：-E[log p(x|z)]
    # 对于伯努利分布 = 二元交叉熵
    BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')

    # KL 散度：KL(q(z|x)||p(z))
    # 闭式解：-0.5 * Σ(1 + log(σ²) - μ² - σ²)
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return BCE + KLD

# 训练循环
model = VAE()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(num_epochs):
    for batch_x in dataloader:
        optimizer.zero_grad()

        # 前向传播
        recon_x, mu, logvar = model(batch_x)

        # 计算损失（负 ELBO）
        loss = vae_loss(recon_x, batch_x, mu, logvar)

        # 反向传播
        loss.backward()
        optimizer.step()

关键实现细节

1. 使用 log(σ²) 而非 σ：数值稳定性更好
2. KL 散度的闭式解：避免采样估计
3. 重参数化：使采样过程可微
4. 单样本估计：重构项通常只采样一次

---

正则化的作用

什么是正则化

正则化是防止模型过拟合的技术，通过约束模型复杂度：

$$\text{损失函数} = \text{数据拟合项} + \lambda \times \text{正则化项}$$

常见类型：

• L2 正则：$\lambda||w||^2$ - 权重衰减
• L1 正则：$\lambda||w||$ - 稀疏性
• Dropout：随机关闭神经元
• Early Stopping：提前停止训练

VAE 中的正则化

ELBO 天然包含正则化：

$$\text{ELBO} = \underbrace{\mathbb{E}[\log p(x|z)]}{\text{拟合数据}} - \underbrace{D{KL}(q(z|x)||p(z))}_{\text{正则化}}$$

KL 散度项的作用：

1. 约束隐空间结构：强制 $q(z|x)$ 接近先验 $p(z) = \mathcal{N}(0,I)$
2. **防止编码器”作弊”**：避免为每个样本学习独特的、互不相关的编码
3. 保证连续性：使隐空间平滑、可插值
4. 实现解耦表示：鼓励学习独立的隐变量

没有 KL 正则化会怎样？

如果只优化重构损失：

loss = reconstruction_loss  # 没有 KL 项

问题：

• 编码器可能���每个样本映射到隐空间的任意位置
• 隐空间结构混乱，无法生成新样本
• 方差 $\sigma$ 可能趋近于 0（退化为确定性编码）

KL 正则化的效果：

• 强制隐变量分布接近标准正态分布
• 保证隐空间的连续性和平滑性
• 使得从先验 $p(z)$ 采样可以生成有意义的样本

---

应用场景

1. 变分自编码器（VAE）

用途：生成模型、表示学习

# 训练后生成新样本
z = torch.randn(batch_size, latent_dim)  # 从先验采样
generated_x = model.decode(z)

2. 主题模型（LDA）

用途：文档主题发现

• 隐变量 $z$：文档的主题分布
• ELBO 用于推断主题

3. 贝叶斯神经网络

用途：不确定性量化

• 隐变量 $z$：网络权重
• ELBO 用于近似权重的后验分布

4. 隐马尔可夫模型（HMM）

用途：序列建模

• 隐变量 $z$：隐藏状态序列
• ELBO 用于推断状态

5. 高斯混合模型（GMM）

用途：聚类分析

• 隐变量 $z$：样本所属的簇
• ELBO 用于推断簇分配

---

最佳实践

1. 平衡重构与 KL 损失

有时需要调整 KL 项的权重：

loss = reconstruction_loss + beta * kl_loss  # β-VAE

• β < 1：更注重重构质量
• β > 1：更注重解耦表示
• β = 1：标准 VAE

2. KL 退火（KL Annealing）

训练初期降低 KL 权重，逐渐增加：

beta = min(1.0, epoch / warmup_epochs)
loss = reconstruction_loss + beta * kl_loss

原因：避免”后验坍缩”（posterior collapse），即 $q(z|x)$ 过早收敛到先验。

3. 自由比特（Free Bits）

为每个隐变量维度设置最小 KL 值：

kl_per_dim = kl_loss / latent_dim
kl_loss = torch.sum(torch.max(kl_per_dim, free_bits_threshold))

作用：防止某些维度被忽略。

4. 数值稳定性

# 使用 log(σ²) 而非 σ
logvar = self.fc_logvar(h)
std = torch.exp(0.5 * logvar)  # 避免负数

# 避免 log(0)
recon_x = torch.clamp(recon_x, min=1e-7, max=1-1e-7)

---

常见问题

问题 1：ELBO 为什么是下界？

解答：根据 Jensen 不等式，对于凹函数 $\log$：

$$\log \mathbb{E}[X] \geq \mathbb{E}[\log X]$$

因此：

$$\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] = \text{ELBO}$$

---

问题 2：为什么叫”证据”下界？

解答：在贝叶斯推断中，$p(x)$ 被称为”证据”（evidence）或”边缘似然”（marginal likelihood），因为它是观测数据 $x$ 的概率。ELBO 是这个证据的下界。

---

问题 3：重参数化技巧为什么有效？

解答：重参数化将随机性从参数中分离出来：

z = μ(x; θ) + σ(x; θ) ⊙ ε,  ε ~ N(0,I)

• $\varepsilon$ 不依赖于 $\theta$，是外部噪声
• $\mu$ 和 $\sigma$ 依赖于 $\theta$，可以求梯度
• 梯度可以通过 $\mu$ 和 $\sigma$ 反向传播到编码器

---

问题 4：VAE 生成的图像为什么模糊？

原因：

1. 高斯似然假设：$p(x|z) = \mathcal{N}(\text{Decoder}(z), I)$ 倾向于生成平均化的结果
2. 重构损失：MSE 或 BCE 惩罚像素级差异，导致模糊

改进方法：

• 使用更复杂的似然分布（如混合高斯）
• 使用感知损失（perceptual loss）
• 结合 GAN（VAE-GAN）

---

参考资源

• Auto-Encoding Variational Bayes (VAE 原论文)
• Tutorial on Variational Autoencoders
• β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework
• Understanding disentangling in β-VAE
• Pattern Recognition and Machine Learning - Christopher Bishop

---