Score matching model

简介

这篇文章主要描述基于得分匹配（Score matching model）的想法，以及之后主要的修改思路。这种思路是生成模型的一种，与GAN、normal-flow等模型具备同样的功能。

本篇文章大量借鉴棒棒生博客，推荐阅读原文博客。本文章在其基础上加入一些作者本人的思考，并且统一符号，增加阅读流畅性。

得分匹配算法（Score matching model）

得分匹配这种思想首先发表于论文—Estimation of Non-Normalized Statistical Models
by Score Matching。

拟合目标

图像、声音、热力学过程等在数学形式上应该是一种概率分布，想要模仿生成相应的事物需要的就是能够得知这种概率分布。但是这种概率分布往往在一个高维空间，难以凭借人类的经验直接获得，因此此时借助神经网络的拟合能力，完成概率密度分布函数的拟合。

利用 $\bf{x}$ 生成一个随机的目标概率密度分布 $p_{\bf{x}}(\cdot)$，接下来尝试拟合该目标函数，拟合函数的概率密度分布记为$p(\cdot;\theta)$，其中$\theta$是一个m维空间的向量。目标是从$\bf{x}$中估计$\theta$。

将拟合函数$p(\cdot;\theta)$进一步分解开来：

$$p(\xi;\theta)=\frac{1}{Z(\theta)}q(\xi;\theta)$$

其中：$Z(\theta)=\int_{\xi\in \mathbb{R}^n}q(\xi;\theta)\mathrm{d}\xi$。这样只需要生成函数$q(\cdot;\theta)$，不需要关心归一化的问题。

通常非归一化采样的方法是马尔科夫链蒙特卡洛模拟。典型的是Ising模型的模拟过程。在得到一系列数据$\{x_1,x_2, x_3…\}$后，使用极大似然估计方法（Maximum Likelihood Estimation, MLE）估计$\theta$：

$$\theta_{MLE}=\text{argmax}_{\theta}\sum_{t=1}^T \text{ln} p(x_t;\theta)$$

然而在实际中这种方案是行不通的，因为配分函数$Z(\theta)$难以计算。计算配分函数是一块难啃的骨头，许多物理学上的困难本质就是来源于如何高效计算配分函数，因此很多方法例如高温展开、信念传播算法等本质就是在解决如何在牺牲一些精确性的条件下，高效计算配分函数。

直接计算存在困难，并且配分函数作为一个仅仅为$\theta$的函数存在，那就转而计算梯度，试图将配分函数在求梯度的过程中磨消。定义：

$$\psi(\xi;\theta)= \left( \begin{matrix} \frac{\partial ln p(\xi;\theta)}{\partial \xi_1} \\ \cdots \\ \frac{\partial ln p(\xi;\theta)}{\partial \xi_n} \end{matrix} \right) = \left( \begin{matrix} \psi_1 (\xi;\theta) \\ \cdots \\ \psi_n (\xi;\theta) \end{matrix} \right) =\nabla_\xi ln p(\xi;\theta)$$

本质上求偏导并不依赖于$Z(\theta)$，可以获得：

$$\begin{equation} \psi(\xi;\theta)=\nabla_\xi lnq(\xi;\theta) \end{equation}$$

上方是一个分布的梯度图像，可以看出如果梯度方向大小均相同的情况下，两个应该是同一种分布。这里可以通过分析下面的函数得到，如果$J(\theta)=0$,由于$p_x(\xi)\geq0$，必然存在$\psi(\xi;\theta)-\psi_x(\xi)\equiv 0$，则说明$\theta=\theta^*$，成功学习到分布。这也证明了解的存在唯一性。

直观的目标函数（Explicit Score Matching，ESM）

将任务目标函数写为考虑两个函数的梯度的均方误差(MSE)：

$$\begin{equation} J_{ESM}(\theta)=\frac{1}{2} \int_{\xi\in\mathbb{R}^n}p_x(\xi)||\psi(\xi;\theta)-\psi_x(\xi)||^2\mathrm{d}\xi \end{equation}$$

其中$\psi_x(\xi)=\nabla_\xi lnp_x(\xi)$，为目标概率密度分布的梯度分布。优化目标为：

$$\hat\theta=\text{argmax}_\theta J(\theta)$$

可以通过重要性采样得到$T$个数据点$\bf{x}(1), \bf{x}(2) \cdots\bf{x}(T)$，从而计算经验期望值：

$$\begin{equation} \tilde{J}_{ESM}(\theta)=\frac{1}{T}\sum_{t=1}^T\sum_{i=1}^n[\psi(\xi;\theta)-\psi_x(\xi)]^2 \end{equation}$$

隐藏的目标函数（Implicit Score Matching，ISM）

在计算均方误差的时候，需要知道$\psi_x(\xi)$，然而知道确切的目标函数形式是不可能的。因此，虽然以上逻辑是合理的，但是本质上是不计算实施的，需要进行一定的调整。

$$\begin{align*} J(\theta)&=\frac{1}{2} \int_{\xi\in\mathbb{R}^n}p_x(\xi)||\psi(\xi;\theta)-\psi_x(\xi)||^2\mathrm{d}\xi \\ &=\int_{\xi\in\mathbb{R}^n}p_x(\xi)[\frac{1}{2} \psi(\xi;\theta)^2+\frac{1}{2} \psi_x(\xi)^2-\psi(\xi;\theta)\psi_x(\xi)]\mathrm{d}\xi \\ \end{align*}$$

针对后一部分结果进行简化，将$\psi$展开为具体的第$i$部分：

$$\begin{align*} &\int_{\xi\in\mathbb{R}^n}-p_x(\xi)\psi(\xi;\theta)\psi_x(\xi)\mathrm{d}\xi \\ &= -\sum_i\int_{\xi\in\mathbb{R}^n}p_x(\xi)\psi_i(\xi;\theta)\psi_{x,i}(\xi)\mathrm{d}\xi \\ &=-\sum_i\int_{\xi\in\mathbb{R}^n}p_x(\xi)\psi_i(\xi;\theta)\frac{\partial\psi_{x}(\xi)}{\partial \xi_i} \mathrm{d}\xi \\ &=-\sum_i\int_{\xi\in\mathbb{R}^n}p_x(\xi)\psi_i(\xi;\theta)\frac{\partial ln p_x(\xi)}{\partial \xi_i} \mathrm{d}\xi \\ &=-\sum_i\int_{\xi\in\mathbb{R}^n}\psi_i(\xi;\theta)\frac{\partial p_x(\xi)}{\partial \xi_i} \mathrm{d}\xi \\ &=-\sum_i\left[ \psi_i(\xi;\theta)p_x(\xi) |_{\xi_i=-\infty}^{\xi_i=\infty} - \int_{\xi\in\mathbb{R}^n} p_x(\xi)\frac{\partial \psi_i(\xi;\theta)}{\partial \xi_i} \mathrm{d}\xi \right]\\ \end{align*}$$

将其中第一部分进行更详细的展开：

$$\begin{align*} \psi_1(\xi;\theta)p_x(\xi) |_{\xi_1=-\infty}^{\xi_1=\infty} &= \lim\limits_{a\to \infty,b\to\infty}[\psi(a,\xi_2\cdots\xi_n;\theta)p_x(a,\xi_2\cdots\xi_n;\theta)-\psi(b,\xi_2\cdots\xi_n;\theta)p_x(b,\xi_2\cdots\xi_n;\theta)] \\ &= 0 \end{align*}$$

后面的原因来源于一个简单的假设，当$||\xi||\to\infty$：

$$p_{x}(\xi)\psi(\xi;\theta)\to0$$

可以这样理解该假设，当采样数量巨大或者分布为连续，单一采样所占概率分布很小，可以视为直接为0。因此得到最后的结果：

$$\begin{align*} \int_{\xi\in\mathbb{R}^n}-p_x(\xi)\psi(\xi;\theta)\psi_x(\xi)\mathrm{d}\xi &=\sum_i \int_{\xi\in\mathbb{R}^n} p_x(\xi)\frac{\partial \psi_i(\xi;\theta)}{\partial \xi_i} \mathrm{d}\xi\\ &=\int_{\xi\in\mathbb{R}^n} p_x(\xi)\partial_{i} \psi_i(\xi;\theta) \mathrm{d}\xi\\ \end{align*}$$ $$\begin{align*} J(\theta)&=\frac{1}{2} \int_{\xi\in\mathbb{R}^n}p_x(\xi)||\psi(\xi;\theta)-\psi_x(\xi)||^2\mathrm{d}\xi \\ &=\int_{\xi\in\mathbb{R}^n}p_x(\xi)[\frac{1}{2} \psi(\xi;\theta)^2+\frac{1}{2} \psi_x(\xi)^2-\psi(\xi;\theta)\psi_x(\xi)]\mathrm{d}\xi \\ &=\int_{\xi\in\mathbb{R}^n}p_x(\xi)[\frac{1}{2} \psi(\xi;\theta)^2+\frac{1}{2} \psi_x(\xi)^2+\partial_{i} \psi_i(\xi;\theta)]\mathrm{d}\xi \\ \end{align*}$$

其中第二项为目标函数的积分，结果应该为一个常数，对于优化目标函数来说一个常数并没有实际意义，因此可以直接消去：

$$\begin{align} J_{ISM}(\theta) &= \int_{\xi\in\mathbb{R}^n}p_x(\xi)[\frac{1}{2} \psi(\xi;\theta)^2+\partial_{i} \psi_i(\xi;\theta)]\mathrm{d}\xi \\ &= \mathbb{E}_{p_x(\xi)}[\frac{1}{2} \psi(\xi;\theta)^2+\text{tr}(\nabla_{\xi} \psi_{\xi}(\xi;\theta))] \\ \tilde{J}_{ISM}(\theta) &= \frac{1}{T}\sum_{t=1}^T\sum_{i=1}^n[\frac{1}{2} \psi(x_t;\theta)^2+\partial_{i} \psi_i(x_t;\theta)] \end{align}$$

总结

首先通过概率密度分布拟合明确了任务目标，通过调节参数$\theta$使得$p_{\bf X}(\cdot)$与$p(\cdot;\theta)$分布尽可能相同。但是直接计算$p(\cdot;\theta)$存在困难，因此将任务目标进行改变，转而寻求概率分布密度函数的梯度尽可能相似。然后求解公式中包含目标函数，无法直接求解，通过分布积分的方法转化为只需要拟合函数便可以求解。

然而在计算ISM的时候存在两点问题：

• 需要计算二次导数，这会导致计算图大小升高
• 当输入数据维度很大的时候很难处理

噪声得分匹配（Denoising Score Matching，DSM）

本算法的思想来源于A Connection Between Score Matching and Denoising Autoencoders。

DSM的提出，可以很好的解决ISM存在的两个问题。

噪声

将目标函数加入噪声 $p_{ x_\sigma }$，其中$\sigma$表示噪声，那么ESM与ISM可以改写为：

$$\begin{align} J_{ESM,p_{\sigma}}(\theta) &=\mathbb{E}_{p_\sigma(\xi)}\left[ \frac{1}{2} ||\psi(\xi;\theta)-\psi_{x_\sigma}(\xi)||^2\right] \\ &= \mathbb{E}_{p_\sigma(\xi)}\left[ \frac{1}{2} ||\psi(\xi;\theta)-\frac{\partial ln p_{\sigma}(\xi)}{\partial \xi}||^2\right] \\ J_{ISM,p_\sigma}(\theta) &= \mathbb{E}_{p_\sigma(\xi)}\left[ \frac{1}{2} \psi(\xi;\theta)^2+\text{tr}(\nabla_{\xi} \psi (\xi;\theta))\right] \\ \end{align}$$

这两个式子在变化前后，形式基本一样。区别的地方那个在于学习的目标函数区别，后者学习的目标函数是包含噪音之后的目标函数。而且可以看出，在变化之后，同样计算二阶导数，计算难度依然没有发生改变。

引入联合概率密度分布$p_\sigma(\tilde{\bf x}, {\bf x})=p_\sigma(\tilde{\bf x}|{\bf x})p_0({\bf x})$，其中$\tilde x$是在准确分布$x$基础上的随机取值，定义噪音得分匹配(DSM)：

$$J_{DSM,p_{\sigma}}(\theta)=\mathbb{E}_{p_\sigma(\tilde x,x)}\left[ \frac{1}{2} ||\psi(\tilde x;\theta)-\frac{\partial ln p_{\sigma}(\tilde x|x)}{\partial \tilde x}||^2\right]$$

假设为高斯核函数关系$p_\sigma(\tilde x|x)=e^{-\frac{||x-\tilde x||^2}{2\sigma^2}}$：

$$\frac{\partial p_{\sigma}(\tilde x|x)}{\partial \tilde x} = \frac{x-\tilde x}{\sigma^2}$$

因此：

$$\begin{equation} J_{DSM,p_{\sigma}}(\theta)=\mathbb{E}_{p_\sigma(\tilde x,x)}\left[ \frac{1}{2} ||\psi(\tilde x;\theta)-\frac{x-\tilde x}{\sigma^2}||^2\right] \end{equation}$$

可以发现，在通过引入噪声函数之后，成功避免了二次导数的计算。然而DSM的引入是直接定义，并没有证明ESM的等价性。

等价性证明

$$\begin{align*} J_{ESM,p_{\sigma}}(\theta) &= \mathbb{E}_{p_\sigma(\tilde x)}\left[ \frac{1}{2} ||\psi(\tilde x;\theta)-\frac{\partial ln p_{\sigma}(\tilde x)}{\partial \tilde x}||^2\right] \\ &= \mathbb{E}_{p_\sigma(\tilde x)}\left[ \frac{1}{2} ||\psi(\tilde x;\theta)||^2\right]-S(\theta)+C_2\\ \end{align*}$$

其中$C_2=\mathbb{E}_{p_\sigma(\tilde x)}\left[ \frac{1}{2}||\frac{\partial ln p_{\sigma}(\tilde x)}{\partial \tilde x}||^2 \right]$，并不依赖于参数$\theta$，可以直接忽略。

$$\begin{align*} S(\theta) &= \mathbb{E}_{p_\sigma(\tilde x)} \left[ <\psi(\tilde{\bf x};\theta),\frac{\partial lnp_\sigma(\tilde{\bf x})}{\partial \tilde{\bf x}}> \right] \\ &= \int_{\tilde{\bf x}} p_\sigma(\tilde{\bf x})<\psi(\tilde{\bf x};\theta),\frac{\partial lnp_\sigma(\tilde{\bf x})}{\partial \tilde{\bf x}}> \mathrm{d}\tilde{\bf x} \\ &= \int_{\tilde{\bf x}} p_\sigma(\tilde{\bf x})<\psi(\tilde{\bf x};\theta),\frac{1}{p_\sigma(\tilde{\bf x})} \frac{\partial p_\sigma(\tilde{\bf x})}{\partial \tilde{\bf x}}> \mathrm{d}\tilde{\bf x} \\ &= \int_{\tilde{\bf x}} <\psi(\tilde{\bf x};\theta), \frac{\partial p_\sigma(\tilde{\bf x})}{\partial \tilde{\bf x}}> \mathrm{d}\tilde{\bf x} \\ &= \int_{\tilde{\bf x}} <\psi(\tilde{\bf x};\theta), \frac{\partial \int_{\bf x} p_\sigma(\tilde{\bf x}|{\bf x})p_\sigma({\bf x})\mathrm{d}{\bf x}}{\partial \tilde{\bf x}}> \mathrm{d}\tilde{\bf x} \\ &= \int_{\tilde{\bf x}} \int_{\bf x} p_\sigma({\bf x})<\psi(\tilde{\bf x};\theta), \frac{\partial p_\sigma(\tilde{\bf x}|{\bf x})}{\partial \tilde{\bf x}}> \mathrm{d}{\bf x} \mathrm{d}\tilde{\bf x} \\ &= \int_{\tilde{\bf x}} \int_{\bf x} p_\sigma({\bf x})p_\sigma(\tilde{\bf x}|{\bf x})<\psi(\tilde{\bf x};\theta), \frac{\partial ln p_\sigma(\tilde{\bf x}|{\bf x})}{\partial \tilde{\bf x}}> \mathrm{d}{\bf x} \mathrm{d}\tilde{\bf x} \\ &= \int_{\tilde{\bf x}} \int_{\bf x} p_\sigma(\tilde{\bf x},{\bf x})<\psi(\tilde{\bf x};\theta), \frac{\partial ln p_\sigma(\tilde{\bf x}|{\bf x})}{\partial \tilde{\bf x}}> \mathrm{d}{\bf x} \mathrm{d}\tilde{\bf x} \\ &= \mathbb{E}_{p_\sigma(\tilde x, x)} \left[ <\psi(\tilde{\bf x};\theta),\frac{\partial lnp_\sigma(\tilde{\bf x}|{\bf x})}{\partial \tilde{\bf x}}> \right] \\ \end{align*}$$

从而得到：

$$\begin{align*} J_{ESM,p_{\sigma}}(\theta) &= \mathbb{E}_{p_\sigma(\tilde x)}\left[ \frac{1}{2} ||\psi(\tilde x;\theta)||^2\right]-\mathbb{E}_{p_\sigma(\tilde x, x)} \left[ <\psi(\tilde{\bf x};\theta),\frac{\partial lnp_\sigma(\tilde{\bf x}|{\bf x})}{\partial \tilde{\bf x}}> \right]+C_2\\ \end{align*}$$

另一方面，针对DSM进行分解：

$$\begin{align*} J_{DSM,p_{\sigma}}(\theta)&=\mathbb{E}_{p_\sigma(\tilde x,x)}\left[ \frac{1}{2} ||\psi(\tilde x;\theta)-\frac{\partial ln p_{\sigma}(\tilde x|x)}{\partial \tilde x}||^2\right] \\ &= \mathbb{E}_{p_\sigma(\tilde x)}\left[ \frac{1}{2} ||\psi(\tilde x;\theta)||^2 \right]-\mathbb{E}_{p_\sigma(\tilde x, x)} \left[ <\psi(\tilde{\bf x};\theta),\frac{\partial lnp_\sigma(\tilde{\bf x}|{\bf x})}{\partial \tilde{\bf x}}> \right]+C_3 \end{align*}$$

其中$C_3=\mathbb{E}_{p_\sigma(\tilde x, x)}\left[\frac{1}{2}||\frac{\partial ln p_{\sigma}(\tilde x|x)}{\partial \tilde x}||^2\right]$，同样不依赖于$\theta$参数，可以认为是一个常量。因此可以得到DSM与ESM的关系：

$$\begin{align} J_{DSM,p_{\sigma}}(\theta) = J_{ESM,p_{\sigma}}(\theta)-C_3+C_2 \end{align}$$

其中$C_2, C_3$与优化目标$\theta$无关，因此可以直接忽略。因此可以证明两者是等价的。

总结与分析

相对于ESM方法，引入噪音简化二阶偏导为一个关于方差$\sigma$的函数。并且成功证明，两者等价。

然而这个方法存在一些问题：

• 学习到的是加上噪音之后的分布，不是原分布
• 加入的方差$\sigma$很难控制调整

切片匹配得分（Sliced Score Matching，SSM）

来源于文章——Sliced Score Matching: A Scalable Approach to Density and Score Estimation。

降维思想

直接对梯度进行比较，会因为维度较高产生问题，可以尝试通过一个函数$V$，将高维度映射到低纬度上。

$$\begin{align*} J_{SSM}(\theta)&=\mathbb{E}_{\xi}\left[\frac{1}{2}||{\bf v}^T \psi(\xi;\theta)- {\bf v}^T\psi_x(\xi)||^2\right]\\ &= \mathbb{E}_{p_x(\xi)}[\frac{1}{2} {\bf v}^T \psi(\xi;\theta)^2 {\bf v}+\text{tr}({\bf v}^T\nabla_{\xi} \psi_{\xi}(\xi;\theta){\bf v})] \\ \end{align*}$$

要求${\bf v}\sim p_v, \mathbb{E}_{p_v}[{\bf v}{\bf v}^T]>0, \mathbb{E}_{p_v}[||{\bf v}||^2_2]<\infty$，这是因为需要最后出一个与$\theta$无关的常数。这种分布也是容易找到的，例如正态分布、均匀分布等。