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 表示每一时间步这个扰动独立采样。\epsilon > 0 是系数。至于为什么是 \frac{\epsilon}{2}, \sqrt{\epsilon} ，我也不懂，我只知道这个式子是从连续的 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[\left\|s_{\theta}(\tilde {x})-\frac{\partial\log\left(q_{\sigma}(\tilde{x})\right)}{\partial\tilde{x}}\right \|_{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 ，并采样一个标准高斯噪声 \epsilon ，对原数据样本加噪：\tilde x = x + \sigma_i \epsilon ，然后模型算出 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 取值为上图所示时，确实能使这个值为固定值。

​

‍