score-based model

paper: Generative Modeling by Estimating Gradients of the Data Distribution

1 引入

什么是分数模型，就是这个模型能估计 "对数概率密度函数对样本的梯度"，用数学表达就是能预测 $\frac{\partial \log p(x)}{\partial x}$ ，这个就叫分数。如果我们有了这个，能干什么？想象一下，假设我扔一个样本 x ，然后你一直用这个模型去预测梯度，然后一直朝着这个梯度走，那么最终你会走到一个地方，这个地方会使得原数据的分布似然尽可能大。

所以如果沿着分数一直走，那么就能够走到数据分布的高概率密度区域，最终生成的样本就会符合原数据分布。

但是我们想做生成任务啊，如果这个分数学的很完美的话，那么我们生成的东西就跟原数据长的几乎差不多了。

所以为了保证生成结果的多样性，我们需要在跟随着梯度移动的过程中，带有随机性。

‍

2 朴素 score-based model

以上是直观理解，下面要具体的展开讲讲数学。

分数模型需要估计分数，它的训练目标可以表示为：

$s_{\theta }$ 就是我们模型预测的结果(分数)，$x$ 是样本的随机变量，要求期望。

不过，$p_{\text{data}}(x)$ 我们不知道啊，但是有一种叫 score matching 的方法，可以让我们在不了解概率密度的情况下，也可以成功训练出模型来估计分数，这种方法利用了分部积分，从而将上述训练目标转换为：

其中，$\frac{\partial s_{\theta }(x)}{\partial x}$ 代表 $s_{\theta }(x)$ 的雅可比矩阵，tr 是迹。

似乎可求了对吧，但是 $\frac{\partial s_{\theta }(x)}{\partial x}$ 这玩意非常难算，因为要对每一个分量求偏导，看下图：

​

而一张图像往往有成白上千个像素，所以，这不太现实。

怎么办呢？（其实可以用 sliced score matching 方法解决）但是还有另一种方法 denoising score matching，就并没有继续前进了，而是回退到了公式一。

怎么办呢？我给随机变量 x 加上一个方差为 ${\sigma }^2$ 的高斯噪声，得到 $\tilde{x}$ 。那么随机变量 $\tilde{x}$ 就是服从均值为 x , 协方差矩阵为 ${\sigma }^{2}I$ 的多维高斯分布。

ok，不妨规定一下符号：

那么，下面这个式子等号可以一直推导（第二个等号原因是 x 服从的分布跟 $\tilde{x}$ 无关）：

所以这个 $q_{\sigma }(\tilde{x}|x)$ 怎么算呢？它表示已知 x 后 $\tilde{x}$ 的概率值，我们又知道 $\tilde{x}$ 是服从均值为 x, 协方差矩阵为 ${\sigma }^{2}I$ 的多维高斯分布。这里我们就可以把 x 当作一个常量，那么，由正态分布的公式，可得：

所以 $q_{\sigma }(\tilde{x}|x)$ 是可计算的。所以上面这个东西 $\frac{\partial \log (q_{\sigma }(\tilde{x}))}{\partial \tilde{x}}$ 是可计算的。

那假如，$p_{\text{data}}(x)$ 和 $q_{\sigma }(\tilde{x})$ 是一样的，由公式 (i)，我们需要计算的目标就变为：

这玩意可求啊，取出样本 x, 加上方差为 ${\sigma }^2$ 的高斯噪声，就可得到 $\tilde{x}$ ，所以上面这个目标是完全可求的。

所以现在要解决的问题只有一个了，如何让 $p_{\text{data}}(x)$ 和 $q_{\sigma }(\tilde{x})$ 尽可能一样？因为随机变量 $\tilde{x}$ 是由 $x$ 加高斯噪声得到，那么我们把方差强度 ${\sigma }^2$ 设置的足够小就好了。

很好，训练问题解决。

模型训练好后，就可以做生成了，假设我们想要生成的内容尽可能跟原数据贴合，那么可以跟随模型预测的分数的方向，迭代更新：

${\tilde{x}}_0$ 随便取，视具体任务而定。

但是这样做，你会发现假如从同一起点出发，最终采样生成都会得到相同的结果。因为回顾整个训练过程，随机过程只体现在对 $x$ 加噪得到 $\tilde{x}$ ，但是我们加的方差 ${\sigma }^2$ 很小，所以导致没加多少噪声其实。所以相当于训练和推理几乎是确定性的。

所以，我们要用朗之万动力学采样（其实就是采样过程中加了随机项），如下：

$z_t\sim N(0,I)$，带下标 $t$ 表示每一时间步这个扰动独立采样。$ϵ>0$ 是系数。至于为什么是 $\frac{ϵ}{2},\sqrt{ϵ}$ ，我也不懂，我只知道这个式子是从连续的 Langevin SDE(Stochastic Differential Equation)变换到离散形式后得到的。

至此，一个朴素的 score-based model 搭建完毕。

‍

3 困局

上述推导很优美，很简单。但是实验发现，会有一些问题：

1. loss 不收敛

目前对此的主流解释是流形假说，即生活中的大部分数据都是在整个编码空间里的低锥流形里的。所以有些方向上的数据几乎不存在 “分布” 这一说法，即该方向的值可能都是差不多一样的。所以，这个方向上就不存在梯度或者梯度信号非常微弱。从而导致这个方向无法正常更新。

举个例子，假设有一个三维空间，考虑其中的 z 轴，假设数据样本的 z 轴取值服从正态分布，那么假设有个样本点，那么考虑其 z 轴上的分数 $\frac{\partial \log p(x)}{\partial x}$ ，因为 z 轴上的概率密度函数解析式已知，所以 z 轴上的分数解析式很容易算出来。

但是假设 z 轴取值全部为 1 ，那么 z 轴上的分数就没有定义域了，因为概率密度函数就是一个点，铺不满整个轴，所以定义域只有 z = 1 ，压根谈不了这个概率密度函数还对 z 求导，所以分数这个概率在 z 轴上就不存在了。

我们 score-based model 的大前提就是基于分数的，如果在很多方向分数这个概念都不存在时，自然这个模型的所有推导就化为泡影了，这就是为什么 loss 起起伏伏的原因。

2. 模型预测(分数)效果差

原数据分布是客观存在的，所以如果起始点或者迭代过程中的点落到了低概率区域，由于这部分低概率区域的数据得到的监督信号很少，所以模型在这里预测的分数就不准。

3. 生成与原分布偏差大

由第二点可知，在低概率区域，模型预测分数不准，从而导致跑偏，从而导致生成与原分布偏差大。

‍

4 NCSN

NCSN, Noise Conditional Score Networks, 来源于宋彪的 paper: Generative Modeling by Estimating Gradients of the Data Distribution。是一种 score-based model ，但是解决了上述遇到的几个痛点。

先把具体做法说出来：

与朴素 score-based model $\frac{1}{2}{E}_{\tilde{x}\sim q_{\sigma }(\tilde{x})}\left[{\Vert s_{\theta }(\tilde{x})-\frac{\partial \log \left(q_{\sigma }(\tilde{x})\right)}{\partial \tilde{x}}\Vert }_2^2\right]$ 的训练不同的是，这篇 paper 会使用不同强度的高斯噪声 $N(0,{\sigma }_i^2)$ 来对数据进行扰动，训练网络去联合估计加噪后的数据对应的分数。

（“联合”代表对于不同噪声级别的分数，都使用同一个网络去估计。）

其实用公式表达很简介，原本是 $s_{\theta }(x,\sigma )$，现在是 $s_{\theta }(x,{\sigma }_i)$ 。

训练完开始采样，采样用的是分层式采样（也可以叫退火朗之万动力学）。具体来说，就是首先从最高强度的噪声级别开始，进行一定步数的采样，这个强度最终的样本会作为下一个噪声级别的初始点，继续进行相同步数的朗之万动力学采样，以此类推。直到最小噪声级别也完成了朗之万动力学采样，就得到了最终生成的样本。

其实比起朴素 score-based ，就是多尺度扰动就吧，其余啥都没改。为什么这么改进就能解决痛点了呢？

其实这有个 trade-off 的问题，增加的噪声级别越高，那么加噪数据就能更多的漂到其它地方，从而填充原分布的低概率密度区域，这样就有更多的监督信号了。而且因为是高斯噪，所以能保证能飘到所有编码空间中，这样前面就可以缓解前面提到的 loss 不收敛问题。

所以噪声级别提高有很多好处就吧，但是问题来了，之前朴素 score-based model 训练目标其实是 $q_{\sigma }(\tilde{x})$ 的分数，但我们实际要求的是 $p_{\text{data}}(x)$ 的分数，之前这么做是因为我们默认加噪方差 ${\sigma }^2$ 很小，这俩近似相等。可是当我们噪声级别提高的话，它们俩就不相等了。

怎么办？多尺度加噪。

具体来说，作者设计了一个加噪序列 $\{{\sigma }_i{\}}_{i=1}^L$，满足 $\frac{{\sigma }_1}{{\sigma }_2}=\frac{{\sigma }_2}{{\sigma }_3}=\cdots =\frac{{\sigma }_{L-1}}{{\sigma }_L}>1$，同时 ${\sigma }_1$ 足够大以能够漂到整个编码空间同时填充低概率区域，${\sigma }_L$ 足够小，以获得对原数据分布良好的近似。

至于 ${\sigma }_i$ 具体数值，去炼丹吧。

好，那么训练目标是什么呢？我们来推导一下：

因为是多尺度噪声，所以总的训练目标是：

其中 $\lambda ({\sigma }_i)$ 代表不同噪声级别损失的权重。所以如何设置 $\lambda ({\sigma }_i)$ 呢？

作者的出发点是，这个权重设计的原则，要使得加权后所有噪声级别的损失都在同一量级里，不受 ${\sigma }_i$ 大小的影响，从而保证公平性。

所以咋设置呢？

作者根据实验，发现当网络训练到收敛时，$s_{\theta }(\tilde{x},\sigma )$ 大致在 $\frac{1}{\sigma }$ 这个量级。

好，那么我就将 $\lambda ({\sigma }_i)$ 设置为 ${\sigma }_i^2$ ，为什么，算一下你就知道了：

好玩的事情发生了，$s_{\theta }(\tilde{x},{\sigma }_i){\sigma }_i$ 此时的量级在 1 ，而且 $\frac{\tilde{x}-x}{{\sigma }_i}$ 这玩意，不就是标准化噪声吗？（因为 $\tilde{x}$ 是均值为 $x$ ，方差为 ${\sigma }^2$ 的高斯分布）此时这一整项不会受到 ${\sigma }_i$ 的影响。

总结一下，在 NCSN 中，只要预先设定好 $\{{\sigma }_i{\}}_{i=1}^L$ ，然后每次迭代时随机选取一个噪声级别 ${\sigma }_i$ ，并采样一个标准高斯噪声 $ϵ$ ，对原数据样本加噪：$\tilde{x}=x+{\sigma }_iϵ$ ，然后模型算出 $s_{\theta }(\tilde{x},{\sigma }_i)$ 后就可以算 loss 然后反向传播了。非常清晰。

训练过程原论文没给伪代码，不过我觉得已经非常清晰了，直接可以自己手搓了。

至于采样过程，上面已经讲过了，就是分层式采样，直接给出伪代码：

​

这个伪代码很好懂，我们来看看。有 L 个级别的噪声，lv1 噪声级别最高，所以从 lv1 开始。${\alpha }_i$ 是步长，为啥要这么设，等会再说。然后迭代 T 次，迭代的公式就是郎之万动力学采样公式，最后迭代的结果作为下一次的起始点。最后得到最终采样。非常清晰。

好了，所以步长为什么那么设，作者给出的解释是：“固定 Langevin dynamics 的信噪比”。信噪比这东西我知道，可以理解为信号与噪声的比值，如果远 < 1 ，说明加噪太大，去噪难度大。但是，我不理解，什么叫 "Langevin dynamics 的信噪比"，这玩意具体是多少？原文中作者说这个信噪比的解析式是 $\frac{{\alpha }_i{s}_{\theta }(x,{\sigma }_i)}{2\sqrt{{\alpha }_i}z}$ ，为什么是这个，我也不知道。不过假定它就是这个，那么，${\alpha }_i$ 取值为上图所示时，确实能使这个值为固定值。

​

‍