变分扩散模型 ELBO 重构推导详解

在变分扩散模型（Variational Diffusion Model）中，证据下界（Evidence Lower Bound, ELBO）的形式通过优化正向和逆向分布的匹配来实现数据生成。初始 ELBO （变分扩散模型中的 Evidence Lower Bound (ELBO) 详解）存在采样复杂性，尤其是过渡块中需要联合分布 ( $q_{\phi }(x_{t-1},x_{t+1}|x_0)$ ) 的样本，这引发了重新设计的动机。后面提出了一种等价的 ELBO 形式，通过贝叶斯定理和条件调整简化了计算。本文将详细推导这一重构过程，解释这种转变，面向具备概率论和深度学习基础的读者。

参考：https://arxiv.org/pdf/2403.18103

---

初始 ELBO 的问题

原始 ELBO

原本定义的 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)}\left[D_{KL}(q_{\phi }(x_T|x_{T-1})\Vert p(x_T))\right]-\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]$$

• 初始块：重构项 ( ${\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}))]$ )。

问题所在

过渡块需要从联合分布 ( $q_{\phi }(x_{t-1},x_{t+1}|x_0)$ ) 抽样，这涉及未来状态 ( $x_{t+1}$ ) 和过去状态 ( $x_{t-1}$ ) 的耦合。直接采样 ( $(x_{t-1},x_{t+1})$ ) 复杂，因为 ( $q_{\phi }(x_{t+1}|x_0)$ ) 依赖多步正向过程，且正向 ( $q_{\phi }(x_t|x_{t-1})$ ) 和逆向 ( $p_{\theta }(x_t|x_{t+1})$ ) 方向相反，增加了计算负担。

---

重构动机与贝叶斯调整

一致性项的挑战

• ( $q_{\phi }(x_t|x_{t-1})$ ) 是正向过渡，( $p_{\theta }(x_t|x_{t+1})$ ) 是逆向过渡，两者方向相反，导致需要同时处理 ( $x_{t-1}$ ) 和 ( $x_{t+1}$ ) 的样本。
• 目标是简化一致性检查，避免“反向”依赖。

贝叶斯定理的引入

通过贝叶斯定理调整条件分布：

$$q(x_t|x_{t-1})=\frac{q(x_{t-1}|x_t)q(x_t)}{q(x_{t-1})}$$

条件于 ( $x_0$ )：

$$q(x_t|x_{t-1},x_0)=\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$$

• 这一变换将正向 ( $q(x_t|x_{t-1},x_0)$ ) 转化为逆向形式的 ( $q(x_{t-1}|x_t,x_0)$ )，方向与 ( $p_{\theta }(x_{t-1}|x_t)$ ) 一致。
• ( $x_0$ ) 的条件确保分布依赖初始状态，避免无限制采样。

---

重构 ELBO 的推导

步骤 1：从 Jensen 不等式开始

从之前的基础推导（变分扩散模型 ELBO 的推导过程详解）出发：

$$\log p(x)\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]$$

代入联合分布：

$$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|x_0)\prod _{t=2}^T{q}_{\phi }(x_t|x_{t-1},x_0)$$

（注意：这里 ( $q_{\phi }(x_t|x_{t-1},x_0)$ ) 因马尔可夫性简化为 ( $q_{\phi }(x_t|x_{t-1})$)，但为一致性保留条件。）

步骤 2：展开对数项

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

分离：

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

步骤 3：应用贝叶斯调整

对第二项，使用贝叶斯定理：

$$\frac{p(x_{t-1}|x_t)}{q_{\phi }(x_t|x_{t-1},x_0)}=\frac{p(x_{t-1}|x_t)}{q_{\phi }(x_{t-1}|x_t,x_0)\frac{q_{\phi }(x_t|x_0)}{q_{\phi }(x_{t-1}|x_0)}}$$

$$=\frac{q_{\phi }(x_{t-1}|x_t,x_0)q_{\phi }(x_t|x_0)}{q_{\phi }(x_{t-1}|x_0)}\cdot \frac{p(x_{t-1}|x_t)}{q_{\phi }(x_{t-1}|x_t,x_0)}$$

整理乘积：

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

步骤 4：期望分离

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

• 第一项：

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

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

• 第二项：

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

使用贝叶斯调整：

$$\frac{p(x_{t-1}|x_t)}{q_{\phi }(x_t|x_{t-1},x_0)}=\frac{p(x_{t-1}|x_t)q_{\phi }(x_{t-1}|x_0)}{q_{\phi }(x_{t-1}|x_t,x_0)q_{\phi }(x_t|x_0)}$$

$$\log \frac{p(x_{t-1}|x_t)}{q_{\phi }(x_t|x_{t-1},x_0)}=\log \frac{p(x_{t-1}|x_t)}{q_{\phi }(x_{t-1}|x_t,x_0)}+\log \frac{q_{\phi }(x_{t-1}|x_0)}{q_{\phi }(x_t|x_0)}$$

步骤 5：简化期望

• 重构项：

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

• 先验匹配项：

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

• 一致性项：

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

第二项的和为：

$$\log \frac{q_{\phi }(x_1|x_0)}{q_{\phi }(x_T|x_0)}=\log q_{\phi }(x_1|x_0)-\log q_{\phi }(x_T|x_0)$$

但重点是第一项：

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

步骤 6：范围调整

从 ( $t=2$ ) 到 ( $t=T$ ) 对应 ( $x_{t-1}$ ) 从 ( $x_1$ ) 到 ( $x_{T-1}$ )，与过渡块 ( $t=1$ ) 到 ( $T-1$ ) 一致，调整索引。

最终 ELBO

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

---

推导总结

• 贝叶斯定理将 ( $q_{\phi }(x_t|x_{t-1},x_0)$ ) 转化为 ( $q_{\phi }(x_{t-1}|x_t,x_0)$ )，与 ( $p_{\theta }(x_{t-1}|x_t)$ ) 方向一致。
• 期望从联合分布简化为单变量，消除了 ( $x_{t+1}$ ) 的依赖。
• 新的 ELBO 保持优化目标，简化了采样复杂性。

---

代码实现片段（伪代码）

def elbo_loss_new(x0, model, T, alpha_schedule):elbo = 0.0x1 = forward_transition(x0, alpha_schedule[1])elbo += torch.mean(model.log_prob_x0_given_x1(x0, x1))  # ReconstructionxT = forward_multi_step(x0, alpha_schedule)kl_prior = kl_divergence(xT, torch.zeros_like(xT), torch.ones_like(xT))elbo -= kl_prior  # Prior matchingfor t in range(2, T + 1):xt = forward_step(x0, t, alpha_schedule)xt_minus_1 = forward_step(x0, t - 1, alpha_schedule)kl_cons = kl_divergence(xt_minus_1, model.reverse_mean(xt, t), model.reverse_cov(xt, t))elbo -= torch.mean(kl_cons)  # Consistencyreturn elbo

---

总结

重构后的 ELBO 通过贝叶斯调整消除了联合采样的复杂性，保持了模型的优化能力。这一设计体现了扩散模型的灵活性，为高效训练提供了可能。

希望这篇推导帮助你理解！

后记

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