线性高斯扩散

前置

• index

1. 线性 SDE 的显式解

先看一维或逐坐标共享标量系数的线性 Itô SDE：

$$\mathrm{d}{X}_t=[\alpha (t)X_t+\gamma (t)]\mathrm{d}t+[\beta (t)X_t+\delta (t)]\mathrm{d}{W}_t.$$

定义随机基本解

$${\Phi }_{t,t_0}:=\exp \left({\int }_{t_0}^t\left[\alpha (s)-\frac{1}{2}\beta (s{)}^2\right]\mathrm{d}s+{\int }_{t_0}^t\beta (s)\mathrm{d}{W}_s\right).$$

由 Itô 乘积公式，

$$\begin{array}{rl}{X}_t={\Phi }_{t,t_0}[X_{t_0}& +{\int }_{t_0}^t{\Phi }_{s,t_0}^{-1}[\gamma (s)-\beta (s)\delta (s)]\mathrm{d}s\\ & +{\int }_{t_0}^t{\Phi }_{s,t_0}^{-1}\delta (s)\mathrm{d}{W}_s].\end{array}$$

当 $\beta \equiv \gamma \equiv 0$ 时，随机乘性项消失，过程成为扩散模型最常用的加性线性高斯 SDE。下文统一讨论

$$\boxed{\mathrm{d}{\mathbf{X}}_t=f(t){\mathbf{X}}_t\mathrm{d}t+g(t)\mathrm{d}{\mathbf{W}}_t}.$$

2. 从 SDE 到扰动核

对 $0\le s<t$，定义

$$a_{t\mid s}:=\exp \left({\int }_s^{t}f(r)\mathrm{d}r\right).$$

变常数公式给出

$${\mathbf{X}}_t=a_{t\mid s}{\mathbf{X}}_s+a_{t\mid s}{\int }_s^t{a}_{r\mid s}^{-1}g(r)\mathrm{d}{\mathbf{W}}_r.$$

Itô 等距说明第二项是零均值高斯，其方差系数为

$$b_{t\mid s}^2:=a_{t\mid s}^2{\int }_s^t\frac{g(r{)}^2}{a_{r\mid s}^2}\mathrm{d}r.$$

因此

$$\boxed{{\mathbf{X}}_t\mid {\mathbf{X}}_s\sim \mathcal{N}\left(a_{t\mid s}{\mathbf{X}}_s,b_{t\mid s}^2\mathbf{I}\right)}.$$

特别地，令

$$a_t:=a_{t\mid 0},b_t:=b_{t\mid 0},$$

则可重参数化为

$${\mathbf{X}}_t=a_t{\mathbf{X}}_0+b_t\mathbf{ϵ},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}).$$

注意方差积分上限是当前时刻 $t$：

$$b_t^2=a_t^2{\int }_0^t\frac{g(s{)}^2}{a_s^2}\mathrm{d}s,$$

不是固定的 $1$。

3. 从扰动核到 SDE

反过来，给定 $a_0=1$、$b_0=0$ 以及 $a_t>0$、$b_t\ge 0$，希望

$$p_{0t}({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;a_t{\mathbf{x}}_0,b_t^2\mathbf{I})$$

来自从 ${\mathbf{X}}_0$ 出发的某个加性线性 SDE。一个稳妥的充分条件是
$a$ 局部绝对连续，并且

$$r_t:=\frac{b_t^2}{a_t^2}$$

局部绝对连续且非减。此时下列等式几乎处处成立；在 $a,b$ 可微时可按普通
导数理解。由

$$a_t=\exp \left({\int }_0^{t}f(s)\mathrm{d}s\right)$$

得到

$$\boxed{f(t)=\frac{a_t^{\prime }}{a_t}}.$$

又由

$$b_t^2=a_t^2{\int }_0^t\frac{g(s{)}^2}{a_s^2}\mathrm{d}s$$

求导可得

$$(b_t^2{)}^{\prime }=2\frac{a_t^{\prime }}{a_t}{b}_t^2+g(t{)}^2,$$

即

$$\boxed{g(t{)}^2=(b_t^2{)}^{\prime }-2\frac{a_t^{\prime }}{a_t}{b}_t^2=a_t^2{\left(\frac{b_t^2}{a_t^2}\right)}^{\prime }}.$$

因此，不是任意一对 $a_t,b_t$ 都定义合法 SDE；条件
$r_t^{\prime }\ge 0$ 几乎处处保证可取实值扩散系数
$g(t)=a_t\sqrt{r_t^{\prime }}$。若初始条件 $a_0=1,b_0=0$ 不成立，该核也不是从
${\mathbf{X}}_0$ 无噪声起步的转移核。

若还要由这些从 $0$ 出发的边缘核构造相容的 Markov 转移族，则对 $s<t$ 必须有

$$a_{t\mid s}=\frac{a_t}{a_s},b_{t\mid s}^2=b_t^2-\frac{a_t^2}{a_s^2}{b}_s^2\ge 0.$$

这些关系保证 Chapman—Kolmogorov 方程成立。

4. 完美遗忘与有限终点

若希望有限 $T$ 时精确满足

$${\mathbf{X}}_T\mid {\mathbf{X}}_0\sim \mathcal{N}(0,{\sigma }_{\mathrm{prior}}^2\mathbf{I})$$

且与 ${\mathbf{X}}_0$ 无关，就需要

$$a_T=0,b_T^2={\sigma }_{\mathrm{prior}}^2.$$

但若 $f\in L^1([0,T])$，则

$$a_T=\exp \left({\int }_0^{T}f(s)\mathrm{d}s\right)>0,$$

所以常规非奇异线性 SDE 在有限终点只能近似遗忘。要使 $a_t\to 0$，必须允许

$${\int }_0^{T}f(s)\mathrm{d}s=-\infty ,$$

即终点处出现不可积收缩。

还需谨慎区分“积分发散”和“系数逐点发散”。仅由
${\int }_0^{T}f=-\infty$ 不能推出 $f(t)\to -\infty$，也不能无条件推出
$g(t)\to \infty$。若进一步假设

$$f(t)\to -\infty ,b_t^2\to {\sigma }_{\mathrm{prior}}^2>0,(b_t^2{)}^{\prime }=O(1),$$

则由

$$g(t{)}^2=(b_t^2{)}^{\prime }-2f(t)b_t^2$$

才可推出 $g(t{)}^2\to \infty$。这就是“完美遗忘导致扩散系数爆炸”所需的附加条件。VP-SDE/DDPM 的实践做法是取 $a_T$ 足够小，而不是在有限时间强制 $a_T=0$。

5. Probability-flow ODE 与 EDM 参数化

加性线性 SDE 的 probability-flow ODE 为

$$\mathrm{d}{\mathbf{X}}_t=\left[\frac{a_t^{\prime }}{a_t}{\mathbf{X}}_t-\frac{1}{2}g(t{)}^2{\nabla }_{\mathbf{x}}\log p_t({\mathbf{X}}_t)\right]\mathrm{d}t.$$

令

$${\sigma }_t:=\frac{b_t}{a_t},u:=\frac{\mathbf{x}}{a_t},$$

并定义加噪数据密度

$$p(u;\sigma ):=\left[p_0\ast \mathcal{N}(0,{\sigma }^2\mathbf{I})\right](u).$$

由变量替换，

$$p_t(\mathbf{x})=a_t^{-d}p\left(\frac{\mathbf{x}}{a_t};{\sigma }_t\right).$$

对 $\mathbf{x}$ 求梯度时必须保留链式因子：

$$\boxed{{\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})=\frac{1}{a_t}{\nabla }_u\log p(u;{\sigma }_t){|}_{u=\mathbf{x}/a_t}}.$$

又因为

$$g(t{)}^2=a_t^2({\sigma }_t^2{)}^{\prime }=2a_t^2{\sigma }_t{\sigma }_t^{\prime },$$

所以 ODE 化为

$$\boxed{\mathrm{d}{\mathbf{X}}_t=\left[\frac{a_t^{\prime }}{a_t}{\mathbf{X}}_t-a_t{\sigma }_t{\sigma }_t^{\prime }{\nabla }_u\log p(u;{\sigma }_t){|}_{u={\mathbf{X}}_t/a_t}\right]\mathrm{d}t}.$$

这正是 EDM 记号 $s(t)=a_t$、$\sigma (t)={\sigma }_t$ 的来源。若去噪器

$$D(u,\sigma )\approx \mathbb{E}[{\mathbf{X}}_0\mid {\mathbf{X}}_0+\sigma \mathbf{ϵ}=u],$$

则 Tweedie 公式给出

$${\nabla }_u\log p(u;\sigma )=\frac{D(u,\sigma )-u}{{\sigma }^2}.$$

代回可得

$$\mathrm{d}{\mathbf{X}}_t=\left[\left(\frac{a_t^{\prime }}{a_t}+\frac{{\sigma }_t^{\prime }}{{\sigma }_t}\right){\mathbf{X}}_t-a_t\frac{{\sigma }_t^{\prime }}{{\sigma }_t}D\left(\frac{{\mathbf{X}}_t}{a_t},{\sigma }_t\right)\right]\mathrm{d}t.$$

完整 EDM 预条件与日程见 EDM。

6. 任意两时刻之间的逆向高斯核

设

$$p({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;a_t{\mathbf{x}}_0,b_t^2\mathbf{I})$$

来自上述相容 Markov 过程。先假设 $0<t<s\le T$ 且
$b_t>0$、$b_{s\mid t}>0$，从而下面的普通 Lebesgue 密度存在；方差退化的
情形应按相应 Dirac 测度或极限处理。记

$$a_{s\mid t}:=\frac{a_s}{a_t},b_{s\mid t}^2:=b_s^2-a_{s\mid t}^2{b}_t^2.$$

于是

$$p({\mathbf{x}}_s\mid {\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_s;a_{s\mid t}{\mathbf{x}}_t,b_{s\mid t}^2\mathbf{I}).$$

由马尔可夫性和贝叶斯公式，

$$p({\mathbf{x}}_t\mid {\mathbf{x}}_s,{\mathbf{x}}_0)\propto p({\mathbf{x}}_s\mid {\mathbf{x}}_t)p({\mathbf{x}}_t\mid {\mathbf{x}}_0).$$

剔除与 ${\mathbf{x}}_t$ 无关的常数后，负二倍对数密度中的二次型为

$$\frac{\Vert {\mathbf{x}}_s-a_{s\mid t}{\mathbf{x}}_t{\Vert }^2}{b_{s\mid t}^2}+\frac{\Vert {\mathbf{x}}_t-a_t{\mathbf{x}}_0{\Vert }^2}{b_t^2}.$$

精度矩阵和一次项分别为

$$\left(\frac{a_{s\mid t}^2}{b_{s\mid t}^2}+\frac{1}{b_t^2}\right)\mathbf{I}=\frac{b_s^2}{b_{s\mid t}^2{b}_t^2}\mathbf{I},$$ $$\frac{a_{s\mid t}}{b_{s\mid t}^2}{\mathbf{x}}_s+\frac{a_t}{b_t^2}{\mathbf{x}}_0.$$

故

$$\boxed{p({\mathbf{x}}_t\mid {\mathbf{x}}_s,{\mathbf{x}}_0)=\mathcal{N}\left({\mathbf{x}}_t;{\mathbf{\mu }}_{t\mid s},b_{t\mid s}^2\mathbf{I}\right)}$$

其中

$$b_{t\mid s}^2=\frac{b_{s\mid t}^2{b}_t^2}{b_s^2},$$ $${\mathbf{\mu }}_{t\mid s}=\frac{a_{s\mid t}{b}_t^2}{b_s^2}{\mathbf{x}}_s+\frac{a_t{b}_{s\mid t}^2}{b_s^2}{\mathbf{x}}_0.$$

该结论不要求 $a_t^2+b_t^2=1$。当 $s\downarrow t$ 时，
$b_{s\mid t}^2\downarrow 0$，条件分布退化为
${\delta }_{{\mathbf{x}}_s}$；它不是 Tweedie 后验均值。Tweedie 公式讨论的是在对数据先验积分之后的
$\mathbb{E}[{\mathbf{X}}_0\mid {\mathbf{X}}_t]$，条件对象不同。

7. Tweedie 后验均值

7.1 一般指数族恒等式

设观测 $Y$ 在给定自然参数 $\eta$ 后具有指数族密度

$$p(y\mid \eta )=p_0(y)\exp \left({\eta }^{\top }T(y)-\varphi (\eta )\right),$$

其中 $T:{\mathbb{R}}^m\to {\mathbb{R}}^k$ 可微，并在所讨论开集上假设
$p_0\in C^1$ 且 $p_0(y)>0$。定义 Jacobian

$$J_T(y{)}_{ij}:=\frac{\partial T_i(y)}{\partial y_j},J_T(y)\in {\mathbb{R}}^{k\times m}.$$

给 $\eta \in {\mathbb{R}}^k$ 指定先验后，
边缘密度为

$$p(y)=p_0(y)\int \exp \left({\eta }^{\top }T(y)-\varphi (\eta )\right)\pi (\mathrm{d}\eta ).$$

若 $p(y)>0$，可以交换对 $y$ 的梯度与积分，且后验一阶矩有限，则

$$\boxed{J_T(y{)}^{\top }\mathbb{E}[\eta \mid Y=y]={\nabla }_y\log p(y)-{\nabla }_y\log p_0(y)}.$$

证明。 对边缘密度求导并除以 $p(y)$：

$$\begin{array}{rl}{\nabla }_y\log p(y)& ={\nabla }_y\log p_0(y)\\ & +\frac{\int J_T(y{)}^{\top }\eta \exp ({\eta }^{\top }T(y)-\varphi (\eta ))\pi (\mathrm{d}\eta )}{\int \exp ({\eta }^{\top }T(y)-\varphi (\eta ))\pi (\mathrm{d}\eta )}\\ & ={\nabla }_y\log p_0(y)+J_T(y{)}^{\top }\mathbb{E}[\eta \mid Y=y].\end{array}$$

整理即得结论。这是一般 Tweedie 公式。若
$J_T(y{)}^{\top }:{\mathbb{R}}^k\to {\mathbb{R}}^m$ 单射，即
$\mathrm{rank}{J}_T(y)=k$，则后验均值唯一确定为

$$\mathbb{E}[\eta \mid Y=y]=(J_T{J}_T^{\top }{)}^{-1}{J}_T\left[{\nabla }_y\log p(y)-{\nabla }_y\log p_0(y)\right].$$

若 $J_T(y{)}^{\top }$ 非单射，恒等式只能确定后验均值关于
$\ker J_T(y{)}^{\top }$ 的等价类。

7.2 线性高斯特例

在线性高斯扩散及其后验采样应用中 (Chung et al., 2022)，考虑

$${\mathbf{X}}_t=a_t{\mathbf{X}}_0+b_t\mathbf{ϵ},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}),a_t,b_t>0.$$

扰动核为

$$p({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\frac{1}{(2\pi b_t^2{)}^{d/2}}\exp \left(-\frac{\Vert {\mathbf{x}}_t-a_t{\mathbf{x}}_0{\Vert }^2}{2b_t^2}\right).$$

展开指数：

$$-\frac{\Vert {\mathbf{x}}_t-a_t{\mathbf{x}}_0{\Vert }^2}{2b_t^2}=-\frac{\Vert {\mathbf{x}}_t{\Vert }^2}{2b_t^2}+\frac{a_t}{b_t^2}{\mathbf{x}}_t^{\top }{\mathbf{x}}_0-\frac{a_t^2}{2b_t^2}\Vert {\mathbf{x}}_0{\Vert }^2.$$

这显示它是以 ${\mathbf{x}}_0$ 为自然参数对象的指数族，其关于观测的充分统计量是

$$T({\mathbf{x}}_t)=\frac{a_t}{b_t^2}{\mathbf{x}}_t,{\nabla }_{{\mathbf{x}}_t}T=\frac{a_t}{b_t^2}\mathbf{I}.$$

也可直接对混合密度

$$p_t({\mathbf{x}}_t)=\int p({\mathbf{x}}_t\mid {\mathbf{x}}_0)p_0({\mathbf{x}}_0)\mathrm{d}{\mathbf{x}}_0$$

求导。在允许梯度与积分交换的正则条件下，

$$\begin{array}{rl}{\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)& =\mathbb{E}\left[{\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t\mid {\mathbf{X}}_0)\mid {\mathbf{X}}_t={\mathbf{x}}_t\right]\\ & =\mathbb{E}\left[-\frac{{\mathbf{x}}_t-a_t{\mathbf{X}}_0}{b_t^2}\mid {\mathbf{X}}_t={\mathbf{x}}_t\right].\end{array}$$

整理得到

$$\boxed{\mathbb{E}[{\mathbf{X}}_0\mid {\mathbf{X}}_t={\mathbf{x}}_t]=\frac{1}{a_t}{\mathbf{x}}_t+\frac{b_t^2}{a_t}{\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)}.$$

这也验证了指数族 Tweedie 公式中的所有链式因子。

7.3 VP/DDPM

令

$$a_t=\sqrt{{\bar{\alpha }}_t},b_t=\sqrt{1-{\bar{\alpha }}_t},$$

则

$$\mathbb{E}[{\mathbf{X}}_0\mid {\mathbf{X}}_t]=\frac{1}{\sqrt{{\bar{\alpha }}_t}}\left[{\mathbf{X}}_t+(1-{\bar{\alpha }}_t)\nabla \log p_t({\mathbf{X}}_t)\right].$$

7.4 VE

令 $a_t=1$、$b_t={\sigma }_t$，则

$$\mathbb{E}[{\mathbf{X}}_0\mid {\mathbf{X}}_t]={\mathbf{X}}_t+{\sigma }_t^2\nabla \log p_t({\mathbf{X}}_t).$$

该关系把分数函数、去噪器和后验均值统一起来，也是条件引导、DDIM
(Song et al., 2020) 等确定性采样器的基础。

Chung, H., Kim, J., McCann, M. T., Klasky, M. L., & Ye, J. C. (2022). Diffusion Posterior Sampling for General Noisy Inverse Problems. arXiv Preprint arXiv:2209.14687.

Song, J., Meng, C., & Ermon, S. (2020). Denoising Diffusion Implicit Models. arXiv Preprint arXiv:2010.02502.