e2hang/Notes

Fork 0

Files

e2hang 2f5211d8fe New Notes

2026-05-16 17:16:51 +08:00

15 KiB

Raw Permalink Blame History

title, draft, tags

title

draft

Score-Based Generative Modeling through SDEs：扩散模型的统一数学理论

一、从 ODE 到 SDE：引入随机性的必然性

1.1 确定性的局限

在 Neural ODE / CNF 中，变换是完全确定性的：

\frac{d x(t)}{d t} = f_\theta(x(t), t)

然而，真实的扩散过程（如物理中的布朗运动、热传导）本质上具有随机性。扩散模型的前向过程也是随机的——每一步都注入独立的高斯噪声。

为了统一确定性 ODE 和随机扩散过程，我们需要将常微分方程（ODE）扩展为随机微分方程（SDE）。

1.2 维纳过程与布朗运动

维纳过程 $W(t)$（也称布朗运动）是 SDE 理论的基础：

定义：W(t) 是一个连续时间随机过程，满足：

$W(0) = 0$（几乎处处）
独立增量：W(t+s) - W(t) \sim \mathcal{N}(0, s I)
路径连续但处处不可导

物理直觉：维纳过程描述了粒子受到大量微小随机碰撞后的位移累积。每一步的位移服从均值 0、方差为时间间隔的高斯分布。

二、前向 SDE（Forward Process）

2.1 形式化定义

设前向过程满足 Itô SDE：

d x(t) = f_t(x(t)) d t + g_t(x(t)) d W(t)

其中：

W(t) 是维纳过程（标准布朗运动）
f_t: \mathbb{R}^D \to \mathbb{R}^D 是漂移系数（drift coefficient）
g_t: \mathbb{R}^D \to \mathbb{R}^{D \times D} 是扩散系数（diffusion coefficient）

2.2 标准前向 SDE（DDPM 类）

DDPM 前向过程对应的 SDE（Variance Preserving, VP-SDE）：

d x(t) = -\frac{1}{2} \beta(t) x(t) d t + \sqrt{\beta(t)} d W(t)

其中 \beta(t) 是噪声调度函数。此为**方差保持（VP）**类 SDE，x(t) 始终保持单位方差附近的扰动。

补充：另有**方差爆炸（Variance Exploding, VE-SDE）**形式：$dx = \sqrt{2} \nabla_x \sigma(t) d W$，其方差随时间增长，适用于需要极端多尺度噪声的场景。两者统称主流 SDE 框架。

对应的离散化形式（Euler-Maruyama 方法）：

x_{t+1} = x_t - \frac{1}{2} \beta_t x_t + \sqrt{\beta_t} \epsilon_t, \quad \epsilon_t \sim \mathcal{N}(0, I)

验证（重参数化推导）：DDPM 离散加噪 x_t = \sqrt{1-\beta_t} x_{t-1} + \sqrt{\beta_t} \epsilon_t 可写为：

x_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0,I)

其中 $\bar\alpha_t = \prod_{s=1}^t (1-\beta_s)$。

连续化时令 $\Delta t = 1/T$，对 x_{t+1} = \sqrt{1-\beta_t} x_t + \sqrt{\beta_t} \epsilon 做 Taylor 展开（$\sqrt{1-\beta} \approx 1 - \frac{\beta}{2}$）：

x_{t+1} \approx x_t - \frac{\beta_t}{2} x_t + \sqrt{\beta_t} \epsilon_t

这正是 Euler-Maruyama 离散形式 $x_{t+1} = x_t - \frac{1}{2}\beta_t x_t + \sqrt{\beta_t} z_t$（$z_t \sim \mathcal{N}(0,I)$），故连续 SDE 欧拉离散与 DDPM 加噪公式等价。\square

2.3 SDE 的解：Fokker-Planck 方程

SDE 的解 x(t) 不再是确定性轨迹，而是一个随机过程。其概率分布由 Fokker-Planck 方程（也称 Kolmogorov 前进方程）描述：

\frac{\partial p_t(x)}{\partial t} = -\nabla \cdot (f_t(x) p_t(x)) + \sum_{i,j} \frac{\partial^2}{\partial x_i \partial x_j} \left[ (g_t g_t^T)_{ij} p_t(x) \right]

物理意义：第一项是漂移项贡献的分布变化（类似流体力学中的对流项），第二项是扩散项贡献的散布（类似热传导的扩散项）。

三、反向时间 SDE：反向生成的数学推导

3.1 核心定理

定理：对于任意前向 Itô SDE：

d x(t) = f(x(t), t) d t + g(x(t), t) d W(t)

对应的反向时间 SDE 为：

d x(t) = \left[ f(x(t), t) - g(x(t), t) g^T(x(t), t) \nabla_x \log p_t(x(t)) \right] d t + g(x(t), t) d \bar{W}(t)

其中：

\bar{W}(t) 是反向时间的维纳过程
\nabla_x \log p_t(x(t)) 是 score function（对数概率密度的梯度）
g g^T \nabla_x \log p_t 是 "逆风"项，抵消前向扩散的效应

3.2 推导要点

设前向过程从 x(0) \sim p_0 演化到 $x(T) \sim p_T$。我们需要找到一个反向过程，从 p_T 回到 $p_0$。

由贝叶斯定理，在反向时间：

p_{T-t}(x(t)) \propto p_T(x(t)) \cdot \frac{p_0(x(0))}{p_T(x(T))}

对 t 求导并利用前向 SDE 的 Fokker-Planck 方程，可以推导出反向 SDE 的形式。

物理直觉：score function \nabla_x \log p_t 指向概率密度增长最快的方向。前向 SDE 添加噪声（增加熵），反向 SDE 需要减去熵（沿 score 方向去噪）。

3.3 Score Function 的直观理解

Score Function：\nabla_x \log p_t(x) 是对数概率密度关于 x 的梯度。

直观理解：

考虑一个概率密度 $p(x)$，在 x 处取对数：\log p(x)
梯度 \nabla_x \log p(x) 指向 p(x) 增长最快的方向
这就像是一个"力"，推动 x 向高概率区域移动

在扩散模型中的意义：

前向过程：向数据添加噪声，\log p_t 变得越来越平坦
反向过程：需要沿着 \nabla \log p_t 的方向去噪，逐步恢复数据结构

四、Probability Flow ODE：确定性等价

4.1 核心发现

关键发现：对于任意 SDE，存在一个等价的确定性 ODE，它产生相同的边际分布 ${p_t}$。

Probability Flow ODE（Song et al., 2021）：

d x(t) = \left[ f(x(t), t) - \frac{1}{2} g(x(t), t) g^T(x(t), t) \nabla_x \log p_t(x(t)) \right] d t

核心物理内涵：概率流 ODE 是所有随机 SDE 轨迹的均值确定性流——它在每点 x(t) 给出的是随机前向轨迹的条件期望，而非单条 SDE 轨迹的真实演化。仅保证边际分布 \{p_t(x)\} 与 SDE 一致，单条轨迹与 SDE 轨迹存在差异。

4.2 与 Reverse-Time SDE 的对比

特性	Reverse-Time SDE	Probability Flow ODE
类型	随机（包含 `d\bar{W}` 项）	确定性（无随机项）
采样	需要 SDE 求解器	可用 ODE 求解器
边际分布	`p_t`	$p_t$（相同）
数值稳定性	较低（有噪声）	较高（确定性）

为什么它们等价？

两者对应的 Fokker-Planck 方程给出相同的解 $p_t$。概率流 ODE 通过将随机项替换为"drift correction"项来消除随机性。

4.3 与 DDPM 的联系

对于 DDPM 风格的 SDE $d x = -\frac{1}{2} \beta(t) x d t + \sqrt{\beta(t)} d W$：

Probability Flow ODE 为：

d x = -\frac{1}{2} \beta(t) \left[ x + \nabla_x \log p_t(x) \right] d t

DDPM 反向过程（连续时间极限）：

d x = -\frac{1}{2} \beta(t) \left[ x + \nabla_x \log p_t(x) \right] d t + \sqrt{\beta(t)} d \bar{W}

两者在形式上是一致的！区别仅在于概率流 ODE 去掉了随机项 $d\bar{W}$。

五、Score Matching：训练目标的推导

5.1 三种等价的 Score Matching 训练目标

5.1.1 直接得分匹配（Direct Score Matching）

\mathcal{L}_{DSM} = \mathbb{E}_{t, x \sim p_t} \left[ \frac{1}{2} \| \nabla_x \log p_t(x) - s_\theta(x, t) \|^2 \right]

其中 s_\theta(x, t) 是网络预测的 score function。

缺点（致命问题）：必须获取真实 $\nabla\log p_t(x)$，实际中无法直接得到——这是 DSM 无法实用的根本原因。去噪得分匹配（DSM-Denoise）正是为规避此问题而诞生。优点：无实际意义，DSM 无法在现实场景中使用。

5.1.2 切片得分匹配（Slice Score Matching）

核心思想：利用随机投影矩阵 $R \in \mathbb{R}^{D \times K}$（列向量独立同分布 $N(0,I)$），通过 R^\top \nabla_x \log p_t(x) 估计 score 的方向，规避显式计算 \nabla_x \log p_t 的二阶导数（散度项）。

设 r_k \sim \mathcal{N}(0, I) 为随机向量，构造估计量：

\hat{s}(x) = \mathbb{E}_{r_k}\left[ \nabla_x \log p_t(x) r_k r_k^\top \right] / \mathbb{E}[r_k r_k^\top] = \mathbb{E}_{r_k}\left[ \nabla_x \log p_t(x) r_k r_k^\top \right]

实际实现时通过下式估计 score：

\mathbb{E}_{v}\left[ (v^\top \nabla_x \log p_t(x)) \cdot v \right] = \mathbb{E}_{v}\left[ \nabla_x^2 \log p_t(x) \cdot vv^\top \right] = \text{tr}(\nabla_x^2 \log p_t(x)) \cdot I

（详细推导见 Song et al., 2019）

5.1.3 去噪得分匹配（Denoising Score Matching）

核心思想：不去直接估计 score，而是在加噪后的数据上做条件期望。

对于加噪过程 $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon$：

\mathcal{L}_{DSM-denoise} = \mathbb{E}_{t, x_0 \sim p_{data}, \epsilon \sim \mathcal{N}(0,I)} \left[ \| s_\theta(x_t, t) + \frac{\epsilon}{\sqrt{1-\bar{\alpha}_t}} \|^2 \right]

物理含义：

\epsilon 是注入的噪声
s_\theta(x_t, t) 应该预测 $-\epsilon / \sigma_t$（因为 $\nabla_{x_t} \log p_{0|t}(x_t | x_0) = -\epsilon / \sigma_t^2$）

5.2 损失权重的设计

实践中常使用加权损失：

\mathcal{L} = \mathbb{E}_t \left[ \lambda(t) \cdot \mathbb{E}_{x_0, x_t} \left[ \| s_\theta(x_t, t) + \frac{x_t - \sqrt{\bar{\alpha}_t} x_0}{1-\bar{\alpha}_t} \|^2 \right] \right]

常用调度：

$\lambda(t) = 1$（均匀权重）
$\lambda(t) = 1 - \bar{\alpha}_t$（与信噪比成反比）
$\lambda(t) = (1 - \bar{\alpha}_t)^2$（更高阶）

六、采样过程

6.1 三种采样范式

给定训练好的 score network $s_\theta(x, t) \approx \nabla_x \log p_t(x)$，我们有三种采样方法：

6.1.1 Reverse-Time SDE 采样

使用数值 SDE 求解器（如 Euler-Maruyama）沿反向时间求解。

Algorithm: SDE Sampling

x_T ~ p_T (纯噪声)
for t = T, T-Δ, ..., 0:
    x_{t-Δ} = x_t - [f_t(x_t) - g_t g_t^T s_θ(x_t,t)] Δ + g_t √Δ · z
    其中 z ~ N(0,I)

6.1.2 Probability Flow ODE 采样

使用确定性 ODE 求解器（可使用更少步数）。

Algorithm: ODE Sampling

x_T ~ p_T (纯噪声)
for t = T, T-Δ, ..., 0:
    x_{t-Δ} = x_t - [f_t(x_t) - ½ g_t g_t^T s_θ(x_t,t)] Δ

6.1.3 Predictor-Corrector 方法

结合两者：先用 ODE predictor 预测，再用 SDE corrector 校正。

Algorithm: PC Sampling

x_T ~ p_T
for t = T, T-Δ, ..., 0:
    x_pred = x_t - [f_t - ½ g g^T s_θ] Δ  # Predictor
    x_{t-Δ} = x_pred + g √Δ · z          # Corrector (Langevin, 步长 g√Δ)

6.2 采样步数的权衡

方法	典型步数	质量	速度
SDE Sampling	1000-10000	最高	慢
ODE Sampling	50-500	高	中
PC Sampling	50-500	最高	中

概率流 ODE 因为没有随机性，可以使用更激进的步长（更大的 $\Delta t$），从而在更少步数内达到满意的质量。

七、数学推导速查

7.1 核心公式

前向 SDE：

d x(t) = f(x(t), t) d t + g(x(t), t) d W(t)

反向时间 SDE：

d x(t) = \left[ f - g g^T \nabla_x \log p_t \right] d t + g d \bar{W}(t)

Probability Flow ODE：

d x(t) = \left[ f - \frac{1}{2} g g^T \nabla_x \log p_t \right] d t

Score Function（DDPM 情形）：

\nabla_{x_t} \log p_{0|t}(x_t | x_0) = -\frac{\epsilon}{\sigma_t^2}

其中 \epsilon 是注入的噪声，$\sigma_t^2 = 1 - \bar{\alpha}_t$。

7.2 维度与复杂度

操作	复杂度	备注
Score 网络前向	`O(D)`	与网络计算量相同
Score 网络反向	`O(D)`	一次反向传播
SDE 采样（`N` 步）	`O(N \cdot D)`	每步需计算 score
ODE 采样（`N` 步）	`O(N \cdot D)`	确定性，无随机开销

八、训练中的潜在问题

8.1 Score 估计的方差问题

问题：当 $t \to 0$（低噪声水平）时，score function \nabla_x \log p_t(x) 的幅度很大，估计方差较高。

原因：低噪声时，数据分布支撑集较小，密度梯度变化剧烈。

解决方案：

使用加权损失 \lambda(t) 降低 t \approx 0 时的权重
采用截断技巧（truncated score）
使用混合损失（结合去噪目标）

8.2 时间 embedding 的重要性

问题：Score network s_\theta(x, t) 需要同时处理空间输入 x 和时间 $t$。

常见方法：

Sinusoidal embedding（原始 DDPM）
Gaussian Fourier embedding
Time MLP / Attention

8.3 数值稳定性

问题：长时间积分可能累积数值误差。