[翻译+笔记] Score-based generation: 通过数据分布的梯度进行生成建模

---

本次翻译+笔记的是Yang Song博士的博客https://yang-song.net/blog/2021/score/. Yang Song博士提出了一系列的score-based的生成模型, 并对后续Diffusion model也产生了很深远的影响.

---

1. 引言

生成的基本任务是从某种程度上来估计或采样原始的数据分布. 现有的估计分布的方法大概可以分为两种:

1. 基于似然的模型. 这些模型尝试去通过最大似然的方式直接学习概率分布函数, 例如自回归模型, 流模型, VAE等等.
2. 隐式的生成模型, 也就是原始的分布通过采样过程来隐式的表示, 主要代表的就是GAN. GAN是学习了从一个简单分布(高斯分布)到一个复杂分布(数据原始分布)的采样映射(也就是Generator的作用). 当然, Discriminator的作用主要是让Generator生成的样本越来越符合原始分布(最好的结果是分辨不出来).

有关VAE和GAN的其他理解, 请参照南开大学李翔老师的"从分布到生成"系列(https://www.bilibili.com/video/BV1cx4y147om/?spm_id_from=333.337.search-card.all.click). 在这个视频系列的观点中, GAN的Discriminator可以理解为一个"学习到"的广义的原始数据分布, 而Generator为Discriminator生成了丰富的负样本, 从而让Discriminator可以越来越逼近真实的分布. VAE的Decoder也是学习的从高斯分布到样本分布的映射, 而Encoder负责的是将数据映射到高斯分布, 从而可以产生监督信号.

然而, 这两种模型都有一些缺点. 基于似然的模型需要一个精心设计(需要比较复杂的网络估计)的归一化常数来进行似然的估计(后面说为什么). 对于隐式生成模型, 往往需要对抗训练的方式, 也有些不稳定, 以及会出现模式坍缩的情况.

这篇博客, 主要是介绍score-based模型, 其最核心的点就是在于不是直接建模数据分布, 而是建模数据分布的log的梯度, 缓解了上面的两个问题.

2. score function, score-based model和score matching

首先来解释一下为什么似然模型需要一个精心控制的归一化常数, 以及引出score函数.

假设数据集的样本为 { x 1 , … , x N } \{\mathbf{x}_1, \dots, \mathbf{x}_N\} {x1​,…,xN​}, 假设数据真实的分布是$p(\mathbf{x})$. 生成模型的根本任务是对$p(\mathbf{x})$进行某种方式的估计, 这样的话对估计的分布进行采样就可以生成了.

对于基于似然的模型, 我们energy-based model来估计分布$p_{\theta }(\mathbf{x})$:

$$p_{\theta }(\mathbf{x})=\frac{e^{-f_{\theta }(\mathbf{x})}}{Z_{\theta }}$$

其中$f_{\theta }(\mathbf{x})$是要训练的函数, 参数为$\theta$, 而$Z_{\theta }$是一个归一化常数, 它是用来保证$\int p_{\theta }(\mathbf{x})dx=1$, $f_{\theta }(\mathbf{x})$叫做energy-based model.

energy-based model的具体解释:
VAE这种模型是直接建模数据分布, 而EBM是隐式建模分布. 其不真正求解分布, 而是定义一个能量函数, 该能量函数在真实样本处(正样本)较小, 而在负样本处较大, 从而使得上式中的$p$在正样本处较大, 负样本处较小, 相当于一个代理的作用. 当然代理求解的代价就是需要$Z$来配平(积分为1), 保证学到的分布是合理的.

因此我们的优化目标是最大化在真实样本处的似然:

$$\underset{\theta }{\max }\sum _{i=1}^N\log p_{\theta }\left({\mathbf{x}}_i\right)$$

但为了保证概率和为1的约束, 对于不同的$f_{\theta }(\mathbf{x})$, 分母$Z$都要变化, 并且即使对于特定的$f_{\theta }(\mathbf{x})$, 我们也很难计算积分. 这样一来, 基于似然的模型就必须来限制模型架构让$Z$满足要求, 或者就要像VAE中的变分推断, 或对比散度中的MCMC(Markov链蒙特卡罗方法)采样来计算这个参数, 但是这样会很复杂.

那么有没有什么办法去掉这个$Z$呢? 答案是有的. 观察$p_{\theta }(\mathbf{x})$的形式, 分子是个指数, 分母是个常数(对于KaTeX parse error: Undefined control sequence: \nathbf at position 1: \̲n̲a̲t̲h̲b̲f̲{x}), 因此我们取log, 不就可以去掉分母了么? 也就是说, 我们不是预测$p$, 而是预测$p$的log的梯度:

$${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$$

这就是作者提出的score function, 对应的模型就称作score-based model, 记作${\mathbf{s}}_{\theta }(\mathbf{x})$. 这样一来, 有:

$${\mathbf{s}}_{\theta }(\mathbf{x})={\nabla }_{\mathbf{x}}\log p_{\theta }(\mathbf{x})=-{\nabla }_{\mathbf{x}}\log f_{\theta }(\mathbf{x})$$

在训练的时候, 和其他模型相似, 采用的也是Fisher散度(差的二范数的期望), 来衡量估计的量和真实量的差别,

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

当然直接计算是很难的, 因为我们根本不知道真实的分布$p(\mathbf{x})$. 好在有一类方法, 就是score matching, 可以用来在不知道真实分布的情况下来最小化这个散度. score matching的目标可以直接在数据集上估计，并通过随机梯度下降进行优化, 类似于训练基于似然的模型的对数似然目标.

此外, 用Fisher散度的好处是, 我们确实也不需要真正学到真实的score function, 而是只要求它们的梯度相似即可. 因为梯度指引了方向, 因此只要梯度相似了, 我们随机初始化一个点, 它就可以沿着梯度到达概率值最大(最符合真实样本)的地方.

score matching的具体说明:

为什么score matching可以绕过真实的概率分布? 绕过后的损失函数又是什么? 由于推导过程比较复杂, 就手写了:

3. 朗之万动力学

在我们得到训练好的score-based model ${\mathbf{s}}_{\theta }(\mathbf{x})$后, 我们就可以通过朗之万动力学的方式来采样.

因为朗之万动力学描述的是一个通过马尔科夫链蒙特卡洛方法, 用score函数(-log的梯度)来进行采样的过程. 其迭代方式如下:

$${\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$$

其中$ϵ$服从标准正态分布. 其中$\log p(\mathbf{x})$就是我们学习到的${\mathbf{s}}_{\theta }(\mathbf{x})$.

关于这部分的细节就先略去了, 我猜测作者是先看到的这种采样方式, 然后去思考的score function~

4. 朴素的基于score的生成模型, 以及它的陷阱

其实在整个数据空间中, 真实的数据分布的值比较大的地方, 是比较少的, 大部分区域都是无意义的样本. 这就导致一个问题, 在pdf(概率密度函数)值小的地方, 我们不能很好地去学习score function:

这样就导致在低密度区域, 学习的就不够准确, 我们学习的是梯度, 因此梯度不准确, 就导致我们无法很好地进行采样生成了, 如下图所示:

5. 用多个噪声干扰来优化score-based生成模型

那么针对上面这个问题, 怎么解决呢? 一个办法是用噪声去干扰数据样本, 这样增加的扰动就可以扩充高概率值部分所占的区域. 然后我们在加噪的样本上训练即可. 下图是一个对两个高斯分布混合的分布进行加噪的例子, 可以看到加噪之后预测的梯度就大大的准确了:

那么问题又出现了: 我们怎么知道我们该加多大的噪声呢? 幅值较大的噪声可以让高概率区域占据更大的面积, 但这也会导致加噪的样本根本就不像原来的样本了. 但幅值小的噪声呢? 它的覆盖效果又不那么好.

为了结合大噪声和小噪声的有点, 我们可以都用! 假设我们用各向同性的高斯噪声来扰动样本, 而且选用$L$种尺度的噪声, 其标准差分别为${\sigma }_1,\dots ,{\sigma }_L$. 对第$i$个噪声, 干扰后的分布为:

在实现中, 对于样本$\mathbf{x}$, 按照$\mathbf{x}+{\sigma }_i(\mathbf{z})$计算即可. 这样我们对每一级加噪的样本, 都训练一个score-based model, 现在是以加噪为先决条件的, 所以叫noise-conditioned score-based model.

这样, 总的损失函数就是各级score function和真值的Fisher散度的加权和, 权重一般选为对应噪声的方差. 我猜测, 噪声越大, 权重越大的原因是大噪声下的分布趋于平滑, 梯度下降, 因此为了平衡各个级别噪声的scale才这么选用权重.

那么训练好之后, 怎么恢复的? 既然我们是逐级加噪, 那么逐级去噪就可以了! 所以按照$i=L,\dots ,1$的顺序运行朗之万动力学, 如下图:

当然, 在实际应用中, 还有一些技巧:

1. 对${\sigma }_1<{\sigma }_2<\dots$的选择, 采用一个几何过程, 也就是${\sigma }_1$相对较小, 而${\sigma }_L$几乎和数据样本点的距离最大值相当. $L$一般成百上千.
2. 预测${\mathbf{s}}_{\theta }$的网络采用UNet.
3. 在测试阶段, 对模型权重, 采用EMA的方式更新.

5. Score-based生成模型和随机微分方程(SDE)

根据前面的讨论, 给原本的数据加噪有利于扩充概率分布中高概率值的区域, 从而可以更好地训练. 当我们将噪声的个数加到无穷的时候, 不仅会得到高质量的样本, 而且可以精准地计算对数似然, 以及可以通过逆问题求解的方式可控地进行生成. 为什么呢? 本节进行解释.

5.1 通过SDE描述干扰数据(加噪)过程

当我们把噪声的数量加到无穷之后, 这个噪声干扰的过程可以视作一个连续时间的随机过程. 我们怎么去表示这个随机过程呢? 很多随机过程, 具体来说, 扩散过程, 是通过随机微分方程表示的, 形式如下:

直观上理解, 随机微分方程表明, $x$的变化是由时间变化和布朗运动决定的, 二者的权重分别是$f(x,t)$和$g(t)$. 注意$w$表示布朗运动, 而$dw$可以近似看作小的白噪声. SDE解的形式为 { x ( t ) } t ∈ [ 0 , T ] \{\mathbf{x}(t)\}_{t\in [0, T]} {x(t)}t∈[0,T]​. 也就是说, 解是一簇时间相关的随机变量的集合. 理解上, 注意随机过程的时间随机性和变量随机性.

假设$\mathbf{x}(t)$对应的pdf为$p_t(\mathbf{x})$, 则$p_0(\mathbf{x})$就是数据的原始分布, 而经过足够长的加噪过程之后, $p_T(\mathbf{x})$实际上和噪声的分布就很像了, 称作先验分布. 其实我们上面进行多级加噪, 和这里按照时间进行的随机加噪过程是等价的.

比如说, SDE是如下形式的话:

$$d\mathbf{x}=e^{t}d\mathbf{w}$$

它表示的过程就是, 在每一步$\mathbf{x}$的变化中, 差异都是均值为0, 方差为指数增长的一个高斯分布 ($dw$近似是高斯白噪声, 由高斯分布的性质得到), 也就等价于我们上面说的加入多级噪声.

5.2 求SDE的逆过程, 用于去噪生成

这里直接给出结论, 一个SDE对应了一个逆形式, 即:

所以加噪和去噪过程, 示意图如下:

在去噪过程中, 我们就可以用我们估计的score function了.

5.3 通过score-based模型和score matching来估计逆SDE

要求解这个逆过程, 我们需要知道最后一步数据的分布$p_T(\mathbf{x})$以及每一步的score function${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})$. 我们已经知道$p_T(\mathbf{x})$近似于噪声分布, 所以这个已知了. 那么问题就在于如何求解score function, 当然我们发现score function是和时间有关的, 因此就需要训练一个时间依赖的score-based模型${\mathbf{s}}_{\theta }(\mathbf{x},t)$.

我们训练的目标依旧是最小化各个时间步的Fisher散度:

训练方式还是通过score matching. 这样去噪过程变为:

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

后面的概率流常微分方程和条件生成部分暂略, 有时间回来补充.

6. 和Diffusion model的关系

直接说结论: Score-based model和Diffusion model都有逐级加噪去噪的过程, 不过它们是对一类模型的不同描述, 一个是通过朗之万动力学, 来准确地建模概率分布的梯度, 并且可以自然地求解逆过程; 而另一个和VAE紧密相连, 是通过ELBO进行训练的.

在ICLR21的一篇工作中, 作者证明了Score-based model和Diffusion model可以视为是SDE的不同离散化的方式, 所以是殊途同归的.