扩散模型中的分数匹配与随机微分方程（Score Matching & Stochastic Differential Equations）
题目描述
扩散模型（Diffusion Models）的经典形式（如DDPM）通过一个离散的前向加噪和反向去噪过程来生成数据。然而，其反向过程可以通过分数匹配（Score Matching） 的理论框架，并结合随机微分方程（SDE） 的视角，得到一个更一般、更统一的连续时间建模方法。本题目将深入解析如何利用分数匹配来估计数据分布的梯度（即分数函数），并通过求解一个反向SDE来实现连续时间的数据生成过程。我们将循序渐进地理解分数匹配的目标、基于分数的生成式建模思想，以及如何从SDE的视角推导出扩散模型的反向生成过程。
解题过程
步骤1：理解“分数”在概率分布中的定义
- 在概率统计中，对于一个概率密度函数 \(p(\mathbf{x})\)，其分数（score） 定义为对数概率密度的梯度：
\[ \nabla_{\mathbf{x}} \log p(\mathbf{x}) \]
- 直观理解：分数函数指向数据空间中概率密度增加最快的方向。例如，在图像空间中，分数函数指示了如何微调当前像素值，使其更像来自真实数据分布。
- 在扩散模型中，我们关心的不仅是原始数据分布 \(p_{\text{data}}(\mathbf{x})\)，更是加噪后的条件分布 \(p_t(\mathbf{x}_t | \mathbf{x}_0)\) 的分数，其中 \(t\) 表示加噪的时间步（或噪声水平）。
步骤2：分数匹配（Score Matching）的基本原理
- 目标：在不知道真实分布 \(p(\mathbf{x})\) 的情况下，通过样本数据 \(\{ \mathbf{x}^{(i)} \}\) 来估计分数函数 \(\nabla_{\mathbf{x}} \log p(\mathbf{x})\)。
- 方法：定义一个参数化的分数网络 \(s_{\boldsymbol{\theta}}(\mathbf{x})\)（例如神经网络），其输出应与真实分数尽量一致。
- 分数匹配的原始目标函数是分数匹配损失，由Hyvärinen (2005) 提出：
\[ J_{\text{SM}}(\boldsymbol{\theta}) = \mathbb{E}_{p(\mathbf{x})} \left[ \frac{1}{2} \| s_{\boldsymbol{\theta}}(\mathbf{x}) - \nabla_{\mathbf{x}} \log p(\mathbf{x}) \|^2 \right] \]
但此式依赖真实分数 \(\nabla_{\mathbf{x}} \log p(\mathbf{x})\)，无法直接计算。
4. 关键推导：通过分部积分，可以得到一个等价的目标函数，仅依赖于样本和分数网络的雅可比矩阵：
\[ J_{\text{DSM}}(\boldsymbol{\theta}) = \mathbb{E}_{p(\mathbf{x})} \left[ \text{tr}(\nabla_{\mathbf{x}} s_{\boldsymbol{\theta}}(\mathbf{x})) + \frac{1}{2} \| s_{\boldsymbol{\theta}}(\mathbf{x}) \|^2 \right] \]
其中 \(\text{tr}(\nabla_{\mathbf{x}} s_{\boldsymbol{\theta}}(\mathbf{x}))\) 是分数网络输出对输入 \(\mathbf{x}\) 的雅可比矩阵的迹。这个形式称为显式分数匹配。
5. 实际问题：计算雅可比矩阵的迹在大规模数据和高维空间中计算成本很高，因此需要更高效的方法。
步骤3：去噪分数匹配（Denoising Score Matching, DSM）
- 为了解决显式分数匹配的计算难题，Vincent (2011) 提出了去噪分数匹配。
- 核心思想：不对原始分布 \(p(\mathbf{x})\) 直接匹配分数，而是对一个加噪后的分布 \(q_{\sigma}(\tilde{\mathbf{x}})\) 进行分数匹配，其中 \(\tilde{\mathbf{x}} = \mathbf{x} + \sigma \boldsymbol{\epsilon}\)，\(\boldsymbol{\epsilon} \sim \mathcal{N}(0, I)\)。
- 加噪分布为：\(q_{\sigma}(\tilde{\mathbf{x}}) = \int p_{\text{data}}(\mathbf{x}) \mathcal{N}(\tilde{\mathbf{x}}; \mathbf{x}, \sigma^2 I) d\mathbf{x}\)。
- DSM 的目标函数为：
\[ J_{\text{DSM}}(\boldsymbol{\theta}; \sigma) = \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{\tilde{\mathbf{x}} \sim \mathcal{N}(\mathbf{x}, \sigma^2 I)} \left[ \frac{1}{2} \| s_{\boldsymbol{\theta}}(\tilde{\mathbf{x}}) - \nabla_{\tilde{\mathbf{x}}} \log p_{\sigma}(\tilde{\mathbf{x}} | \mathbf{x}) \|^2 \right] \]
其中 \(p_{\sigma}(\tilde{\mathbf{x}} | \mathbf{x}) = \mathcal{N}(\tilde{\mathbf{x}}; \mathbf{x}, \sigma^2 I)\)，其分数函数是 \(\nabla_{\tilde{\mathbf{x}}} \log p_{\sigma}(\tilde{\mathbf{x}} | \mathbf{x}) = -(\tilde{\mathbf{x}} - \mathbf{x}) / \sigma^2\)。
5. 因此，DSM 损失简化为：
\[ J_{\text{DSM}}(\boldsymbol{\theta}; \sigma) = \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{\boldsymbol{\epsilon} \sim \mathcal{N}(0, I)} \left[ \frac{1}{2} \| s_{\boldsymbol{\theta}}(\mathbf{x} + \sigma \boldsymbol{\epsilon}) + \boldsymbol{\epsilon} / \sigma \|^2 \right] \]
- 直观解释：训练网络 \(s_{\boldsymbol{\theta}}\) 去预测加噪样本 \(\tilde{\mathbf{x}}\) 对应的噪声方向（即 \(-\boldsymbol{\epsilon} / \sigma\)）。这与扩散模型中的去噪目标有紧密联系。
步骤4：从离散扩散到连续SDE的视角
- 扩散模型的前向过程可以看作是一个连续的随机过程，由一个随机微分方程（SDE） 描述：
\[ d\mathbf{x} = f(\mathbf{x}, t) dt + g(t) d\mathbf{w} \]
其中：
- \(\mathbf{w}\) 是标准维纳过程（布朗运动）。
- \(f(\mathbf{x}, t)\) 是漂移系数，通常设为 \(f(\mathbf{x}, t) = -\frac{1}{2} \beta(t) \mathbf{x}\)（如方差保持SDE）。
- \(g(t)\) 是扩散系数，通常设为 \(g(t) = \sqrt{\beta(t)}\)。
- \(\beta(t)\) 是噪声调度函数，控制噪声随时间增加的速度。
- 前向SDE从 \(t=0\)（数据分布）演化到 \(t=T\)（近似纯噪声分布，如高斯分布）。
- 关键结果（Anderson, 1982）：存在一个对应的反向SDE，可以从噪声分布生成数据：
\[ d\mathbf{x} = [f(\mathbf{x}, t) - g(t)^2 \nabla_{\mathbf{x}} \log p_t(\mathbf{x})] dt + g(t) d\bar{\mathbf{w}} \]
其中 \(\bar{\mathbf{w}}\) 是反向时间的维纳过程，\(p_t(\mathbf{x})\) 是前向过程在时间 \(t\) 的边缘分布。
4. 反向SDE的漂移项中出现了分数函数 \(\nabla_{\mathbf{x}} \log p_t(\mathbf{x})\)。因此，如果我们能估计每个时间 \(t\) 的分数，就可以通过数值求解这个反向SDE来从噪声中采样生成数据。
步骤5：结合分数匹配训练分数网络
- 我们需要一个网络 \(s_{\boldsymbol{\theta}}(\mathbf{x}, t)\) 来估计 \(\nabla_{\mathbf{x}} \log p_t(\mathbf{x})\)，其中 \(t \in [0, T]\)。
- 训练目标：扩展DSM到连续时间。使用一个加噪调度 \(\sigma(t)\)（通常与 \(g(t)\) 相关），定义损失函数：
\[ J_{\text{DSM-C}}(\boldsymbol{\theta}) = \mathbb{E}_{t \sim \mathcal{U}(0,T)} \mathbb{E}_{\mathbf{x}_0 \sim p_{\text{data}}} \mathbb{E}_{\boldsymbol{\epsilon} \sim \mathcal{N}(0, I)} \left[ \lambda(t) \cdot \| s_{\boldsymbol{\theta}}(\mathbf{x}_t, t) + \boldsymbol{\epsilon} / \sigma(t) \|^2 \right] \]
其中 \(\mathbf{x}_t = \mathbf{x}_0 + \sigma(t) \boldsymbol{\epsilon}\)，\(\lambda(t)\) 是正权重函数，用于平衡不同时间步的损失贡献。
3. 在实际训练中，通常选择 \(\lambda(t) = g(t)^2\) 或 \(\lambda(t) = 1\) 的简化形式。
4. 网络 \(s_{\boldsymbol{\theta}}(\mathbf{x}, t)\) 通常是一个U-Net类的结构，输入为加噪样本 \(\mathbf{x}_t\) 和时间嵌入 \(t\)，输出为同尺寸的分数估计。
步骤6：通过反向SDE生成样本
- 训练好分数网络 \(s_{\boldsymbol{\theta}}(\mathbf{x}, t)\) 后，用其近似真实分数 \(\nabla_{\mathbf{x}} \log p_t(\mathbf{x})\)。
- 反向SDE变为：
\[ d\mathbf{x} = [f(\mathbf{x}, t) - g(t)^2 s_{\boldsymbol{\theta}}(\mathbf{x}, t)] dt + g(t) d\bar{\mathbf{w}} \]
- 生成过程：从 \(t = T\) 开始，采样初始噪声 \(\mathbf{x}_T \sim \mathcal{N}(0, I)\)，然后数值求解上述反向SDE（例如使用欧拉-丸山法、预测-校正法等）从 \(t=T\) 到 \(t=0\)，得到生成样本 \(\mathbf{x}_0\)。
- 求解反向SDE的过程是一个连续时间的去噪过程，网络 \(s_{\boldsymbol{\theta}}\) 在每个时间步提供去噪方向。
步骤7：概率流常微分方程（Probability Flow ODE）
- 同一前向SDE对应一个确定性的概率流常微分方程（ODE）：
\[ d\mathbf{x} = \left[ f(\mathbf{x}, t) - \frac{1}{2} g(t)^2 \nabla_{\mathbf{x}} \log p_t(\mathbf{x}) \right] dt \]
- 这个ODE与反向SDE有相同的边缘分布 \(p_t(\mathbf{x})\)。因此，我们也可以通过求解这个ODE来生成样本，其优势是确定性采样、可精确计算似然等。
- 在生成时，用训练好的分数网络 \(s_{\boldsymbol{\theta}}(\mathbf{x}, t)\) 替换 \(\nabla_{\mathbf{x}} \log p_t(\mathbf{x})\)，然后使用ODE求解器（如Runge-Kutta方法）从噪声 \(\mathbf{x}_T\) 求解到 \(\mathbf{x}_0\)。
总结
通过分数匹配，我们可以训练一个网络来估计数据分布（及加噪分布）的梯度（分数）。结合SDE的视角，扩散模型的前向加噪过程可表述为一个前向SDE，其反向生成过程是一个依赖于分数函数的反向SDE。训练分数网络后，通过数值求解反向SDE（或对应的概率流ODE），即可实现连续时间的、灵活的数据生成。这一框架统一了多种扩散模型变体，并允许更高效的采样器和更好的生成质量。