Aapo Hyvärinen, Estimation of Non-Normalized Statistical Models by Score Matching (2005)
本文收录于 泛 AI 论文选读。
TL;DR
给定若干样本，拟合其概率分布的 PDF 的时候，归一化因子通常是难以计算的，这导致直接的最大似然估计难以进行。此时可以通过优化一个能用采样计算的目标 $\ref{J-function}$ 拟合该 PDF 的 score function。
引入
考虑拟合一个随机向量 $\boldsymbol{x}\in \mathbb{R}^n$ 的 PDF（概率密度函数）$p_{\boldsymbol{x}}(.)$。我们需要求出一个参数化的 PDF $p(.; \boldsymbol{\theta})$。
这里问题是，我们的模型通常是非归一化的。也就是说我们只能够计算某种 $q(\boldsymbol{\xi}, \boldsymbol{\theta})$，但是实际上
$$
p(\boldsymbol{\xi}; \boldsymbol{\theta}) = \frac{1}{Z(\boldsymbol{\theta})}q(\boldsymbol{\xi}; \boldsymbol{\theta})
$$
其中归一化因子 $Z(\boldsymbol{\theta})$ 在高维空间下通常没有解析解，数值上也难以计算：
$$
Z(\boldsymbol{\theta}) = \int_{\boldsymbol{\xi}\in \mathbb{R}^n} q(\boldsymbol{\xi}; \boldsymbol{\theta})\mathrm{d}\boldsymbol{\xi}
$$
所以这时，给定一系列样本，用随机梯度下降之类的方法来优化这些样本的最大似然就有些困难。文章提出了一种新的方法。
Score Matching
直觉上，如果我们让 $\log p(.; \boldsymbol{\theta})$ 对齐 $\log p_{\boldsymbol{x}}(.)$ 的梯度，那么参数化的分布也将与欲拟合之分布对齐。现在定义 score function：
$$
\boldsymbol{\psi}(\boldsymbol{\xi}; \boldsymbol{\theta}) = \nabla_{\boldsymbol{\xi}}\log p(\boldsymbol{\xi}; \boldsymbol{\theta}), \quad \boldsymbol{\psi}_{\boldsymbol{x}} = \nabla_{\boldsymbol{\xi}} \log p_{\boldsymbol{x}}(\boldsymbol{\xi})
$$
注意 $Z(\boldsymbol{\theta})$ 对上述梯度的贡献为 $0$。现在定义一个易于用采样计算的损失
$$
J(\boldsymbol{\theta}) = \frac{1}{2}\int_{\boldsymbol{\xi}\in\mathbb{R}^n} p_{\boldsymbol{x}}(\boldsymbol{\xi}) \lVert \boldsymbol{\psi}(\boldsymbol{\xi}; \boldsymbol{\theta}) - \boldsymbol{\psi}_{\boldsymbol{x}}(\boldsymbol{\xi})\rVert^2 \mathrm{d}\boldsymbol{\xi} \label{J-function}
$$
在能够知道 $\boldsymbol{\psi}_{\boldsymbol{x}}$ 的情况下直接优化上述函数，称为 Explicit Score Matching。但注意现在我们无法准确计算 $\boldsymbol{\psi}_{\boldsymbol{x}}$。
可以用核密度估计之类的方法来近似 $\boldsymbol{\psi}_{\boldsymbol{x}}$。但是直觉上，在高维空间这种做法效果会非常差。
幸运的是：
定理 1. 若 score function 和欲拟合的分布性质足够好，则有
$$
J(\boldsymbol{\theta}) = {\color{lightgrey} const} + \int_{\boldsymbol{\xi}\in \mathbb{R}^n} p_{\boldsymbol{x}}(\boldsymbol{\xi})\sum_{i=1}^n \left[\partial_i \psi_i(\boldsymbol{\xi}; \boldsymbol{\theta}) + \frac 12\psi_i(\boldsymbol{\xi}; \boldsymbol{\theta})^2\right]\mathrm{d}\boldsymbol{\xi}
$$
其中 $const$ 是一个和 $\boldsymbol{\theta}$ 无关的常数。
性质要求细节
必要的条件包括：
- $\boldsymbol{\psi}(\boldsymbol{\xi}; \boldsymbol{\theta})$ 可微。
- $\mathbb{E}[\lVert \boldsymbol{\psi}(\boldsymbol{\xi}; \boldsymbol{\theta})\rVert^2]$ 有限，$\mathbb{E}[\lVert \boldsymbol{\psi}(\boldsymbol{\xi}; \boldsymbol{\theta})\rVert^2]$ 有限。
- 当 $\lVert \boldsymbol{\xi}\rVert \rightarrow \infty$ 时 $p_{\boldsymbol{x}}(\boldsymbol{\xi})\boldsymbol{\psi}(\boldsymbol{\xi}; \boldsymbol{\theta}) \rightarrow 0$。
- 可能需要一些一致连续性之类的东西，但是都来做 AI 了，当然这些都是默认的（逃
证明.
首先将 $\ell_2$ 范数展开，以下两种项是无需考虑的：
- 和 $\boldsymbol{\theta}$ 无关的项（积分后为常数）。
- 和 $p_{\boldsymbol{x}}(\boldsymbol{\xi})$ 无关的项（可以直接采样）。
只需要考虑计算：
$$
\int_{\boldsymbol{\xi}\in \mathbb{R}^n}p_{\boldsymbol{x}}(\boldsymbol{\xi})\sum_{i=1}^n \psi_i(\boldsymbol{\xi}; \boldsymbol{\theta})\psi_{\boldsymbol{x}, i}(\boldsymbol{\xi})\mathrm{d}\boldsymbol{\xi}
$$
求和是有限的，所以不妨单独考察求和中的一项，不妨设就是第 $n$ 项：
$$
\begin{align}
&\int_{\boldsymbol{\xi}\in \mathbb{R}^n}p_{\boldsymbol{x}}(\boldsymbol{\xi})\psi_n(\boldsymbol{\xi}; \boldsymbol{\theta})\psi_{\boldsymbol{x}, n}(\boldsymbol{\xi})\mathrm{d}\boldsymbol{\xi} \\
&= \int_{\boldsymbol{\xi}\in \mathbb{R}^n}p_{\boldsymbol{x}}(\boldsymbol{\xi})\psi_n(\boldsymbol{\xi}; \boldsymbol{\theta})\frac{\partial \log p_{\boldsymbol{x}}(\boldsymbol{\xi})}{\partial \xi_n}\mathrm{d}\boldsymbol{\xi} \\
&= \int_{\boldsymbol{\xi}\in \mathbb{R}^n}\psi_n(\boldsymbol{\xi}; \boldsymbol{\theta})\frac{\partial p_{\boldsymbol{x}}(\boldsymbol{\xi})}{\partial \xi_n}\mathrm{d}\boldsymbol{\xi} \\
&= \iint\cdots\int \mathrm{d}\xi_1…\xi_{n - 1} \int_{-\infty}^{\infty} \psi_n(\boldsymbol{\xi}; \boldsymbol{\theta})\frac{\partial p_{\boldsymbol{x}}(\boldsymbol{\xi})}{\partial \xi_n}\mathrm{d}\xi_n \\
&= \iint\cdots\int \mathrm{d}\xi_1…\xi_{n - 1} \left(\left[p_\boldsymbol{x}(\boldsymbol{\xi})\boldsymbol{\psi}(\boldsymbol{\xi}; \boldsymbol{\theta})\right]_{\xi_n = -\infty}^{\xi_n = \infty} - \int_{\infty}^\infty p_{\boldsymbol{x}}(\boldsymbol{\xi})\partial_n\psi_n(\boldsymbol{\xi}; \boldsymbol{\theta})\right)\mathrm{d}\xi_n \\
&= \int_{\boldsymbol{\xi}\in\mathbb{R}^n} p_\boldsymbol{x}(\boldsymbol{\xi})\partial_i\psi_i(\boldsymbol{\boldsymbol{\xi}; \boldsymbol{\theta}})\mathrm{d}\boldsymbol{\xi}
\end{align}
$$
代回整理即可。
撇去和优化无关的部分，此函数容易使用采样估计：
$$
\tilde{J}(\boldsymbol{\theta}) = \frac 1T\sum_{t = 1}^T\sum_{i=1}^n\left[\partial_i \psi_i(\boldsymbol{x}_t; \boldsymbol{\theta}) + \frac 12 \psi_i(\boldsymbol{x}(t); \boldsymbol{\theta})^2\right]
$$
关于当参数化模型的表达范围包含 $p_\boldsymbol{x}$，且仅有一个参数 $\boldsymbol{\theta}^*$ 能使得参数化模型等于 $p_\boldsymbol{x}$，则 $\arg \min_{\boldsymbol{\theta}} J(\boldsymbol{\theta}) = \boldsymbol{\theta}^*$ 的结论，非常显然，此处不再重复。
这种做法被后人称为 Implicite Score Matching。
*Denoising Score Matching
这部分原文中没有。
注意 score matching 的关键难点是 $\boldsymbol{\psi}_{\boldsymbol{x}}$ 无法计算。上述方法通过类似于分部积分的办法成功推出原来的优化目标等价于一个不含 $\boldsymbol{\psi}_{\boldsymbol{x}}$，但是推出的东西需要零阶优化或者高阶导数计算，计算复杂度非常糟糕。
另外一种方法是做一些妥协，给原分布混上一个无偏的小方差噪声 $q(\boldsymbol{\tilde{x}} | \boldsymbol{x})$（比如 $\mathcal{N}(0, 10^{-6}\mathbf{I})$），然后优化
$$
\frac 12\mathbb{E}_{\boldsymbol{x}\sim p_{data}}\mathbb{E}_{\boldsymbol{\tilde{x}}\sim q(\boldsymbol{\tilde{x}} | \boldsymbol{x})} [\lVert\boldsymbol{\psi}(\boldsymbol{\tilde{x}}; \boldsymbol{\theta}) - \nabla_{\boldsymbol{\tilde x}}\log q(\boldsymbol{\tilde{x}} | \boldsymbol{x})\rVert^2]
$$
只需要注意到
$$
\log q(\boldsymbol{\tilde{x}}, \boldsymbol{x}) = \log q(\boldsymbol{x}) + \log q(\boldsymbol{\tilde{x}} | \boldsymbol{x})
$$
而前者是没有梯度的。另外，回忆我们有全期望公式
$$
\begin{aligned}
\mathbb{E}_{x\sim p(x)}[f(x)] &= \int p(x) f(x) \mathrm{d}x \\
&= \iint p(x, y) f(x) \mathrm{d}x\mathrm{d}y \\
&= \int p(y)\mathrm{d}y \int p(x | y)f(x)\mathrm{d}x \\
&= \mathbb{E}_{y\sim p(y)}[\mathbb{E}_{x\sim p(x | y)}[f(x)]]
\end{aligned}
$$
就知道优化此函数实际上是在优化
$$
\frac 12\mathbb{E}_{\boldsymbol{\tilde{x}}\sim q(\boldsymbol{\tilde{x}})}[\lVert\boldsymbol{\psi}(\boldsymbol{\tilde{x}}; \boldsymbol{\theta}) - \nabla_{\boldsymbol{\tilde x}}\log q(\boldsymbol{\tilde{x}})\rVert^2]
$$
这正是 $q$ 的 ESM。
*Sliced Score Matching
这部分原文中没有。
可以用
$$
\mathrm{tr}[\mathbf{A}] = \mathbb{E}_{\boldsymbol{v}\sim \mathcal{N}(0, \mathbf{I})}[\boldsymbol{v}^\top\mathbf{A}\boldsymbol{v}]
$$
来无偏估计矩阵的迹。注意
$$
\boldsymbol{v}^\top \nabla_{\boldsymbol{x}}\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}) = \nabla_{\boldsymbol{x}}\boldsymbol{v}^\top \boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x})
$$
后者是一个标量，这似乎优化了使用自动微分求解梯度的次数（原来为了求迹，需要求 $m$ 次梯度）