从扩散模型开始的生成模型范式演变--SDE(2)

上一篇SDE文章中，从去噪分数匹配、分数网络出发，逐渐将SDE引入到扩散过程表征中，主要集中在原理讲解和推导中，本文则进一步补充能在扩散过程中实际使用的SDE公式实例。

一种SDE定义

相信很多读者自己学习或在看完上一篇文章后，对SDE实际如何使用情况还是不清楚。其实，因为SDE是随机微分方程，所以其与DDPM中加噪的方差其实都是人为设定的，相当于超参数。所以在训练前就需要将SDE具体形式定义好，在此将SDE定义为
$$dx={\sigma }^{t}d\omega t\in [0,1]\text{\hspace{1em}(1)}$$
公式(1)中丢弃了SDE定义中的漂移部分，只设置了扩散部分。此时已将时间区间正则化，即范围为[0,1]，也知道了$dx$随时间变化的公式，因为当前加噪是连续过程，故
$$x(t)=x(0)+{\int }_0^{t}dx=x(0)+{\int }_0^t{\sigma }^{s}d\omega (s)\text{\hspace{1em}(2)}$$
我们知道$\omega (s)$是标准布朗运动，其符合高斯分布$N(0,tI)$；而${\int }_0^t{\sigma }^{s}d\omega (s)$是一个随机积分，在时间区间[0,t]范围内的一个加权和；根据随机微分方程性质，${\int }_0^t{\sigma }^{s}d\omega (s)$也服从高斯分布，其均值为0，方差计算如下：

$$\begin{array}{rl}Var[{\int }_0^t{\sigma }^{s}d\omega (s)]& ={\int }_0^t{\sigma }^{2s}ds\\ & =\frac{1}{2\log \sigma }{\sigma }^{2s}{|}_0^t\\ & =\frac{1}{2\log \sigma }({\sigma }^{2t}-1)\end{array}\text{\hspace{1em}(3)}$$
即${\int }_0^t{\sigma }^{s}d\omega (s)\sim N(0,\frac{1}{2\log \sigma }({\sigma }^{2t}-1)I)$。
因为$\omega (s)$符合高斯分布，$p_{0t}(x(t)|x(0))$也符合高斯分布，整合上述推导有
$$p_{0t}(x(t)|x(0))=N(x(t);x(0),\frac{1}{2\log \sigma }({\sigma }^{2t}-1)I)\text{\hspace{1em}(4)}$$
在上一篇SDE文档中提到过，当$p_{0t}(x(t)|x(0))$为高斯分布时，$\lambda (t)=方差=\frac{1}{2\log \sigma }({\sigma }^{2t}-1)$。

$p_{0t}$是一个联合分布，$p_{t=0}$和$p_{t=1}$属于边缘分布，因为是连续概率，$p_{t=1}$概率等于$p_{t=0}$分布下所有值的概率积分，重新用$y,x$分别表示$x(0),x(t)$，$p_{t=1}$的计算公式如下
$$p_{t=1}=\int p_0(y)N(x;y,\frac{1}{2\log \sigma }({\sigma }^2-1)I)dy\text{\hspace{1em}(5)}$$
因为$t=1$，所以公式(5)中的分布为$N(x;y,\frac{1}{2\log \sigma }({\sigma }^2-1)I)$；当$\sigma$很大时，该分布中的均值y也可以将其看是趋近于0。故$N(x;y,\frac{1}{2\log \sigma }({\sigma }^2-1)I)\approx N(x;0,\frac{1}{2\log \sigma }({\sigma }^2-1)I)$。将其带入公式(5)有
$$\begin{array}{rl}{p}_{t=1}& \approx \int p_0(y)N(x;0,\frac{1}{2\log \sigma }({\sigma }^2-1)I)dy\\ & =N(x;0,\frac{1}{2\log \sigma }({\sigma }^2-1)I)\int p_0(y)dy\\ & =N(x;0,\frac{1}{2\log \sigma }({\sigma }^2-1)I)\end{array}\text{\hspace{1em}(6)}$$
上述大费周章推动出公式(6)的原因是$p_{t=1}$是采样起始的先验分布，并且公式(6)中的$p_{t=1}$符合高斯分布，也与数据分布无关。在进行逆SDE过程求解时，可基于公式(6)采样初始值。

欧拉数值采样

常规SDE公式定义为$dx=f(x,t)dt+g(t)dw$，其对应的逆SDE公式为$dx=[f(x,t)-g^2(t){\nabla }_x\log p_t(x)]dt+g(t)d\bar{w}$，当前人为定义的SDE为公式(1)，则其对应的逆SDE公式如下
$$dx=-{\sigma }^{2t}{\nabla }_x\log p_t(x)dt+{\sigma }^{t}d\bar{w}\text{\hspace{1em}(7)}$$
训练过程使得分数网络$s_{\theta }(x,t)$近似分数函数，即$s_{\theta }(x,t)={\nabla }_x\log p_t(x)dt$，公式(7)进而转换为
$$dx=-{\sigma }^{2t}{s}_{\theta }(x,t)dt+{\sigma }^{t}d\bar{w}\text{\hspace{1em}(8)}$$
在加噪过程中$p_{t=0}$对应原始数据分布，$p_{t=1}$对应加噪结束后的先验分布，前文已推导出该先验分布，即公式(6)。逆SDE过程先从先验分布中随机采样一个样本，然后使用Euler-Maruyama等数值方法求解逆SDE。该方法基于 SDE 的简单离散化，用${\Delta }_t$替代$dt$，用$z\sim N(0,g^2(t)\Delta tI)$替代$d\omega$。当应用于公式(7)中可得到以下迭代采样公式
$$x_{t-\Delta t}=x_t+{\sigma }^{2t}{s}_{\theta }(x,t)\Delta t+{\sigma }^t\sqrt{\Delta t}{z}_t\text{\hspace{1em}(9)}$$
其中$x_t\in N(0,I)$。

Predictor-Corrector Sampler

上述欧拉采样速度较快，但是质量不高。既然已通过分数匹配方法训练得到了可估计分数的分数网络，故也可以运用一些基于分数模拟采样的方法从分布中采样，如朗之万动力学采样。将逆SDE的数值解法与模拟采样相结合，可以进一步提高采样质量。具体步骤是使用数值解法–欧拉采样先获得一个预测值，然后使用模拟采样–朗之万动力学采样进行校正；数值解法称为Predictor，模拟采样方法称为Corrector，此种融合方式称为Predictor-Corrector Sampler，或简称为PC-Sampler。

当分数已知时，基于分数的朗之万MCMC采样方法可从$p(x)$中生成样本，具体公式如下
$$x_{i+1}=x_i+ϵ{\nabla }_x\log p(x_i)+\sqrt{2ϵ}{z}_i,\text{\hspace{1em}(10)}$$
其中$z_i\in N(0,I)$，$ϵ>0$是步长。

PC-Sampler的伪代码如下所示，采样过程还是很直观的。首先应用逆SDE的数值解法从$x_t$中获取$x_{t-\Delta t}$，此过程称为predictor步骤。再使用朗之万动力学采样精炼$x_{t-\Delta t}$，使得$x_{t-\Delta t}$成为分布$p_{t-\Delta t}(x)$更准确的样本，此步骤称为corrector步骤，能降低数值解法中的误差。

回忆上一篇SDE文章中的退火朗之万动力学采样过程，也是一个嵌套for循环形式，外层循环是从大到小的噪声等级，内层循环是与公式(10)一致的朗之万动力学采样。想到此时，又是一阵豁然开朗。因为逆SDE的数值采样方法本质就是通过离散化的方式求解连续数值，而退火朗之万动力学采样中的不同加噪等级就是有限的、离散的加噪过程。

后续研究表明，DDPM和Score-based Generative Model均可融入到基于分数的SDE中，并且采样方法可以整合到PC-Sampler框架中

基于ODE数值采样

对于所有扩散过程，在求解逆SDE时，存在一个相应的确定性过程，其轨迹与SDE具有相同的边缘概率密度，即逆向求解的结果和逆SDE求解出来的样本服从相同分布。这个确定性过程如下，称为概率流常微分方程：
$$dx=[f(x,t)-\frac{1}{2}g(t{)}^2{\nabla }_x\log p_t(x)]dt\text{\hspace{1em}(11)}$$

以上示意图显示了概率流 ODE 的轨迹与 SDE 轨迹的不同，但仍然从相同的分布中采样。因此，也可以从$p_{t=1}$分布中的样本出发，以逆时间方向对ODE进行积分，最终可获得符合$p_{t=0}$的样本。

按照前文公式(1)中定义的SDE过程，其对应的ODE方程如下
$$dx=-\frac{1}{2}{\sigma }^{2t}{s}_{\theta }(x,t)dt\text{\hspace{1em}(12)}$$
基于公式(12)，使用一些数值解法就能进行样本采样。