扩散模型的SDE、ODE视角以及Fokker-Planck方程

Holi2025-03-18

本文参考《Tutorial on Diffusion Models for Imaging and Vision》【该文中对DSM的${\bf x}’$替换为${\bf x}$的部分的叙述貌似有点小问题以及最后对fokker-Planck方程的部分的一个例子也有点小问题】以及《 SCORE-BASED GENERATIVE MODELING THROUGH STOCHASTIC DIFFERENTIAL EQUATIONS》。

得分匹配模型

宋飏等人在DDPM发表之前就提出了多轮噪声尺度的得分匹配模型，这一类模型的本质后续被证明和DDPM是在同一体系下的，不过重新梳理其理论推导有助于我们后续从另一个视角理解扩散模型。

我们已有一些数据$\{ {\bf x}_i\in\mathbb{R}^D\}_{i=1}^N$，数据的分布为$q({\bf x})$，但是现在我们面临的问题不再是拟合这个分布，而是找到这个分布中概率最大的数据点(其实就是学习已有的数据分布，便于生成一些看起来最常见的数据)，那么其实本质上就是找到

$\max_{\bf x} q({\bf x})$

直观来看，我们可以通过梯度上升来找到这个数据点$\nabla_{\bf x} q({\bf x})$，但是由于$q({\bf x})$中常常存在归一化系数，对这个系数的拟合会耗费额外的不必要资源(毕竟我们只关心相对的最大值)，因此不妨将最优化问题转变为

$\max_{\bf x}\log q({\bf x})$

这样的话，我们再求梯度$\nabla_{\bf x}\log q({\bf x})$，归一化系数就直接被干掉了，这就是$q({\bf x})$的分数【有文献提到，这里的分数实际上和数学意义上常见的分数的定义不同】。那么我们可以用一个网络${\bf s}_\theta:\mathbb{R}^D\to\mathbb{R}^D$来拟合该分数，拟合完成后，就可以使用近似的梯度上升法来找到概率最大的数据点

${\bf x}_{t+1}={\bf x}_t+\tau\nabla_{\bf x}\log q({ {\bf x}_t })$

这时我们可以引入离散朗之万动力学方程

${\bf x}_{t+1}={\bf x}_t+\tau\nabla_{\bf x}\log q({ {\bf x}_t })+\sqrt{2\tau}{\bf z},{\bf z}\sim\mathcal{N}(0,{\bf I})$

可以看到，朗之万动力学方程实际上就是加入了噪声项的梯度上升。但是为什么要引入这样的一个噪声项呢？一个直观理解是，我们实际上不是希望完全求解最开始的优化问题，也就是不是希望单单找到那个最优点，而是希望最终收敛的结果具有一定多样性，建立对高概率区域附近的分布的认知。【但是实际上这个噪声的引入是有严格的数学背景的，这和Fokker-Planck方程有密切的关系，该方程可以给出任何一种随机微分方程在任一时刻的分布$p({\bf x},t)$】

上面就是得分匹配函数的核心理念【先训练${\bf s}_\theta({\bf x})\approx \nabla_{\bf x}\log q({\bf x})$，再应用离散朗之万动力学方程进行采样】，但是在实际应用上述算法时，存在一个核心问题在于，真值$\nabla_{\bf x}\log q({\bf x})$如何得到？【毕竟我们的目标就是拟合这个玩意，要是我们已经有这个东西的简单计算方式就不需要再拟合这个玩意了(我们收集到的数据是离散的，而真实的分布空间是连续的)】

【明确得分匹配ESM】该方法认为数据真实分布是多个简单的核分布的叠加【在真实样本点处放置核函数，使得整个空间变为较为平滑的连续分布空间】，即

$q({\bf x})\approx q_h({\bf x})=\frac{1}{M}\sum_{m=1}^M\frac{1}{h}K(\frac{ {\bf x}-{\bf x}^{(m)} }{h})$

那么拟合目标就变为

$\begin{align} J_{ESM}(\theta)\\ :=\mathbb{E}_{q({\bf x})}||{\bf s}_\theta({\bf x})-\nabla_{\bf x}\log q({\bf x})||^2\\ \approx \mathbb{E}_{q_h({\bf x})}||{\bf s}_\theta({\bf x})-\nabla_{\bf x}\log q_h({\bf x})||^2\\ =\int||{\bf s}_\theta({\bf x})-\nabla_{\bf x}\log q_h({\bf x})||^2 q_h({\bf x})d{\bf x}\\ \approx\frac{1}{M}\sum_{m=1}^M\int||{\bf s}_\theta({\bf x})-\nabla_{\bf x}\log q_h({\bf x})||^2\frac{1}{h}K\left(\frac{ {\bf x}-{\bf x}^{(m)} }{h}\right)d{\bf x} \end{align}$

由于核函数一般是一些简单的分布形式，这样对其求梯度的过程就比较容易计算出来了【优化目标中的积分实际上是通过蒙特卡洛方法的近似实现的，每次从相应的核函数的定义中随机抽取几个样本，求解损失函数】(容易处理的，tractable)。【此处可以进一步详细介绍】

【模糊得分匹配ISM】在该方法中，优化目标为

$J_{ISM}(\theta):=\mathbb{E}_{q({\bf x})}\left[\text{Tr}(\nabla_{\bf x}{\bf s}_{\theta}({\bf x}))+\frac{1}{2}||{\bf s}_\theta({\bf x})||^2\right]$

其中$\nabla_{\bf x}{\bf s}_\theta({\bf x})$为${\bf s}_\theta({\bf x})$的雅可比矩阵，而该方法的优化目标是可以被蒙特卡洛方法估计的

$J_{ISM}(\theta)\approx \frac{1}{M}\sum_{m=1}^M\sum_i(\partial_i{\bf s}_\theta({\bf x}^{(m)})+\frac{1}{2}|[{\bf s}_\theta({\bf x}^{(m)})]_i|^2)$

其中$\partial_i{\bf s}_\theta({\bf x}^{(m)})=\frac{\partial }{\partial x_i}[{\bf s}_\theta({\bf x})]_i=\frac{\partial^2}{\partial x_i^2}\log p_\theta({\bf x})$，但是如果分数是由深度神经网络拟合得到的，那么迹是比较难以计算的。

【去噪得分匹配DSM】这种方法在2011年就被提出了，而在深度学习时代由宋飏等人进一步发扬光大。由于前述的两种方法存在诸多局限，因此该方法被提出，损失函数如下所示

$J_{DSM}(\theta):=\mathbb{E}_{q({\bf x},{\bf x}')}\left[\frac{1}{2}||{\bf s}_\theta({\bf x})-\nabla_{\bf x}\log q({\bf x|{\bf x}}')||^2\right]$

通过理论计算，实际上可以验证$J_{DSM}(\theta)=J_{ESM}(\theta)+C$，也就说这样训练出来的最优${\bf s}_\theta({\bf x})$和ESM的拟合目标拟合出来的最优结果是一致的，下面是详细推导：

$\begin{align} J_{ESM}(\theta)\\ =\mathbb{E}_{q({\bf x})}\left[\frac{1}{2}||{\bf s}_\theta({\bf x})-\nabla_{\bf x}\log q({\bf x})||^2\right]\\ =\mathbb{E}_{q({\bf x})}\left[\frac{1}{2}||{\bf s}_\theta({\bf x})||^2-{\bf s}_\theta({\bf x})^T\nabla_{\bf x}\log q({\bf x})+\frac{1}{2}||\nabla_{\bf x}\log q({\bf x})||^2\right]\\ =\mathbb{E}_{q({\bf x})}\left[\frac{1}{2}||{\bf s}_\theta({\bf x})||^2-{\bf s}_\theta({\bf x})^T\nabla_{\bf x}\log q({\bf x})\right]+C\\ =\mathbb{E}_{q({\bf x})}\left[\frac{1}{2}||{\bf s}_\theta({\bf x})||^2\right]-\int q({\bf x}){\bf s}_\theta({\bf x})^T\nabla_{\bf x}\log q({\bf x})d{\bf x}+C\\ =\mathbb{E}_{q({\bf x})}\left[\frac{1}{2}||{\bf s}_\theta({\bf x})||^2\right]-\int {\bf s}_\theta({\bf x})^T\nabla_{\bf x}q({\bf x})d{\bf x}+C \end{align}$

对于第二项，我们可以利用隐变量条件概率积分来替换$q({\bf x})$，即$q({\bf x})=\int q({\bf x}|{\bf x}’)q({\bf x}’)d{\bf x}’$，那么进一步推导得到

$\begin{align} \int {\bf s}_\theta({\bf x})^T\nabla_{\bf x}q({\bf x})d{\bf x}\\ =\int {\bf s}_\theta({\bf x})^T\nabla_{\bf x}\left(\int q({\bf x}|{\bf x}')q({\bf x}')d{\bf x}'\right)d{\bf x}\\ =\int {\bf s}_\theta({\bf x})^T\left(\int \nabla_{\bf x}q({\bf x}|{\bf x}')q({\bf x}')d{\bf x}'\right)d{\bf x}\\ =\int {\bf s}_\theta({\bf x})^T\left(\int \nabla_{\bf x}q({\bf x}|{\bf x}')\frac{q({\bf x}|{\bf x}')}{q({\bf x}|{\bf x}')}q({\bf x}')d{\bf x}'\right)d{\bf x}\\ =\int {\bf s}_\theta({\bf x})^T\left(\int \nabla_{\bf x}\log q({\bf x}|{\bf x}')q({\bf x}|{\bf x}')q({\bf x}')d{\bf x}'\right)d{\bf x}\\ =\int\int {\bf s}_\theta({\bf x})^T\nabla_{\bf x}\log q({\bf x}|{\bf x}')q({\bf x},{\bf x}')d{\bf x}'d{\bf x}\\ =\mathbb{E}_{q({\bf x},{\bf x}')}\left[{\bf s}_\theta({\bf x})^T\nabla_{\bf x}\log q({\bf x}|{\bf x}')\right] \end{align}$

将这个结果代入回$J_{ESM}(\theta)$，得到

$\begin{align} J_{ESM}(\theta)\\ =\mathbb{E}_{q({\bf x})}\left[\frac{1}{2}||{\bf s}_\theta({\bf x})||^2\right]-\mathbb{E}_{q({\bf x},{\bf x}')}\left[{\bf s}_\theta({\bf x})^T\nabla_{\bf x}\log q({\bf x}|{\bf x}')\right]+C\\ =\mathbb{E}_{q({\bf x},{\bf x}')}\left[\frac{1}{2}||{\bf s}_\theta({\bf x})||^2\right]-\mathbb{E}_{q({\bf x},{\bf x}')}\left[{\bf s}_\theta({\bf x})^T\nabla_{\bf x}\log q({\bf x}|{\bf x}')\right]+C\\ =J_{DSM}(\theta)+C' \end{align}$

其实对比该损失函数和前述的损失函数，不难发现其就是把拟合目标由$\nabla_{\bf x}\log q({\bf x})$变为了$\nabla_{\bf x}\log q({\bf x|{\bf x}}’)$，通过这样的替换，我们实际上就可以将难以计算的拟合目标变得容易求解。比如我们可以令$q({\bf x}|{\bf x}’)=\mathcal{N}({\bf x}|{\bf x}’,\sigma^2)$，即${\bf x}={\bf x}’+\sigma{\bf z}$，那么这样我们就容易求解出

$\begin{align} \nabla_{\bf x}\log q({\bf x}|{\bf x}')\\ =\nabla_{\bf x}\left\{-\frac{||{\bf x}-{\bf x}'||^2}{2\sigma^2}\right\}\\ =-\frac{ {\bf x}-{\bf x}'}{\sigma^2}\\ =-\frac{\bf z}{\sigma} \end{align}$

假如我们把${\bf x}’$当作真实的数据空间，那么要拟合的${\bf x}$分布空间实际上就变成了被高斯噪声扰动后的数据空间，此时我们可以重新写出DSM的拟合目标

$J_{DSM}(\theta)=\mathbb{E}_{ {\bf x}\sim q({\bf x}),{\bf z}\sim{\mathcal{N}(0,{\bf I})} }\left[||{\bf s}_\theta({\bf x}+\sigma {\bf z})+\frac{\bf z}{\sigma}||^2\right]$

这一拟合目标是十分直观的，实际上就是学习如何消去这个噪声${\bf z}$，也就是对于每一个输入的带噪样本点${\bf x}’$，要找到其指向最近的真实样本点${\bf x}$的向量，但是要格外注意的是，只有当$\sigma$趋近于零，我们得到的才是对真实数据空间$q({\bf x})$的拟合，否则实际上得到的是对被高斯噪声扰动后的数据空间的拟合。

在实际拟合时，我们是从数据集$\left\{ {\bf x}^{(m)}\right\}_{m=1}^M$出发，完成如下优化问题

$\begin{align} \theta^*=\mathop{\arg\min}\limits_{\theta}\mathbb{E}_{ {\bf x}\sim q({\bf x}),{\bf z}\sim{\mathcal{N}(0,{\bf I})} }\left[\left\Vert{\bf s}_\theta({\bf x}+\sigma {\bf z})+\frac{\bf z}{\sigma}\right\Vert^2\right]\\ \approx\mathop{\arg\min}\limits_{\theta}\frac{1}{M}\sum_{m=1}^M\frac{1}{2}\left\Vert{\bf s}_\theta({\bf x}^{(m)}+\sigma{\bf z^{(m)} })+\frac{ {\bf z}^{(m)} }{\sigma}\right\Vert^2 \end{align}$

如上的训练过程是在一个固定$\sigma$尺度下训练得到分数拟合函数，那么自然会想到将其扩展到多个噪声尺度【即宋飏等人提出的noise conditioned score network (NCSN)】：

$J_{NCSN}(\theta)=\frac{1}{L}\sum_{i=1}^L\lambda(\sigma_i)\mathcal{l}(\theta;\sigma_i)$

其中$\mathcal{l}(\theta;\sigma_i)=\mathbb{E}_{ {\bf x}\sim q({\bf x}),{\bf z}\sim{\mathcal{N}(0,{\bf I})} }\left[||{\bf s}_\theta({\bf x}+\sigma_i {\bf z})+\frac{\bf z}{\sigma_i}||^2\right]$，而$\lambda(\sigma)=\sigma^2$【这是经验公式】，且一般$\sigma_i$是随着$i$增大而逐级递减的。在推理时，则采取如下退火重要性采样朗之万方程：

${\bf x}_{t+1}={\bf x}_t+\frac{\alpha_i}{2}{\bf s}_\theta({\bf x}_t,\sigma_i)+\sqrt{\alpha_i}{\bf z}_t,\ \ \ \ \ {\bf z}_t\sim\mathcal{N}(0,{\bf I})$

其中$\alpha_i=\sigma_i^2/\sigma_L^2$。对于每个噪音尺度$\sigma_i$，都会进行$T$次关于$t$的迭代，最终得到的${\bf x}_T$会作为下一次噪声尺度迭代的开始${\bf x}_0$。

Sliced Score Matching (SSM)

此处是宋飏等人在2019年论文提出的新的得分匹配函数的优化目标，该优化目标通过将对得分函数的高维空间的完全匹配替换为对随机方向的得分函数投影的匹配，从而避免了计算$\nabla_{\bf x}\log q({\bf x})$的麻烦，这实际上使得后续为找到非仿射系数的前向SDE的反向SDE成为可能。

我们通常需要知道过渡核 $p_{0t}(x(t) \,|\, x(0))$ 才能有效地拟合得分模型。当漂移项$f(t)$是仿射的时，过渡核始终是高斯分布，其均值和方差通常具有闭式形式，可以通过标准方法获得。对于更一般的随机微分方程 (SDE)，我们可以通过求解Kolmogorov 正向方程来获得$p_{0t}(x(t) \,|\, x(0))$。另外，我们还可以直接对 SDE 进行模拟采样，以从$p_{0t}(x(t) \,|\, x(0))$中生成样本，并在模型训练中将方程 (7) 中的去噪分数匹配替换为切片分数匹配 (SSM)。这种做法能够绕过显式计算$ \nabla_x \log p_{0t}(x(t) \,|\, x(0))$的需求。

这里直接摆上原文的证明(偷懒)，要注意的是SSM优化目标的成立需要满足下面证明提到的几个假设(在某些场景下，假设可能是不成立的)【该证明的核心步骤是分部积分】

SDE与ODE视角下的扩散模型

DDPM与得分匹配模型本质上都是对连续随机微分方程的离散化拟合。因此从连续随机微分方程的视角出发，我们可以更进一步理解扩散模型的分布特性(引入fokker-Planck方程)。

【常微分方程简单例子】设有一个离散迭代方程${\bf x}_i=\left(1-\frac{\beta}{2}\right){\bf x}_{i-1}$，那么我们可以将其转化为在时间$[0,1]$内的连续微分方程，只需令${\bf x}(\frac{i}{N})={\bf x}_i$，以及$\beta=\frac{\hat\beta}{N}$【这个$\beta$的代换是十分重要的，因为假如没有这个代换，我们就无法进行后续取极限的操作】，那么就可以得到${\bf x}(\frac{i}{N})=\left(1-\frac{\hat\beta}{2N}\right){\bf x}(\frac{i-1}{N})$，不妨令$\Delta t=\frac{1}{N}$，那么就可以得到${\bf x}(t+\Delta t)-{\bf x}(t)=-\frac{\hat\beta}{2}{\bf x}(t)\Delta t$，令分割无限细化，就可以得到$\frac{d{\bf x}(t)}{dt}=-\frac{\hat\beta}{2}{\bf x}(t)$，该微分方程的解为${\bf x}(t)=\exp(-\frac{\hat\beta}{2}t)$，这样我们就得到了对上述离散迭代行为的一个简单认识。【从离散微分方程到连续微分方程的映射中，很喜欢做的一个处理就是把无限长的序列压缩到单位时间长度中，这样的话单次离散微分方程的计算实际上就是发生在一个极小时间内了，这时再在差分等式右侧凑出来一个$\Delta t$，也就可以顺理成章地转化为微分运算了】

简而言之，一阶常微分方程的形式为如下所示：

$\frac{d{\bf x}(t)}{dt}={\bf f}(t,{\bf x})$

一般容易计算其解析形式为

${\bf x}(t)={\bf x}_0+\int_0^t{\bf f}(s,{\bf x}(s))ds$

而对于一阶随机微分方程，其形式其实和ODE非常接近，只不过增加了噪声项【此处表达式的正确性存疑】

$\frac{d{\bf x}(t)}{dt}={\bf f}(t,{\bf x})+{\bf g}(t,{\bf x}){\bf \xi}(t),\ \ \ \ \ \ \ \ {\bf \xi}(t)\sim\mathcal{N}(0,{\bf I})$

常见的形式是包含维纳过程$d{\bf w}$的，$d{\bf w}\sim\mathcal{N}(0,dt{\bf I})$

$d{\bf x}={\bf f}(t,{\bf x})dt+{\bf g}(t,{\bf x})d{\bf w}$

我们也可以将其表达为积分形式

${\bf x}(t,\omega)={\bf x}_0+\int_0^t{\bf f}(s,{\bf x}(s,\omega))ds+\int_0^t{\bf g}(s,{\bf x}(s,\omega))d{\bf w}(s,\omega)$

其中$\omega$为${\bf x}$在概率空间中的状态【在该问题中即为随机过程变量${\bf x}$的某种轨迹，整个概率空间由多种多样的轨迹组成】，为了求解得到一个随机过程的一条轨迹就需要先选定一种随机状态$\omega$。

对于形如如下所示的连续的前向扩散过程(在$t$时刻具有分布$p_t({\bf x})$)

$d{\bf x}={\bf f}({\bf x},t)dt+g(t)d{\bf w}$

那么其对应的反向扩散过程为【该部分证明需要参考1982年《Reverse-time diffusion equation models》中的证明】【这个反向扩散方程可以保证反向轨迹在$t$时刻的分布为$q_t({\bf x})$】

$d{\bf x}=[{\bf f}({\bf x},t)-g(t)^2\nabla_{\bf x}\log q_t({\bf x})]dt+g(t)d\bar{\bf w}\tag{*}\label{*}$

其中$\bar{\bf w}$为反向扩散过程中的维纳过程(独立于过去的维纳过程，不独立于未来，这一性质是与${\bf w}$相反的)

根据宋飏等人论文的论述(基于1982年反向扩散过程物理学论文)，更一般的扩散形式为【${\bf f}(\cdot,t):\mathbb{R}^d\to\mathbb{R}^d$为向量函数，${\bf G}(\cdot,t):\mathbb{R}^d\to\mathbb{R}^{d\times d}$为矩阵函数，该矩阵函数意味着可以对各个通道施加与状态${\bf x}$相关的且存在一定相关性的噪声，比$g(t)$具有更强的灵活性】
$d{\bf x}={\bf f}({\bf x},t)dt+{\bf G}({\bf x},t)d{\bf w}$
相应地，其对应的反向扩散过程为
$\begin{align} 【d{\bf x}=\{ {\bf f}({\bf x},t)-\frac{1}{q({\bf x},t)}\nabla_{\bf x}\cdot (q({\bf x},t){\bf G}({\bf x},t){\bf G}({\bf x},t)^T)\}dt+{\bf G}({\bf x},t)d\bar{\bf w}】\\ d{\bf x}=\{ {\bf f}({\bf x},t)-\nabla_{\bf x}\cdot[{\bf G}({\bf x},t){\bf G}({\bf x},t)^T]-{\bf G}({\bf x},t){\bf G}({\bf x},t)^T\nabla_{\bf x}\log q({\bf x},t)\}dt+{\bf G}({\bf x},t)d\bar{\bf w} \end{align}$
【上下两种表述是一样的】其中$\nabla\cdot F({\bf x}):=[\nabla\cdot {\bf f}^1({\bf x}),…,\nabla\cdot{\bf f}^d({\bf x})]^T$，${\bf F}({\bf x})=[{\bf f}^1({\bf x}),…,{\bf f}^d({\bf x})]^T$。在1982年中的原文其实是表述为
$\bar{\bf f}^i({\bf x},t)={\bf f}^i({\bf x},t)-\frac{1}{q({\bf x},t)}\sum_{j,k}\frac{\partial}{\partial x_j}[q({\bf x},t)\cdot g^{jk}({\bf x},t)g^{ik}({\bf x},t)]$
对于离散SDE，那么前向过程写为
${\bf x}_{i+1}={\bf x}_{i}+{\bf f}_i({\bf x}_i)+{\bf G}_i({\bf z}_i)$
其对应的反向ODE过程为
${\bf x}_i={\bf x}_{i+1}-{\bf f}_{i+1}({\bf x}_{i+1})+\frac{1}{2}{\bf G}_{i+1}{\bf G}_{i+1}^T{\bf s}_{\theta}({\bf x}_{i+1},i+1)$

对于上述SDE，还有一个很有趣的期望性质，令${\bf M}(t)$为分布均值，${\bf\Sigma}(t)$为分布的方差矩阵，那么其满足如下ODE【参考《Applied Stochastic Differential Equations》 by Simo Särkkä & Arno Solin】：
$\frac{d{\bf M} }{dt}=\mathbb{E}[{\bf f}({\bf x},t)]$ $\frac{d{\bf\Sigma} }{dt}=\mathbb{E}[{\bf f}({\bf x},t)({\bf x}-{\bf M})^T]+\mathbb{E}[({\bf x}-{\bf M}){\bf f}({\bf x},t)^T]+\mathbb{E}[{\bf G}({\bf x},t){\bf Q}{\bf G}({\bf x},t)^T]$
其中$d{\bf w}d{\bf w}^T={\bf Q}dt$。

据此，宋飏等人给出了DDPM的SDE对应的方差矩阵的演变ODE
$\frac{d{\bf \Sigma}_{VP}(t) }{dt}=\beta(t)({\bf I}-{\bf\Sigma}_{VP}(t))$
事实上这个方程很好求解，可以计算得到
${\bf\Sigma}_{VP}(t)={\bf I}+\exp\left(-\int_0^t\beta(s)ds\right)({\bf\Sigma}_{VP}(0)-{\bf I})$
由此可见，DDPM分布随时间演变的方差是有界的(受限于${\bf\Sigma}_{VP}(0)$，因此称为VP)。宋飏等人进一步提出了sub-VP对应的前向SDE
$d{\bf x}=-\frac{1}{2}\beta(t){\bf x}dt+\sqrt{\beta(t)\left(1-\exp\left(-2\int_0^t\beta(s)ds\right)\right)}d{\bf w}$
该SDE对应的方差矩阵的演化为
${\bf\Sigma}_{subVP}(t)={\bf I}+\exp\left(-2\int_0^t\beta(s)ds\right){\bf I}+\exp\left(-\int_0^t\beta(s)ds\right)({\bf\Sigma}_{subVP}(0)-2{\bf I})$
其满足的重要性质是${\bf\Sigma}_{subVP}(t) \preceq{\bf\Sigma}_{VP}(t)$。

此外，有文献显示，上述SDE实际上可以由另外一个ODE来替代
$d{\bf x}=\left\{ {\bf f}({\bf x},t)-\frac{1}{2}\nabla\cdot[{\bf G}({\bf x},t){\bf G}({\bf x},t)^T]-\frac{1}{2}{\bf G}({\bf x},t){\bf G}({\bf x},t)^T\nabla_{\bf x}\log q_t({\bf x})\right\}dt$
宋飏等人给出了一个基于fokker-planck方程的简易证明，此处直接粘贴原文证明。【这部分的证明核心也是函数乘积的求导，其推导的思想就是fokker-planck方程和前向马尔可夫SDE存在对应关系，那么我们可以将fokker-planck方程中与原SDE的$d{\bf w}$前的系数相关的项归并到与${\bf f}({\bf x},t)$相关的项中，从而实现将${\bf G}({\bf x},t)$置零，这样就等价于消去了维纳过程项，原本的SDE就被转化为ODE了】

将上述ODE重新写为
$d{\bf x}=\bar{\bf f}({\bf x},t)dt$
其中$\bar{\bf f}({\bf x},t)$一般为模型拟合得到的，为$\bar{\bf f}_\theta({\bf x},t)$，那么套用Fokker-Planck方程，可以解得如下关系
$\log (p_\theta)_0({\bf x}(0))=\log (p_{\theta})_T({\bf x}(T))+\int_0^T\nabla\cdot\bar{\bf f}_\theta({\bf x}(t),t)dt$
通常情况下，计算$\nabla\cdot\bar{\bf f}_\theta({\bf x}(t),t)$耗费比较大，因此宋飏等人提出采用Skilling-Hutchinson估计器
$\nabla\cdot\bar{\bf f}_\theta({\bf x}(t),t)=\mathbb{E}_{p({\bf\epsilon})}[\epsilon^T\nabla\bar{\bf f}_\theta({\bf x}(t),t)\epsilon]$
其中，$\nabla\bar{\bf f}_\theta({\bf x}(t),t)$表示$\bar{\bf f}_\theta({\bf x}(t),t)$的Jacobian矩阵，随机变量$\epsilon$满足：$\mathbb{E}_{p(\epsilon)}[\epsilon]=0$且$\text{Cov}_{p(\epsilon)}={\bf I}$，$\epsilon^T\nabla\bar{\bf f}_\theta({\bf x}(t),t)$可以通过反向模式自动微分(reverse mode automatic differentiation)高效计算，其计算成本与评估$\bar{\bf f}_\theta({\bf x}(t),t)$相当。

由此，我们实际上可以通过上述模拟来精确计算拟合出的分布在任意${\bf x}$处的对数值$\log (p_\theta)_0({\bf x})$，由此便容易得到负对数似然。【我们自然希望真实的图片的平均负对数似然尽可能小】

DDPM的随机微分方程

对于DDPM模型，有如下离散随机微分方程形式

${\bf x}_i=\sqrt{1-\beta_i}{\bf x}_{i-1}+\sqrt{\beta_i}{\bf z}_{i-1},\ \ \ \ \ {\bf z}_{i}\sim\mathcal{N}(0,{\bf I})$

如果我们同样令${\bf x}_i={\bf x}(\frac{i}{N})$【即$t(i)=\frac{i}{N}$】以及$\Delta t=\frac{1}{N}$，$f_1(t(i))=\frac{\sqrt{1-\beta(t+\Delta t)})-1}{\Delta t}=N(\sqrt{1-\beta_{i+1}}-1)$，$f_2(t)=\frac{\beta(t+\Delta t)}{\Delta t}=N\beta_{i+1}$，【此处凑出来的$f_1(t)$和$f_2(t)$都是为了让$\Delta t$出现在它该出现的地方】那么有

$\begin{align} {\bf x}(t+\Delta t)=\sqrt{1-\beta_i}{\bf x}(t)+\sqrt{\beta_i}{\bf z}(t)\\ \Rightarrow {\bf x}(t+\Delta t)-{\bf x}(t)=(\sqrt{1-\beta_i}-1){\bf x}(t)+\sqrt{\beta_i}{\bf z}(t)\\ \Rightarrow {\bf x}(t+\Delta t)-{\bf x}(t)=(f_1(t)\Delta t){\bf x}(t)+\sqrt{\beta(t+\Delta t)}\frac{ {\bf w}(t+\Delta t)-{\bf w}(t)}{\sqrt{\Delta t} }\\ \Rightarrow d{\bf x}=f_1(t){\bf x}(t)dt+\sqrt{f_2(t)}d{\bf w} \end{align}$

在一些文献中，常常认为$\beta$足够小，那么就有$f_1(t(i))\approx-\frac{1}{2}N\beta_{i+1}$，该形式和$f_2(t(i))$一致，因此不妨直接令$\hat\beta(t(i))=N\beta_{i+1}$，这样就得到了常见的一种随机微分方程形式

$d{\bf x}=-\frac{1}{2}\hat\beta(t){\bf x}(t)dt+\sqrt{\hat\beta(t)}d{\bf w}$

当$\Delta t$越接近于零，那么该连续的随机微分方程对离散微分方程的拟合越接近，当二者足够接近时，我们实际上就可以使用拟合SDE的方法来拟合离散随机微分方程的过程。

实际上，我们也可以反过来，认为DDPM给出的离散随机方程实际上是在离散地近似求解上述连续SDE。

而对于上述连续前向过程的连续反向过程，参考$(\ref{*})$式子，就应该是如下形式(此处为了书写方便，令$\beta(t):=\hat\beta(t)$)【注意此处的$dt$是反向的，即为负，$\bar{\bf w}$为反向的维纳过程】

$d{\bf x}=-\beta(t)\left[\frac{\bf x}{2}+\nabla_{\bf x}\log q_t({\bf x})\right]dt+\sqrt{\beta(t)}d\bar {\bf w}$

下面可以证明，DDPM的离散反向过程本质上就是对该连续反向过程的离散模拟。

$\begin{align} {\bf x}(t-\Delta t)-{\bf x}(t)=-\beta(t)\left[\frac{ {\bf x}(t)}{2}+\nabla_{\bf x}\log q_t({\bf x})\right](-\Delta t)+\sqrt{\beta(t)}({\bar{\bf w}(t-\Delta t)}-{\bar{\bf w}(t)})\\ \Rightarrow {\bf x}(t-\Delta t)=\left(1+\frac{\beta(t)\Delta t}{2}\right){\bf x}(t)+\beta(t)\nabla_{\bf x}\log q_t({\bf x})\Delta t+\sqrt{\beta(t)}({\bar{\bf w}(t-\Delta t)}-{\bar{\bf w}(t)})\\ \Rightarrow {\bf x}_{i-1}=\left(1+\frac{\beta(t)\Delta t}{2}\right){\bf x}_i+\beta(t)\nabla_{\bf x}\log q_t({\bf x})\Delta t+\sqrt{\beta(t)}\sqrt{\Delta t}{\bf z}_i\\ \Rightarrow {\bf x}_{i-1}=\left(1+\frac{\beta(\frac{i}{N})\frac{1}{N}}{2}\right){\bf x}_i+\beta\left(\frac{i}{N}\right)\nabla_{\bf x}\log q_t({\bf x})\frac{1}{N}+\sqrt{\frac{\beta\left(\frac{i}{N}\right)}{N}}{\bf z}_i\\ \Rightarrow {\bf x}_{i-1}\approx\left(1+\frac{\beta(\frac{i}{N})\frac{1}{N}}{2}\right){\bf x}_i+\beta\left(\frac{i}{N}\right)\nabla_{\bf x}\log q_t({\bf x})\frac{1}{N}+\frac{1}{2}\left(\beta\left(\frac{i}{N}\right)\frac{1}{N}\right)^2\nabla_{\bf x}\log q_t({\bf x})+\sqrt{\frac{\beta(\frac{i}{N})}{N}}{\bf z}_i\\ \Rightarrow {\bf x}_{i-1}\approx\left(1+\frac{\beta(\frac{i}{N})\frac{1}{N}}{2}\right)\left({\bf x}_i+\beta\left(\frac{i}{N}\right)\frac{1}{N}\nabla_{\bf x}\log q_t({\bf x})\right)+\sqrt{\frac{\beta(\frac{i}{N})}{N}}{\bf z}_i\\ \Rightarrow {\bf x}_{i-1}\approx\left(1+\frac{\hat\beta_i}{2}\right)\left({\bf x}_i+\hat\beta_i\nabla_{\bf x}\log q_t({\bf x})\right)+\sqrt{\hat\beta_i}{\bf z}_i\\ \Rightarrow {\bf x}_{i-1}\approx\frac{1}{\sqrt{1-\hat\beta_i}}\left({\bf x}_i+\hat\beta_i\nabla_{\bf x}\log q({\bf x}_i,i)\right)+\sqrt{\hat\beta_i}{\bf z}_i\\ \end{align}$

其中做了一次代换$\hat\beta_i=\beta(\frac{i}{N})\frac{1}{N}$【且为了凑出DDPM的离散形式，实际上引入了一个二阶小量】，当$N$非常大的时候，我们可以认为这种离散化方程可以很好地模拟原本的连续随机过程。事实上宋飏等人的论文中给出了另一种可行的离散方程近似形式【该形式是由SDE离散化推导得到的，而不是通过ELBO优化得到的，原本的DDPM的离散方程被称为祖先采样，但是不是所有形式的扩散过程都是可以通过ELBO优化得到祖先采样形式的】

${\bf x}_{i-1}=\left(2-\sqrt{1-\hat\beta_i}\right){\bf x}_i+\hat\beta_i\nabla_{\bf x}\log q({\bf x}_i,i)+\sqrt{\hat\beta_i}{\bf z}_i$

SMLD(得分匹配模型)的随机微分方程

其前向离散随机微分方程为

${\bf x}_i={\bf x}_{i-1}+\sqrt{\sigma_i^2-\sigma_{i-1}^2}{\bf z}_{i-1}$

可以类似于DDPM的过程将其转化为连续形式

$\begin{align} {\bf x}(t+\Delta t)={\bf x}(t)+\sqrt{\sigma^2(t+\Delta t)-\sigma^2(t)}{\bf z}(t)\\ \Rightarrow d{\bf x}=\sqrt{\frac{\sigma^2(t+\Delta t)-\sigma^2(t)}{\Delta t}}d{\bf w}\\ \Rightarrow d{\bf x}=\sqrt{\frac{d[\sigma^2(t)]}{d t}}d{\bf w} \end{align}$

相应地，容易写出其对应的反向SDE为

$d{\bf x}=-\left(\frac{d[\sigma^2(t)]}{dt}\nabla_{\bf x}\log q_t({\bf x})\right)dt+\sqrt{\frac{d[\sigma^2(t)]}{d t}}d\bar{\bf w}$

同样通过离散化，可以得到

${\bf x}_{i-1}={\bf x}_i+(\sigma_i^2-\sigma_{i-1}^2)\nabla_{\bf x}\log q_i({\bf x}_i)+\sqrt{(\sigma_i^2-\sigma_{i-1}^2)}{\bf z}_i$

Predictor-Corrector算法

此处借鉴了求解ODE常用的数值方法，先进行初步预测，然后对初步预测进行多轮次纠正，使得结果更为准确。

文中还给出了实验中采用的Corrector算法，但是效果不佳，这里就不放了。

条件生成

此处参考网站1。

如果我们希望进行基于$\nabla_{\bf x}q_t({\bf x}|{\bf y})$的条件生成，那么实际上我们很容易计算得到

$\begin{align} \nabla_{\bf x}q_t({\bf x}|{\bf y})\\ =\nabla_{\bf x}q_t({\bf y}|{\bf x})+\nabla_{\bf x}q_t({\bf x})-\nabla_{\bf x}q_t({\bf y})\\ =\nabla_{\bf x}q_t({\bf y}|{\bf x})+\nabla_{\bf x}q_t({\bf x}) \end{align}$

这样的话，我们看起来需要额外训练一个模型来拟合$\nabla_{\bf x}q_t({\bf y}|{\bf x})$，如果在拟合出的$\nabla_{\bf x}(p_\phi)_t({\bf y}|{\bf x})$前加上权重$\gamma$，我们实际上就可以控制遵循条件引导的强弱

$\gamma\nabla_{\bf x}(p_\phi)_t({\bf y}|{\bf x})+\nabla_{\bf x}(p_\theta)_t({\bf x})$

如果不希望训练额外的$\nabla_{\bf x}(p_\phi)_t({\bf y}|{\bf x})$，那么实际上可以再把$\nabla_{\bf x}q_t({\bf y}|{\bf x})$变回去，得到

$\begin{align} \gamma\nabla_{\bf x}(p_\phi)_t({\bf y}|{\bf x})+\nabla_{\bf x}(p_\theta)_t({\bf x})\\ =\gamma\nabla_{\bf x}(p_\theta)_t({\bf x}|{\bf y})+(1-\gamma)\nabla_{\bf x}(p_\theta)_t({\bf x}) \end{align}$

这样的话本质上只需要训练$\nabla_{\bf x}(p_\theta)_t({\bf x}|{\bf y})$，这时只需要引入所谓的空条件，那么也就相当于同时训练了$\nabla_{\bf x}(p_\theta)_t({\bf x})$。

Fokker-Planck方程

通过将Kramers-Moyal展开截断到2阶，就可以得到fokker-planck方程(或者称为Kolmogorov’s前向方程)，该方程可基于SDE推理出各个时刻的分布情况【关于这部分的更细致内容可以参考《Stochastic Differential Equations》 by Bernt Øksendal】

(此处是1维形式)对于任意的马尔可夫过程$\xi (t)$，其在$t$时刻的分布满足$p(x,t)$【$\xi(t)=x$】，那么该随时间变化的分布必然满足
$\frac{\partial }{\partial t}p(x,t)=-\frac{\partial}{\partial x}D^{(1)}(x,t)p(x,t)+\frac{\partial^2}{\partial x^2}D^{(2)}(x,t)p(x,t)$
其中
$D^{(m)}(x,t)=\frac{1}{m!}\lim\limits_{\Delta t\to 0}\left[\frac{1}{\Delta t}\mathbb{E}[(\xi(t+\Delta t)-x)^m]\Big|_{\xi(t)=x}\right]$

对于一般的像前面提到的那样的随机微分方程(${\bf x}:\mathbb{R}^d$)：

$d{\bf x}={\bf f}({\bf x},t)dt+{\bf G}({\bf x},t)d{\bf w}$

则fokker-planck方程形式为

$\frac{\partial p({\bf x},t)}{\partial t}=-\sum_{i=1}^d\frac{\partial}{\partial x_i}[f_i({\bf x},t)p({\bf x},t)]+\frac{1}{2}\sum_{i=1}^d\sum_{j=1}^d\frac{\partial^2}{\partial x_i\partial x_j}\left[\sum_{k=1}^dG_{ik}({\bf x},t)G_{jk}({\bf x},t)p({\bf x},t)\right]$

【此处留待进一步解析】