变分扩散模型 ELBO 的推导过程详解

变分扩散模型 ELBO 的推导过程详解

变分扩散模型（Variational Diffusion Model）通过证据下界（Evidence Lower Bound, ELBO）优化模型参数，实现从噪声到数据的生成。ELBO 的具体形式，其推导基于概率论中的边缘化、Jensen 不等式以及条件独立的性质。本文将详细推导 ELBO 表达式 ( ${\text{ELBO}}_{\phi ,\theta }(x)$ )，从基本的对数似然分解开始，逐步展开每一步计算，面向具备概率论和深度学习基础的读者。

---

ELBO 的目标与初始设定

问题背景

目标是最大化 ( $\log p(x)$ )，其中 ( $x=x_0$ ) 是输入数据，( $p(x_0)$ ) 是数据的真实分布。扩散模型通过多步正向过程 ( $q_{\phi }(x_{1:T}|x_0)$ ) 和逆向过程 ( $p_{\theta }(x_{0:T})$ ) 建模这一分布。ELBO 提供了一个可优化的下界：

$${\text{ELBO}}_{\phi ,\theta }(x)={\mathbb{E}}_{q_{\phi }(x_1|x_0)}[\log p_{\theta }(x_0|x_1)]-{\mathbb{E}}_{q_{\phi }(x_{T-1}|x_0)}[D_{KL}(q_{\phi }(x_T|x_{T-1})\Vert p(x_T))]-\sum _{t=1}^{T-1}{\mathbb{E}}_{q_{\phi }(x_{t-1},x_{t+1}|x_0)}[D_{KL}(q_{\phi }(x_t|x_{t-1})\Vert p_{\theta }(x_t|x_{t+1}))]$$

其中 ( $x_T\sim \mathcal{N}(0,I)$ )。
直观的解释可以参考笔者的另一篇博客：变分扩散模型中的 Evidence Lower Bound (ELBO) 详解
下面我们推导这一表达式的来源。

---

推导步骤

步骤 1：对数似然的边缘化

从贝叶斯定理出发，( $\log p(x_0)$ )（即 ( $\log p(x)$ )）可以通过对所有中间状态 ( $x_{1:T}$ ) 积分获得：

$$\log p(x_0)=\log \int p(x_{0:T})d{x}_{1:T}$$

这里 ( x 0 : T = { x 0 , x 1 , … , x T } x_{0:T} = \{x_0, x_1, \dots, x_T\} x0:T​={x0​,x1​,…,xT​} ) 是从 ( $t=0$ ) 到 ( $t=T$ ) 的所有状态，( $p(x_{0:T})$ ) 是联合逆向分布。

步骤 2：引入变分分布

为了便于优化，引入正向过程 ( $q_{\phi }(x_{1:T}|x_0)$ ) 作为辅助分布。利用分母-分子技巧：

$$\log p(x_0)=\log \int \frac{p(x_{0:T})}{q_{\phi }(x_{1:T}|x_0)}{q}_{\phi }(x_{1:T}|x_0)d{x}_{1:T}$$

重排：

$$=\log {\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\frac{p(x_{0:T})}{q_{\phi }(x_{1:T}|x_0)}\right]$$

步骤 3：应用 Jensen 不等式

( $\log$ ) 是一个凹函数，根据 Jensen 不等式：

$$\log {\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\frac{p(x_{0:T})}{q_{\phi }(x_{1:T}|x_0)}\right]\ge {\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log \frac{p(x_{0:T})}{q_{\phi }(x_{1:T}|x_0)}\right]$$

等号成立当且仅当 ( $\frac{p(x_{0:T})}{q_{\phi }(x_{1:T}|x_0)}$ ) 恒定（即两者完全匹配）。因此：

$$\log p(x_0)\ge {\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log \frac{p(x_{0:T})}{q_{\phi }(x_{1:T}|x_0)}\right]$$

右边即为 ELBO。

步骤 4：分解联合分布

逆向分布 ( $p(x_{0:T})$ )

根据马尔可夫性质，逆向过程是条件独立的，可知：

$$p(x_{0:T})=p(x_T)\prod _{t=1}^{T}p(x_{t-1}|x_t)$$

• ( $p(x_T)$ ) 是先验分布，通常设为 ( $\mathcal{N}(0,I)$ )。
• ( $p(x_{t-1}|x_t)$ ) 是逆向过渡分布，参数化为 ( $p_{\theta }(x_{t-1}|x_t)$ )。

展开：

$$p(x_{0:T})=p(x_T)p(x_0|x_1)\prod _{t=2}^{T}p(x_{t-1}|x_t)$$

正向分布 ( $q_{\phi }(x_{1:T}|x_0)$ )

同样，( $q_{\phi }(x_{1:T}|x_0)$ ) 是马尔可夫链：

$$q_{\phi }(x_{1:T}|x_0)=\prod _{t=1}^T{q}_{\phi }(x_t|x_{t-1})$$

• ( $q_{\phi }(x_1|x_0)$) 是从 ( $x_0$ ) 到 ( $x_1$ ) 的过渡。
• ( $q_{\phi }(x_T|x_{T-1})$ ) 是最后一步。

因此：

$$q_{\phi }(x_{1:T}|x_0)=q_{\phi }(x_T|x_{T-1})\prod _{t=1}^{T-1}{q}_{\phi }(x_t|x_{t-1})$$

步骤 5：代入 ELBO 表达式

将 ( $p(x_{0:T})$ ) 和 ( $q_{\phi }(x_{1:T}|x_0)$ ) 代入：

$$\log p(x_0)\ge {\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log \frac{p(x_T)p(x_0|x_1){\prod }_{t=2}^{T}p(x_{t-1}|x_t)}{q_{\phi }(x_T|x_{T-1}){\prod }_{t=1}^{T-1}{q}_{\phi }(x_t|x_{t-1})}\right]$$

分母-分子展开：

$$={\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log [\frac{p(x_T)}{q_{\phi }(x_T|x_{T-1})}\cdot \frac{p(x_0|x_1)}{{\prod }_{t=1}^{T-1}{q}_{\phi }(x_t|x_{t-1})}\cdot \prod _{t=2}^{T}p(x_{t-1}|x_t)]\right]$$

步骤 6：分离期望

由于期望是线性的，可分离为：

$$={\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log \frac{p(x_T)}{q_{\phi }(x_T|x_{T-1})}\right]+{\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log p(x_0|x_1)\right]+{\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log \frac{{\prod }_{t=2}^{T}p(x_{t-1}|x_t)}{{\prod }_{t=1}^{T-1}{q}_{\phi }(x_t|x_{t-1})}\right]$$

调整索引

将 ( ${\prod }_{t=2}^{T}p(x_{t-1}|x_t)$ ) 的 ( $t$ ) 替换为 ( $t+1$ )：

$$\prod _{t=2}^{T}p(x_{t-1}|x_t)=\prod _{t=1}^{T-1}p(x_t|x_{t+1})$$

因此：

$$\frac{{\prod }_{t=2}^{T}p(x_{t-1}|x_t)}{{\prod }_{t=1}^{T-1}{q}_{\phi }(x_t|x_{t-1})}=\prod _{t=1}^{T-1}\frac{p(x_t|x_{t+1})}{q_{\phi }(x_t|x_{t-1})}$$

步骤 7：简化期望

• 重构项：

$${\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log p(x_0|x_1)\right]$$

因为 ( $\log p(x_0|x_1)$ ) 只依赖 ( $x_0$ ) 和 ( $x_1$ )，期望可简化为：

$$={\mathbb{E}}_{q_{\phi }(x_1|x_0)}\left[\log p_{\theta }(x_0|x_1)\right]$$

• 先验匹配项：

$${\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log \frac{p(x_T)}{q_{\phi }(x_T|x_{T-1})}\right]$$

( $\log \frac{p(x_T)}{q_{\phi }(x_T|x_{T-1})}$ ) 只依赖 ( $x_T$ ) 和 ( $x_{T-1}$ )。由于 ( $q_{\phi }(x_{1:T}|x_0)$ ) 是马尔可夫链，( $q_{\phi }(x_T|x_{T-1},x_0)=q_{\phi }(x_T|x_{T-1})$ )，期望可写为：

$$={\mathbb{E}}_{q_{\phi }(x_{T-1}|x_0)}\left[{\mathbb{E}}_{q_{\phi }(x_T|x_{T-1})}\left[\log \frac{p(x_T)}{q_{\phi }(x_T|x_{T-1})}\right]\right]$$

因为 ( $\log \frac{p(x_T)}{q_{\phi }(x_T|x_{T-1})}$ ) 是 KL 散度：

$$=-{\mathbb{E}}_{q_{\phi }(x_{T-1}|x_0)}\left[D_{KL}(q_{\phi }(x_T|x_{T-1})\Vert p(x_T))\right]$$

• 过渡项：

$${\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log \prod _{t=1}^{T-1}\frac{p(x_t|x_{t+1})}{q_{\phi }(x_t|x_{t-1})}\right]$$

利用对数的性质：

$$=\sum _{t=1}^{T-1}{\mathbb{E}}_{q_{\phi }(x_{1:T}|x_0)}\left[\log \frac{p(x_t|x_{t+1})}{q_{\phi }(x_t|x_{t-1})}\right]$$

因为 ( $\log \frac{p(x_t|x_{t+1})}{q_{\phi }(x_t|x_{t-1})}$ ) 只依赖 ( $x_{t-1},x_t,x_{t+1}$ )：

$$=\sum _{t=1}^{T-1}{\mathbb{E}}_{q_{\phi }(x_{t-1},x_t,x_{t+1}|x_0)}\left[\log \frac{p(x_t|x_{t+1})}{q_{\phi }(x_t|x_{t-1})}\right]$$

由于 ( $q_{\phi }(x_t|x_0)=q_{\phi }(x_t|x_{t-1})$ )（马尔可夫性）：

$$=\sum _{t=1}^{T-1}{\mathbb{E}}_{q_{\phi }(x_{t-1},x_{t+1}|x_0)}\left[{\mathbb{E}}_{q_{\phi }(x_t|x_{t-1})}\left[\log \frac{p(x_t|x_{t+1})}{q_{\phi }(x_t|x_{t-1})}\right]\right]$$

这即 KL 散度：

$$=-\sum _{t=1}^{T-1}{\mathbb{E}}_{q_{\phi }(x_{t-1},x_{t+1}|x_0)}\left[D_{KL}(q_{\phi }(x_t|x_{t-1})\Vert p_{\theta }(x_t|x_{t+1}))\right]$$

步骤 8：参数化

将 ( $p(x_0|x_1)$ ) 和 ( $p(x_t|x_{t+1})$) 替换为可学习的 ( $p_{\theta }(x_0|x_1)$ ) 和 ( $p_{\theta }(x_t|x_{t+1})$ )，得到最终 ELBO。

---

推导总结

• 从 ( $\log p(x_0)$ ) 的积分形式出发，引入 ( $q_{\phi }(x_{1:T}|x_0)$ ) 并应用 Jensen 不等式，得到 ELBO 下界。
• 通过分解 ( $p(x_{0:T})$ ) 和 ( $q_{\phi }(x_{1:T}|x_0)$ ) 的马尔可夫结构，分离出初始重构、先验匹配和过渡一致性项。
• 期望的简化利用了条件独立性，确保每项只依赖相关变量。

---

代码实现片段（伪代码）

def elbo_loss(x0, model, T, alpha_schedule):
    elbo = 0.0
    x1 = forward_transition(x0, alpha_schedule[1])
    elbo += torch.mean(model.log_prob_x0_given_x1(x0, x1))  # Initial block
    
    xT_minus_1 = forward_multi_step(x0, alpha_schedule[:T])
    xT = forward_transition(xT_minus_1, alpha_schedule[T])
    kl_final = kl_divergence(xT, torch.zeros_like(xT), torch.ones_like(xT))
    elbo -= torch.mean(kl_final)  # Final block
    
    for t in range(1, T):
        x_t_minus_1, x_t_plus_1 = sample_pair(x0, t, alpha_schedule)
        x_t = forward_transition(x_t_minus_1, alpha_schedule[t])
        x_t_reverse = model.reverse(x_t_plus_1, t)
        kl_trans = kl_divergence(x_t, x_t_reverse.mean, x_t_reverse.cov)
        elbo -= torch.mean(kl_trans)  # Transition blocks
    
    return elbo

---

总结

ELBO 的推导从对数似然的下界出发，通过 Jensen 不等式和马尔可夫性质，分解为三个优化目标：初始重构、先验匹配和过渡一致性。这一过程体现了扩散模型的概率建模本质，为训练逆向去噪过程提供了理论基础。

希望这篇推导加深了你的理解！

后记

2025年3月5日15点32分于上海，在grok 3大模型辅助下完成。