连续状态、连续时间扩散（Score-SDE）

前置

• 分数匹配
• 时间反演与概率流

1. 记号与前向 SDE

对标量函数 $\varphi :{\mathbb{R}}^d\to \mathbb{R}$，$\nabla \varphi$ 是梯度；对向量场
$\varphi :{\mathbb{R}}^d\to {\mathbb{R}}^d$，

$$\nabla \cdot \varphi (\mathbf{x})=\sum _{i=1}^d\frac{\partial {\varphi }_i}{\partial x_i}$$

是散度。对矩阵场
$A:{\mathbb{R}}^d\to {\mathbb{R}}^{d\times d}$，本文按行定义

$$(\nabla \cdot A{)}_i:=\sum _{j=1}^d\frac{\partial A_{ij}}{\partial x_j}.$$

Score-SDE (Song et al., 2020) 以如下正向 Itô SDE 统一连续状态扩散：

$$\mathrm{d}{\mathbf{X}}_t=\mathbf{f}({\mathbf{X}}_t,t)\mathrm{d}t+G({\mathbf{X}}_t,t)\mathrm{d}{\mathbf{W}}_t,0\le t\le T,$$

并记

$$A(\mathbf{x},t):=G(\mathbf{x},t)G(\mathbf{x},t{)}^{\top }.$$

假设系数、密度 $p_t$ 及边界衰减满足反向时间公式所需的正则条件。前向 Fokker—Planck 方程为

$${\partial }_t{p}_t=-\nabla \cdot (\mathbf{f}{p}_t)+\frac{1}{2}\sum _{i,j}{\partial }_{x_i}{\partial }_{x_j}(A_{ij}{p}_t).$$

2. 逆向时间 SDE

2.1 沿原时间坐标反向积分

若仍用原时间变量 $t$，但从 $T$ 积分到 $0$，则

$$\boxed{\mathrm{d}{\mathbf{X}}_t=\left[\mathbf{f}-\nabla \cdot A-A{\nabla }_{\mathbf{x}}\log p_t\right]({\mathbf{X}}_t,t)\mathrm{d}t+G({\mathbf{X}}_t,t)\mathrm{d}{\overline{\mathbf{W}}}_t}$$

其中积分方向满足 $\mathrm{d}t<0$。$\overline{\mathbf{W}}$ 是相对于逆向滤过的 reverse-time Brownian motion；其反向增量在长度
$|\mathrm{d}t|$ 的区间上协方差为 $|\mathrm{d}t|\mathbf{I}$。它不是把同一条前向 Brownian 路径简单倒放。

2.2 改用递增的逆向时钟

令 $\tau =T-t$、${\mathbf{Y}}_{\tau }={\mathbf{X}}_{T-\tau }$。此时 $\tau$ 从 $0$ 增加到 $T$，公式等价写为

$$\mathrm{d}{\mathbf{Y}}_{\tau }=\left[-\mathbf{f}+\nabla \cdot A+A{\nabla }_{\mathbf{x}}\log p_{T-\tau }\right]({\mathbf{Y}}_{\tau },T-\tau )\mathrm{d}\tau +G({\mathbf{Y}}_{\tau },T-\tau )\mathrm{d}{\mathbf{B}}_{\tau }.$$

两种约定只能选一种贯穿推导；若把第一式的漂移与正的 $\mathrm{d}t$ 混用，所有符号都会反转。

当 $G(\mathbf{x},t)=g(t)\mathbf{I}$ 时，$A=g(t{)}^2\mathbf{I}$ 且
$\nabla \cdot A=0$，沿递减 $t$ 的逆向 SDE 简化为

$$\mathrm{d}{\mathbf{X}}_t=\left[\mathbf{f}({\mathbf{X}}_t,t)-g(t{)}^2\nabla \log p_t({\mathbf{X}}_t)\right]\mathrm{d}t+g(t)\mathrm{d}{\overline{\mathbf{W}}}_t.$$

3. Probability-flow ODE

与正向 SDE 具有相同一时边缘密度 $(p_t{)}_{t\in [0,T]}$ 的确定性 ODE 是

$$\boxed{\mathrm{d}{\mathbf{X}}_t=\left[\mathbf{f}-\frac{1}{2}\nabla \cdot A-\frac{1}{2}A\nabla \log p_t\right]({\mathbf{X}}_t,t)\mathrm{d}t}.$$

从 $T$ 向 $0$ 解该 ODE 可将先验样本确定性地输运到数据分布。这里的“等价”只表示它与 SDE 共享每个时刻的边缘分布；随机 SDE 与确定性 ODE 不共享样本轨迹，也通常不共享多时联合分布或路径测度。

4. 分数网络的训练

训练时间依赖网络

$${\mathbf{s}}_{\theta }(\mathbf{x},t)\approx {\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}).$$

可操作的主要方法是去噪分数匹配；目标、等价性证明、权重和训练伪代码统一见
分数匹配。

4.1 Sliced score matching

若 $\mathbf{v}$ 独立于 $\mathbf{x}$ 且
$\mathbb{E}[\mathbf{v}{\mathbf{v}}^{\top }]=\mathbf{I}$，sliced score matching 的单样本项为

$${\mathbf{v}}^{\top }({\nabla }_{\mathbf{x}}{\mathbf{s}}_{\theta }(\mathbf{x},t))\mathbf{v}+\frac{1}{2}({\mathbf{v}}^{\top }{\mathbf{s}}_{\theta }(\mathbf{x},t){)}^2.$$

第一项含的是分数向量场的 Jacobian
${\nabla }_{\mathbf{x}}{\mathbf{s}}_{\theta }$，不能写成
${\mathbf{v}}^{\top }{\mathbf{s}}_{\theta }\mathbf{v}$。对 $t$、$\mathbf{x}\sim p_t$ 和
$\mathbf{v}$ 取期望，并乘时间权重 $\lambda (t)$，即可用随机向量估计经典 score matching 中的散度项。

5. 条件逆过程

给定条件 $\mathbf{y}$，把无条件分数替换为
${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}\mid \mathbf{y})$：

$$\mathrm{d}{\mathbf{X}}_t=\left[\mathbf{f}-\nabla \cdot A-A{\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}\mid \mathbf{y})\right]\mathrm{d}t+G\mathrm{d}{\overline{\mathbf{W}}}_t.$$

贝叶斯恒等式给出

$${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}\mid \mathbf{y})={\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})+{\nabla }_{\mathbf{x}}\log p_t(\mathbf{y}\mid \mathbf{x}).$$

分类器指导、classifier-free guidance、loss guidance 与 model guidance 的训练和近似边界统一见
条件扩散。

6. VP-SDE：DDPM 的连续极限

离散 DDPM 一步为

$${\mathbf{X}}_{n+1}=\sqrt{1-{\beta }_n}{\mathbf{X}}_n+\sqrt{{\beta }_n}{\mathbf{ϵ}}_n.$$

要取得非退化的连续极限，离散方差必须随步长缩放：

$$\boxed{{\beta }_n=\beta (t_n)\Delta t+o(\Delta t)}.$$

利用 $\sqrt{1-z}=1-\frac{1}{2}z+o(z)$，

$${\mathbf{X}}_{n+1}-{\mathbf{X}}_n=-\frac{1}{2}\beta (t_n){\mathbf{X}}_n\Delta t+\sqrt{\beta (t_n)}\sqrt{\Delta t}{\mathbf{ϵ}}_n+o(\Delta t){\mathbf{X}}_n+o(\sqrt{\Delta t}){\mathbf{ϵ}}_n,$$

其中第二个余项来自 $\sqrt{{\beta }_n}$，一般不能并入 $o(\Delta t)$。
在 Euler—Maruyama 意义下，漂移与条件协方差分别收敛，故极限 SDE 为

$$\boxed{\mathrm{d}{\mathbf{X}}_t=-\frac{1}{2}\beta (t){\mathbf{X}}_t\mathrm{d}t+\sqrt{\beta (t)}\mathrm{d}{\mathbf{W}}_t}.$$

它的扰动核是

$${\mathbf{X}}_t\mid {\mathbf{X}}_0\sim \mathcal{N}\left(a(t){\mathbf{X}}_0,[1-a(t{)}^2]\mathbf{I}\right),a(t)=\exp \left(-\frac{1}{2}{\int }_0^t\beta (s)\mathrm{d}s\right).$$

对应的逆向 SDE（递减 $t$）为

$$\mathrm{d}{\mathbf{X}}_t=\left[-\frac{1}{2}\beta (t){\mathbf{X}}_t-\beta (t)\nabla \log p_t({\mathbf{X}}_t)\right]\mathrm{d}t+\sqrt{\beta (t)}\mathrm{d}{\overline{\mathbf{W}}}_t.$$

因此，DDPM 表中的单步 ${\beta }_n$ 与 VP-SDE 的速率 $\beta (t)$ 数值范围不可直接比较：前者约等于后者乘 $\Delta t$。

7. VE-SDE

VE-SDE 无漂移：

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

Itô 积分解为

$$\boxed{{\mathbf{X}}_t={\mathbf{X}}_0+{\int }_0^{t}g(s)\mathrm{d}{\mathbf{W}}_s}.$$

注意积分必须含 $\mathrm{d}{\mathbf{W}}_s$，不能写成普通积分
${\int }_0^{t}g(s)\mathrm{d}s$。条件方差为

$$\mathrm{Var}({\mathbf{X}}_t\mid {\mathbf{X}}_0)=\left({\int }_0^{t}g(s{)}^2\mathrm{d}s\right)\mathbf{I}.$$

若 ${\sigma }^2$ 绝对连续且非减，并令

$$g(t)=\sqrt{\frac{\mathrm{d}}{\mathrm{d}t}\sigma (t{)}^2},$$

即 $({\sigma }^2{)}^{\prime }(t)\ge 0$ 几乎处处，则

$${\mathbf{X}}_t\mid {\mathbf{X}}_0\sim \mathcal{N}\left({\mathbf{X}}_0,[\sigma (t{)}^2-\sigma (0{)}^2]\mathbf{I}\right).$$

只有在 $\sigma (0)=0$ 时，该方差才等于 $\sigma (t{)}^2$。经典 VE 实现常令

$$\sigma (t)={\sigma }_{\min }{\left(\frac{{\sigma }_{\max }}{{\sigma }_{\min }}\right)}^t$$

的几何形式；此时若
$\sigma (0)={\sigma }_{\min }>0$，严格的 SDE 增量方差是
$\sigma (t{)}^2-{\sigma }_{\min }^2$。若直接把扰动核定义成
$\mathcal{N}({\mathbf{X}}_0,\sigma (t{)}^2\mathbf{I})$，则起点已经含
${\sigma }_{\min }$ 级噪声，这是另一种边界约定，必须明确。对上述几何日程，

$$g(t)=\sigma (t)\sqrt{2\log \frac{{\sigma }_{\max }}{{\sigma }_{\min }}}.$$

在离散网格上，与该 SDE 增量一致的一步为

$${\mathbf{X}}_{t_i}={\mathbf{X}}_{t_{i-1}}+\sqrt{\sigma (t_i{)}^2-\sigma (t_{i-1}{)}^2}{\mathbf{ϵ}}_i.$$

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2020). Score-Based Generative Modeling through Stochastic Differential Equations. arXiv Preprint arXiv:2011.13456.