【理论推导】基于分数的生成模型Score-based Model

生成模型分类

已有的生成模型大致可以分为两类

• likelihood-based model：以 VAE 和 normalizing flow model 为主，直接拟合数据分布，但通常对网络模型的结构设计提出了较大限制
• implicit generative model：以 GAN 为主，使用判别器间接判断输出的数据是否符合数据分布，存在训练困难，容易训练失败以及模式崩塌等问题

Score Matching & Non-Normalized Distribution

Score Matching 最早在 2005 年《Estimation of Non-Normalized Statistical Models by Score Matching》提出，主要是用于解决概率分布的归一化的问题。在生成模型中，我们常希望使用含参模型 $p_{\theta }(x)$ 给出的概率密度函数能够拟合真实的概率密度函数 $p(x)$，$p_{\theta }(x)$ 作为一个概率密度函数须满足归一化性质 ${\int }_x{p}_{\theta }(x)dx=1$，因此，我们通常使用模型 $q_{\theta }(x)$ 给出一个未归一化的概率密度函数，然后使用归一化项 $Z(\theta )$ 来保证 $p_{\theta }(x)$ 的性质，即
$$p_{\theta }(x)=\frac{1}{Z(\theta )}{q}_{\theta }(x)$$
其中 $Z(\theta )$ 为与样本无关的一个常量。由于 $Z(\theta )$ 的存在，无论是基于梯度的优化过程还是正向推理都变得很难计算，因此，考虑使 $p_{\theta }(x)$ 关于输入的梯度逼近 $p(x)$ 关于输入的梯度

$${▽}_x{q}_{\theta }(x)={▽}_x{p}_{\theta }(x)\approx {▽}_{x}p(x)$$
形式化地，我们定义分数函数 $\psi :{\mathbb{R}}^n\to {\mathbb{R}}^n$ 如下
$$\psi (x)={▽}_{x}p(x)$$
我们利用MSE损失函数找到最优参数 $\theta$ 的过程称作是 Score Matching，损失函数定义为
$$\begin{array}{rl}{J}_{\text{ESM}}(\theta )& =\frac{1}{2}{\int }_{x}p(x)||{\psi }_{\theta }(x)-\psi (x)|{|}_2^{2}dx\\ & ={\mathbb{E}}_{x\sim p(x)}[\frac{1}{2}||{\psi }_{\theta }(x)-\psi (x)|{|}_2^2]\end{array}$$
这种损失函数的表达形式也被称为显式分数匹配 (Explicit Score Matching)

对 $J_{\text{ESM}}$ 进行展开，有
$$\begin{array}{rl}J(\theta )& =\frac{1}{2}{\int }_{x}p(x)||{\psi }_{\theta }^2(x)|{|}_2^{2}dx-{\int }_{x}p(x){\psi }_{\theta }^T(x)\psi (x)dx+\frac{1}{2}{\int }_{x}p(x)||{\psi }^2(x)|{|}_2^{2}dx\\ & =\frac{1}{2}{\int }_{x}p(x)||{\psi }_{\theta }^2(x)|{|}_2^{2}dx-{\int }_{x}p(x){\psi }_{\theta }^T(x)\psi (x)dx+C\end{array}$$
对于第二项进行维度展开，有
$$\begin{array}{rl}{\int }_{x}p(x){\psi }_{\theta }^T(x)\psi (x)dx& =\sum _{i=1}^n{\int }_{x}p(x){\psi }_{\theta }^{(i)}(x){\psi }^{(i)}(x)dx\\ & =\sum _{i=1}^n{\int }_{x}p(x){\psi }_{\theta }^{(i)}(x)\frac{\partial \log p(x)}{\partial x^{(i)}}dx\\ & =\sum _{i=1}^n{\int }_x{\psi }_{\theta }^{(i)}(x)\frac{\partial p(x)}{\partial x^{(i)}}dx\end{array}$$
我们考虑 $i=1$ 时的情况，由分部积分公式
$$\begin{gathered} \underset{a\to \infty ,b\to -\infty }{\lim }f(a,x^{(2)},...,x^{(n)})g(a,x^{(2)},...,x^{(n)})-f(b,x^{(2)},...,x^{(n)})g(b,x^{(2)},...,x^{(n)}) \\ ={\int }_{-\infty }^{\infty }f(x)\frac{\partial g(x)}{\partial x^{(1)}}d{x}^{(1)}+{\int }_{-\infty }^{\infty }g(x)\frac{\partial f(x)}{\partial x^{(1)}}d{x}^{(1)} \end{gathered}$$
因此，有下式
$$\begin{array}{rl}{\int }_x{\psi }_{\theta }^{(1)}(x)\frac{\partial p(x)}{\partial x^{(1)}}dx& ={\int }_{x^{(2)}...x^{(n)}}{\int }_{x^{(1)}}{\psi }_{\theta }^{(1)}(x)\frac{\partial p(x)}{\partial x^{(1)}}d{x}^{(1)}d(x^{(2)}...x^{(n)})\\ & ={\int }_{x^{(2)}...x^{(n)}}\left[\underset{a\to \infty ,b\to -\infty }{\lim }(p(a,x^{(2)},...){\psi }_{\theta }^{(1)}(a,x^{(2)},...)-p(b,x^{(2)},...){\psi }_{\theta }^{(1)}(b,x^{(2)},...))-{\int }_{x^{(1)}}p(x)\frac{\partial {\psi }_{\theta }^{(1)}(x)}{\partial x^{(1)}}d{x}^{(1)}\right]d(x^{(2)}...x^{(n)})\end{array}$$
若假定 ${\lim }_{||x||\to \infty }p(x){\psi }_{\theta }(x)=0$，上式可化简如下所示
$$\begin{array}{r}{\int }_x{\psi }_{\theta }^{(1)}(x)\frac{\partial p(x)}{\partial x^{(1)}}dx=-{\int }_x\frac{\partial {\psi }_{\theta }^{(1)}(x)}{\partial x^{(1)}}p(x)dx\end{array}$$
将 (3)(4) 式代回到 (2) 式，得到等价的隐式函数优化目标 $J_{\text{ISM}}$
$$\begin{array}{rl}{J}_{\text{ISM}}(\theta )& ={\int }_{x}p(x)\sum _{i=1}^n\left(\frac{\partial {\psi }_{\theta }^{(i)}(x)}{\partial x_i}+\frac{1}{2}{\psi }_{\theta }^{(i)}(x{)}^2\right)dx\\ & ={\int }_{x}p(x)\left(\text{tr}({▽}_x{\psi }_{\theta }(x))+\frac{1}{2}||{\psi }_{\theta }(x)|{|}_2^2\right)dx\\ & ={\mathbb{E}}_{x\sim p(x)}[\text{tr}({▽}_x{\psi }_{\theta }(x))+\frac{1}{2}||{\psi }_{\theta }(x)|{|}_2^2]\end{array}$$

如果我们得到了一个最优的分数函数 ${\psi }_{\theta }(x)$，那么我们可以引入郎之万动力学方程 (Langevin Dynamics) 进行数据的生成任务，具体来说，给定一个固定的步长 $ϵ>0$，初始值 $x_0\sim \pi (x)$ 采样自某先验分布之中，$z_t\sim \mathcal{N}(0,I)$，有如下迭代公式
$$\begin{array}{rl}{x}_t& =x_{t-1}+\frac{ϵ}{2}{▽}_x\log p(x_{t-1})+\sqrt{ϵ}{z}_t\\ & \approx x_{t-1}+\frac{ϵ}{2}{\psi }_{\theta }(x_{t-1})+\sqrt{ϵ}{z}_t\end{array}$$
当满足 $t\to \infty$ 和 $ϵ\to 0$ 时，$x_t$ 等价于从 $p(x)$ 中进行采样

但以上方法存在以下几个问题

• 在搭建神经网络时，需要使用反向传播优化模型参数，(5)式中存在关于输出的一阶导数，这表明反向传播需要计算关于输出的二阶导数
• 对于尺寸较大的输入，以上建模方式并不生效

问题2的一种直观解释：图像分布可以看作是高维空间中的一个流形，因此对于空间中的大多位置都是不存在合理的训练样本进行训练的，对于采样自先验分布的初值 $x_0$，会因为不准确的分数函数而移动到某局部最优值上，得到较差的生成结果

Sliced Score Matching

Score Matching 最主要的问题是效率问题，为了避免计算Score函数相对于输入的二阶 Hessian 矩阵，作者提出使用如下式子替换 $J_{\text{ESM}}(\theta )$
$$\begin{array}{rl}{J}_{\text{ESSM}}(\theta )& ={\mathbb{E}}_{v\sim p_v(v)}{\mathbb{E}}_{x\sim p(x)}[\frac{1}{2}(v^T{\psi }_{\theta }(x)-v^T\psi (x){)}^2]\end{array}$$
其中 $v$ 是抽样自 $p_v$ 的随机方向向量，满足 ${\mathbb{E}}_{p_v}[v{v}^T]\succ 0$ （矩阵是正定的），且 ${\mathbb{E}}_{p_v}[||v|{|}_2^2]<\infty$ ，例如标准正态分布、多元Rademacher分布和超球面上的均匀分布。

对 (6) 式采用类似地推导，可以得到 $J_{\text{ISSM}}(\theta )$
$$J_{\text{ISSM}}(\theta )={\mathbb{E}}_{v\sim p_v(v)}{\mathbb{E}}_{x\sim p(x)}[v^T{▽}_x{\psi }_{\theta }(x)v+\frac{1}{2}(v^T{\psi }_{\theta }(x){)}^2]$$

证明过程如下所示

如何加速SM：对于${▽}_x{\psi }_{\theta }(x)$，原本需要计算一个 Hessian 矩阵，现在损失函数只需先计算标量 $v^T{\psi }_{\theta }(x)$，然后相对于标量求导 ${▽}_x(v^T{\psi }_{\theta }(x))=v^T{▽}_x{\psi }_{\theta }(x)$，即可得到损失函数

Denoising Score Matching

SSM提供了一个无偏的分数函数拟合方法，但是仍需要较大的计算量进行图像的生成。类比于VAE的思路，如果每个样本不是表示单独的样本点，而是表示一个局部的概率分布，即在训练样本 $x$ 上施加一个噪声 $\tilde{x}|x$，可以有效地缓解Score Matching中的问题2，假定 $q_{\sigma }(\tilde{x}|x)$ 表示噪声的概率分布，其中 $\sigma$ 为噪声分布中的参数，那么噪声样本的概率分布为
$$q_{\sigma }(\tilde{x})={\int }_x{q}_{\sigma }(\tilde{x}|x)p(x)dx$$
那么，学习的目标即为使用模型拟合该噪声分布的分数函数，例如 $J_{\text{ESM}}$ 表示转化为
$$\begin{array}{rl}{J}_{\text{ESM},q_{\sigma }}& ={\mathbb{E}}_{\tilde{x}\sim q_{\sigma }(\tilde{x})}[\frac{1}{2}||{\psi }_{\theta }(\tilde{x})-\psi (\tilde{x})|{|}_2^2]\\ & ={\mathbb{E}}_{\tilde{x}\sim q_{\sigma }(\tilde{x})}[\frac{1}{2}||{\psi }_{\theta }(\tilde{x})-\frac{\partial \log q_{\sigma }(\tilde{x})}{\partial \tilde{x}}|{|}_2^2]\end{array}$$
另一方面，定义 $J_{\text{DSM},q_{\sigma }}(\theta )$
$$J_{\text{DSM},q_{\sigma }}(\theta )={\mathbb{E}}_{x,\tilde{x}\sim q_{\sigma }(x,\tilde{x})}[\frac{1}{2}||{\psi }_{\theta }(\tilde{x})-\frac{\partial \log q_{\sigma }(\tilde{x}|x)}{\partial \tilde{x}}|{|}_2^2]$$
下面证明，在Score Matching对 $\theta$ 进行优化的目标下，$J_{\text{DSM},q_{\sigma }}(\theta )$ 与 $J_{\text{ESM},q_{\sigma }}(\theta )$ 等价
$$\begin{array}{rl}{J}_{\text{ESM},q_{\sigma }}& ={\mathbb{E}}_{\tilde{x}\sim q_{\sigma }(\tilde{x})}[\frac{1}{2}||{\psi }_{\theta }(\tilde{x})-\frac{\partial \log q_{\sigma }(\tilde{x})}{\partial \tilde{x}}|{|}_2^2]\\ & ={\mathbb{E}}_{x,\tilde{x}\sim q_{\sigma }(x,\tilde{x})}[\frac{1}{2}||{\psi }_{\theta }(\tilde{x})|{|}_2^2]-S(\theta )+C\end{array}$$
其中 $C$ 表示不包含 $\theta$ 的部分，而 $S(\theta )$ 推导如下所示
$$\begin{array}{rl}S(\theta )& ={\mathbb{E}}_{\tilde{x}\sim q_{\sigma }(\tilde{x})}[<{\psi }_{\theta }(\tilde{x}),\frac{\partial \log q_{\sigma }(\tilde{x})}{\partial \tilde{x}}>]\\ & ={\int }_{\tilde{x}}{q}_{\sigma }(\tilde{x})<{\psi }_{\theta }(\tilde{x}),\frac{1}{q_{\sigma }(\tilde{x})}\cdot \frac{\partial q_{\sigma }(\tilde{x})}{\partial \tilde{x}}>d\tilde{x}\\ & ={\int }_{\tilde{x}}<{\psi }_{\theta }(\tilde{x}),\frac{\partial q_{\sigma }(\tilde{x})}{\partial \tilde{x}}>d\tilde{x}\\ & ={\int }_{\tilde{x}}<{\psi }_{\theta }(\tilde{x}),\frac{\partial }{\partial \tilde{x}}{\int }_x{q}_{\sigma }(\tilde{x}|x)p(x)dx>d\tilde{x}\\ & ={\int }_{\tilde{x}}<{\psi }_{\theta }(\tilde{x}),{\int }_{x}p(x)\frac{\partial q_{\sigma }(\tilde{x}|x)}{\partial \tilde{x}}dx>d\tilde{x}\\ & ={\int }_{\tilde{x}}<{\psi }_{\theta }(\tilde{x}),{\int }_{x}p(x)q_{\sigma }(\tilde{x}|x)\frac{\partial \log q_{\sigma }(\tilde{x}|x)}{\partial \tilde{x}}dx>d\tilde{x}\\ & ={\int }_{\tilde{x}}{\int }_{x}p(x)q_{\sigma }(\tilde{x}|x)<{\psi }_{\theta }(\tilde{x}),\frac{\partial \log q_{\sigma }(\tilde{x}|x)}{\partial \tilde{x}}>dxd\tilde{x}\\ & ={\mathbb{E}}_{x,\tilde{x}\sim q_{\sigma }(x,\tilde{x})}[<{\psi }_{\theta }(\tilde{x}),\frac{\partial \log q_{\sigma }(\tilde{x}|x)}{\partial \tilde{x}}>]\end{array}$$
对 $J_{\text{DSM},q_{\sigma }}$ 进行展开，即可得证 $J_{\text{DSM},q_{\sigma }}$等价于 $J_{\text{ESM},q_{\sigma }}$，注意到如果施加的噪声为 $\mathcal{N}(0,{\sigma }^{2}I)$，那么有
$$\frac{\partial \log q_{\sigma }(\tilde{x}|x)}{\partial \tilde{x}}=\frac{1}{{\sigma }^2}(x-\tilde{x})$$
因此，优化 $J_{\text{DSM},q_{\sigma }}$ 目标函数如下所示
$$\begin{array}{r}{J}_{\text{DSM},q_{\sigma }}(\theta )={\mathbb{E}}_{x,\tilde{x}\sim q_{\sigma }(x,\tilde{x})}[\frac{1}{2}||{\psi }_{\theta }(\tilde{x})-\frac{1}{{\sigma }^2}(x-\tilde{x})|{|}_2^2]\end{array}$$
使用噪声的方式可以巧妙地避开了二阶导，提高了训练的效率，观察函数表示，我们发现DSM本质上，在使用分数函数拟合一个降噪器，这种方法存在以下两种弊端

• 拟合出来的分布是加了噪声之后的分布
• 噪声的级别 $\sigma$ 难以设置，较强的噪声损害了拟合出的数据分布，而较小的噪声无法覆盖到足够大的区域使得对于大部分初值无法优化到理想的位置

NCSN给出的方法是，在推理的不同阶段从大到小给出噪声，即

超参数的选择：噪声级别，步长与迭代次数的选择可见 Improved Techniques for Training Score-Based Generative Models

参考资料

Generative Modeling by Estimating Gradients of the Data Distribution
Sliced Score Matching: A Scalable Approach to Density and Score Estimation
Improved Techniques for Training Score-Based Generative Models