Score-Based Generative Models

$$

Brief Introduction

在从VAE到DDPM一文中，我们从 VAE 出发得到了 DDPM 的理论框架，发现 DDPM 可以看作是具有人为指定的后验分布（即前向过程）的层次化 VAE。这一角度的代表论文是 2015 年的 Sohl-Dickstein et al.^[1] 和 2020 年的 DDPM^[2]。与此同时，宋飏博士从 score matching + Langevin dynamics (SMLD)^[3] 的角度出发得到了与 DDPM 本质相同的模型，称为 score-based generative models. 两个角度出发点截然不同，但殊途同归，也是非常有意思了。

Score Function

考虑一个以 $$ 为参数的能量模型 (EBM) $$，依以下形式定义一个概率分布： $$ 其中 $$ 是归一化因子。训练生成模型无非就是希望极大化样本的对数似然：

$$ 我们发现一个问题：$$ 涉及对整个数据空间的积分，因而是 intractable 的。为了解决这个问题，score-based generative models 使用 score function 巧妙地回避开了 $$. 对于任意一个概率分布 $$，score function 定义为对数概率密度函数对数据的梯度：

$$ 直观上，它是一个向量场，指示了从当前点去到概率密度更大的地方的方向，如下图所示：

source:https://yang-song.net/blog/2021/score/

具体到能量模型，将 $$ 式代入可得其 score function 为： $$ 可以看见由于 $$ 与 $$ 无关，对 $$ 求梯度后就直接变成零了！这启发我们直接对 score function 建模，即构建一个模型 $$，用它去近似真实数据分布的 score function： $$ 此过程称作 score matching. 如果模型学习到了数据的 score function，也就相当于学到了数据分布。

Score Matching

为了训练模型 $$，我们最小化它与 $$ 的均方距离（也即是模型给出的分布和数据分布之间的 Fisher Divergence）： $$ 下标 ESM 是 Explicit Score Matching 的意思，explicit 是指该目标函数要求我们显式地知道 $$，但这恰恰是我们不知道的。围绕这个问题，我们下面介绍三种方法。

Implicit Score Matching

论文^[5]的作者指出，我们可以使用一个分部积分去掉 $$. 具体而言，将 $$ 式展开： $$ 第一项与 $$ 无关可以略去，最后一项可以直接优化，而中间项还需要再处理一下： $$

一元函数分部积分

设 $$ 是关于 $$ 的一元函数且具有连续导数，且下文出现的积分都存在，由于： $$ 两边同时在 $$ 上积分，得到： $$ 即： $$ 扩展到多元函数

设 $$ 是关于 $$ 的多元标量函数，固定除 $$ 外的其他变元，所有求导均为对 $$ 的偏导，那么上述推导依旧使用，因此有结论： $$

应用分部积分，并且假设 $$，那么中间项变成： $$ 于是最终优化目标为： $$ 下标 ISM 指 Implicit Score Matching，与 ESM 相对应。如果用 $$ 的 Jacobian 矩阵 $$ 来表示，就是： $$ 当然，由于 $$ 本身是 $$ 的梯度，所以前者的 Jacobian 矩阵，也是后者的 Hessian 矩阵，所以也可以写成： $$ 只是记号上的不同罢了。

可以看见，现在的优化目标期望里面已经没有 $$ 或 $$ 了，而期望本身通过训练过程的采样来近似。

当然，上述推导过程要成立也有一些条件，比如用分部积分要求 $$ 和 $$ 可导，出现的所有期望要存在，以及上文已经提到过的 $$. 好在这些条件都容易满足。

Denoising Score Matching

上文中，我们使用分部积分巧妙地去掉了 $$，但是代价是引入了 $$，意味着我们还要对 score function 求导（即求 $$ 的 Hessian 矩阵）。虽然这个操作可以通过 torch.autograd.grad() 实现，但反向传播次数增加了 $$ 倍（$$ 为数据维数），在深度网络和高维数据的场景下效率很差。另一种直接的想法是，既然我们不知道 $$，那就对数据做非参估计呗。考虑使用带高斯核的 Parzen 窗做非参估计（也即核密度估计）： $$ 其中，当 $$ 时，高斯核趋向于 Dirac delta 函数，$$. 更一般地，当 Parzen 窗的数量趋于无穷时，其本质就是用高斯噪声扰动数据：

$$ 考虑在扰动后的数据分布上做 score matching，即将 $$ 式中的 $$ 换成 $$： $$ 可是 $$ 还是无法处理。注意到由于 $$ 是高斯核，$$ 是可以直接计算的： $$ 所以我们希望能把期望里面的 $$ 变换成 $$. 把平方打开，第一个平方项与 $$ 无关扔掉，另一个平方项保留，着重考虑中间项： $$ 这样我们就得到了变换后的目标函数： $$ 看上去有点丑，不妨配凑回平方的形式，即是论文中提出的 Denoising Score Matching (DSM) 优化目标： $$

直观上，$$ 指出了加噪数据 $$ 回到无噪声数据的移动方向，这正与 score function 的定义相符。

Sliced Score Matching

本节其实与以去噪为基础的扩散模型关系不大，但为了完整性，同时考虑到它也是宋飏大佬的工作，就顺便读了读论文。

Denoising score matching 虽然解决了 implicit score matching 无法 scale 到深度网络和高维数据的问题，但是本身也有两个问题：只能学习带有噪声的分布、性能对 noise level 很敏感。而宋飏等人在 2020 年提出的 sliced score matching^[7] 是另一种解决方法，其思想是把 $$ 随机投影到向量 $$ 上，只在投影的方向上做匹配。具体而言，$$ 式改写作：

$$ 并要求 $$ 且 $$. 常用的 $$ 包括标准多元正态分布 $$、多元 Rademacher 分布（$$ 上的均匀分布）、超球面 $$ 上的均匀分布等。

与原始 score matching 类似，我们可以用分部积分去掉 $$. 同样只需着重考虑中间项： $$ 利用多元函数分部积分结论，有： $$ 后面那项是我们想要的，现在需要证明前面那项等于 0. 在原始 score matching 中，我们是通过假设 $$ 完成的，但是现在多了个 $$ 来搅局。不过好在最后要对 $$ 求期望，所以我们依旧能够期待前项可以等于 0. 利用绝对值不等式把 $$ 拆出来，并利用柯西不等式，我们有： $$ 其中柯西不等式的作用只是为了用上 $$ 的假设条件，进而有限项乘以无穷小依旧是无穷小。

于是我们就得到了 SSM 的 implicit 形式： $$ 特别地，当 $$ 是标准多元正态分布或多元 Rademacher 分布时，后面一项可以直接积出来：$$，即： $$ 下标 VR 是 Variance Reduction 的意思，因为直接计算相比采样近似减小了方差，效果更好。

说了这么多，$$ 式中不还是有 $$ 吗，怎么就减少了计算量呢？注意 $$，所以我们可以先计算出标量 $$，然后对 $$ 求梯度，最后乘上 $$，这样反向传播次数就是直接计算 $$ 的 $$ 了。

另外，每次采样的投影向量个数也是影响效率的重要因素。采样越多，近似越准确，计算量也越大，所以这是一个 trade-off. 论文作者实验发现每次采样 1 个向量效果就不错了。

Langevin Dynamics

在模型学到了 score function 之后，我们怎么从中采样呢？直观上，score function $$ 是当前点指向密度更大处的方向，所以我们可以先随机初始化一个位置，然后不断沿着 score function 的方向走，最后就能走到密度的极大值点——这个过程很像梯度下降，只不过梯度下降是在参数空间中进行，而采样是在数据空间中进行： $$ 然而这样做每次都到达 $$ 的极大值点，并不是我们希望的从 $$ 采样。不过 Langevin dynamics 告诉我们，只需要在上式的基础上添加一个高斯噪声项，就可以化优化为采样： $$ 其中 $$. 可以证明，当 $$ 且 $$ 时，$$ 的分布收敛到 $$. 相关理论参见这篇文章。

source:https://yang-song.net/blog/2021/score/

Noise Perturbations

截至目前，我们用 score matching 训练模型预测 score function，然后通过 Langevin dynamics 采样，看起来似乎已经很完美了。但这个做法在实际中很难 work. 究其原因，罪魁祸首在于 manifold hypothesis，即数据分布在高维空间中的低维流形上。换句话说，空间中绝大部分都是低密度区域。这将导致两个问题：

首先，从训练层面上讲，低密度区域数据量少（甚至没有），无法较好地训练模型，导致模型预测的 score function 不准确：

source:https://yang-song.net/blog/2021/score/

又由于我们的初始化点极有可能落在低密度区域中，所以不准确的 score function 将对采样过程产生极大的影响。

其次，即便 score function 能被准确预测，但对于被低密度区域隔开的两个分布来说，Langevin dynamics 很难在有限步数内收敛到正确的分布（尽管 Langevin dynamics 在理论上能收敛到正确的分布，但可能需要很小的步长和很多的步数），如图所示：

source:Generative modeling by estimating gradients of the data distribution.

真实情况是右上的正态分布权重更大，但中图的采样结果显示两个分布被予以了几乎相同的权重。

宋飏等人提出解决这两个问题的办法是：给数据添加高斯噪声，并且随着 Langevin dynamics 的进行逐渐减小 noise level（退火）。添加噪声后的数据遍布整个空间，填充了低密度区域，因此模型能够得到更多的监督信号，Langevin dynamics 也能更快的收敛到高密度区域；同时，逐渐减小的噪声使得带噪分布最终收敛到真实分布，因此也能保证采样质量不受影响。

source:https://yang-song.net/blog/2021/score/

形式化地说，我们确定一个逐渐减小的噪声序列 $$，那么对应 $$ 的噪声扰动后的数据分布为： $$ 我们用一个神经网络（称作 Noise Conditional Score Networks，NCSN）同时预测不同 noise level 的 score function，即： $$ 于是训练目标就是对每一个 noise level 做 score matching： $$ 其中 $$ 是给不同 noise level 加的权重。考虑到 denoising score matching 正好符合现在的加噪去噪设定，所以我们选择它而非 sliced score matching. 根据前文的叙述，训练目标变换为：

$$ 其中 $$.

怎么选取权重 $$ 呢？根据作者经验，一般有 $$，所以为了各个 noise level 的损失大致相同，取 $$.

训练完成后，我们就可以随机初始化采样点，逐渐减小 noise level，在每个 level 下通过 Langevin dynamics 以步长 $$ 更新若干步，最终生成新的数据。算法流程图如下所示：

source:Generative modeling by estimating gradients of the data distribution.

其中步长设置为 $$ 的动机是保证不同 noise level 下 Langevin dynamics 的信噪比 $$ 的模长固定。

这一代的 NCSN 已然能与 GAN 媲美，但局限在 32x32 的大小。在后续的工作^[4]中，宋飏等人探索了更多的技巧，进一步地提升了图像生成质量，并能稳定生成 256x256 的图片。这些技巧包括：

• 选择 $$ 与训练数据中两两之间的最大欧氏距离差不多大。

• 选择 $$ 为公比为 $$ 的等比序列，且 $$

• 将 NCSN 参数化为 $$，其中 $$ 是一个无条件的网络。

• 选择 $$（每个 noise level 下 Langevin dynamics 的步数）在可承受范围尽可能大，并选择 $$（Langevin dynamics 里与步长有关的参数）使得一个复杂的式子（懒得抄了）约等于 $$.

• 在采样的时候使用 EMA model.

具体内容和分析本文不再叙述，感兴趣的读者可以去看论文。

Connection to DDPM

在引言里，我们说 SMLD 和 DDPM 本质是相同的，现在来看看为什么。

Forward Process

DDPM 与 SMLD 的加噪过程分别为： $$ 可见它们都是选择高斯分布的转移形式，但对于均值处是否加权略有分歧。DDPM 为均值加权，效果就是随着扩散过程的进行，数据的均值往 $$ 移动，并且一直维持着有限的方差且越来越接近 $$，最后服从 $$；而 SMLD 均值没有加权，所以无论怎么加噪数据的中心始终在原本的位置，因此加噪过程其实是依靠增大 $$ 来让数据弥散到整个空间。正因如此，在之后的文章中我们可以看到，宋飏等人将前者称为 Variance Preserving，而后者称为 Variance Exploding. 但其实这个差别不是很重要，如果把 SMLD 的均值也加上权重，并不影响任何推导过程。

为进一步探索二者的关系，考虑对 $$ 求 score function： $$ 我们发现，DDPM 中添加的噪声 $$，其实就是 score function 的反方向。因此，DDPM 的噪声预测模型和 SMLD 的 score 模型之间有如下关系： $$ 为了看得更明显，将 $$ 式代入 DDPM 的损失函数得： $$ 为了方便比较，将其改写一下： $$ 对比 $$ 式： $$ 可以看到，$$ 式和 $$ 式完全相同！甚至连系数都保持了一致：

• 根据 $$ 式有：$$.
• 根据 $$ 式有：$$.

Reverse Process

虽然 DDPM 和 SMLD 在加噪过程和优化目标上保持了惊人的一致，但是在生成过程上还是有所区别。对于 DDPM，我们构建逆向的马尔可夫链，每一步从 $$ 中采样： $$ 或写作： $$ 而对于 SMLD，我们在每一个 noise level $$ 下都进行 $$ 步的 Langevin dynamics 游走： $$ 上述过程重复 $$ 次。根据 Predictor-Corrector 的思想，可以认为 DDPM 只做了相邻时间步之间的 prediction，没有同一时间步内的 correction；SMLD 则只有同一时间步内的 correction，而没有相邻时间步之间的 prediction.

References

1. Sohl-Dickstein, Jascha, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2256-2265. PMLR, 2015. ↩︎
2. Ho, Jonathan, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems 33 (2020): 6840-6851. ↩︎
3. Song, Yang, and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems 32 (2019). ↩︎
4. Song, Yang, and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems 33 (2020): 12438-12448. ↩︎
5. Hyvärinen, Aapo, and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research 6, no. 4 (2005). ↩︎
6. Vincent, Pascal. A connection between score matching and denoising autoencoders. Neural computation 23, no. 7 (2011): 1661-1674. ↩︎
7. Song, Yang, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pp. 574-584. PMLR, 2020. ↩︎
8. Yang Song. Generative Modeling by Estimating Gradients of the Data Distribution. https://yang-song.net/blog/2021/score/ ↩︎
9. Denoising Diffusion-based Generative Modeling: Foundations and Applications. https://cvpr2022-tutorial-diffusion-models.github.io ↩︎
10. Chih-Sheng Chen. Score Matching 系列 (一) Non-normalized 模型估計. https://bobondemon.github.io/2022/01/08/Estimation-of-Non-Normalized-Statistical-Models-by-Score-Matching/ ↩︎
11. Chih-Sheng Chen. Score Matching 系列 (二) Denoising Score Matching (DSM) 改善效率並可 Scalable. https://bobondemon.github.io/2022/03/06/A-Connection-Between-Score-Matching-and-Denoising-Autoencoders/ ↩︎
12. Chih-Sheng Chen. Score Matching 系列 (三) Sliced Score Matching (SSM) 同時保持效率和效果. https://bobondemon.github.io/2022/03/06/Sliced-Score-Matching-A-Scalable-Approach-to-Density-and-Score-Estimation/ ↩︎
13. 扩散模型与能量模型，Score-Matching和SDE，ODE的关系 - 中森的文章 - 知乎 https://zhuanlan.zhihu.com/p/576779879 ↩︎
14. Yang Song. Sliced Score Matching: A Scalable Approach to Density and Score Estimation. https://yang-song.net/blog/2019/ssm/ ↩︎