流匹配与扩散模型-03：分数函数与分数匹配

分数函数

Why Score Function?

流匹配通过学习向量场训练流模型；扩散模型通常使用 分数函数 Score Function 描述分布局部结构。

对分类器而言，可用 Softmax 将非归一化分数归一化为概率。一般可写作 $p(x) = \frac{\tilde{p}(x)}{Z}$，其中 $\tilde{p}(x)$ 是非归一化模型，$Z = \int \tilde{p}(x) dx$ 是配分函数。

扩散模型需要处理整张图像的高维连续分布，直接计算配分函数 $Z$ 等价于对高维状态空间积分，通常不可行。

分数函数用于绕过配分函数计算。

分数函数的定义

对于任何概率分布 $q(x)$，其分数函数被定义为该分布对数似然函数对空间坐标的梯度，即：

$$s(x) = \nabla_x \log q(x)$$

它是 $x$ 处对数似然的梯度。由于配分函数 $Z$ 与 $x$ 无关，$\nabla_x \log p(x) = \nabla_x \log \tilde{p}(x) - \nabla_x \log Z = \nabla_x \log \tilde{p}(x)$，因此无需显式计算 $Z$。

分数函数在数据空间中定义向量场，方向指向概率密度上升最快处。扩散模型学习该向量场，用于指导噪声样本向数据分布演化。

条件分数函数与边缘分数函数

扩散模型和流匹配都使用时间参数 $t$ 控制概率路径。路径从 $p_{init}$ 演化到 $p_{data}$，对应两类分数：条件分数函数和边缘分数函数。

定义与演化

令 $\delta_z$ 表示集中在真实数据点 $z \in \mathbb{R}^d$ 的狄拉克分布。条件概率路径 $p_t(x|z)$ 描述单个数据点的加噪/去噪分布，满足 $p_0(\cdot|z) = p_{init}$ 与 $p_1(\cdot|z) = \delta_z$。对应的条件分数函数为 $\nabla \log p_t(x|z)$。

边缘概率路径 $p_t(x)$ 是条件路径对 $p_{data}(z)$ 的边缘化结果，即 $p_t(x) = \int p_t(x|z)p_{data}(z) dz$。边缘分数函数为 $\nabla \log p_t(x)$。

边缘分布函数的导出

目标是用可计算的条件分数表示边缘分数。推导如下：

由对数求导法则：

$$\nabla \log p_t(x) = \frac{\nabla p_t(x)}{p_t(x)}$$

代入边缘分布 $p_t(x) = \int p_t(x|z)p_{data}(z) dz$，并将对 $x$ 的梯度移入积分：

$$\nabla \log p_t(x) = \frac{\nabla \int p_t(x|z)p_{data}(z) dz}{p_t(x)} = \frac{\int \nabla p_t(x|z) p_{data}(z) dz}{p_t(x)}$$

利用 $\nabla p_t(x|z) = p_t(x|z) \nabla \log p_t(x|z)$：

$$\nabla \log p_t(x) = \frac{\int \left[ p_t(x|z) \nabla \log p_t(x|z) \right] p_{data}(z) dz}{p_t(x)}$$

重排得：

$$\nabla \log p_t(x) = \int \nabla \log p_t(x|z) \frac{p_t(x|z)p_{data}(z)}{p_t(x)} dz$$

权重 $\frac{p_t(x|z)p_{data}(z)}{p_t(x)}$ 是后验分布 $p_t(z|x)$。因此，边缘分数是条件分数在后验分布下的期望。

高斯路径

条件分数函数

设噪声调度函数 $\alpha_t,\beta_t$ 满足 $\alpha_0=0, \beta_0=1$ 与 $\alpha_1=1, \beta_1=0$。高斯条件路径为：

$$p_t(x|z) = \mathcal{N}(x; \alpha_t z, \beta_t^2 I_d)$$

对多维高斯密度取对数并对 $x$ 求梯度，得到条件分数：

$$\nabla \log p_t(x|z) = \nabla_x \left[ -\frac{d}{2} \log(2\pi\beta_t^2) - \frac{1}{2\beta_t^2} \|x - \alpha_t z\|^2 \right] = -\frac{x - \alpha_t z}{\beta_t^2}$$

条件分数关于 $x$ 与 $z$ 线性，方向指向 $\alpha_t z$。方差 $\beta_t^2$ 越大，分数幅度越小；方差趋近于 0 时，分数幅度增大。

向量场与分数函数的数学等价性：高斯路径下，流匹配的目标向量场与扩散模型的分数函数存在线性转换。条件向量场为：

$$u_t^{target}(x|z) = \left( \dot{\alpha}_t - \frac{\dot{\beta}_t}{\beta_t} \alpha_t \right) z + \frac{\dot{\beta}_t}{\beta_t} x$$

定义时间系数：

$$a_t = \beta_t^2 \frac{\dot{\alpha}_t}{\alpha_t} - \dot{\beta}_t \beta_t$$

$$b_t = \frac{\dot{\alpha}_t}{\alpha_t}$$

代入条件分数后：

$$u_t^{target}(x|z) = \left( \beta_t^2 \frac{\dot{\alpha}_t}{\alpha_t} - \dot{\beta}_t \beta_t \right) \left( \frac{\alpha_t z - x}{\beta_t^2} \right) + \frac{\dot{\alpha}_t}{\alpha_t} x$$

$$u_t^{target}(x|z) = a_t \nabla \log p_t(x|z) + b_t x$$

由于 $a_t,b_t$ 只依赖时间，边缘化后有：

$$u_t^{target}(x) = a_t \nabla \log p_t(x) + b_t x$$

因此，在高斯路径下，学习边缘目标向量场与学习边缘分数函数等价。

Denoiser 重参数化

由上述线性关系可得到 去噪器 Denoiser 视角。定义边缘去噪器为后验均值：

$$D_t(x) = \int z \frac{p_t(x|z)p_{data}(z)}{p_t(x)} dz = \mathbb{E}_{z|x}[z]$$

预测 $D_t(x)$ 后即可恢复无条件向量场和分数函数。工程上常预测原始图像或噪声残差，而不是直接预测分数，因为纯噪声阶段的分数模长可能很大，数值稳定性较差。

基于 SDE 的概率采样

扩散模型使用 SDE 采样，在确定性向量场外加入随机噪声，并通过分数项控制分布演化。

SDE Extension Trick

SDE Extension Trick：在确定性 ODE 轨迹中加入随机扰动，同时保持任意时刻的边缘分布仍为 $p_t(x)$。

对任意时间依赖扩散系数 $\sigma_t \ge 0$，构造 SDE：

$$dX_t = \left[ u_t^{target}(X_t) + \frac{\sigma_t^2}{2} \nabla \log p_t(X_t) \right] dt + \sigma_t dW_t$$

初始条件为 $X_0 \sim p_{init}$。其中 $\sigma_t dW_t$ 注入布朗噪声，会扩散分布；$\frac{\sigma_t^2}{2} \nabla \log p_t(X_t)$ 则沿密度上升方向补偿该扩散，使宏观边缘分布仍沿目标路径演化。

Fokker-Planck 方程

可用 Fokker-Planck 方程证明上述 SDE 的边缘分布确实为 $p_t(x)$。FPE 描述 Ito SDE 所诱导的概率密度随时间演化。

对一般 Ito SDE：$dX_t = \mu_t(X_t) dt + \sigma_t dW_t$，对应 FPE 为：

$$\partial_t p_t(x) = -\nabla \cdot (p_t(x) \mu_t(x)) + \frac{\sigma_t^2}{2} \Delta p_t(x)$$

它包含两项：

1. 对流/漂移分量 Drift Term [$-\nabla \cdot (p_t \mu_t)$]：对应确定性向量场对概率质量的搬运。
2. 扩散分量 Diffusion Term [$\frac{\sigma_t^2}{2} \Delta p_t(x)$]：对应布朗噪声引起的密度扩散，形式类似热方程。

将 SDE Extension Trick 中的漂移项代入 FPE：

$$\mu_t(x) = u_t^{target}(x) + \frac{\sigma_t^2}{2} \nabla \log p_t(x)$$

计算对流项：

$$-\nabla \cdot \left( p_t \left( u_t^{target} + \frac{\sigma_t^2}{2} \nabla \log p_t \right) \right)$$

展开散度：

$$= -\nabla \cdot (p_t u_t^{target}) - \frac{\sigma_t^2}{2} \nabla \cdot (p_t \nabla \log p_t)$$

由 $\nabla \log p_t = \frac{\nabla p_t}{p_t}$，有 $p_t \nabla \log p_t = \nabla p_t$，因此第二项为：

$$-\frac{\sigma_t^2}{2} \nabla \cdot (\nabla p_t) = -\frac{\sigma_t^2}{2} \Delta p_t$$

漂移项贡献为：

$$-\nabla \cdot (p_t u_t^{target}) - \frac{\sigma_t^2}{2} \Delta p_t$$

代回 FPE：

$$\partial_t p_t = \left( -\nabla \cdot (p_t u_t^{target}) - \frac{\sigma_t^2}{2} \Delta p_t \right) + \frac{\sigma_t^2}{2} \Delta p_t$$

扩散项与分数补偿项抵消，剩下：

$$\partial_t p_t = -\nabla \cdot (p_t u_t^{target})$$

这就是纯 ODE 流模型的连续性方程。因此，该 SDE 与目标 ODE 具有相同的边缘分布演化；随机性只改变样本路径，不改变宏观概率路径。

Langevin Dynamics

若目标分布不随时间变化，即 $p_t=p$ 且 $\partial_t p_t=0$，则目标向量场可取 $u_t^{target}=0$。

代入 SDE Extension Trick：

$$dX_t = \frac{\sigma_t^2}{2} \nabla \log p(X_t) dt + \sigma_t dW_t$$

该式为 过阻尼朗之万方程 Overdamped Langevin Equation。在 MCMC 和分子动力学中，它用于采样复杂平稳分布。若满足遍历性条件，过程最终收敛到目标分布 $p$。早期基于分数的生成模型也利用朗之万步骤进行降噪采样。

分数匹配

ODE/SDE 采样都需要边缘分数 $\nabla \log p_t(x)$，但该量依赖未知的 $p_{data}$。因此需要训练分数网络 $s_t^\theta: \mathbb{R}^d \times \rightarrow \mathbb{R}^d$ 来近似它。

显式分数匹配

最直接的方法是最小化预测分数与真实边缘分数的距离，称为 显式分数匹配 Explicit Score Matching：

$$\mathcal{L}_{SM}(\theta) = \mathbb{E}_{t \sim Unif, x \sim p_t} \left[ \| s_t^\theta(x) - \nabla \log p_t(x) \|^2 \right]$$

该损失的问题是：一方面需要未知的 $\nabla \log p_t(x)$；另一方面，即便通过积分分部消去显式依赖，低密度区域的分数估计仍不稳定。

去噪分数匹配

去噪分数匹配 Denoising Score Matching 将目标从未知边缘分数转为已知条件分数。

定义条件去噪分数匹配损失：

$$\mathcal{L}_{CSM}(\theta) = \mathbb{E}_{t \sim Unif, z \sim p_{data}, x \sim p_t(\cdot|z)} \left[ \| s_t^\theta(x) - \nabla \log p_t(x|z) \|^2 \right]$$

这里先采样真实数据 $z$ 和时间 $t$，再由已知条件路径采样 $x$。目标 $\nabla \log p_t(x|z)$ 可计算，高斯路径下有闭式解。

核心结论：

边缘分数匹配损失等同于去噪分数匹配损失加上一个常数项，即 $\mathcal{L}_{SM}(\theta) = \mathcal{L}_{CSM}(\theta) + C$。

两者对 $\theta$ 的梯度相同，因此最小化 $\mathcal{L}_{CSM}$ 等价于优化不可直接计算的 $\mathcal{L}_{SM}$。关键在于处理 $\mathcal{L}_{SM}$ 展开后的交叉项。

利用贝叶斯重排公式 $\nabla \log p_t(x) = \int \nabla \log p_t(x|z) \frac{p_t(x|z)p_{data}(z)}{p_t(x)} dz$：

$$\int_0^1 \int_x p_t(x) \cdot s_t^\theta(x)^T \cdot \left[ \int_z \nabla \log p_t(x|z) \frac{p_t(x|z)p_{data}(z)}{p_t(x)} dz \right] dx dt$$

外层 $p_t(x)$ 与后验表达式分母抵消，得到：

$$\int_0^1 \int_x \int_z s_t^\theta(x)^T \nabla \log p_t(x|z) p_t(x|z) p_{data}(z) dz dx dt$$

这正是联合分布下的期望：$\mathbb{E}_{t \sim Unif, z \sim p_{data}, x \sim p_t(\cdot|z)}$

展开平方项后，剩余差异只是不依赖 $\theta$ 的常数 $C$。因此，用条件分数训练即可得到边缘分数的最优估计。

DDPM

DDPM 是去噪分数匹配的工程化形式。考虑高斯条件路径 $p_t(x|z) = \mathcal{N}(x; \alpha_t z, \beta_t^2 I_d)$。

条件分数为 $-\frac{x - \alpha_t z}{\beta_t^2}$。含噪样本通过重参数化得到：$x = \alpha_t z + \beta_t \epsilon$，其中 $\epsilon \sim \mathcal{N}(0, I_d)$。

代入条件分数：

$$\nabla \log p_t(x|z) = -\frac{(\alpha_t z + \beta_t \epsilon) - \alpha_t z}{\beta_t^2} = -\frac{\beta_t \epsilon}{\beta_t^2} = -\frac{\epsilon}{\beta_t}$$

代入 DSM 损失：

$$\mathcal{L}_{CSM}(\theta) = \mathbb{E}_{t, z, \epsilon} \left[ \left\| s_t^\theta(x) + \frac{\epsilon}{\beta_t} \right\|^2 \right] = \mathbb{E}_{t, z, \epsilon} \left[ \frac{1}{\beta_t^2} \left\| \beta_t s_t^\theta(x) + \epsilon \right\|^2 \right]$$

直接预测 $s_t^\theta(x)$ 会在 $\beta_t \to 0$ 时产生数值不稳定。Ho et al. (2020) 将网络改为预测噪声 $\epsilon_t^\theta(x)$。

令 $\epsilon_t^\theta(x) = -\beta_t s_t^\theta(x)$，并去掉全局权重 $\frac{1}{\beta_t^2}$，得到 DDPM 降噪损失：

$$\mathcal{L}_{DDPM}(\theta) = \mathbb{E}_{t \sim Unif, z \sim p_{data}, \epsilon \sim \mathcal{N}(0, I_d)} \left[ \left\| \epsilon_t^\theta(\alpha_t z + \beta_t \epsilon) - \epsilon \right\|^2 \right]$$

算法流程：

DDPM 通过预测噪声间接学习分数函数，使训练目标稳定且易实现。