Langevin方程和Score Matching

朗之万方程

Langevin方程描述了微观粒子运动（Langevin读音类似朗之万，因为是法国人）。大概是下面这种形式

$$m\ddot{x} = -\gamma \dot{x} + F_{\text{ext}} + dW_t$$

$\ddot x$是加速度，$\dot x$是速度。这个公式本质上就是牛顿第二定律，右边第一项表示粒子受到的黏滞阻力（和速度成反比），$F_{\text{ext}}$表示外力（比如重力，这里我们总是认为外力来自某个保守场），$dW_t$是维纳过程的增量。

$W_t$维纳过程，其实就是布朗运动，或者说随机游走。我们认为它的均值是0、关于$t$是独立的，并且是平稳的，服从高斯分布。也就是

$$dW_t \sim \mathcal N(0, \sigma^2)$$

因为每个时刻都是独立的，所以自相关函数形式如下，$\delta$表示狄拉克函数，也就是说只有$s=t$的时候有个无穷大的相关性，其他时候不相关，$\langle \cdot \rangle$ 的含义是取期望/取平均。

$$\langle dW_s \cdot dW_t \rangle = 2\sigma \cdot \delta(s-t)$$

所以总的来看，朗之万方程大概就是在说，粒子受到的力来自三个部分，除了阻力和外力，还有一个来自涨落的随机力，这个随机力均值为0，并且每一时刻的随机力之间是不相关的。

玻尔兹曼分布

在统计物理当中有个结论。当系统处于平衡态的时候，粒子在空间上服从玻尔兹曼分布，也就是说，某一点$x$处的密度$n(x)$和势能$E_p$有关。

$$n(x) \propto \exp(-E_p/kT)$$

势能越大，概率越小。温度越低，系统的能量也越低，越趋近于低势能的点；温度越高，分布就越平均。$k$是玻尔兹曼常数，出现在这里的原因是为了把内能和开氏温标对齐。

简单起见，考虑最简单的重力场中的大气密度分布，这是个经典例子。假设重力场中的大气处处等温，几乎无限高。我们任选一个参考平面为$h=0$，那么在高度$h$处，任取一个小小的高度微元$dh$

上下的压强差 $dP = n(h) mg\cdot dh$

压强来自粒子的碰撞，$P \propto n \langle E_k \rangle$，$\langle E_k \rangle$表示平均动能，类似于求期望，平均动能实际上就是内能，正比于$kT$。因为温度处处相同，所以$P \propto n/(kT)$，代入上面的方程

$$\frac{dn}{dh} = C nmgh/(kT)$$

这是个可以分离变量的微分方程，很好求，求解出来就会有上面的指数形式。这个结论不难推广到一般的保守场。

Langevin方程从微观的角度描述了粒子的运动，而玻尔兹曼定律从宏观的角度描述了粒子密度的分布。这也就意味着我们只要按照Langevin方程模拟粒子的运动，就能从对应的玻尔兹曼分布中采样。至于这个分布长什么样子，取决于势能，也就是保守场。

Score Matching

概率密度就像是气体密度，可以想象成有很多粒子按照概率密度弥散在空间当中，只不过这个空间可能是高维的，并且需要归一化。玻尔兹曼常数$k$在这里有点碍眼，可以直接去掉。在样本点$x$（或者叫做状态$x$）的概率密度：

$$p(x) = \frac{1}{Z} \exp (-E(x)/T)$$

概率密度需要归一化，所以有个有个$Z = \int \exp(-E(x)/T)dx$在分母（如果状态是离散的，积分就变成求和，概率密度改成概率分布）。这个形式和softmax一模一样。这个积分在一般情况下是没法算的，intractable。

原则上，我们可以训练一个模型去学习这个势能$E_{\theta} (x)$。拿到一个样本，模型输出其势能，然后最大化对数似然函数，记模型参数是$\theta$。不失一般性，我们直接假设$T=1$，那就有似然函数

$$J(\theta|x) = \log p(x|\theta) = -E_\theta(x) - \log Z$$

使用梯度上升的方法最大化这个函数会碰到问题。$Z$和$\theta$有关，不能看成常数，因此

$$\nabla_\theta\log p(x|\theta) = -\nabla_\theta E_\theta(x) - \nabla_\theta Z$$

等式右边第一项好求，反向传播一下就能算出来。第二项归一化因子涉及高维度的积分，没法求。因此在训练的时候，基于势能的方法没法进行。

有没有方法可以彻底绕开这个归一化因子呢？我们可以注意到$Z$和$x$无关（积分积掉了），只要两边都对$x$求梯度，$Z$就不见了。换句话说，我们不需要学习势能，转而学习势能的负梯度——也就是力场！。

$$\nabla_x \log p(x|\theta) = -\nabla_xE_\theta(x) = s_\theta(x)$$

$s_\theta(x)$就是score函数。我们的优化目标转而变成把score函数和真实分布的梯度匹配上，最小化二者的平方损失

$$\text{MSE}(s_\theta(x), \nabla_x \log p_{\text{data}}(x))$$

也就是最小化

$$J(\theta) = \mathbb E_{x_\sim p_\text{data}(x)} ∥ s_\theta(x) - \nabla_x\log p_{\text{data}}(x)∥^2$$

但是我们并不知道真实分布对样本$x$的梯度。我们把期望写回积分的形式。根据期望的定义，有

$$J(\theta) = \int_\Omega p_\text{data}(x) ∥s_\theta(x) - \nabla_x\log p_\text{data}(x)∥^2dx$$

其中$\Omega$是样本空间。

我们简化一下记号，$p_\text{data}(x)$记作$p(x)$，$\nabla_x \log p_\text{data}(x)$记作$s(x)$，把平方损失展开：

$$∥s_\theta(x) - s(x)∥^2 = ∥s_\theta(x)∥^2 + ∥s(x)∥^2 - 2s_\theta(x)\cdot s(x)$$

第二项和$\theta$无关，直接写成常数$C$，代回到积分当中

$$\int_\Omega p(x) \left[∥s_\theta(x)∥^2 - 2s_\theta(x)\cdot s(x) + C\right]dx$$

把常数提出去，变成$\int_\Omega p(x)Cdx$，和$\theta$无关，在优化的时候直接不用管。那么剩下的项就是

$$\int_\Omega p(x)∥s_\theta(x)∥^2dx - 2\int_\Omega p(x)s_\theta(x)\cdot s(x)dx$$

第一项就是$\mathbb E_{x_\sim p_\text{data}(x)}\left[∥s_\theta(x)∥^2\right]$，这个好算，就是从数据集里面采样，然后算score函数的平方。

第二项可以作分部积分

$$\begin{aligned} 2\int_\Omega p(x)s_\theta(x)\cdot \nabla_x \log p(x)dx =& 2\int_\Omega s_\theta(x) \cdot \nabla_x p(x) dx \\ =& 2\int_\Omega s_\theta(x) \cdot d\nabla_x p(x) \\ = & 2\int_{\partial \Omega}s_\theta(x) \cdot \nabla_x p(x) \cdot n(x) dS - 2\int_\Omega p(x) \nabla_x s_\theta(x) dx \end{aligned}$$

如果$\Omega = \mathbb R^d$，并且概率密度在无穷远处衰减为0，那么第一项会是0。只剩下第二项，也就是$s_\theta(x)$对$x$的梯度的期望。

所以最后优化目标变成了

$$J(\theta) = \mathbb E_{x \sim p(x)} \left[∥s_\theta(x)∥^2 - 2\nabla_x s_\theta(x)\right]$$

不需要真实PDF的梯度了。因为我们采用神经网络来拟合$s(x)$，因此$s_\theta(x)$对$x$的梯度也可以用反向传播算法求出来。这个优化目标里面每一项都可以算。数据集如果很大，期望不好求，可以minibatch + 随机梯度下降。

前面说到，score函数其实就是力场，有了力场之后，我们可以通过模拟Langevin方程来采样。甚至还可以引入Metropolis-Hasting来确保是平稳的（想象一下那个拒绝概率，分子分母里面的势能都在指数上面，相除会变成做差，做差实际上就是对$x$的梯度，也就是score函数）。