PhoenixPeng's blog.

扩散模型——4.分数模型(SDEs)

2023/07/26

一、前言

  你是否想过,当我们在玩生成模型的时候,实际上是在玩什么?其实就是想让它表示概率分布,从而能够通过其实现采样生成。

NCSN 基于噪声的分数模型

  至于如何表示,不同的模型有不同的方法,但它们大致都可被归纳为两种范式,分别是:对数据的采样过程建模对概率密度进行建模。前者不纠结于数据分布的概率密度,而是通过其它方法达到表示概率分布的目的,因此也被称为隐式(implicit)生成模型;而后者则直接让模型去估计概率密度,不绕圈子,于是被称为显式(explicit)生成模型。

  而在显示生成模型中,有两类概率生成模型,它们都涉及到了如何缓慢的增加噪声从而迭代的破坏训练数据,然后学习扭转这种破坏以形成数据的生成模型。它们分别是DDPM以及SMLD(朗之万动力学),DDPM我们再熟悉不过了,但是对于SMLD而言似乎还有些许陌生,所以这一节主要是介绍score-based model的一些核心思想。

1.什么是分数?🤖

  像VAE等生成模型是直接学习数据分布的概率\(\textcolor{blue}{p_\theta(x)}\),从下面的示意图可以看到,其颜色越深的地方代表数据分布的概率越高,密度越高;相反越浅的地方则代表没有什么目标数据。

  Score-based model其核心就是分数(score),它不是直接学习数据分布的概率,而是学习score。所以首先要解释一下score这一概念。

  当我们能够从未知数据分布\(\textcolor{blue}{p(x)}\)中采样出独立同分布的样本\(\textcolor{blue}{\lbrace x_i\in R^D\rbrace^N_{i=1}}\),我们定义分数为概率密度函数\(\textcolor{blue}{p(x)}\)的对数似然关于x的梯度即对数概率的梯度\(\textcolor{blue}{\nabla_x \log p(x)}\)。而分数模型,就是用于估计(预测)分数

  数据往往是多维的。由分数的定义以及从数学的角度出发来看,它应当是一个“矢(向)量场”(vector field)。既然是向量,那么就有方向,这个方向就是:对于输入数据(样本)来说,其对数概率密度增长最快的方向,其箭头代表分数,圆弧描绘了概率密度的情况。然后我们可以看到这些梯度的整体趋势是指向数据密度越高的位置。

  在采样过程中沿着分数的方向走,就能够走到数据分布的高概率密度区域,最终生成的样本就会符合原数据分布。这也就是为什么我们想要估计分数的原因:

  那么可以进一步定义分数网络为\(\textcolor{blue}{s_\theta:R^D \rightarrow R^D}\),它的作用是对分数进行建模,即该网络是由参数\(\textcolor{blue}{\theta}\)所构成,训练的目的是使得网络的输出能够逼近我们真实分布\(\textcolor{blue}{p_{data}(x)}\)的分数。基于分数的生成模型的框架涉及到两个要素:Score matching(分数匹配)以及Langevin dynamic(朗之万动力学)

2.朗之万动力学采样(Sampling with Langevin Dynamics)🦊

  概括地来说,该方法首先从先验分布随机采样一个初始样本,然后利用模型估计出来的分数\(\textcolor{blue}{\nabla_x \log p(x)}\)逐渐将样本向数据分布的高概率密度区域靠近。为保证生成结果的多样性,我们需要采样过程带有随机性。正好!布朗运动就是带有随机性的,朗之万动力学方程因此也带有随机项。

  朗之万动力学:它是一种采样算法,仅仅只用到分数函数\(\textcolor{blue}{\nabla_x \log p(x)}\)就能从概率密度\(\textcolor{blue}{p(x)}\)中生成样本。当给定一个步长\(\textcolor{blue}{\epsilon>0}\)并给定一个初始值\(\textcolor{blue}{\tilde{x}_0 \backsim \pi(x)}\)(从先验分布\(\textcolor{blue}{\pi}\)采样的初始值),那么朗之万算法就能通过下面这个公式去迭代:

\[\huge \tilde{x}_t = \tilde{x}_{t-1}+\frac{\epsilon}{2}\nabla_x\log p(\tilde{x}_{t-1})+\sqrt{\epsilon}z_t-(1)\]

  其中\(\textcolor{blue}{z \backsim N(0,I)}\),\(\textcolor{blue}{\epsilon}\)是预定义的步长,当\(\textcolor{blue}{\epsilon\rightarrow0}\)并且\(\textcolor{blue}{T\rightarrow\infty}\)时,我们就可以近似的认为上式得到\(\textcolor{blue}{\tilde{x}_T}\)的分布等于\(\textcolor{blue}{p(x)}\),也就是说\(\textcolor{blue}{\tilde{x}_T}\)能够从\(\textcolor{blue}{p(x)}\)分布中进行采样。由此我们可以看出朗之万算法仅仅依赖于分数函数\(\textcolor{blue}{\nabla_x \log p(x)}\),因此我们希望能够用朗之万算法从一个分布\(\textcolor{blue}{p_{data}(x)}\)中获取样本的话,我们就首先需要训练分数网络使得\(\textcolor{blue}{s_\theta(x) \approx \nabla_x \log p_{data}(x)}\),这样上面的公式中的分数就能够通过该分数网络去替代,由此,我们就能够从随机采样噪声出发一步步逼近目标数据分布,最后从分布中采样出样本。这就是基于分数模型生成样本的一个主要思想。

3.分数匹配🦚

  由上面的朗之万算法,我们可以直接训练分数网络\(\textcolor{blue}{s_\theta(x)}\)来去估计\(\textcolor{blue}{\nabla_x \log p_{data}(x)}\),并不需要训练网络去估计\(\textcolor{blue}{p_{data}(x)}\),那么我们就需要一个目标函数去对其算出损失,然后进一步去优化,而该训练目标可表示为:

\[\huge min\frac{1}{2}E_{p_{data}}[||s_\theta(x)-\nabla_x \log p_{data}(x)||_2^2]-(2)\]

  然而,这就要求我们知道概率密度\(\textcolor{blue}{p_{data}(x)}\),但通常我们是无法知道原数据所服从的分布。并且,如果提前知道了数据分布,那么直接从其中采样生成样本就好了,还搞那么多飞机来炼丹干嘛。

  那么分数匹配就是用来干这个的,它就是为了让我们在不了解概率分布的情况下,也可以让分数网络输出的分数能够匹配\(\textcolor{blue}{p_{data}(x)}\)的分数,这种方法利用了分部积分,将式1转化为了如下的函数形式:

\[\huge E_{p_{data}(x)}[tr(\nabla_x s_\theta(x))+\frac{1}{2}||s_\theta(x)||_2^2]-(3) \]

  (2)转成(3)的推导过程如下:

\[\huge \frac{1}{2}E_{p_{data}}[||s_{\theta}(x)-\nabla_x \log p_{data}(x)||_2^2] \]

  把上面loss的期望写开,二次项打开,可以得到:

\[\Large \frac{1}{2}\int p_{data}(x)[||s_{\theta}(x)||_2^2-2(\nabla_x \log p_{data}(x))^Ts_{\theta}(x)+||\nabla_x \log p_{data}(x)||_2^2]dx-(4) \]

  其中,[] 内的最后一项对于模型参数\(\textcolor{blue}{\theta}\)来说是常数项,因此我们可以忽略掉;而第一项,通过模型的输出可以直接计算出来,于是可以先不管它,在这里我们主要关注第二项:

\[\Large \begin{aligned} & \frac{1}{2}\int p_{data}(x)[-2(\nabla_x \log p_{data}(x))^Ts_{\theta}(x)]dx\\ &=-\int p(x)[(\frac{\partial \log(p(x))}{\partial x})^Ts_{\theta}(x)]dx \\ &= -\int p(x)[\sum_{i=1}^N\frac{\partial \log(p(x))}{\partial x_i}s_{\theta_i}(x)]dx \\&=-\sum_{i=1}^N\int p(x)[(\frac{1}{p(x)}\frac{\partial p(x)}{\partial x_i})s_{\theta_i}(x)]dx \\ &=-\sum_{i=1}^N \int[\frac{\partial p(x)}{\partial x_i}s_{\theta_i}(x)]dx \end{aligned}\]

  根据分部积分法:\(\textcolor{blue}{u^{\prime}v=(uv)^{\prime}-uv^{\prime}}\),则\(\textcolor{blue}{\int u^{\prime}v dx=uv-\int uv^{\prime}dx}\),我们可以得到:

\[\Large \begin{aligned} & -\sum_{i=1}^N \int[\frac{\partial p(x)}{\partial x_i}s_{\theta_i}(x)]dx \\ &=-\sum_{i=1}^N \int[\frac{\partial p(x)s_{\theta_i}(x)}{\partial x_i}-p(x)\frac{\partial s_{\theta_i}(x)}{\partial x_i}]dx \\ &=-\sum_{i=1}^N\int p(x)s_{\theta_i}(x) |_{- \infty}^{+\infty}dx_{/i}+\int p(x)\frac{\partial s_{\theta_i}(x)}{\partial x_i}dx \end{aligned}\]

  第一项为0可忽略,故:

\[\Large \begin{aligned} &\sum_{i=1}^N \int p(x)\frac{\partial s_{\theta_i}(x)}{\partial x_i}dx\\&=\int p(x) \sum_{i=1}^N \frac{\partial s_{\theta_i}(x)}{\partial x_i} dx \\ &= \int p(x) tr(\frac{\partial s_{\theta}(x)}{\partial x})dx \\ &= \int p(x) tr(\nabla_x s_{\theta}(x))dx\end{aligned}\]

  最后将上式与式(4)的第一项相加即可得到:

\[\Large \begin{aligned} & \int [p(x)tr(\nabla_x s_{\theta}(x))+ \frac{1}{2}p(x)||s_{\theta}(x)||_2^2]dx\\ & =\int p_{data}(x)[tr(\nabla_x s_\theta(x))+\frac{1}{2}||s_\theta(x)||_2^2]dx \\ &= E_{p_{data}(x)}[tr(\nabla_x s_\theta(x))+\frac{1}{2}||s_\theta(x)||_2^2]\end{aligned}\]

  其中\(\textcolor{blue}{\nabla_x s_\theta(x)}\)为雅可比矩阵(雅可比矩阵是函数的一阶偏导数以一定方式排列成的矩阵),而\(\textcolor{blue}{tr(\cdot)}\)为矩阵的迹,其定义为方阵上对角线所有元素之和。如果我们数据分布很高的情况下,上式并不太好能够泛化,计算复杂度比较高(当遇到高维数据时,其图像特征的通道数常都成百上千,那么就需要经历多次反向传播来求偏微分),那么后续就有人研究如何绕过雅可比矩阵的求解,就引出了去噪分数匹配。

其实也有的方法是对其降维,但后续的NCSN采用的是去噪分数匹配方法,所以这里不作介绍,有兴趣可以自查。

4.去噪分数匹配🐳

  这个做法其实是绕过了上面计算雅可比矩阵的方式,回归到了式(2)的做法。那么我们就回到了最初的问题上,我们是不知道\(\textcolor{blue}{p_{data}(x)}\)的分布的,但我们可以自己搞一个出来。首先对原始数据样本\(\textcolor{blue}{x}\)按照预定好的噪声分布\(\textcolor{blue}{q_\sigma(\tilde{x}|x)}\)去加噪,使得其结果满足我们预先定义好的分布,比如高斯分布。这样,我们就知道了现在的概率密度,接下来就可以使用式(2)的方法来去估计加噪后的数据的分数。

  记训练样本为\(\textcolor{blue}{x}\),加噪后的数据为\(\textcolor{blue}{\tilde{x}}\),预定义分布为\(\textcolor{blue}{q_\sigma(\tilde{x}|x) \sim \mathcal{N}(\tilde{x};x,\sigma^2I)}\),加噪后的数据分布表示为:

\[\huge q_\sigma(\tilde{x}_t) = \int{q_\sigma(\tilde{x}|x)p_{data}(x)}dx\]

  我们在这个数据分布上计算的分数表示为:

\[\huge \frac{\partial \log q_\sigma(\tilde{x}_t)}{\partial\tilde{x}} \equiv \frac{\partial \log (q_\sigma(\tilde{x}|x)p_{data}(x))}{\partial\tilde{x}}=\frac{\partial \log (q_\sigma(\tilde{x}|x))}{\partial\tilde{x}}\]

  由于\(\textcolor{blue}{p_{data}(x)}\)\(\textcolor{blue}{\tilde{x}}\)不相干,于是导数为0,因此最右边那项没有了它的身影。那么其目标函数我们可以进一步写出:

\[\huge\frac{1}{2} E_{q_\sigma(\tilde{x}|x)p_{data}(x)}[||s_\theta(\tilde{x})-\nabla_{\tilde{x}} \log q_\sigma(\tilde{x}|x)||_2^2]\]

  一旦\(\textcolor{blue}{s_\theta(\tilde{x})}\)收敛之后,我们就有理由相信:\(\textcolor{blue}{s_\theta^*(x) = \nabla_x \log q_\sigma(x) }\),此时的\(\textcolor{blue}{q_\sigma(x)}\)仅仅只是我们加噪后的分布,并不是初始的数据分布\(\textcolor{blue}{p_{data}(x)}\)。当我们加的噪声方差\(\textcolor{blue}{\sigma}\)很小的时候,我们可以近似的认为:\(\textcolor{blue}{q_\sigma(x) \approx p_{data}(x)}\),如果当噪声方差\(\textcolor{blue}{\sigma}\)足够大的时候,\(\textcolor{blue}{q_\sigma(x) \approx \mathcal{N}(x_t;0,\sigma_{max}^2I)}\)

  在这种方法下,我们只要从标准高斯分布\(\textcolor{blue}{N(0,I)}\)中采样随机噪声\(\textcolor{blue}{\epsilon}\),然后乘上我们预定义的\(\textcolor{blue}{\sigma}\),接着加到样本中\(\textcolor{blue}{x}\),从而就能得到加噪后的样本\(\textcolor{blue}{\tilde{x}=x+\sigma \epsilon}\)。根据高斯分布的性质,这样加噪后\(\textcolor{blue}{q_\sigma(\tilde{x}|x)}\)就会满足分布\(\textcolor{blue}{\mathcal{N}(\tilde{x};x,\sigma^2I)}\)估计分数时,记得要将加噪后的样本\(\textcolor{blue}{\tilde{x}}\) 喂给模型,而非原数据样本\(\textcolor{blue}{x}\)

5.分数模型的困境🦅

  分数模型在实际应用时会面临一些困难。概括起来,主要是以下三点:

1️⃣loss 不收敛

2️⃣模型估计的分数不准

3️⃣生成结果与原数据分布偏差大

  接下来将逐步分析这三点

5.1 loss不收敛

  如果你根据前文介绍的 score matching 去训练分数模型,会发现模型被“玩坏”了:loss 不断震荡,可谓“一秒天堂,一秒地狱”。

  为何会这样呢?明明前面分析得头头是道。这不得不引出先一种叫 “流形假设”(mainfold hypotheis) 的理论。

  流行假设认为,生活中的真实数据大部分都倾向于分布在低维空间中,这也是大名鼎鼎的 “流形学习”(mainfold learning) 的大前提。也就是说,我们的编码空间(通过神经网络等一系列方法)的维度可能很广泛,但是数据经过编码后在许多维度上都存在冗余,实际上用一部分维度就足以表示它,说明实质上它仅分布在整个编码空间中的一部分(低维流形),并没有“占满”整个编码空间。

  基于流形假设,就可以来解释 loss 不收敛的问题了:

  首先,分数这个东西:\(\textcolor{blue}{\nabla_x \log p(x)}\) ,是针对整个编码空间定义的,然而当我们的数据仅分布在编码空间的低维流形时,分数就会出现“没有定义” (undefined) 的情况。

  另外,score matching 的训练目标是基于数据“占满”了整个编码空间这个前提下而推出来的。当数据仅分布在低维流形中时,score matching 的方法就不适用了。

  综合以上两点,在基于 score matching 的方式下,训练分数模型就会出现 loss 震荡而不收敛的情况。

5.2 分数估计不准确

  在数据密度较低的区域,会由于没有足够的样本来供模型训练学习,从而导致这部分的分数估计不准,实质上就是训练不充分

在低数据密度区域,由于缺乏数据样本,分数匹配可能没有足够的证据来准确估计分数函数。为了了解这一点,回顾前文提到的分数匹配中最小化分数估计平方误差的期望\(\textcolor{blue}{E_{p_{data}}[||s_\theta(x)-\nabla_x \log p_{data}(x)||_2^2]}\),即实际上,关于数据分布的期望总是由独立同分布的样本\(\textcolor{blue}{\lbrace x_i\rbrace^N_{i=1} \backsim p_{data}(x)}\)估计的。考虑任务使低密度区域\(\textcolor{blue}{R\in R^D}\),\(\textcolor{blue}{p_{data}(R) \approx 0}\)。大多数情况下\(\textcolor{blue}{\lbrace x_i\rbrace^N_{i=1} \bigcap R =\emptyset}\),分数匹配将没有足够的数据样本来准确估计出对于\(\textcolor{blue}{x \in R}\)\(\textcolor{blue}{\nabla_x \log p_{data}(x)}\)

  由上图我们可以看出来,在概率密度比较高的地方数据分数以及估计分数比较匹配,但是在概率密度比较低的区域,就开始估计的不太精确。其中右图颜色(橘红色)深浅代表数据概率密度的大小(越深越大,对应的样本越可能出现),箭头代表分数。这个问题会导致在早期推理的时候,模型很容易根据错误的梯度而”脱轨“,不容易获得较好的结果。

5.3 生成结果偏差大

  分数估计不准,带来的一个结果就是生成效果不佳,采样生成出来的样本不太符合原数据分布。

  由于分数模型的生成是根据模型估计的分数并且基于朗之万动力学采样去实现的。当数据分布是由多个分布按一定比例混合而成时,在朗之万动力学采样的方法下,即使用真实的分数(对应式 (1) ),采样生成出来的结果也可能不能反应出各个分布之间的比例关系。

  这是因为我们在根据数据梯度行走时,比如做选择是向左高密度区域走还是向右高密度区域走,往往可能是随机的,这也是朗之万采样使用贪心思想的问题。如下图,在实际的数据分布中(a),右上部分的分布对应的样本比较多,左下部分的比较少;而生成的结果(b)确是两种分布对应的样本一样多,并没有反映出原始各分布之间的关系。

  看到这你可能觉得,这分数模型根本用不下去呀,面临的问题重重,但是后来有工作提出了NCSN网络。成功的将这三个问题完美解决,那么接下来我们就来看它是如何做到的。

二、Noise Conditional Score Network(NCSN)

  为了解决上述所带来的问题,这里引入了新的分数网络NCSN,实际上就是用高斯噪声去扰动数据

  针对问题1️⃣:NCSN 用了高斯噪声去扰动数据,而高斯噪声是分布在整个编码空间的,因此,扰动后的数据就不会仅“驻扎”在低维流形,而是能够有机会出现在整个编码空间的任意部分,从而避免了分数(score)没有定义(undefined)的情况,同时使得 score matching适用,这就破解了“loss 不收敛”的困局。

  针对问题2️⃣:对数据加噪声,以此扩大数据范围,从而让原本低密度的数据区域膨胀,这样能够比较准确估计score的区域就会增大,从而提升了生成样本的准确性,这破解了“分数估计不准”的困局。

  右图所展示的是我们的数据分布被加噪后的样子,我们可以由此发现,加噪后的区域其低密度区域已经很少了,而加噪后的分数已经估计的分数都比较准确。

  针对问题3️⃣:事实上在加噪扰动后,原本“稀疏”的低概率密度区域就能够被“填满”,并且两个分布中混合比例高的那个分布所占这部分(“填满”后的区域)的比例会更高扰动后的分数\(\textcolor{blue}{\nabla_x \log p(x)}\) 将更多地由混合比例高的那个分布的分数“演变”而来。这就是说,分数的方向将更多地指向混合比例高的那个分布。最终,基于分数去进行采样生成的时候,将更多地生成混合比例高的那个分布的样本,从而生成的整体分布就能够成功反映出原数据中两个分布的相对关系了。解决了问题!

  那么我们该如何进行加噪呢?我们知道当加噪的噪声足够小的时候,会有\(\textcolor{blue}{q_\sigma(x) \approx p_{data}(x)}\);但问题是加的噪声太小,我们仍然不是很好的能够解决低密度的问题。那么如果我们把噪声加的太大的话,就会导致数据分布被破坏掉,并且当我们进行生成的时候,最后我们也只能走到带有噪声的数据分布,并不是干净的。

  🔥较强的噪声——更多区域可以准确估计score——更严重损坏原本的数据分布

  💧较弱的噪声——避免损坏太多原本的数据分布——无法在大多区域准确估计score

  为了平衡这个加噪方式带来的后果,最终的解决思路是在不同的阶段加不同强度的噪声,噪声从大到小。在一开始的时候加较大的噪声,因为这个时候当前位置离目标数据分布还比较远,那么在推理早期的时候就能够在当前位置估计出准确的score,也就是准确的梯度方向。当推理进行到中期或者后期的时候,这个时候,已经离目标分布很接近了,这个时候就不需要加很强的噪声,也能够准确的估计出score,并且被加噪后的分布也可以近似是看作成干净的目标分布,这样数据的整个生成就完成了。

注:介绍到这,如果你只关注SDEs那么直接跳到第三章既可以,下面是对分数模型的一些扩展,如果你感兴趣也可以继续看下去。

1.噪声设计原则🚩

  NCSN 在不同的阶段加了不同强度的高斯噪声,每种噪声强度所对应的分布可表示为\(\textcolor{blue}{\mathcal{N}(0,\sigma^2_i)}\),那么到底应该如何选取噪声强度呢\(\textcolor{blue}{\sigma_i}\)?作者的想法很简单:我们先加噪大一些,让数据先移动到高密度区域那块,再加噪小一些,能让数据在高密度区域更精准地到更符合数据分布的区域。

  具体地,作者采取了一个噪声序列\(\textcolor{blue}{\{\sigma_i\}^L_{i=1}}\),并满足\(\textcolor{blue}{\frac{\sigma_1}{\sigma_2}=...=\frac{\sigma_{L-1}}{\sigma_L}>1}\),(同时\(\textcolor{blue}{\sigma_1}\)需要足够大,以能够充分“填充”低概率密度区域;而\(\textcolor{blue}{\sigma_L}\)得足够小,以获得对原数据分布良好的近似,避免过度扰动),我们有:

\[\huge q_\sigma(x)=\int p_{data}(t)\mathcal{N}(x|t,\sigma^2I)\]

  其中\(\textcolor{blue}{q_\sigma(x)}\)就是被加噪之后的数据分布,那我们的score估计函数即为:

\[\huge s_\theta(x,\sigma) \approx \nabla_x \log q_\sigma(x)\]

2.网络结构设计🌈

  由于生成的图片和原图一样大小,也就是每个像素点都需要由朗之万动力学采样生成,因此模型对于每个像素点都要估计其对应的分数。也就是说,从张量的角度来看,网络的输出(分数)要和输入图像的形状(shape)一致

  作者的网络结构采用了以噪声为条件的UNet,其中加入了 空洞卷积(dilated/atrous convolution) 和 以噪声为条件的实例归一化(conditional instance normalization++)

  另外,即使对于同一个像素点,在不同的噪声强度下也要对应估计出不同的分数,于是,模型还需要将噪声强度\(\textcolor{blue}{\sigma_i}\)也作为输入

3.模型训练方法🎨

  对于训练方法,作者采用了去噪分数匹配。通过前文介绍的去噪分数匹配,我们已经知道其训练目标为:

\[\huge\frac{1}{2} E_{q_\sigma(\tilde{x}|x)p_{data}(x)}[||s_\theta(\tilde{x})-\nabla_{\tilde{x}} \log q_\sigma(\tilde{x}|x)||_2^2]-(4)\]

  由于我们有\(\textcolor{blue}{q_\sigma(\tilde{x}|x) \sim \mathcal{N}(\tilde{x};x,\sigma^2I)}\),故有\(\textcolor{blue}{q_\sigma(\tilde{x}|x) =\frac{1}{\sqrt{2\pi \sigma}}e^{-\frac{(\tilde{x}-x)^2}{2\sigma^2}}}\),将其取对数并对\(\textcolor{blue}{\tilde{x}}\)求导,就可获得其分数的表示形式:

\[\huge \frac{\partial \log q_\sigma(\tilde{x}|x)}{\partial \tilde{x}}=\frac{\partial \log \frac{1}{\sqrt{2\pi \sigma}}+\partial \log e^{\frac{-(\tilde{x}-x)^2}{2\sigma^2}}}{\partial \tilde{x}}=\frac{\partial \frac{-(\tilde{x}-x)^2}{2\sigma^2}}{\partial \tilde{x}}=-\frac{(\tilde{x}-x)}{\sigma^2}\]

  我们可以将其代入到式4的目标函数当中,就能得到在某个噪声级别\(\textcolor{blue}{\sigma}\)对应的损失函数:

\[\huge\frac{1}{2} E_{p_{data}(x)}E_{\tilde{x} \sim \mathcal{N}(x,\sigma^2I)}[||s_\theta(\tilde{x},\sigma)+\frac{(\tilde{x}-x)}{\sigma^2}||_2^2]-(5)\]

  其中\(\textcolor{blue}{s_\theta(\tilde{x},\sigma)}\)代表网络在特定噪声级别\(\textcolor{blue}{\sigma}\)下估计的分数。

  而 NCSN 使用了多个噪声级别,于是,分别对它们的损失加权求和后再求均值,就得到了联合的损失函数,它表示为:

\[\huge L(\theta;\{\sigma_i\}^L_{i=1})=\frac{1}{L}\sum_{i=1}^L \lambda(\sigma_i)l(\theta;\sigma_i)-(6)\]

  其中\(\textcolor{blue}{\lambda( \sigma_i)>0}\)代表不同噪声级别的损失的权重。

  目前一个关键的地方就在于如何设置\(\textcolor{blue}{\lambda( \sigma_i)}\),作者的出发点是:加权后所有噪声级别的损失都在同一量级,不受\(\textcolor{blue}{\sigma_i}\)的大小影响。这样也保证了“公平性”,即不会有哪个噪声级别的损失非常大(小),从而使得模型过份“重视”(“忽略”)这个噪声级别所要学习的内容。

  最终,这个经验性的结论启发了作者\(\textcolor{blue}{\lambda( \sigma_i)}\)设置为\(\textcolor{blue}{\sigma^2_i}\) ,具体原因可以参考原文。如果将\(\textcolor{blue}{\lambda( \sigma_i)=\sigma^2_i}\)代入式6并结合式5,我们可以发现

\[\huge \lambda(\sigma_i)l(\theta;\sigma_i)=\sigma_i^2l(\theta;\sigma_i)=\frac{1}{2}E[||\sigma_i s_{\theta}(\tilde{x},\sigma_i)+\sigma_i \frac{(\tilde{x}-x)}{\sigma^2_i}||_2^2]\]

\[\huge = \frac{1}{2}E[||\sigma_i s_{\theta}(\tilde{x},\sigma_i)+\frac{(\tilde{x}-x)}{\sigma_i}||_2^2]-(7)\]

  前文讲述去噪分数匹配时有提到,由于\(\textcolor{blue}{q_\sigma(\tilde{x}|x) \sim \mathcal{N}(\tilde{x};x,\sigma^2I)}\),因此有\(\textcolor{blue}{\tilde{x}=x+\sigma \epsilon}\),那么其实\(\textcolor{blue}{\frac{(\tilde{x}-x)}{\sigma_i}}\)就是噪声\(\textcolor{blue}{N(0,I)}\),这是与\(\textcolor{blue}{\sigma_i}\)无关的量,即不受后者影响。

  其次,作者发现当网络训练到收敛时,估计出的分数的量级\(\textcolor{blue}{||s_\theta(\tilde{x},\sigma_i)||_2}\)即 L2 norm)在\(\textcolor{blue}{\frac{1}{\sigma}}\)这个水平,于是\(\textcolor{blue}{||\sigma_i s_\theta(\tilde{x},\sigma_i)||_2}\)的量级就会是1。

  综上,\(\textcolor{blue}{\lambda( \sigma)=\sigma^2}\)的设置下,加权后的损失\(\textcolor{blue}{\lambda(\sigma)l(\theta;\sigma)}\)的量级就会与\(\textcolor{blue}{\sigma}\)的大小无关,即能够保证所有噪声级别的损失都在同一量级,使得模型能够“一视同仁”。

  在这种训练方式下,只要预先设定好\(\textcolor{blue}{\{\sigma_i\}^L_{i=1}}\),然后在每次迭代时随机选取其中一个噪声级别\(\textcolor{blue}{\sigma_i}\),并采样一个标准高斯噪声\(\textcolor{blue}{\epsilon}\),对原始数据样本加噪: \(\textcolor{blue}{\tilde{x}=x+\sigma_i \epsilon}\),接着将加噪后的样本\(\textcolor{blue}{\tilde{x}}\)与对应的噪声级别\(\textcolor{blue}{\sigma_i}\)一并喂给网络食用,待其吐出预测的分数\(\textcolor{blue}{s_\theta(\tilde{x},\sigma_i)}\)后,就可以使用式7来计算损失了。

4.采样生成方法🚀

  当然还有一个超参数:在噪声强度为\(\textcolor{blue}{\sigma_i}\)下的步长\(\textcolor{blue}{\alpha_i}\)如何设置呢?作者选择了 \(\textcolor{blue}{\alpha_i∝\sigma_i^2}\) 。这是为了固定信噪比的量级(具体计算可以参考原文)。也就是说噪声越大,步长越大,很合理吧!对于具体数值的设置,作者选择是:

\[\huge \alpha_i=\epsilon \cdot \sigma_i^2/\sigma_L^2\]

  万事俱备,我们还需要给这样改进的朗之万采样一个名字,那作者选择“退火”(anneal)这个词,因为这样加噪的思想和模拟退火算法很像。最后我们的退火朗之万采样算法即为:

  首先,初始样本\(\textcolor{blue}{\tilde{x}_0}\)从某种固定的先验分布去采样就好(作者默认采用均匀分布);然后,从最大的噪声级别\(\textcolor{blue}{\sigma_1}\)开始使用朗之万动力学采样,直至最小的噪声级别\(\textcolor{blue}{\sigma_L}\) 。在每个噪声级别进行采样前,会先设定步长(step size) \(\textcolor{blue}{\alpha_i}\);接着,在每个噪声级别下,基于一定的步数\(\textcolor{blue}{T}\)去迭代进行朗之万动力学采样。该噪声级别最后一步采样生成的样本会作为下一个噪声级别的初始样本\(\textcolor{blue}{\tilde{x}_0}\);最后,待所有噪声级别的朗之万动力学采样过程均完成时,就得到了最终的生成结果。

该算法有两层循环,第一层循环是\(\textcolor{blue}{i:1 \sim L}\)其代表是噪声的不同级别,在不同级别下都有预先设置好的步长\(\textcolor{blue}{\alpha_i}\)以及噪声\(\textcolor{blue}{\sigma_i}\) ,第二层循环是\(\textcolor{blue}{t:1 \sim T}\)。即在每个级别会采样\(\textcolor{blue}{T}\)步,当\(\textcolor{blue}{L\cdot T=1000,L=1000,T=1}\)相当于我们有1000个不同级别的噪声强度,并在不同强度下只采样一步,就跟DDPM是完全一样的了

  由于对数据进行了扰动,因此能够避免数据仅“驻扎”在低维流形,于是模型估计出来的分数就会比较准确。那么沿着分数(梯度)的方向,某个噪声级别下最终的采样生成结果就会聚敛到其对应的高概率密度区域。而由于相邻噪声级别之间的差异比较小,因此它们的高概率密度区域是非常近似的,相当于某噪声级别下样本最终聚敛到的区域也是下个噪声级别的高概率密度区域,而概率密度高,分数估计通常更准。这不是捡了个便宜嘛!对于下一个噪声级别来说,以当前噪声级别最终生成的样本作初始化无疑是白嫖一波~ 每个噪声级别都白嫖“前辈”们的,层层推波助澜,这就是所谓的站在巨人的肩膀上。

5.与扩散模型的联系🚅

  我们再看看文章给出的“Intermediate samples of annealed Langevin dynamics”:

  可以看到这不是和DDPM效果一样嘛!所以看openreview上,DDPM被诘问的最多的问题就是:你这和NCSN有什么区别

  事实上,DDPM是一种显式的去噪,NCSN也可以看成是一个去噪的过程,它是基于梯度来实现去噪生成的。那么根据\(\textcolor{blue}{\tilde{x}=x+\sigma \epsilon}\),我们有:

\[\huge \frac{\partial \log q_\sigma(\tilde{x}|x)}{\partial \tilde{x}}=\frac{-(\tilde{x}-x)}{\sigma^2}=-\frac{\epsilon}{\sigma}\]

  这说明分数的方向与所加噪声的方向是相反的!因此,在采样生成时沿着分数的方向走,就是往噪声的反方向走,这样就会走回原来数据样本的方向。这,就是实打实的去噪过程!

  进一步,加噪后\(\textcolor{blue}{q_\sigma(\tilde{x}|x)}\)所服从的分布,其均值是原本数据\(\textcolor{blue}{x}\) ,由于它是高斯分布,因此\(\textcolor{blue}{x}\)处就是概率密度最高的地方。而分数\(\textcolor{blue}{\nabla_{\tilde{x}} \log q_{\sigma}(x)}\)的方向就是指向概率密度增长最快的方向,所以沿着它走就会走回原数据样本\(\textcolor{blue}{x}\)之处。

6.条件生成🌕

  分数模型的核心是基于梯度来玩生成,这带来的一个便利之处就是能够简单地进一步实现条件生成。条件生成就是在基于指定的条件下去生成样本。打个比方:你输入描述“白色的猫”,并且将其编码后输入模型,模型就会在基于这个描述的条件下给你输出一张白色的猫的图片,而非输出任意颜色的猫或者根本不是猫。

  顺着分数模型基于梯度的思想,假设我们现在要求模型在基于\(\textcolor{blue}{y}\)的条件下去做生成,那么我就要让模型去估计该条件下的分数\(\textcolor{blue}{\nabla_{x} \log p(x|y)}\) 。进一步,通过贝叶斯定理,可以推导出:

\[\huge \frac{\partial \log p(x|y)}{\partial x}= \frac{\partial \frac{\log p(y|x)p(x)}{p(y)}}{\partial x}=\frac{\partial \log p(y|x)+\partial \log p(x)-\partial \log p(y)}{\partial x}\]

\[\huge =\frac{\partial \log p(y|x)}{\partial x} +\frac{\partial \log p(x)}{\partial x}\]

  最后面那一项就是本文一直讨论的无条件下的分数,而前面那一项代表什么呢?拿最简单的例子来说,\(\textcolor{blue}{y}\)代表物体类别时, \(\textcolor{blue}{p(y|x)}\)实际上就是一个分类器的输出。这就是说,我们只要在原来的分数模型的基础上额外加一个分类器,就可以完成该条件生成。并且,这个分类器与分数模型解耦,没有架构限制。

  进一步,如果说我们希望生成的结果更加符合指定的条件,那么就应当让模型更倚重去估计前面那一项。于是,我们可以为它加上一个权重,从而令训练目标变为:

\[\huge \frac{\partial \log p(x|y)}{\partial x}=\gamma \frac{\partial \log p(y|x)}{\partial x} +\frac{\partial \log p(x)}{\partial x}-(7)\]

   \(\textcolor{blue}{\gamma}\)越大,在模型的训练目标中,前面那项所占的比重就越大。特别地, \(\textcolor{blue}{\gamma=0}\)时,等价于无条件生成。

  而NCSN 是基于扰动后的数据去做分数匹配的,也就是式7中的\(\textcolor{blue}{x}\)其实是加噪后的数据样本,这就要求我们要用加噪后的数据去训练分类器。而通常 pretrained 的分类器都是在原始数据集(比如 MNIST、ImageNet) 上训的,这就不适用了

  可是,假如我们就是不想额外训一个分类器,与其要多训一个模型,还不如想想办法专心训好分数模型本身。那么有什么办法可以去掉分类器呢?那么就是要让\(\textcolor{blue}{\nabla_{x}\log p(y|x)}\)这一项消失。同样地,从这个式子,我们可以将它用以下代替:

\[\huge \frac{\partial \log p(y|x)}{\partial x}=\frac{\partial \log p(x|y)}{\partial x} -\frac{\partial \log p(x)}{\partial x}\]

  接着代入到式7,就可以得到:

\[\huge \frac{\partial \log p(x|y)}{\partial x}=\gamma \frac{\partial \log p(x|y)}{\partial x} +(1-\gamma)\frac{\partial \log p(x)}{\partial x}\]

  由此我们就不需要分类器了,只要将条件\(\textcolor{blue}{y}\)编码后一并输入到分数模型,再去估计其对应的“条件分数”即可!

  在实现时,可以将\(\textcolor{blue}{y}\)\(\textcolor{blue}{x}\)拼接(concat)在一起喂给模型。另外,为了避免让条件输入与无条件输入分别喂给模型导致需要经历两次独立的前向过程,我们可以为无条件的输入也赋予一个“无意义”的条件\(\textcolor{blue}{y^\prime}\) (比如编码为某个特定值,这个值专门预留给无条件的情况)。这样一来,两种情况的输入形式就对齐了,于是可以一并输入到模型通过一次前向过程就能完成条件与无条件的分数估计。

这里需要特别注意一点,就是前面需要分类器的那种做法中,\(\textcolor{blue}{\nabla_{x}\log p(y|x)}\)这一项并不是条件分数!因为它是\(\textcolor{blue}{y}\)的条件分布对\(\textcolor{blue}{x}\)求导,而想要成为分数,则必须是关于\(\textcolor{blue}{x}\)的分布对\(\textcolor{blue}{x}\)求导,如式7前面那一项\(\textcolor{blue}{\nabla_{x}\log p(x|y)}\)才是名副其实的条件分数。

  最后,给基于分类器与无分类器两种玩法冠个名头吧。其实它们早已名声在外,分别叫作:classifier-guidance & classifier-free guidance,这两个内容也会在后续的文章当中陆续介绍。

  光是前面的内容就足够一篇文章的内容了,然而这仅仅只是开始,下面才正式到SDE的介绍。

三、SDEs

  在扩散模型和NCSN当中我们都能看到,Diffustion的关键在于在多个不同规模上加入噪声扰动,通过将这些不同规模的噪声的数量扩展到无穷大,我们不仅可以得到高质量的生成样本,而且可以得到准确的Log-likelihood和可控的生成过程(比如给定生成标签)。

  注意到,不论是DDPM,还是NCSN而言,它们都存在一个共同的特性:加噪链和去噪过程都是离散化的,即他们按照时间步\(\textcolor{blue}{1 \sim T}\)被打散,我们已经强调过,当然时间\(\textcolor{blue}{T}\)需要足够大,才可以保证我们最终加噪后的样本趋近于标准正态分布。事实上,这是一个随机的过程,如果把上述有限步\(\textcolor{blue}{T}\)拓展到无限步会怎么样?因为经过实验验证, \(\textcolor{blue}{T}\)越大,可以得到更准确的似然估计,以及质量更好的结果。而这样我们对数据进行连续时间的扰动就可以被建模成一个随机微分方程 (stochastic differential equation, SDE)SDE的思想是想将这个离散的随机控制过程连续化,那么就要涉及随机微分方程(随机控制)的理论和原理了。Song Yang详细的对这个过程进行了可视化的描述(从图中我们可以看出想要通过一个连续的过程来描述该问题)。

  由此我们可以看出,作者提出了一种随机微分方程(SDE),它通过缓慢注入噪声将复杂的数据分布平滑地转换为已知的先验分布,以及相对应的反向 SDE,通过缓慢消除噪声将先验分布转换回数据分布。反向 SDE 仅取决于扰动数据分布的时间相关梯度场(分数)。通过利用基于分数的生成模型的优化,我们可以使用神经网络准确地估计这些分数,并使用数值 SDE 求解器来生成样本。

1.使用SDE扰动数据🌒

  为了生成样本,我们需要考虑一个扩散过程,去不断对数据加噪使得数据分布变成一个噪声分布。那么当\(\textcolor{blue}{\lbrace x_{(t)}\rbrace_{t=0}^T}\)是一个连续变量时,我们可以通过SDE去表达这个过程:

\[\huge dx=f(x,t)dt+g(t)dw-(8)\]

  其中\(\textcolor{blue}{f(\cdot,t):R^d \rightarrow R^d}\)是个向量函数,被称为确定性漂移系数\(\textcolor{blue}{g(t)\in R^d}\)称为扩散系数,w是随机布朗运动的一个量,\(\textcolor{blue}{dw}\)是无限小的白噪声。根据SDE移动的粒子不仅遵循确定性漂移\(\textcolor{blue}{f(x,t)}\),而且还受到来自\(\textcolor{blue}{g(t)dw}\)的随机噪声的影响。

  这个SDE公式表达一个扩散过程,相当于DDPM当中的前向扩散过程,只不过在DDPM中,时间\(\textcolor{blue}{t}\)是离散的,而在SDE中,时间是连续的。因此SDE可以看做是DDPM在连续时间上的扩展。我们可以把这个式子看成是下述离散形式在\(\textcolor{blue}{Δt→0}\)时的极限:

\[\huge x_{t+Δt}−x_t=f_t(x_t)Δt+g_t\sqrtΔtε,ε∼N(0,I)\]

  有没有感觉这个式子很眼熟,和朗之万迭代采样的式子很像。在这里越小的\(\textcolor{blue}{Δt}\)意味着是对原始SDE越好的近似,如果取\(\textcolor{blue}{Δt=0.001}\),那就对应于原来的\(\textcolor{blue}{T=1000}\),如果是\(\textcolor{blue}{Δt=0.01}\)则对应于\(\textcolor{blue}{T=100}\)。也就是说,在连续时间的SDE视角之下,不同的\(\textcolor{blue}{T}\)是SDE不同的离散化程度的体现,它们会自动地导致相似的结果,我们不需要事先指定\(\textcolor{blue}{T}\),而是根据实际情况下的精确度来取适当的\(\textcolor{blue}{T}\)进行数值计算。该SDE的解是一个随机变量连续的集合\(\textcolor{blue}{\lbrace x_{(t)}\rbrace_{0:T}}\)

  在\(\textcolor{blue}{t=0}\)时,\(\textcolor{blue}{p_0(x)=p(x)}\),表示初始时刻没有任何噪声混入的原始数据分布;在经过足够长的时间\(\textcolor{blue}{T}\)后,随着规模不断增大的噪声的混入,\(\textcolor{blue}{p_T(x)}\)变成了一个可以处理(tractable)的噪声分布(比如高斯分布),记为\(\textcolor{blue}{\pi(x)}\),被称为先验分布(prior distrubution)。

根据DDPM推导上述式8:

  和DDPM一样,这个前向过程不需要模型去学习,我们关注的是它的逆过程,利用逆过程从数据分布\(\textcolor{blue}{p_T}\)中去采样新的样本。

2.SDE的逆向过程🌖

  在噪声规模数量有限的情况下(DDPMs, NCSNs),我们通过逆向过程逐渐减少噪声来生成样本,NCSNs称之为退火朗之万动力学,即使用朗之万动力学从每个混入噪声的分布中按照逆序采样;类似的,对于无限数量的噪声尺度,我们使用逆向SDE来对混入噪声的过程进行反向来进行样本生成。

  对于任何的SDE,其逆向过程仍为一个扩散过程,其封闭数学表达式可以表示为:

  其中\(\textcolor{blue}{dt}\)表示一个无穷小的负向时间步长,\(\textcolor{blue}{\bar{w}}\)表示一个从时间\(\textcolor{blue}{t=T}\)\(\textcolor{blue}{t=0}\)的逆时的布朗运动过程,\(\textcolor{blue}{\nabla_{x}\log p_t(x)}\)表示概率函数的分数。一旦知道了正向SDE的漂移和扩散系数,以及每个\(\textcolor{blue}{t \in [0,T]}\)\(\textcolor{blue}{p_t(x)}\)的分数,就可以计算出逆向SDE。逆向SDE扩散过程即利用一个训练的神经网络\(\textcolor{blue}{s_\theta(x,t)}\)来估计分数\(\textcolor{blue}{\nabla_{x}\log p_t(x)}\)

  我们可以通过SDE将数据中混入噪声,最终得到一个简单的噪声分布;这个过程可以逆转,如果我们知道每个时间步的score分布,那么我们就可以从噪声中生成样本。实际上,对逆向SDE的求解的过程,就是产生一个基于Score的生成模型的过程。

逆向SDE微分方程式的推导:

3.使用基于Score的模型和score matching来估计reverse SDE🌌

  想要求解逆向SDE不仅需要分数函数\(\textcolor{blue}{\nabla_{x}\log p_t(x)}\),也需要终端分布\(\textcolor{blue}{p_T(x)}\),因为那身正向SDE的终点,也是逆向SDE的起点。前面提到,\(\textcolor{blue}{p_T(x)}\)接近于先验分布\(\textcolor{blue}{\pi(x)}\),通常是一个完全是噪声组成的可解(tractable)的简单噪声分布。

  为了估计分数函数\(\textcolor{blue}{\nabla_{x}\log p_t(x)}\),可以训练一个时间依赖的基于score的模型\(\textcolor{blue}{s_\theta(x,t)}\),使得\(\textcolor{blue}{s_\theta(x,t)\approx \nabla_{x}\log p_t(x)}\)

  这样,就可以将一个训练好的分数模型与上述的逆向SDE式子相结合就能得到:

\[\huge dx = [f(x,t)-g^2(t)s_\theta(x,t)]dt+g(t)dw\]

  当我们得到已经训练好的基于分数的模型\(\textcolor{blue}{s_\theta(x,t)}\),就可以从\(\textcolor{blue}{p_T(x)}\)出发逐步”去噪“来获得没有噪声的数据样本\(\textcolor{blue}{p_0(x)}\),也就是生成的数据。

  其优化目标函数,也是NCSN的目标函数的连续化结果:

  其中\(\textcolor{blue}{t \sim Uniform[0,T]}\)表示在时间间隔内的均匀分布\(\textcolor{blue}{[0,T]}\)。而\(\textcolor{blue}{\lambda(t):R \rightarrow R_{>0}}\)是一个正加权函数。通常使用\(\textcolor{blue}{\lambda(t) \propto 1/E[||\nabla_{x(t)}\log p(x(t)|x(0))||^2_2]}\)以平衡不同分数匹配损失随时间变化的幅度。

4.求解逆向SDE🌠

  样本生成过程就是逆向SDE的解。所以如何求解反向SDE呢?本文作者提供了几种方法,而第一种就是通过用数值SDE求解器求解估计的逆向SDE。

4.1 通用数值 SDE 求解器

  其最简单的数值SDE求解器是Euler-Maruyama方法。它使用有限的时间步长和小的高斯噪声对SDE进行离散化。具体来说,它选择一个小的负时间步长\(\textcolor{blue}{Δt\approx 0}\)初始化\(\textcolor{blue}{t\leftarrow T}\),并迭代以下过程,直到\(\textcolor{blue}{t\approx 0}\)

  这里\(\textcolor{blue}{z_t \sim N(0,I)}\),Euler-Maruyama 方法在性质上与 Langevin 动力学相似,两者都进行了更新\(\textcolor{blue}{x}\)通过遵循受高斯噪声干扰的分数函数。

  除了 Euler-Maruyama 方法外,还可以直接使用其他数值 SDE 求解器来求解用于样本生成的反向 SDE,例如 Milstein 方法和随机 Runge-Kutta 方法。后者的作者介绍了自适应步长 SDE 求解器,可以更快地生成样本,质量更好。

4.2 预测-校正采样器

  值得一提的就是逆向SDE有两特殊的特性,可实现更灵活的采样方法

  1️⃣估计的分数\(\textcolor{blue}{\nabla_{x}\log p_t(x)}\)可以通过训练一个基于分数的模型\(\textcolor{blue}{s_\theta(x,t)}\)来获得。

  2️⃣在不同时间步长获得的样本可以具有任意相关性,而不必形成从反向 SDE 采样的特定轨迹。

  由于这两个特性,我们可以应用MCMC(马尔科夫链蒙特卡洛)方法来微调从数值SDE求解器获得的轨迹。具体来说,宋博士提出了预测器-校正采样器

  🎯预测器可以是任何数值SDE求解器,它能够从现有样本\(\textcolor{blue}{x(t) \sim p_{t}(x)}\)中预测\(\textcolor{blue}{x(t+Δt) \sim p_{t+Δt}(x)}\)

  🎳校正器可以是任何仅依赖于分数函数的 MCMC 过程,例如 Langevin 动力学和哈密顿蒙特卡洛。

  在预测器-校正器采样器的每个步骤中,首先使用预测器来选择适当的步长\(\textcolor{blue}{Δt< 0}\),然后根据当前样本\(\textcolor{blue}{x(t)}\)预测\(\textcolor{blue}{x(t+Δt)}\)。接下来,根据我们基于分数的模型\(\textcolor{blue}{s_\theta(x,t+Δt)}\)运行几个校正步骤来改进样本\(\textcolor{blue}{x(t+Δt)}\),使\(\textcolor{blue}{x(t+Δt)}\)成为更高质量的样本。

4.3 ODE(常微分方程)

  尽管能够生成高质量的样本,但基于 Langevin MCMC 和 SDE 求解器的采样器并未能提供计算基于分数的生成模型的精确对数似然的方法。下面将介绍一个基于常微分方程 (ODE) 的采样器,该采样器允许精确似然计算。

  对于所有扩散过程,都存在一个相应的确定性过程,其轨迹与 SDE 具有相同的边际概率密度\(\textcolor{blue}{\{p_t(x)\}_{t=0}^T}\)扩散过程的SDE有确定的ODE对应解代表着我们可以借用业界一直在研究的ODE求解器,快速求解(在几十步内)从而避开了扩散所需几千步的庞大计算量。同时,将扩散由离散向连续统一,使得条件控制生成变得特别自然。ODE的形式如下:

  接下来就一步步推导,如何从SDE转换到ODE:

  我们已知SDE的正向加噪公式为:

\[\huge dx=f(x,t)dt+g(t)dw-(8)\]

  其描述的是一个随机运动过程中粒子的位置变化情况,其包含着运动的随机性。而粒子整体的分布则可以通过Fokker-Planck方程进行描述,其分布则是确定的,因此我们可以写出它对应的F-P方程(F-P方程的形式是如何推导出来的,具体可以参考文章《生成扩散模型漫谈(六):一般框架之ODE篇》,它是描述边际分布的偏微分方程:

\[\huge \frac{\partial}{\partial t}p(x,t)=- \nabla_{x}[f(x,t)p(x,t)]+\frac{1}{2}g_t^2\nabla_{x}^2p(x,t)-(9)\]

  对于任意满足\(\textcolor{blue}{σ^2_t≤g^2_t}\)的函数\(\textcolor{blue}{σ_t}\),我们有如下推导:

\[\huge \begin{aligned}\frac{\partial}{\partial t}p(x,t) &=- \nabla_{x}[f(x,t)p(x,t)]+\frac{1}{2}g_t^2\nabla_{x}^2p(x,t)\\&=- \nabla_{x}[f(x,t)p(x,t)-\frac{1}{2}(g_t^2-\sigma_t^2)\nabla_{x}p(x,t)]\\&+\frac{1}{2}\sigma_t^2\nabla_{x}^2p(x,t) \end{aligned}\]

  其中,\(\textcolor{blue}{\frac{\partial f(x)}{\partial x}=\frac{\partial \log f(x)}{\partial x} \cdot f(x),\frac{\partial \log f(x)}{\partial x}=\frac{\partial f(x)}{\partial x} \cdot \frac{1}{f(x)}}\)则有:

\[\huge \begin{aligned}\frac{\partial p(x,t)}{\partial t}=&- \nabla_{x}[(f(x,t)-\frac{1}{2}(g_t^2-\sigma_t^2)\nabla_{x} \log p(x,t))p(x,t)]\\&+\frac{1}{2}\sigma_t^2\nabla_{x}^2p(x,t)-(10)\end{aligned}\]

  形式上该F-P方程又相当于原来的F-P的\(\textcolor{blue}{f(x,t)}\)换成了\(\textcolor{blue}{(f(x,t)+\frac{1}{2}(g_t^2-\sigma_t^2)\nabla_{x} \log p(x,t))}\)\(\textcolor{blue}{g_t}\)换成了\(\textcolor{blue}{\sigma_t}\)根据式9对应于式8,上式则对应于:

\[\huge dx =(f(x,t)-\frac{1}{2}(g_t^2-\sigma_t^2)\nabla_{x} \log p(x,t))dt + \sigma_tdw-(11)\]

  别忘了式10跟式9是完全等价的,所以这意味着式8和式11这两个随机微分方程所对应的边际分布\(\textcolor{blue}{p_t(x)}\)是完全等价的!这个结果告诉我们存在不同方差的前向过程,它们产生的边际分布是一样的。

  特别地,我们可以写出式(11)对应的反向SDE:

\[\huge dx =(f(x,t)-\frac{1}{2}(g_t^2+\sigma_t^2)\nabla_{x} \log p(x,t))dt + \sigma_tdw-(12)\]

  式11允许我们改变采样过程的方差,这里我们特别考虑\(\textcolor{blue}{\sigma_t=0}\)的极端情形,此时SDE退化为ODE(常微分方程):

  这个ODE称为“概率流ODE(Probability flow ODE)”,由于实践中的\(\textcolor{blue}{\nabla_{x}\log p_t(x)}\)需要用神经网络\(\textcolor{blue}{s_\theta(x,t)}\)近似,所以上式也对应一个“神经ODE”。同时,这一方程的解,是一个确定性的轨迹,其从纯噪音分布\(\textcolor{blue}{\log p_t(x)}\)指向原始图像分布\(\textcolor{blue}{\log p_0(x)}\)。并且它的解是可逆的,因此其既可以用于描述前向传播过程,也可以用于描述逆向过程

  下图就展示了SDE和ODE解的过程,可以看到ODE的轨迹(白色线)是确定光滑的,而SDE的轨迹是随机的。

  我们可以通过求解式13来采样\(\textcolor{blue}{x(0) \sim p_{0}}\)。 使用 ODE 求解器不仅可以产生高质量的样本,而且还允许我们明确地权衡精度和效率。误差容限较大,在不影响样本视觉质量的情况下,功能评估次数可减少90%以上,如下图所示。左图为概率流 ODE 能够随着数值精度的变化而采用自适应步长进行快速采样,并在不损害质量的情况下减少评分函数评估的数量(NFE)。

  作者表明尽管概率流 ODE 求解器允许快速采样,但如果不使用校正器,它们的样本通常比 SDE 求解器具有更高(更差)的 FID 分数。

四、总结

  本文首先介绍了基于得分匹配和朗之万方程的Score-based model模型和Diffusion的模型联系十分紧密,从另一个视角揭示了Diffusion模型的本质,方便我们更加深入理解和优化Diffusion模型。并且还介绍了SDE将分数模型与扩散模型联系起来

  

  

  

  

  paper:Score-Based Generative Modeling through Stochastic Differential Equations

  参考文章:生成扩散模型漫谈(五):一般框架之SDE篇

  图像生成别只知道扩散模型(Diffusion Models),还有基于梯度去噪的分数模型:NCSN(Noise Conditional Score Networks) - 知乎 (zhihu.com)

  Generative Modeling by Estimating Gradients of the Data Distribution(score-based model,NCSN) - 知乎 (zhihu.com)

  Generative Modeling by Estimating Gradients of the Data Distribution | Yang Song

CATALOG
  1. 1. 一、前言
    1. 1.1. 1.什么是分数?🤖
    2. 1.2. 2.朗之万动力学采样(Sampling with Langevin Dynamics)🦊
    3. 1.3. 3.分数匹配🦚
    4. 1.4. 4.去噪分数匹配🐳
    5. 1.5. 5.分数模型的困境🦅
      1. 1.5.1. 5.1 loss不收敛
      2. 1.5.2. 5.2 分数估计不准确
      3. 1.5.3. 5.3 生成结果偏差大
  2. 2. 二、Noise Conditional Score Network(NCSN)
    1. 2.1. 1.噪声设计原则🚩
    2. 2.2. 2.网络结构设计🌈
    3. 2.3. 3.模型训练方法🎨
    4. 2.4. 4.采样生成方法🚀
    5. 2.5. 5.与扩散模型的联系🚅
    6. 2.6. 6.条件生成🌕
  3. 3. 三、SDEs
    1. 3.1. 1.使用SDE扰动数据🌒
    2. 3.2. 2.SDE的逆向过程🌖
    3. 3.3. 3.使用基于Score的模型和score matching来估计reverse SDE🌌
    4. 3.4. 4.求解逆向SDE🌠
      1. 3.4.1. 4.1 通用数值 SDE 求解器
      2. 3.4.2. 4.2 预测-校正采样器
      3. 3.4.3. 4.3 ODE(常微分方程)
  4. 4. 四、总结