变分扩散模型中的ELBO推导：从概率论到信息论的深度解析

1. 变分扩散模型与ELBO的核心思想

变分扩散模型（Variational Diffusion Models, VDM）作为当前生成式AI领域的前沿技术，其核心在于通过一个精心设计的马尔可夫链过程，逐步将数据分布转化为可处理的高斯噪声，再学习逆向的去噪过程。这个框架的成功很大程度上依赖于证据下界（Evidence Lower Bound, ELBO）的巧妙构造与优化。

理解ELBO需要把握三个关键视角：

1. 变分推断视角：ELBO提供了对数似然的一个可优化下界，使得我们可以通过最大化这个下界来间接最大化数据的对数似然。这与变分自编码器（VAE）的思想一脉相承，但扩散模型通过时间维度进行了扩展。

2. 信息论视角：ELBO的推导过程中，KL散度的出现揭示了模型如何在信息压缩（噪声化）与信息恢复（去噪）之间建立平衡。每一步的KL散度实际上度量了信息损失的程度。

3. 热力学视角：扩散过程可以类比于热力学中的熵增过程，而ELBO优化则对应于系统寻找最低能量状态的过程。这种类比为理解模型的收敛性提供了物理直觉。

在数学表达上，ELBO将原始的最大化对数似然问题分解为三个关键部分：

ELBO = 重构项 - 先验匹配项 - 过渡一致性项

这种分解不仅使优化变得可行，还为模型训练提供了明确的指导方向。接下来我们将深入每个部分的数学细节。

2. 从边缘化到Jensen不等式：ELBO的初步构建

2.1 问题形式化

给定观测数据x₀，我们的目标是最大化其对数似然log p(x₀)。在扩散模型中，我们引入隐变量x₁:T表示中间状态，通过边缘化可以得到：

log p(x₀) = log ∫ p(x₀:T) dx₁:T

这个积分直接计算是不可行的，因此需要引入变分分布q(x₁:T|x₀)来构建下界。

2.2 变分下界的推导

应用Jensen不等式是构建下界的关键步骤。我们首先通过引入变分分布构造期望：

log p(x₀) = log E_q[p(x₀:T)/q(x₁:T|x₀)]

然后利用log函数的凹性，应用Jensen不等式：

log E_q[p(x₀:T)/q(x₁:T|x₀)] ≥ E_q[log p(x₀:T)/q(x₁:T|x₀)] = ELBO

这个不等式成立的原因是log函数在定义域内是凹函数，根据Jensen不等式，对于凹函数f，有f(E[X]) ≥ E[f(X)]。

2.3 马尔可夫假设的作用

扩散模型的核心假设是前向和逆向过程都具有马尔可夫性。这使得联合分布可以分解为：

p(x₀:T) = p(x_T) ∏ p(x_{t-1}|x_t)
q(x₁:T|x₀) = ∏ q(x_t|x_{t-1})

这种分解极大地简化了ELBO的计算，使其可以表示为时间步的累加形式。

3. ELBO的分解与物理意义

3.1 三项分解的详细推导

将ELBO展开并利用马尔可夫性质，可以得到三个关键部分：

1. 重构项：

   E_q[log p(x₀|x₁)]
   

   这项鼓励模型在第一步去噪时就能较好地重建原始数据。

2. 先验匹配项：

   D_KL(q(x_T|x_{T-1}) || p(x_T))
   

   确保最终状态的分布与预设的先验分布（通常是标准高斯）一致。

3. 过渡一致性项：

   ∑ D_KL(q(x_{t-1}|x_t,x₀) || p(x_{t-1}|x_t))
   

   使学习的逆向过程与真实的反向过渡保持一致。

3.2 信息论解释

从信息论角度看，这三项分别对应：

• 重构项：保留原始数据的信息量
• 先验匹配项：约束最终状态的熵
• 过渡一致性项：控制每一步的信息传输效率

KL散度的使用使得模型在每一步都能自动平衡信息保留与噪声引入的trade-off。

3.3 实现中的简化

在实际实现中，通常会做以下简化：

1. 将方差设为固定值，只让网络预测均值
2. 用预测噪声而非直接预测均值来参数化模型
3. 使用重参数化技巧使训练过程可导

这些技巧在保持理论严谨性的同时大大简化了实现复杂度。以下是PyTorch风格的伪代码：

def elbo_loss(x0, model, T, betas):
    # 前向过程加噪
    x1 = sqrt(1-betas[1])*x0 + sqrt(betas[1])*torch.randn_like(x0)
    
    # 重构项
    recon_loss = F.mse_loss(model.predict_x0(x1,1), x0)
    
    # 先验匹配项
    xT = forward_until_T(x0, betas)
    prior_loss = D_KL(N(xT,1), N(0,1))
    
    # 过渡一致性项
    trans_loss = 0
    for t in range(2,T):
        xt = forward_to_t(x0, betas, t)
        pred_noise = model(xt, t)
        true_noise = ... # 计算真实噪声
        trans_loss += F.mse_loss(pred_noise, true_noise)
    
    return recon_loss + prior_loss + trans_loss

4. 数学细节与高级主题

4.1 条件高斯分布的解析解

在推导过程中，关键的一步是得到q(x_{t-1}|x_t,x₀)的解析表达式。通过贝叶斯定理和马尔可夫性质，可以证明：

q(x_{t-1}|x_t,x₀) = N(x_{t-1}; μ̃(x_t,x₀), β̃_t I)

其中：

μ̃ = [√α_t(1-̄α_{t-1})x_t + √̄α_{t-1}(1-α_t)x₀]/(1-̄α_t)
β̃_t = (1-α_t)(1-̄α_{t-1})/(1-̄α_t)

这个解析解为计算KL散度提供了便利。

4.2 噪声预测的参数化

实践中，我们通常不直接预测均值，而是预测噪声ε。这是因为可以建立以下等价关系：

μ_θ = 1/√α_t (x_t - (1-α_t)/√(1-̄α_t) ε_θ)

这种参数化不仅数值稳定，还与得分匹配（score matching）有深刻联系。

4.3 连续时间极限

当时间步T→∞时，扩散过程可以表示为随机微分方程（SDE）。在这种视角下：

• 前向过程对应扩散SDE
• 逆向过程对应逆时间SDE
• ELBO优化对应于求解Fokker-Planck方程

这种连续视角为理解扩散模型提供了更统一的框架。

5. 实践中的优化技巧

5.1 噪声调度策略

噪声调度（即β_t的选择）对模型性能至关重要。常见策略包括：

调度类型	公式	特点
线性调度	β_t = β_min + (β_max-β_min)t/T	简单但可能不够平滑
余弦调度	β_t = (1-cos(πt/T))/2 * β_max	更平滑的过渡，适合图像
平方根调度	β_t = (√α_t - √α_{t-1})²	理论性质更好但计算复杂

5.2 训练目标的重加权

原始ELBO中各项权重相同，但实践中可以对不同时间步赋予不同权重：

L = ∑ w_t E[||ε - ε_θ(x_t,t)||²]

常见加权策略包括：

• 均匀加权：w_t = 1
• SNR加权：w_t ∝ 1/SNR(t)
• 单调加权：给中间时间步更高权重

5.3 架构设计要点

成功的扩散模型实现通常包含以下设计元素：

1. U-Net架构：保持空间分辨率的同时捕获多尺度特征
2. 位置编码：将时间步t编码为模型可以理解的信号
3. 注意力机制：在关键层加入自注意力以捕获长程依赖
4. 条件归一化：使用自适应归一化层注入时间信息

6. 理论联系与扩展

6.1 与VAE的关系

扩散模型可以视为一种特殊的时间扩展VAE，两者对比：

特性	VAE	扩散模型
隐变量维度	固定	与数据同维×时间步
推断过程	一步编码	多步前向扩散
生成过程	一步解码	多步逆向去噪
目标函数	ELBO	时间扩展ELBO

6.2 与得分匹配的联系

通过特定的参数化方式，扩散模型的训练目标可以解释为：

E[||s_θ(x_t,t) - ∇log p(x_t)||²]

其中s_θ是得分函数。这种视角揭示了扩散模型与基于得分的生成模型的深刻联系。

6.3 最新进展方向

当前研究前沿包括：

1. 快速采样：通过改进SDE求解器或蒸馏技术减少采样步数
2. 条件生成：开发更灵活的条件控制机制
3. 多模态扩展：构建统一的跨模态扩散框架
4. 理论理解：深入分析扩散模型的泛化行为和收敛特性

理解ELBO的推导不仅对掌握扩散模型至关重要，也为开发新模型提供了理论基础。通过概率论和信息论的透镜，我们可以更清晰地看到这些强大模型背后的设计原理。