无需 Path Measure，也能轻松推出 Diffusion ELBO

TL; DR: 仅靠求导计算和配方，就能导出 Diffusion 的 ELBO，不需要对连续路径变分。
每个局部的 Diffusion Loss 都控制了当前时间下 Cross Entropy 的增长速度。

所有的 Gaussian diffusion path 本质相同，不妨考虑最简单的：从数据 \(X_0\sim p_0\) 出发，加上标准 Brownian noise，

\[X_t=X_0+\sqrt t\,\epsilon,\qquad \epsilon\sim\mathcal N(0,I). \]

它对应

\[{\rm d}X_t={\rm d}W_t, \]

因此边缘密度 \(p_t\) 满足热方程

\[\partial_t p_t=\frac12\Delta p_t. \]

设 denoiser \(D_t(x)\) 估计后验均值 \(\mathbb E[X_0\mid X_t=x]\)。由 Tweedie 公式，它对应的 score 估计为

\[s_t^D(x)=\frac{D_t(x)-x}{t}. \]

以此定义模型的 backward SDE：

\[{\rm d}X_t=-s_t^D(X_t)\,{\rm d}t+{\rm d}\bar W_t, \qquad t:T\to0. \]

记该反向过程在时刻 \(t\) 的边缘密度为 \(q_t\)。由于时间沿 \(T\to0\) 运行，\(q_t\) 满足 backward Fokker–Planck 方程

\[\partial_t q_t=\nabla\cdot(s_t^Dq_t)-\frac12\Delta q_t. \]

下面不比较整条 path measure，只考察同一时刻的交叉熵

\[\mathcal C(t)=\mathbb E_{X_t\sim p_t}[-\log q_t(X_t)]. \]

记

\[s_t^p=\nabla\log p_t, \qquad s_t^q=\nabla\log q_t. \]

将 \(p_t,q_t\) 的 Fokker–Planck 方程代入，对 \(\mathcal C(t)\) 求导并分部积分，二阶项相消，得到

\[-\frac{{\rm d}\mathcal C}{{\rm d}t}= \mathbb E_{p_t}[\nabla\cdot s_t^D] +\mathbb E_{p_t}[s_t^D\cdot s_t^q] -\frac12\mathbb E_{p_t}\|s_t^q\|^2. \]

再用 Stein 恒等式，

\[\mathbb E_{p_t}[\nabla\cdot s_t^D]=-\mathbb E_{p_t}[s_t^D\cdot s_t^p], \]

并直接配方：

\[-\frac{{\rm d}\mathcal C}{{\rm d}t}= \frac12\mathbb E_{p_t} \left[ -\|s_t^q-s_t^D\|^2 +\|s_t^p-s_t^D\|^2 -\|s_t^p\|^2 \right] \le \frac12 \mathbb E_{p_t} \left[ \|s_t^p-s_t^D\|^2-\|s_t^p\|^2 \right]. \]

由 denoising score matching 恒等式以及 Tweedie 公式，

\[\mathbb E_{p_t} \left[ \|s_t^p-s_t^D\|^2-\|s_t^p\|^2 \right] =\frac{\mathbb E\|D_t(X_t)-X_0\|^2-Nt}{t^2}, \]

其中 \(N\) 是维数。

因此

\[\mathbb E_{p_0}[-\log q_0(X_0)] \le \mathbb E_{p_T}[-\log q_T(X_T)] + \frac12\int_0^T\frac{\mathbb E\|D_t(X_t)-X_0\|^2-Nt}{t^2}{\rm d}t. \]

忽略常数后，训练目标即

\[\frac12\int_0^T \frac{1}{t^2}\mathbb E\|D_t(X_t)-X_0\|^2\,{\rm d}t. \]

带 Codec 的情形

为了控制截断行为，在 \(t=\tau\) 时刻提前解码，相当于以 \(q_\tau\) 为 prior、\(p_\tau(\cdot | x)\) 为 encoder、\(q(x\mid x_\tau)\) 为 decoder 的 VAE:

\[\begin{align*} -\log q(x) & = -\log \mathbb E_{p_\tau}\left[ \frac{q_\tau(x_\tau)}{p_\tau(x_\tau)} q(x \mid x_\tau) \right]\ \\ & \le -\mathbb E_{p_\tau}[\log q(x\mid x_\tau) + \log q_\tau(x_\tau) - \log p_\tau(x_\tau)] \\ & = -\mathbb E_{p_\tau}[ \log q(x\mid x_\tau) + \log q_\tau(x_\tau) ] - \mathcal H\big(p_\tau(\cdot|x)\big) \\ \mathbb E[-\log q(x)] & \le \mathbb E_{p_\tau}[ -\log q(x\mid x_\tau) ] + \mathbb E_{p_T}[-\log q_T(X_T)] + \frac12\int_\tau^T \frac{\mathbb E\|D_t(X_t)-X_0\|^2-Nt}{t^2}{\rm d}t - \mathcal H\big(p_\tau(\cdot|x)\big) \\ & = \mathbb E_{p_\tau}[ -\log q(x\mid x_\tau) ] + \mathcal D_{\rm KL}(p_T \| q_T) + \mathcal H\big(p_T(\cdot|x)\big) + \frac12\int_\tau^T \frac{\mathbb E\|D_t(X_t)-X_0\|^2-Nt}{t^2}{\rm d}t - \mathcal H\big(p_\tau(\cdot|x)\big) \\ & = \mathbb E_{p_\tau}[ -\log q(x\mid x_\tau) ] + \mathcal D_{\rm KL}(p_T \| q_T) + \frac12\int_\tau^T \frac{\mathbb E\|D_t(X_t)-X_0\|^2}{t^2}{\rm d}t \\ & \xlongequal{T \to +\infty} \mathbb E_{p_\tau}[ -\log q(x\mid x_\tau) ] + \frac12\int_\tau^{+\infty} \frac{\mathbb E\|D_t(X_t)-X_0\|^2}{t^2}{\rm d}t. \end{align*} \]

这里用到了

\[\mathcal H\big(p_T(\cdot|x)\big) - \frac12\int_\tau^T \frac{N}{t}{\rm d}t - \mathcal H\big(p_\tau(\cdot|x)\big) = \frac N2\ln(2\pi{\rm e}T)-\frac N2\ln(T/\tau)-\frac N2\ln(2\pi{\rm e}\tau)=0. \]

结论

以往的连续 Diffusion 需要先定义 Path Measure 并对整个路径的分布变分。如今我们仅靠求导计算和配方，最多再加一层单点的变分，就导出了 Diffusion 的 ELBO。

这还揭示了一个事实：Diffusion 看似对整个 Path Measure 变分，实际上每个局部 MSE Loss 都控制了当前时间下 Cross Entropy 的增长速度。