Chapter 4 DDPM 模型

附录1：为什么要求KL散度有闭式解？

---

一、问题的核心:有两个KL,你混淆了

让我先把问题挑明。在 diffusion 数学里出现的 KL 不是同一个——有两个层次,作用完全不同。一个不需要闭式,一个需要闭式。

KL 类型 A:外层 KL(在恒等式里)

$$\log p_\theta(x_0) = \text{ELBO} + D_{\text{KL}}(q(x_{1:T} \mid x_0) \,\|\, p_\theta(x_{1:T} \mid x_0))$$

这个 KL 出现在”似然 = ELBO + KL”这个恒等式里。它的角色是:

• 理论分析:解释 ELBO 是下界、解释下界的间隙是什么
• 不需要计算:训练时我们根本不算它,只算 ELBO

所以这个 KL 不需要闭式——它只是个理论概念,告诉我们 ELBO 优化的本质。我们无需关心它具体值多少。

KL 类型 B:内层 KL(ELBO 化简后出现)

DDPM 把 ELBO 重写成:

$$-\text{ELBO} = \mathbb{E}_q\left[\sum_{t=2}^T \underbrace{D_{\text{KL}}(q(x_{t-1} \mid x_t, x_0) \,\|\, p_\theta(x_{t-1} \mid x_t))}_{\text{KL 类型 B}} + \cdots\right]$$

这个 KL 出现在 ELBO 的内部——是 ELBO 化简后的求和项。它的角色是:

• 每次训练都要算它(因为它是 loss 的一部分)
• 必须能高效计算

这个 KL 需要闭式——闭式才能算得快、梯度精确。

---

二、所以两者的关系是

让我用一张图表示清楚:

log p(x_0)  =  ELBO  +  D_KL(q‖p)   ← 外层 KL(类型 A,理论用,不算)
              ↓
            ELBO 内部 = Σ E[ D_KL(...) ] + ...   ← 内层 KL(类型 B,训练算,要闭式)
                          ↑
                    这里的 KL 才需要闭式

关键区分:

	外层 KL(类型 A)	内层 KL(类型 B)
出现位置	恒等式 $\log p = \text{ELBO} + \text{KL}$	ELBO 化简后的求和项
是 谁 vs 谁 的 KL	完整变分分布 $q(x_{1:T}\	x_0)$ vs 真实后验 $p_\theta(x_{1:T}\	x_0)$	单步前向后验 $q(x_{t-1}\	x_t,x_0)$ vs 单步逆向 $p_\theta(x_{t-1}\	x_t)$
是否需要计算	不需要	每次训练都要
是否需要闭式	不需要	需要(否则训练慢/方差大)
性质	高维($T$ 维联合分布之间)	低维($d$ 维单步分布之间)

---

三、为什么内层 KL 必须高效计算?

回到训练流程。每次训练一步:

1. 采一个数据点 $x_0$
2. 采一个时刻 $t$
3. 采一个噪声 $\epsilon$,得到 $x_t$
4. 计算 loss(就是某个 $t$ 的内层 KL 项)
5. 反向传播

第 4 步每次迭代都要做几百万次(整个训练过程)。如果它不闭式:

• 要嵌套蒙特卡洛(里面再采样估计)→ 慢
• 估计有方差 → 训练不稳
• 梯度不精确 → 收敛慢

所以内层 KL 必须闭式——这是工程必需。

---

四、那”蒙特卡洛估计 ELBO”到底估的是什么?

回到你的疑问:既然要蒙特卡洛估 ELBO,为什么还要闭式?

答案是:蒙特卡洛和闭式估的是不同维度的随机性。我们两者都在用——把它们组合起来才得到最终的训练 loss。

ELBO 的完整结构

DDPM 改写后:

$$\text{ELBO} = -\mathbb{E}_{q}\left[\sum_{t=2}^T D_{\text{KL}}(q(x_{t-1} \mid x_t, x_0) \| p_\theta(x_{t-1} \mid x_t)) + \cdots\right]$$

注意有两层结构:

• 外层:对 $q$ 的期望 $\mathbb{E}_q[\cdots]$ —— 这层用蒙特卡洛
• 内层:每个 KL —— 这层用闭式

训练时怎么算?

外层(蒙特卡洛):从 $q$ 采样 $x_t$、随机选 $t$,得到一个具体的 KL 项。这是单样本 MC 估计。

内层(闭式):对这个具体的 $(x_t, x_0, t)$,直接用高斯 KL 闭式公式算出 KL 的精确值——不需要再采样。

具体地,对一个 batch 里的样本 $(x_0, t, \epsilon)$:

$$\text{loss} = \underbrace{\frac{\| \tilde\mu_t(x_t, x_0) - \mu_\theta(x_t, t) \|^2}{2\sigma_t^2}}_{\text{KL 闭式公式,无需采样}}$$

整个 loss 对这一组 $(x_0, t, \epsilon)$ 而言是个确定性的、可微的、精确的值。多个样本平均就是对外层期望的 MC 估计。

---

五、对比 Sohl-Dickstein:为什么”内层闭式”是关键

现在你能精确看懂 Sohl-Dickstein 和 DDPM 的差异了:

Sohl-Dickstein 的形式

$$\text{ELBO} = \mathbb{E}_q\left[\sum_t \log \frac{p_\theta(x_{t-1} \mid x_t)}{q(x_t \mid x_{t-1})}\right]$$

每一项 $\log \frac{p_\theta}{q}$ 依赖于 $x_{t-1}$ 和 $x_t$ 的具体值——它本身就是个随机变量。

没有”内层闭式”——每一项都是采样得到的随机数。

外层蒙特卡洛要估这个随机变量的期望,但它方差大,估计困难。

DDPM 的形式

$$\text{ELBO} = \mathbb{E}_q\left[\sum_t D_{\text{KL}}(q(x_{t-1} \mid x_t, x_0) \| p_\theta(x_{t-1} \mid x_t))\right]$$

每一项是个 KL,不依赖于具体的 $x_{t-1}$ 取值——它只依赖 $x_t, x_0$,且有闭式。

有”内层闭式”——每一项是确定性的精确值(给定 $x_t, x_0$)。

外层蒙特卡洛只对 $x_t$ 的随机性估,因为内层已经被解析积掉了——方差大幅降低。

---

六、用一个具体类比讲清楚

想象你要估这个嵌套期望:

$$J = \mathbb{E}_X\left[\mathbb{E}_Y[f(X, Y) \mid X]\right]$$

方法 1:不用闭式,纯蒙特卡洛。采 $X_i$、再采 $Y_{i,j}$,平均所有 $f(X_i, Y_{i,j})$。两层采样,方差大。

方法 2:如果内层期望 $g(X) := \mathbb{E}_Y[f(X, Y) \mid X]$ 有闭式,直接采 $X_i$,平均 $g(X_i)$。一层采样,方差小。

这就是 Rao-Blackwell 化——能解析积的就解析积,不要乱采样。

回到 diffusion:外层期望(对 $x_0, x_t$ 的随机性)用蒙特卡洛——这没办法,只能采;内层期望(每个 KL 项中”对 $x_{t-1}$ 的随机性”)用闭式直接算——这是降方差的关键。

---

七、再用代码看一眼

# 训练一步
x_0 = sample_data()
t   = uniform(1, T)
eps = randn_like(x_0)

# 第一层随机性 → 蒙特卡洛(就这一组采样)
x_t = sqrt(alpha_bar[t]) * x_0 + sqrt(1 - alpha_bar[t]) * eps

# 第二层"随机性"原本要积掉 → 用闭式 KL 公式直接算
mu_q     = posterior_mean(x_0, x_t, t)    # 真实后验均值,闭式
mu_theta = network(x_t, t)                 # 网络输出
loss     = 0.5 * ||mu_q - mu_theta||^2 / sigma**2   # 高斯 KL 闭式

loss.backward()

看到了吗?

• 外层(x_0, t, eps):蒙特卡洛
• 内层(loss):闭式公式,不需要再采样估计

如果内层不闭式,这行 loss = ... 就要变成”再做一次蒙特卡洛估计”,训练会慢得多、噪声大得多。

---

八、总结:你的混淆是怎么解开的

回到你最初的问题:

“我们不是用蒙特卡洛估计去估计 ELBO 的吗?为什么要求 KL 有闭式解?”

答案:

1. ELBO 里有两层结构:外层期望 + 内层每项
2. 外层:用蒙特卡洛估(对 $x_0, t, x_t$ 的随机性)
3. 内层每项:最好是闭式,这样不用再嵌套一层蒙特卡洛
4. 闭式 + 蒙特卡洛是配合关系,不是二选一——把外层的”无法避免的随机性”留给 MC,把内层的”能算的部分”用闭式精确算掉
5. “$\log p = \text{ELBO} + \text{KL}$”中的 KL(外层 KL)只是理论恒等式,不计算,不需要闭式
6. ELBO 化简后出现的 KL(内层 KL)才是训练时真正算的,必须闭式

---

九、要点回顾

• diffusion 数学里有两个 KL,一定要分清:外层(理论用)和内层(训练算)
• 蒙特卡洛 ≠ 闭式的对立面,两者是配合关系
• Rao-Blackwell 化:能闭式就闭式,降方差靠这个
• DDPM 比 Sohl-Dickstein 强的本质:让内层有闭式,把”两层蒙特卡洛”减少到”一层蒙特卡洛”
• 训练 loss 对单个样本而言是确定性、精确、可微的——多样本平均才是 MC 估计

---

清楚了吗?这个区分非常微妙,但理解了它你就抓到了 diffusion 数学优雅性的真正命门。

继续推进的话,正好可以看一下 DDPM 的贝叶斯改写——你会亲眼看到“外层 KL”如何在等式中出现,”内层 KL”如何被造出来,以及高斯 KL 闭式如何把内层化简成 MSE。要不要现在就推这一段?