分数匹配与去噪分数匹配

前置

• 线性高斯扩散

1. 分数学习目标

设 $p_t$ 是前向扩散在时刻 $t$ 的边缘密度，分数为

$${\mathbf{s}}^{\star }(\mathbf{x},t):={\nabla }_{\mathbf{x}}\log p_t(\mathbf{x}).$$

训练网络 ${\mathbf{s}}_{\theta }$ 的理想加权 Fisher divergence 为

$${\mathcal{J}}_{\mathrm{score}}(\theta )=\frac{1}{2}{\mathbb{E}}_{t\sim \pi }\left[\lambda (t){\mathbb{E}}_{{\mathbf{X}}_t\sim p_t}{\Vert {\mathbf{s}}_{\theta }({\mathbf{X}}_t,t)-\nabla \log p_t({\mathbf{X}}_t)\Vert }_2^2\right].$$

这里：

• $\pi$ 是明确指定的训练时间分布，不必是均匀分布；
• $\lambda (t)\ge 0$ 是损失权重；
• 连续时间高斯扰动在 $t=0$ 常有 $b_0=0$，条件分数奇异，因此实践中通常取
  $t\sim \pi$ 支持于 $[\epsilon ,T]$，$\epsilon >0$；
• $p_0=p_{\mathrm{data}}$ 是数据边界分布，$p_T$ 是终端边缘，只有日程设计充分时才近似所选先验 $p_{\mathrm{prior}}$。

真实边缘分数通常不可计算，因为

$$p_t(\mathbf{x})=\int p_{0t}(\mathbf{x}\mid {\mathbf{x}}_0)p_0({\mathbf{x}}_0)\mathrm{d}{\mathbf{x}}_0.$$

2. 去噪分数匹配

若扰动核 $p_{0t}({\mathbf{x}}_t\mid {\mathbf{x}}_0)$ 的条件分数可解，则去噪分数匹配
（DSM）(Vincent, 2011) 目标为

$$\boxed{{\mathcal{J}}_{\mathrm{DSM}}(\theta )=\frac{1}{2}{\mathbb{E}}_{t\sim \pi }\left[\lambda (t){\mathbb{E}}_{\begin{array}{c}{\mathbf{X}}_0\sim p_0,\\ {\mathbf{X}}_t\sim p_{0t}(\cdot \mid {\mathbf{X}}_0)\end{array}}{\Vert {\mathbf{s}}_{\theta }({\mathbf{X}}_t,t)-{\nabla }_{{\mathbf{x}}_t}\log p_{0t}({\mathbf{X}}_t\mid {\mathbf{X}}_0)\Vert }_2^2\right]}.$$

该目标只需从数据和已知扰动核采样。

3. DSM 等价性证明

固定 $t$。在可将梯度移入积分、且所有期望有限的正则条件下，

$$\begin{array}{rl}{\nabla }_{\mathbf{x}}{p}_t(\mathbf{x})& =\int {\nabla }_{\mathbf{x}}{p}_{0t}(\mathbf{x}\mid {\mathbf{x}}_0)p_0({\mathbf{x}}_0)\mathrm{d}{\mathbf{x}}_0\\ & =\int p_{0t}(\mathbf{x}\mid {\mathbf{x}}_0)p_0({\mathbf{x}}_0){\nabla }_{\mathbf{x}}\log p_{0t}(\mathbf{x}\mid {\mathbf{x}}_0)\mathrm{d}{\mathbf{x}}_0.\end{array}$$

除以 $p_t(\mathbf{x})$ 得 Fisher identity：

$$\nabla \log p_t(\mathbf{x})=\mathbb{E}\left[{\nabla }_{\mathbf{x}}\log p_{0t}(\mathbf{x}\mid {\mathbf{X}}_0)\mid {\mathbf{X}}_t=\mathbf{x}\right].$$

因此对任意 ${\mathbf{s}}_{\theta }$，

$$\begin{array}{rl}& {\mathbb{E}}_{{\mathbf{X}}_t\sim p_t}\left[⟨{\mathbf{s}}_{\theta }({\mathbf{X}}_t,t),\nabla \log p_t({\mathbf{X}}_t)⟩\right]\\ & ={\mathbb{E}}_{{\mathbf{X}}_0,{\mathbf{X}}_t}\left[⟨{\mathbf{s}}_{\theta }({\mathbf{X}}_t,t),{\nabla }_{{\mathbf{x}}_t}\log p_{0t}({\mathbf{X}}_t\mid {\mathbf{X}}_0)⟩\right].\end{array}$$

展开两个平方目标：

$$\begin{array}{rl}\frac{1}{2}{\mathbb{E}}_{p_t}\Vert {\mathbf{s}}_{\theta }-\nabla \log p_t{\Vert }^2=& \frac{1}{2}{\mathbb{E}}_{p_t}\Vert {\mathbf{s}}_{\theta }{\Vert }^2-{\mathbb{E}}_{p_t}⟨{\mathbf{s}}_{\theta },\nabla \log p_t⟩\\ & +\frac{1}{2}{\mathbb{E}}_{p_t}\Vert \nabla \log p_t{\Vert }^2,\end{array}$$ $$\begin{array}{rl}\frac{1}{2}{\mathbb{E}}_{p_0{p}_{0t}}\Vert {\mathbf{s}}_{\theta }-\nabla \log p_{0t}{\Vert }^2=& \frac{1}{2}{\mathbb{E}}_{p_t}\Vert {\mathbf{s}}_{\theta }{\Vert }^2\\ & -{\mathbb{E}}_{p_0{p}_{0t}}⟨{\mathbf{s}}_{\theta },\nabla \log p_{0t}⟩\\ & +\frac{1}{2}{\mathbb{E}}_{p_0{p}_{0t}}\Vert \nabla \log p_{0t}{\Vert }^2.\end{array}$$

前两项逐一相同，最后一项都与 $\theta$ 无关。因此

$${\mathcal{J}}_{\mathrm{DSM}}(\theta )={\mathcal{J}}_{\mathrm{score}}(\theta )+C$$

且 $C$ 与 $\theta$ 无关；二者具有相同的参数梯度和最优解。这个证明不声称两个目标的数值相等。

若改用经典 Hyvärinen score matching，把目标化为

$${\mathbb{E}}_{p_t}\left[\nabla \cdot {\mathbf{s}}_{\theta }+\frac{1}{2}\Vert {\mathbf{s}}_{\theta }{\Vert }^2\right],$$

还需额外假设边界通量
$p_t(\mathbf{x}){\mathbf{s}}_{\theta }(\mathbf{x},t)$ 在无穷远消失，才能合法分部积分。

4. 线性高斯核下的参数化

设

$${\mathbf{X}}_t=a_t{\mathbf{X}}_0+b_t\mathbf{ϵ},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}),b_t>0.$$

则

$$p_{0t}({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\mathcal{N}(a_t{\mathbf{x}}_0,b_t^2\mathbf{I})$$

且条件分数为

$$\boxed{{\nabla }_{{\mathbf{x}}_t}\log p_{0t}({\mathbf{x}}_t\mid {\mathbf{x}}_0)=-\frac{{\mathbf{x}}_t-a_t{\mathbf{x}}_0}{b_t^2}=-\frac{\mathbf{ϵ}}{b_t}}.$$

4.1 score 参数化

网络直接输出

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

DSM 单样本损失为

$$\lambda (t){\Vert {\mathbf{s}}_{\theta }({\mathbf{x}}_t,t)+\frac{\mathbf{ϵ}}{b_t}\Vert }_2^2.$$

4.2 epsilon 参数化

定义

$${\mathbf{s}}_{\theta }({\mathbf{x}}_t,t):=-\frac{{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)}{b_t}.$$

则

$$\lambda (t){\Vert {\mathbf{s}}_{\theta }+\frac{\mathbf{ϵ}}{b_t}\Vert }^2=\frac{\lambda (t)}{b_t^2}{\Vert {\mathbf{ϵ}}_{\theta }-\mathbf{ϵ}\Vert }^2.$$

所以“不加权 epsilon MSE”对应的是分数 DSM 权重

$$\lambda (t)=b_t^2,$$

而不是所有权重选择都自动等价于同一个无权 epsilon 目标。

4.3 $x_0$ 参数化

若网络预测 ${\hat{\mathbf{x}}}_{0,\theta }$，则

$${\mathbf{ϵ}}_{\theta }=\frac{{\mathbf{x}}_t-a_t{\hat{\mathbf{x}}}_{0,\theta }}{b_t},$$ $${\mathbf{s}}_{\theta }=\frac{a_t{\hat{\mathbf{x}}}_{0,\theta }-{\mathbf{x}}_t}{b_t^2}.$$

因此 score、epsilon 和 $x_0$ 参数化在 $a_t,b_t\ne 0$ 时可相互转换，但损失权重会随转换改变。

5. 引导目标的 epsilon 形式

若希望在无条件分数上加入一个给定向量场

$$\mathbf{h}({\mathbf{x}}_t,t,\mathbf{y})\approx {\nabla }_{{\mathbf{x}}_t}\log p_t(\mathbf{y}\mid {\mathbf{x}}_t),$$

则目标分数是

$$-\frac{\mathbf{ϵ}}{b_t}+\mathbf{h}.$$

采用 ${\mathbf{s}}_{\theta }=-{\mathbf{ϵ}}_{\theta }/b_t$ 时，

$${\Vert -\frac{{\mathbf{ϵ}}_{\theta }}{b_t}+\frac{\mathbf{ϵ}}{b_t}-\mathbf{h}\Vert }^2=\frac{1}{b_t^2}{\Vert {\mathbf{ϵ}}_{\theta }-\mathbf{ϵ}+b_t\mathbf{h}\Vert }^2.$$

例如启发式 loss guidance 取
$\mathbf{h}={\nabla }_{{\mathbf{x}}_t}[-\ell ({\mathbf{x}}_t,\mathbf{y})]$，则 epsilon 目标为

$$\mathbf{ϵ}-b_t\mathbf{h}=\mathbf{ϵ}+b_t{\nabla }_{{\mathbf{x}}_t}\ell .$$

该代数关系本身是精确的，但 $\mathbf{h}$ 是否真是条件似然分数取决于引导方法的近似；详见
条件扩散。

6. 训练伪代码

算法：无条件扩散模型训练

1. 选定时间分布 $\pi$、权重 $\lambda$ 和扰动核
   $p_{0t}$。
2. 重复直到收敛：
   1. 采样 ${\mathbf{x}}_0\sim p_{\mathrm{data}}$；
   2. 采样 $t\sim \pi$；连续时间通常限制 $t\ge \epsilon$；
   3. 采样 $\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$；
   4. 构造
      ${\mathbf{x}}_t=a_t{\mathbf{x}}_0+b_t\mathbf{ϵ}$；
   5. 计算 $$\lambda (t){\Vert {\mathbf{s}}_{\theta }({\mathbf{x}}_t,t)+\frac{\mathbf{ϵ}}{b_t}\Vert }^2$$ 或其等价的正确加权 epsilon 形式；
   6. 反向传播并更新 $\theta$。

7. 逆向 SDE 采样伪代码

在网格

$$0=t_0<t_1<\cdots <t_N=T$$

上，从终端向起点积分。令

$$\Delta t_i:=t_{i-1}-t_i<0.$$

算法：Euler—Maruyama 逆向采样

1. 采样 ${\mathbf{x}}_{t_N}\sim p_{\mathrm{prior}}$。这里应使用实际终端先验；只有模型定义如此时才写成
   $\mathcal{N}(0,\mathbf{I})$。
2. 对 $i=N,N-1,\dots ,1$：
   1. 计算逆向漂移 $${\mathbf{f}}_{\mathrm{rev}}=\mathbf{f}({\mathbf{x}}_{t_i},t_i)-g(t_i{)}^2{\mathbf{s}}_{\theta }({\mathbf{x}}_{t_i},t_i);$$
   2. 采样 $$\Delta {\overline{\mathbf{W}}}_i\sim \mathcal{N}(0,|\Delta t_i|\mathbf{I});$$
   3. 更新 $${\mathbf{x}}_{t_{i-1}}={\mathbf{x}}_{t_i}+{\mathbf{f}}_{\mathrm{rev}}\Delta t_i+g(t_i)\Delta {\overline{\mathbf{W}}}_i.$$
3. 返回 ${\mathbf{x}}_{t_0}$。

若使用 probability-flow ODE，则删除随机增量，并把分数项系数从 $g^2$ 改为 $\frac{1}{2}{g}^2$。两种采样器只共享边缘路径，不共享单条轨迹。

Vincent, P. (2011). A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7), 1661–1674.