扩散模型的损失推导

定义

扩散模型真正想最大化的是数据对数似然 $\log p_\theta(x_0)$ ，但 $p_\theta(x_0)$ 需要对所有中间潜变量 $x_{1:T}$ 积分，无法直接算。于是用变分推断把它松弛成可优化的下界——ELBO（Evidence Lower Bound，证据下界）：

\log p_\theta(x_0) \geq \mathbb{E}_{q(x_{1:T}\mid x_0)}\left[\log\frac{p_\theta(x_{0:T})}{q(x_{1:T}\mid x_0)}\right] =: \mathcal{L}

DDPM 里那个「硬给」的损失 $L_{simple} = \|\epsilon - \epsilon_\theta(x_t, t)\|^2$ ，正是从这个 ELBO 一步步简化出来的。

训练的本质是：让反向去噪过程（网络学的）尽可能贴近前向加噪过程的逆（已知的）。

关键洞察在于：前向过程虽然整体是「加噪」，但只要给定 $x_0$ 和 $x_t$ ，就能算出 $x_{t-1}$ 最可能长什么样——即后验 $q(x_{t-1}\mid x_t, x_0)$ 有闭式解（因为加噪是线性高斯的）。这个后验就是去噪的「标准答案」，网络照着它学即可。

所以整个推导的线索是：

后验 $q(x_{t-1}\mid x_t,x_0)$ 可解析算 → 它是高斯，有明确均值 → 网络只需学这个均值 → 均值最终只依赖「加进去的噪声 $\epsilon$ 」→ 损失退化成预测噪声的 MSE。

利用马尔可夫性质与贝叶斯公式，把 $\mathcal{L}$ 拆开：

\mathcal{L} = \underbrace{\mathbb{E}_q[\log p_\theta(x_0\mid x_1)]}_{L_0} - \underbrace{\sum_{t=2}^{T} D_{KL}\big(q(x_{t-1}\mid x_t,x_0)\,\|\,p_\theta(x_{t-1}\mid x_t)\big)}_{L_{t-1}} - \underbrace{D_{KL}\big(q(x_T\mid x_0)\,\|\,p(x_T)\big)}_{L_T}

前向的三个分布都是高斯： $q(x_t\mid x_0)$ 、 $q(x_{t-1}\mid x_0)$ 、 $q(x_t\mid x_{t-1})$ 。对高斯套贝叶斯，得后验仍是高斯：

q(x_{t-1}\mid x_t,x_0) = \mathcal{N}\left(\tilde\mu_t(x_t,x_0),\ \tilde\beta_t I\right)

\tilde\mu_t = \frac{\sqrt{\bar\alpha_{t-1}}\,\beta_t}{1-\bar\alpha_t}x_0 + \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}x_t, \qquad \tilde\beta_t = \frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\beta_t

其中 $\alpha_t = 1-\beta_t$ ， $\bar\alpha_t = \prod_{s=1}^t \alpha_s$ （见 diffusion process）。

由重参数化 $x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\epsilon$ ，反解：

x_0 = \frac{x_t - \sqrt{1-\bar\alpha_t}\,\epsilon}{\sqrt{\bar\alpha_t}}

代入 $\tilde\mu_t$ ，消掉 $x_0$ ：

\tilde\mu_t = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\,\epsilon\right)

这是关键一步：后验均值被写成「当前 $x_t$ 减掉一个可预测的噪声」，于是「学后验」等价于「学噪声」。

让网络取同样形式，只是把未知的 $\epsilon$ 换成网络输出 $\epsilon_\theta(x_t,t)$ ：

\mu_\theta(x_t,t) = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\,\epsilon_\theta(x_t,t)\right)

两个方差相同的高斯，其 $D_{KL}$ 退化为均值差的加权平方：

L_{t-1} = \mathbb{E}\left[\frac{1}{2\sigma_t^2}\|\tilde\mu_t - \mu_\theta\|^2\right] \propto \mathbb{E}\left[\|\epsilon - \epsilon_\theta(x_t,t)\|^2\right]

DDPM 实验发现，直接丢掉前面的权重系数、均匀采样时间步 $t$ ，训练反而更稳：

L_{simple} = \mathbb{E}_{t,\,x_0,\,\epsilon}\left[\|\epsilon - \epsilon_\theta(x_t,t)\|^2\right]

[!note] 结论「预测噪声」不是拍脑袋定的，而是「学反向过程」在数学上等价的结果； $L_{simple}$ 是 ELBO 的简化形式，简化掉的是那些 KL 项的权重系数。