score-based method

Generative Modeling by Estimating Gradients of the Data Distribution

https://yang-song.net/blog/2021/score/#the-score-function-score-based-models-and-score-matching

基于最大似然的生成模型，由于配分函数难以解决。所以对于传统的方法，主要有两种解决方法：1.控制模型结构，使配分函数可以求得（如自回归模型，流模型）。2.近似配分函数，典型的方法是变分自编码器（使用mcmc采样）。

隐式的生成模型，代表是GAN，可以直接学习数据分布。但训练困难，不稳定。

建模函数$s_{\theta }(x)={\nabla }_x\log p(x)$，优化目标是 fisher divergence：

$${\mathbb{E}}_{p(\mathbf{x})}\left[{\Vert {\nabla }_{\mathbf{x}}\log p(\mathbf{x})-{\mathbf{s}}_{\theta }(\mathbf{x})\Vert }_2^2\right]$$

由于没有办法求未知的分布${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$，所以使用score matching的方法。

$s_{\theta }(x)$是score-base model。In fact, the only requirement on the score-based model is that it should be a vector-valued function with the same input and output dimensionality, which is easy to satisfy in practice.

summary：

As a brief summary, we can represent a distribution by modeling its score function, which can be estimated by training a score-based model of free-form architectures with score matching.

使用Langevin dynamics的方法进行采样求解

Langevin dynamics

Langevin dynamics是一种类似mcmc采样的方法，可以直接从${\nabla }_x\log p(x)$中采样$x\sim p(x)$，公式为：

${\mathbf{x}}_{i+1}\leftarrow {\mathbf{x}}_i+ϵ{\nabla }_{\mathbf{x}}\log p(\mathbf{x})+\sqrt{2ϵ}{\mathbf{z}}_i,i=0,1,\cdots ,K$

其中$x_0$服从任意先验分布，${\mathbf{z}}_i\sim \mathcal{N}(0,I)$，$ϵ\to 0$ 且 $K\to \infty$，此时$x_K\sim p(x)$。

所以可以将$s_{\theta }(x)\approx {\nabla }_x\log p(x)$带入上面方程，则可以采样得到服从于$p(x)$的样本$x_K$。

Naive score-based generative modeling and its pitfalls

上面方法描述了如何通过langevin dynamic来获得样本和训练模型，但是这里仍然存在一些问题。

优化的目标：

$${\mathbb{E}}_{p(\mathbf{x})}\left[{\Vert {\nabla }_{\mathbf{x}}\log p(\mathbf{x})-{\mathbf{s}}_{\theta }(\mathbf{x})\Vert }_2^2\right]=\int p(\mathbf{x}){\Vert {\nabla }_{\mathbf{x}}\log p(\mathbf{x})-{\mathbf{s}}_{\theta }(\mathbf{x})\Vert }_2^2\mathrm{d}\mathbf{x}$$

当在$p(x)$的低密度区域采样数据会比较少，则估计自然不准。当使用langevin dynamic采样，对于高维数据来说，很容易采样到$p(x)$的低密度区域，使其无法生成代表数据的高质量样本。

Score-based generative modeling with multiple noise perturbations

解决上述问题的一个办法，就是往$p(x)$中添加噪声点填充低密度区域。此时的问题是如何选择合适噪声尺度，过小的话达不到效果，过大的话会破坏原本的数据结构。为了解决这两者，同时使用了多个尺度的噪声扰动。具体操作为：

分别添加$L$个逐渐增大的方差获得$L$个新的分布$p_{{\sigma }_i}(x)$：${\sigma }_1<{\sigma }_2<\cdots <{\sigma }_L$，此时从$p_{{\sigma }_i}(x)$中采样很容易，即$x\sim p(x)$，然后计算$x+{\sigma }_{i}z,{\mathbf{z}}_i\sim \mathcal{N}(0,I)$.

下面对噪声扰动的$p(x)$分布的score function，${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$进行估计，使用score matching方法训练${\mathbf{s}}_{\theta }(\mathbf{x},i)$（Noise Conditional Score-Based Model，又名Noise Conditional Score Network，NCSN）。

训练${\mathbf{s}}_{\theta }(\mathbf{x},i)$的目标函数是：

${\mathbf{s}}_{\theta }(\mathbf{x},i)$训练完成后，则可以使用Langevin dynamics来采样得到样本，$i=L,L-1,\dots ,1$，This method is called annealed Langevin dynamics

Score-based generative modeling with stochastic differential equations (SDEs)

当噪声尺度的数量推广到无穷大时，此时可以不仅可以获得高质量的样本，还能对对数似然精确求解，和对逆过程的可控生成。

当使用无穷多的噪声来扰动数据样本时，就是连续时间上的随机过程。

使用下述形式来表示这个随机过程：

$$dx=f(x,t)dt+g(t)dw$$

其中$f(x,t):{\mathbb{R}}^d\to {\mathbb{R}}^d$，是一个向量函数，被称作drift coefficient。$g(t)\in \mathbb{R}$是一个实值函数，被称作diffusion coefficient。w是布朗运动，dw可以被看作无意义的白噪声。这个随机微分方程的解是随机变量 { x ( t ) } t ∈ [ 0 , T ] \{x(t)\}_{t\in [0,T]} {x(t)}t∈[0,T]​的连续集合。t是从开始时刻0到时刻T中的每一个时刻。

令$p_t(x)$是$x(t)$的概率密度函数，此时$p_t(x)$类似于$p_{{\sigma }_i}(x)$。显然$p_0(x)=p(x)$，当添加足够步长$T$时，$p_T(x)$就会变成可以tractable的先验概率分布$\pi (x)$，在噪声尺度有限时，$p_T(x)$相当于$p_{{\sigma }_L}(x)$。此时相当于对数据添加了极大方差${\sigma }_L$扰动。

添加噪声扰动有很多方式，例如下面的SDE：

$$dx=e^{t}dw$$

这个数据扰动相当于均值为0，方差是指数级增长的噪音。相当于$\mathcal{N}\left(0,{\sigma }_1^{2}I\right),\mathcal{N}\left(0,{\sigma }_2^{2}I\right),\cdots ,\mathcal{N}\left(0,{\sigma }_L^{2}I\right)$，${\sigma }_1<{\sigma }_2<\cdots <{\sigma }_L$且是等比数列。所以，还有其他种类的SDE，如the Variance Exploding SDE (VE SDE), the Variance Preserving SDE (VP SDE), and the sub-VP SDE.

Reversing the SDE for sample generation

对于有限个噪声尺度，可以使用annealed Langevin dynamics方法来获取样本，但是对于无限个噪声尺度，使用reverse SDE方法来采样。

对于任意的SDE，都有其逆SDE，形式如下：

$$\mathrm{d}\mathbf{x}=\left[\mathbf{f}(\mathbf{x},t)-g^2(t){\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})\right]\mathrm{d}t+g(t)\mathrm{d}\mathbf{w}$$

Estimating the reverse SDE with score-based models and score matching

要求解上面的reverse SDE，需要知道最终的概率分布函数$p_T(x)$和score function ${\nabla }_x\log p_t(x)$。其中$p_T(x)=\pi (x)$,是设定的先验分布。为了估计${\nabla }_x\log p_t(x)$，我们训练了一个******Time-Dependent Score-Based Model $s_{\theta }(x,t)\approx {\nabla }_x\log p_t(x)$，**这个函数类似于上面的对于有限的噪声尺度的$s_{\theta }(x,i)$，所以$s_{\theta }(x,i)\approx {\nabla }_x\log p_{{\sigma }_i}(x)$。

最后的训练目标是连续状态下的加权Fisher divergences

$${\mathbb{E}}_{t\in \mathcal{U}(0,T)}{\mathbb{E}}_{p_t(\mathbf{x})}\left[\lambda (t){\Vert {\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})-{\mathbf{s}}_{\theta }(\mathbf{x},t)\Vert }_2^2\right],$$

其中$\mathcal{U}(0,T)$是$[0,T]$之间的均匀分布。$\lambda :\mathbb{R}\to {\mathbb{R}}_{>0}$是一个权值函数。一般使$\lambda (t)\propto 1/\mathbb{E}\left[{\Vert {\nabla }_{\mathbf{x}(t)}\log p(\mathbf{x}(t)\mid \mathbf{x}(0))\Vert }_2^2\right]$来平衡不同时间点上使用score match的损失。

当$s_{\theta }(x,t)$训练到最优的时候，就可以利用估计出 reverse SDE:

$$\mathrm{d}\mathbf{x}=\left[\mathbf{f}(\mathbf{x},t)-g^2(t)s_{\theta }(x,t)\right]\mathrm{d}t+g(t)\mathrm{d}\mathbf{w}$$

从$x(T)\sim \pi$开始，逐步求解上述的reverse SDE公式来获得样本$x(0)$，记最终是从$p_{\theta }$中获得$x(0)$的分布，当$s_{\theta }(x,t)$被训练到最优的时候，$p_{\theta }\approx p_0$,此时可以看作$x(0)\sim p_0$。

当$\lambda (t)=g^2(t)$时，就可以得到$p_{\theta }$和$p_0$之间的Fisher divergence和KL divergence之间重要的关系。

$$\mathrm{KL}\left(p_0(\mathbf{x})\Vert p_{\theta }(\mathbf{x})\right)\le \frac{T}{2}{\mathbb{E}}_{t\in \mathcal{U}(0,T)}{\mathbb{E}}_{p_t(\mathbf{x})}\left[\lambda (t){\Vert {\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})-{\mathbf{s}}_{\theta }(\mathbf{x},t)\Vert }_2^2\right]+\mathrm{KL}\left(p_T\Vert \pi \right).$$

由于极大似然估计和最小化KL散度之间的等价性，$\lambda (t)=g^2(t)$被称作为likelihood weighting function。利用这个加权函数，说明可以利用score base的方法来提高概率分布的极大似然估计。

How to solve the reverse SDE

利用数值SDE求解器来求解reverse SDE，有许多方法，如Euler-Maruyama method、Milstein method, 和 stochastic Runge-Kutta methods。

reverse SDE有两个重要的特点：

1. 可以通过******Time-Dependent Score-Based Model $s_{\theta }(x,t)$来估计$\ast \ast \ast \ast \ast \ast {\nabla }_x\log p_t(x)$。
2. 我们需要求的只有边际分布$p_t(x)$，不同的时间步之间存在任意的相关性。也即是说不需要按reverse SDE的顺序逐步采样来求$p_t(x)$。

由于这两点，我们可以使用MCMC采样来微调从数值 SDE 求解器获得reverse SDE的解。特别的，提出了Predictor-Corrector samplers。predictor可以是任意的numeric SDE 求解器来根据$x(t)\sim p_t(x)$预测$x(t+\Delta t)\sim p_{t+\Delta t}(x)$。corrector是利用score function（例如Langevin dynamics和Hamiltonian Monte Carlo）的MCMC步骤。

详细的，在每一个时间步，首先利用predictor去选择一个合适的时间步长$\Delta t<0$，然后根据样本$x(t)$来预测$x(t+\Delta t)$。然后corrector根据score-based model $s_{\theta }(x,t+\Delta t)$运行几步来提高$x(t+\Delta t)$的质量。最后$x(t+\Delta t)$就会变为来自$p_{t+\Delta t}(x)$的高质量样本。

Probability flow ODE

Langevin MCMC方法并无法精确的计算其中的似然函数。但一种基于常微分方程的求解器可以做到。

在不改变随机微分方程中随机过程的边缘分布的情况下，可以将任何的SDE转换为ODE（常微分方程）。因此，通过求解ODE方程，可以从reverse SDE的边缘分布中采样数据，它的ODE方程如下：

$$\mathrm{d}\mathbf{x}=\left[\mathbf{f}(\mathbf{x},t)-\frac{1}{2}{g}^2(t){\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})\right]\mathrm{d}t$$

下图描述了SDE和Probability flow ODE的轨迹，可以看出求解Probability flow ODE的轨迹所获得边缘分布和SDE相同。