--- title: Score-Based SDE:扩散模型的统一理论 draft: false tags: - Score-Based - SDE - 扩散模型 - 生成模型 - 深度学习 --- # 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)$ 是一个连续时间随机过程,满足: 1. $W(0) = 0$(几乎处处) 2. 独立增量:$W(t+s) - W(t) \sim \mathcal{N}(0, s I)$ 3. 路径连续但处处不可导 **物理直觉**:维纳过程描述了粒子受到大量微小随机碰撞后的位移累积。每一步的位移服从均值 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 数值稳定性 **问题**:长时间积分可能累积数值误差。 **解决方案**: - 使用自适应求解器 - 对极端 $t$ 值使用更小步长 - 监控能量守恒等不变量 --- ## 九、与其他模型的关系 ### 9.1 与 DDPM 的关系 | 特性 | DDPM | Score SDE | |------|------|-----------| | **数学框架** | 变分推断 + ELBO | Score Matching + SDE | | **前向过程** | 离散马尔可夫链 | 连续 SDE | | **反向过程** | 离散条件分布 | 连续 SDE/ODE | | **训练目标** | 噪声预测 $\epsilon$ | Score function $\nabla \log p$ | | **采样** | 多步离散 | 可连续(ODE) | **数学等价性**: - DDPM 的噪声预测目标 $\epsilon_\theta$ 与 score function 的关系:$s_\theta(x_t, t) = -\epsilon_\theta(x_t, t) / \sigma_t^2$ - 当 $\lambda(t) = \sigma_t^2$ 时,两者的损失函数等价 ### 9.2 与 Flow Matching 的关系 Score SDE 和 Flow Matching 提供了两种互补的视角: | 视角 | Score SDE | Flow Matching | |------|-----------|---------------| | **核心量** | $\nabla \log p_t$(score) | $v_t$(速度场) | | **物理类比** | 力场(指向密度增长方向) | 速度场(描述运动) | | **训练目标** | Score Matching | 向量场回归 | | **统一关系** | $v_t = -\frac{1}{2} g g^T \nabla \log p_t$(仅在概率流 ODE 场景下成立,非 Score-SDE 与 Flow Matching 的通用统一公式) | --- ## 十、总结 Score-Based SDE 框架将扩散模型的前向和反向过程统一到随机微分方程的理论体系中: 1. **前向 SDE**:将加噪过程建模为 Itô SDE,描述数据如何逐渐被噪声污染 2. **反向时间 SDE**:严格推导了从噪声恢复到数据所需的随机微分方程 3. **Probability Flow ODE**:给出了一个确定性的等价 ODE,可以更高效地采样 4. **Score Matching**:提供了训练 score network 的三种等价目标 这一理论框架的核心价值在于: - **统一性**:将 DDPM、NF、VAE 等多种生成模型纳入同一框架 - **灵活性**:可以根据精度和速度需求选择不同的采样方法 - **可扩展性**:容易引入新的 SDE 形式(如 sub-VP、VESDE 等) --- **延伸阅读**: 1. Song et al., "Score-Based Generative Modeling through Stochastic Differential Equations" (ICLR 2021) 2. Song et al., "Maximum Likelihood Training of Score-Based Diffusion Models" (NeurIPS 2021) 3. Huang et al., "Stochastic Interpolants: A Unifying Framework for Flow and Diffusion Models" (ICLR 2024)