Reading

SDE和扩散模型

- SMLD 和 DDPM 中使用的噪声扰动可以看作是两个不同 SDE 的离散化
- 扩散模型和评分模型在连续时间极限下完全等价,也就是说将有限次数的加噪过程推广到无穷次, 也就是推广到连续的情况下,可以得到一个更加一般的扩散过程,这个过程可以用SDE来表示,求解更加方便
- 两种方法的目标函数可以互相转换

随机微分

在DDPM中,扩散过程被划分为了固定的T步,还是用DDPM中的类比来说,就是“拆楼”和“建楼”都被事先划分为了T步,这个划分有着相当大的人为性。事实上,真实的“拆”、“建”过程应该是没有刻意划分的步骤的,我们可以将它们理解为一个在时间上连续的变换过程,可以用随机微分方程(Stochastic Differential Equation,SDE)来描述。

为此,我们用下述SDE描述前向过程(“拆楼”):

\[d\boldsymbol{x} = \boldsymbol{f}_t(\boldsymbol{x}) dt + g_t d\boldsymbol{w}\tag{1}\]

相信很多读者都对SDE很陌生,笔者也只是在硕士阶段刚好接触过一段时间,略懂皮毛。不过不懂不要紧,我们只需要将它看成是下述离散形式在\(\Delta t\to 0\)时的极限:

\[\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t = \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t + g_t \sqrt{\Delta t}\boldsymbol{\varepsilon},\quad \boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\tag{2}\]

再直白一点,如果假设拆楼需要1天,那么拆楼就是\(\boldsymbol{x}\)\(t=0\)\(t=1\)的变化过程,每一小步的变化我们可以用上述方程描述。至于时间间隔\(\Delta t\),我们并没有做特殊限制,只是越小的\(\Delta t\)意味着是对原始SDE越好的近似,如果取\(\Delta t=0.001\),那就对应于原来的\(T=1000\),如果是\(\Delta t = 0.01\)则对应于\(T=100\),等等。也就是说,在连续时间的SDE视角之下,不同的\(T\)是SDE不同的离散化程度的体现,它们会自动地导致相似的结果,我们不需要事先指定\(T\),而是根据实际情况下的精确度来取适当的\(T\)进行数值计算。

所以,引入SDE形式来描述扩散模型的本质好处是“将理论分析和代码实现分离开来”,我们可以借助连续性SDE的数学工具对它做分析,而实践的时候,则只需要用任意适当的离散化方案对SDE进行数值计算。

对于(2),读者可能比较有疑惑的是为什么右端第一项是\(\mathcal{O}(\Delta t)\)的,而第二项是\(\mathcal{O}(\sqrt{\Delta t})\)的?也就是说为什么随机项的阶要比确定项的阶要高?这个还真不是那么容易解释,也是SDE比较让人迷惑的地方之一。简单来说,就是\(\boldsymbol{\varepsilon}\)一直服从标准正态分布,如果随机项的权重也是\(\mathcal{O}(\Delta t)\),那么由于标准正态分布的均值为\(\boldsymbol{0}\)、协方差为\(\boldsymbol{I}\),临近的随机效应会相互抵消掉,要放大到\(\mathcal{O}(\sqrt{\Delta t})\)才能在长期结果中体现出随机效应的作用。

\(\boldsymbol{w}_{t}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{t})\)
\(\boldsymbol{w}_{t+\Delta t}-\boldsymbol{w}_{t}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{\Delta t})\)
\(\Delta t \to0\)
\(d\boldsymbol{w}=\sqrt{dt}\boldsymbol{\varepsilon},\quad \boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\)

用概率的语言,(2)意味着条件概率为

\[\begin{aligned} p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t) =&\, \mathcal{N}\left(\boldsymbol{x}_{t+\Delta t};\boldsymbol{x}_t + \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t, g_t^2\Delta t \,\boldsymbol{I}\right)\\ \propto&\, \exp\left(-\frac{\Vert\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t - \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t\Vert^2}{2 g_t^2\Delta t}\right) \end{aligned}\tag{3}\]

逆向方程

简单起见,这里没有写出无关紧要的归一化因子。按照DDPM的思想,我们最终是想要从“拆楼”的过程中学会“建楼”,即得到\(p(\boldsymbol{x}_t|\boldsymbol{x}_{t+\Delta t})\),为此,我们用贝叶斯定理:

\[\begin{aligned} p(\boldsymbol{x}_t|\boldsymbol{x}_{t+\Delta t}) =&\, \frac{p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t)p(\boldsymbol{x}_t)}{p(\boldsymbol{x}_{t+\Delta t})} = p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t) \exp\left(\log p(\boldsymbol{x}_t) - \log p(\boldsymbol{x}_{t+\Delta t})\right)\\ \propto&\, \exp\left(-\frac{\Vert\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t - \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t\Vert^2}{2 g_t^2\Delta t} + \log p(\boldsymbol{x}_t) - \log p(\boldsymbol{x}_{t+\Delta t})\right) \end{aligned}\tag{4}\]

不难发现,当\(\Delta t\)足够小时,只有当\(\boldsymbol{x}_{t+\Delta t}\)\(\boldsymbol{x}_t\)足够接近时,\(p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t)\) 才会明显不等于0,反过来也只有这种情况下\(p(\boldsymbol{x}_t|\boldsymbol{x}_{t+\Delta t})\) 才会明显不等于0。因此,我们只需要对\(\boldsymbol{x}_{t+\Delta t}与\boldsymbol{x}_t\)足够接近时的情形做近似分析,为此,我们可以用泰勒展开:

\[\log p(\boldsymbol{x}_{t+\Delta t})\approx \log p(\boldsymbol{x}_t) + (\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t)\cdot \nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t) + \Delta t \frac{\partial}{\partial t}\log p(\boldsymbol{x}_t)\]

注意不要忽略了\(\frac{\partial}{\partial t}\)项,因为\(p(\boldsymbol{x}_t)\)实际上是“\(t\) 时刻随机变量等于\(\boldsymbol{x}_t\) 的概率密度”,而\(p(\boldsymbol{x}_{t+\Delta t})\)实际上是“\(t+\Delta t\)时刻随机变量等于\(\boldsymbol{x}_{t+\Delta t}\)的概率密度”,也就是说\(p(\boldsymbol{x}_t)\) 实际上同时是\(t\)\(\boldsymbol{x}_t\) 的函数,所以要多一项 \(t\) 的偏导数。代入到(4)后,配方得到

\[p(\boldsymbol{x}_t|\boldsymbol{x}_{t+\Delta t}) \propto \exp\left(-\frac{\Vert\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t - \left[\boldsymbol{f}_t(\boldsymbol{x}_t) - g_t^2\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t) \right]\Delta t\Vert^2}{2 g_t^2\Delta t} + \mathcal{O}(\Delta t)\right)\]

\(\Delta t\to 0\)时,\(\mathcal{O}(\Delta t)\to 0\)不起作用,因此

\[\begin{aligned} p(\boldsymbol{x}_t|\boldsymbol{x}_{t+\Delta t}) \propto&\, \exp\left(-\frac{\Vert\boldsymbol{x}_{t+\Delta t} - \boldsymbol{x}_t - \left[\boldsymbol{f}_t(\boldsymbol{x}_t) - g_t^2\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t) \right]\Delta t\Vert^2}{2 g_t^2\Delta t}\right) \\ \approx&\,\exp\left(-\frac{\Vert \boldsymbol{x}_t - \boldsymbol{x}_{t+\Delta t} + \left[\boldsymbol{f}_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t}) - g_{t+\Delta t}^2\nabla_{\boldsymbol{x}_{t+\Delta t}}\log p(\boldsymbol{x}_{t+\Delta t}) \right]\Delta t\Vert^2}{2 g_{t+\Delta t}^2\Delta t}\right) \end{aligned}\]

\(p(\boldsymbol{x}_t|\boldsymbol{x}_{t+\Delta t})\) 近似一个均值为\(\boldsymbol{x}_{t+\Delta t} - \left[\boldsymbol{f}_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t}) - g_{t+\Delta t}^2\nabla_{\boldsymbol{x}_{t+\Delta t}}\log p(\boldsymbol{x}_{t+\Delta t}) \right]\Delta t\)、协方差为\(g_{t+\Delta t}^2\Delta t\,\boldsymbol{I}\) 的正态分布,取\(\Delta t\to 0\) 的极限,那么对应于SDE:

\[d\boldsymbol{x} = \left[\boldsymbol{f}_t(\boldsymbol{x}) - g_t^2\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x}) \right] dt + g_t d\boldsymbol{w}\tag{5}\]

这就是反向过程对应的SDE,最早出现在《Reverse-Time Diffusion Equation Models》中。这里我们特意在\(p\)处标注了下标\(t\),以突出这是\(t\)时刻的分布。

image

所以相比前向过程,逆向过程的SDE多了一个score fuction \(\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\),需要对score fuction进行求解,就可以完成采样过程。

得分匹配

现在我们已经得到了逆向的SDE为(5),如果进一步知道\(\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\) ,那么就可以通过离散化格式

\[\boldsymbol{x}_t - \boldsymbol{x}_{t+\Delta t} = - \left[\boldsymbol{f}_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t}) - g_{t+\Delta t}^2\nabla_{\boldsymbol{x}_{t+\Delta t}}\log p(\boldsymbol{x}_{t+\Delta t}) \right]\Delta t - g_{t+\Delta t} \sqrt{\Delta t}\boldsymbol{\varepsilon}\tag{6}\]

来逐步完成“建楼”的生成过程【其中\(\boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\)】,从而完成一个生成扩散模型的构建。

那么如何得到\(\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\) 呢?\(t\)时刻的\(p_t(\boldsymbol{x})\)就是前面的\(p(\boldsymbol{x}_t)\),它的含义就是 \(t\) 时刻的边缘分布。在实际使用时,我们一般会设计能找到\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\) 解析解的模型,这意味着

\[\small p(\boldsymbol{x}_t|\boldsymbol{x}_0) = \lim_{\Delta t\to 0}\int\cdots\iint p(\boldsymbol{x}_t|\boldsymbol{x}_{t-\Delta t})p(\boldsymbol{x}_{t-\Delta t}|\boldsymbol{x}_{t-2\Delta t})\cdots p(\boldsymbol{x}_{\Delta t}|\boldsymbol{x}_0) d\boldsymbol{x}_{t-\Delta t} d\boldsymbol{x}_{t-2\Delta t}\cdots d\boldsymbol{x}_{\Delta t}\]

是可以直接求出的,比如当\(\boldsymbol{f}_t(\boldsymbol{x})\) 是关于\(\boldsymbol{x}\) 的线性函数时,\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\) 就可以解析求解。在此前提下,有

\[p(\boldsymbol{x}_t) = \int p(\boldsymbol{x}_t|\boldsymbol{x}_0)\tilde{p}(\boldsymbol{x}_0)d\boldsymbol{x}_0=\mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right]\]

于是

\[\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t) = \frac{\mathbb{E}_{\boldsymbol{x}_0}\left[\nabla_{\boldsymbol{x}_t} p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right]}{\mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right]} = \frac{\mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right]}{\mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right]}\]

可以看到最后的式子具有“\(\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\)的加权平均”的形式,由于假设了\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\)有解析解,因此上式实际上是能够直接估算的,然而它涉及到对全体训练样本 \(\boldsymbol{x}_0\) 的平均,一来计算量大,二来泛化能力也不够好。因此,我们希望用神经网络学一个函数\(\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\),使得它能够直接计算\(\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t)\)

很多读者应该对如下结果并不陌生(或者推导一遍也不困难):

\[\mathbb{E}[\boldsymbol{x}] = \mathop{\text{argmin}}_{\boldsymbol{\mu}}\mathbb{E}_{\boldsymbol{x}}\left[\Vert \boldsymbol{\mu} - \boldsymbol{x}\Vert^2\right]\]

即要让\(\boldsymbol{\mu}\) 等于\(\boldsymbol{x}\) 的均值,只需要最小化\(\Vert \boldsymbol{\mu} - \boldsymbol{x}\Vert^2\) 的均值。同理,要让\(\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\) 等于\(\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\) 的加权平均【即\(\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t)\)】,则只需要最小化\(\left\Vert \boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right\Vert^2\) 的加权平均,即

\[ \frac{\mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\left\Vert \boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right\Vert^2\right]}{\mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right]}\]

分母的\(\mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right]\) 只是起到调节Loss权重的作用,简单起见我们可以直接去掉它,这不会影响最优解的结果。最后我们再对\(\boldsymbol{x}_t\) 积分(相当于对于每一个\(\boldsymbol{x}_t\)都要最小化上述损失),得到最终的损失函数

\[\begin{aligned}&\,\int \mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\left\Vert \boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right\Vert^2\right] d\boldsymbol{x}_t \\ =&\, \mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{x}_t \sim p(\boldsymbol{x}_t|\boldsymbol{x}_0)\tilde{p}(\boldsymbol{x}_0)}\left[\left\Vert \boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right\Vert^2\right] \end{aligned}\tag{7}\]

这就是“(条件)得分匹配”的损失函数。得分匹配的最早出处可以追溯到2005年的论文《Estimation of Non-Normalized Statistical Models by Score Matching》,至于条件得分匹配的最早出处,笔者追溯到的是2011年的论文《A Connection Between Score Matching and Denoising Autoencoders》。不过,虽然该结果跟得分匹配是一样的,但其实在这一节的推导中,我们已经抛开了“得分”的概念了,纯粹是由目标自然地引导出来的答案,笔者认为这样的处理过程更有启发性,希望这一推导能降低大家对得分匹配的理解难度。

结果倒推

至此,我们构建了生成扩散模型的一般流程:

1、通过随机微分方程(1)定义“拆楼”(前向过程);
2、求\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\)的表达式;
3、通过损失函数(7)训练\(\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\)(得分匹配);
4、用\(\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}t, t)\)替换式8的\(\nabla{\boldsymbol{x}}\log p_t(\boldsymbol{x})\),完成“建楼”(反向过程)。

可能大家看到SDE、微分方程等字眼,天然就觉得“恐慌”,但本质上来说,SDE只是个“幌子”,实际上将对SDE的理解转换到(2)(3)上后,完全就可以抛开SDE的概念了,因此概念上其实是没有太大难度的。

不难发现,定义一个随机微分方程(1)是很容易的,但是从(1)求解\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\)却是不容易的。原论文的剩余篇幅,主要是对两个有实用性的例子推导和实验。然而,既然求解\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\)不容易,那么按照笔者的看法,与其先定义(1)再求解\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\),倒不如像DDIM一样,先定义\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\),然后再来反推对应的SDE?

例如,我们先定义

\[p(\boldsymbol{x}_t|\boldsymbol{x}_0) = \mathcal{N}(\boldsymbol{x}_t; \bar{\alpha}_t \boldsymbol{x}_0,\bar{\beta}_t^2 \boldsymbol{I})\]

并且不失一般性假设起点是\(t=0\),终点是 \(t=1\),那么 \(\bar{\alpha}_t,\bar{\beta}_t\) 要满足的边界就是

\[ \bar{\alpha}_0 = 1,\quad \bar{\alpha}_1 = 0,\quad \bar{\beta}_0 = 0,\quad \bar{\beta}_1 = 1\]

当然,上述边界条件理论上足够近似就行,也不一定非要精确相等,比如上一篇文章我们分析过DDPM相当于选择了\(\bar{\alpha}_t = e^{-5t^2}\),当t=1时结果为\(e^{-5}\approx 0\)

有了\(p(\boldsymbol{x}_t|\boldsymbol{x}_0)\),我们去反推(1),本质上就是要求解\(p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t)\),它要满足

\[p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_0) = \int p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t) p(\boldsymbol{x}_t|\boldsymbol{x}_0) d\boldsymbol{x}_t\]

我们考虑线性的解,即

\[d\boldsymbol{x} = f_t\boldsymbol{x} dt + g_t d\boldsymbol{w}\]

《DDIM = 高观点DDPM》一样,我们写出

\[\begin{array}{c|c|c} \hline \text{记号} & \text{含义} & \text{采样}\\ \hline p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_0) & \mathcal{N}(\boldsymbol{x}_t;\bar{\alpha}_{t+\Delta t} \boldsymbol{x}_0,\bar{\beta}_{t+\Delta t}^2 \boldsymbol{I}) & \boldsymbol{x}_{t+\Delta t} = \bar{\alpha}_{t+\Delta t} \boldsymbol{x}_0 + \bar{\beta}_{t+\Delta t} \boldsymbol{\varepsilon} \\ \hline p(\boldsymbol{x}_t|\boldsymbol{x}_0) & \mathcal{N}(\boldsymbol{x}_t;\bar{\alpha}_t \boldsymbol{x}_0,\bar{\beta}_t^2 \boldsymbol{I}) & \boldsymbol{x}_t = \bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon}_1 \\ \hline p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t) & \mathcal{N}(\boldsymbol{x}_{t+\Delta t}; (1 + f_t\Delta t) \boldsymbol{x}_t, g_t^2 \Delta t\, \boldsymbol{I}) & \boldsymbol{x}_{t+\Delta t} = (1 + f_t\Delta t) \boldsymbol{x}_t + g_t\sqrt{\Delta t}\boldsymbol{\varepsilon}_2 \\ \hline {\begin{array}{c}\int p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t) \\ p(\boldsymbol{x}_t|\boldsymbol{x}_0) d\boldsymbol{x}_t\end{array}} & & {\begin{aligned}&\,\boldsymbol{x}_{t+\Delta t} \\ =&\, (1 + f_t\Delta t) \boldsymbol{x}_t + g_t\sqrt{\Delta t} \boldsymbol{\varepsilon}_2 \\ =&\, (1 + f_t\Delta t) (\bar{\alpha}_t \boldsymbol{x}_0 + \bar{\beta}_t \boldsymbol{\varepsilon}_1) + g_t\sqrt{\Delta t} \boldsymbol{\varepsilon}_2 \\ =&\, (1 + f_t\Delta t) \bar{\alpha}_t \boldsymbol{x}_0 + ((1 + f_t\Delta t)\bar{\beta}_t \boldsymbol{\varepsilon}_1 + g_t\sqrt{\Delta t} \boldsymbol{\varepsilon}_2) \\ \end{aligned}} \\ \hline \end{array}\]

由此可得

\[\begin{aligned} \bar{\alpha}_{t+\Delta t} =&\, (1 + f_t\Delta t) \bar{\alpha}_t \\ \bar{\beta}_{t+\Delta t}^2 =&\, (1 + f_t\Delta t)^2\bar{\beta}_t^2 + g_t^2\Delta t \end{aligned}\]

\(\Delta t\to 0\),分别解得

\[ f_t = \frac{d}{dt} \left(\ln \bar{\alpha}_t\right) = \frac{1}{\bar{\alpha}_t}\frac{d\bar{\alpha}_t}{dt}, \quad g_t^2 = \bar{\alpha}_t^2 \frac{d}{dt}\left(\frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\right) = 2\bar{\alpha}_t \bar{\beta}_t \frac{d}{dt}\left(\frac{\bar{\beta}_t}{\bar{\alpha}_t}\right)\]

\(\bar{\alpha}_t\equiv 1\) 时,结果就是论文中的VE-SDE(Variance Exploding SDE);而如果取\(\bar{\alpha}_t^2 + \bar{\beta}_t^2=1\)时,结果就是原论文中的VP-SDE(Variance Preserving SDE)。

至于损失函数,此时我们可以算得

\[\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t|\boldsymbol{x}_0) = -\frac{\boldsymbol{x}_t - \bar{\alpha}_t\boldsymbol{x}_0}{\bar{\beta}_t^2}=-\frac{\boldsymbol{\varepsilon}}{\bar{\beta}_t}\]

第二个等号是因为\(\boldsymbol{x}_t = \bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}\),为了跟以往的结果对齐,我们设\(\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) = -\frac{\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)}{\bar{\beta}_t}\),此时(7)

\[\frac{1}{\bar{\beta}_t^2}\mathbb{E}_{\boldsymbol{x}_0\sim \tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I})}\left[\left\Vert \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\bar{\alpha}_t\boldsymbol{x}_0 + \bar{\beta}_t\boldsymbol{\varepsilon}, t) - \boldsymbol{\varepsilon}\right\Vert^2\right]\tag{8}\]

忽略系数后就是DDPM的损失函数,而用 \(-\frac{\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t+\Delta t}, t+\Delta t)}{\bar{\beta}_{t+\Delta t}}\) 替换掉(6)\(\nabla_{\boldsymbol{x}_{t+\Delta t}}\log p(\boldsymbol{x}_{t+\Delta t})\)后,结果与DDPM的采样过程具有相同的一阶近似(意味着\(\Delta t\to 0\)时两者等价)。

具体实例

上述 \(d{x}={f}({x},t)dt+g(t)d\mathbf{w}\) 给出了一般化的SDE的定义,实际上 \({f}({x},t)\)\(g(t)\) 的不同取值也会得到不同的生成模型。SMLD 和 DDPM 中使用的噪声扰动可以看作是两个不同 SDE 的离散化。

VE-SDE在 \(t \to \infty\) 时总是产生方差爆炸的过程。相反,VP-SDE产生具有有界方差的过程。

VE-SDE(Variance Exploding SDE)

当使用总共\(N\)个噪声尺度时,SMLD的每个扰动核 \(p_{\sigma_{i}}\left({x}\mid{x}_{0}\right)\) 对应于 \(x_i\) 的分布,即可写成一下马尔可夫链形式:

\[\mathbf{x}_i=\mathbf{x}_{i-1}+\sqrt{\sigma_i^2-\sigma_{i-1}^2}\mathbf{z}_{i-1},\quad i=1,\cdots,N\]

在上述表达中,我们假设初始随机变量 \(\mathbf{z}_{i-1}\) 符合均值为 \(\mathbf{0}\) 协方差矩阵为单位矩阵 \(\mathbf{I}\) 的多元正态分布 \(\mathcal{N}(\mathbf{0},\mathbf{I})\) 。为了简化符号,引入了参数 \(\sigma_0=0\) 。在 \(N\to\infty\) 的极限下,随机变量序列 \(\{\sigma_i\}_{i=1}^N\) 变成了连续函数 \(\sigma(t)\) , \(\mathbf{z}_i\) 变成了连续随机过程 \(\mathbf{z}(t)\) 。马尔可夫链 \(\{\mathbf{x}_i\}_{i=1}^N\) 也演变为了连续随机过程 \(\{\mathbf{x}(t)\}_{t=0}^1\) ,其中我们使用连续时间变量 \(t\in[0,1]\) 进行索引,而不再使用整数 \(i\) 。由此可以将上式写为:

\[\mathbf{x}(t+\Delta t)=\mathbf{x}(t)+\sqrt{\sigma^2(t+\Delta t)-\sigma^2(t)}\mathbf{z}(t)\approx\mathbf{x}(t)+\sqrt{\frac{\mathrm{d}\left[\sigma^2(t)\right]}{\mathrm{d}t}\Delta t}\mathbf{z}(t)\]

\(\Delta t\ll1\) 时,近似等式成立。在 \(\Delta t\rightarrow0\) 的极限下,它收敛于以下SDE:

\[\mathrm{d}\mathbf{x}=\sqrt{\frac{\mathrm{d}\left[\sigma^2(t)\right]}{\mathrm{d}t}}\mathrm{d}\mathbf{w}\]

VP-SDE(Variance Preserving SDE)

同样对于 DDPM 的扰动核 \(\text{}\{p_{\alpha_i}(\mathbf{x}\mid\mathbf{x}_0)\}_{i=1}^N\),离散马尔可夫链是

\[\mathbf{x}_i=\sqrt{1-\beta_i}\mathbf{x}_{i-1}+\sqrt{\beta_i}\mathbf{z}_{i-1},\quad i=1,\cdots,N\]

同样简化时间变量到 \(t\in[0,1]\) 的范围,上式可改写为:

\[\mathbf{x}_i=\sqrt{1-\frac{\bar{\beta}_i}{N}}\mathbf{x}_{i-1}+\sqrt{\frac{\bar{\beta}_i}{N}}\mathbf{z}_{i-1},\quad i=1,\cdots,N.\]

\(\Delta t = 1/N\) ,上式可以改写为离散形式下的:

\[\begin{aligned} \mathbf{x}(t+\Delta t)& =\sqrt{1-\beta(t+\Delta t)\Delta t}\mathbf{x}(t)+\sqrt{\beta(t+\Delta t)\Delta t}\mathbf{z}(t) \\ &\approx\mathbf{x}(t)-\frac{1}{2}\beta(t+\Delta t)\Delta t\mathbf{x}(t)+\sqrt{\beta(t+\Delta t)\Delta t}\mathbf{z}(t) \\ &\approx\mathbf{x}(t)-\frac{1}{2}\beta(t)\Delta t\mathbf{x}(t)+\sqrt{\beta(t)\Delta t}\mathbf{z}(t), \end{aligned}\]

\(\Delta t\ll1\) 时,近似等式成立。在 \(\Delta t\rightarrow0\) 的极限下,它收敛于以下SDE:

\[\mathrm{d}\mathbf{x}=-\frac12\beta(t)\mathbf{x}\mathrm{d}t+\sqrt{\beta(t)}\mathrm{d}\mathbf{w}\]

整合联系

DDPM与SMLD的联系

DDPM中我们希望找到一个 \(\epsilon_{\theta}(x_t,t)\) 去近似 \(\frac{x_{t}-\bar{\alpha}_tx_{0}}{\bar{\beta}_t}\) ,而SMLD中则是希望找到一个\(s_{\theta}({x},t) \)来估计 \(\nabla_{x}\log p({x})\)

而在这两个模型中, \(\boldsymbol{x}_t\) 的条件概率都可以写成 \(p(\boldsymbol{x}_t|\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_t;\bar{\alpha}_t\boldsymbol{x}_0,\bar{\beta}_t^2\boldsymbol{I})\) (DDPM要求 \(\bar{\alpha}_{t}^{2}+\bar{\beta}_{t}^{2}=1\) ;而SMLD要求 \(\bar{\alpha}_{t}=1 \)

于是可以得到 \(\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t|\boldsymbol{x}_0)=-\frac{\boldsymbol{x}_t-\bar{\alpha}_t\boldsymbol{x}_0}{\bar{\beta}_t^2}=-\frac\varepsilon{\bar{\beta}_t}\),令 \(s_{\theta}(x_{t},t)=-\frac{\epsilon_{\theta}(x_{t},t)}{\bar{\beta}_{t}}\) ,此时SMLD的损失函数可以写为:

\[\frac1{\bar{\beta}_t^2}\mathbb{E}_{\boldsymbol{x}_0\sim\tilde{p}(\boldsymbol{x}_0),\boldsymbol{\varepsilon}\sim\mathcal{N}(\mathbf{0},\boldsymbol{I})}\left[\left\|\boldsymbol{\varepsilon}_{\theta}(\bar{\alpha}_t\boldsymbol{x}_0+\bar{\beta}_t\boldsymbol{\varepsilon},t)-\boldsymbol{\varepsilon}\right\|^2\right]\]

忽略系数后就是DDPM的损失函数,这说明二者实质上是相通的,也即VE-SDE和VP-SDE是相通的。由此即将DDPM、SMLD和SDE联系了起来。

SDE与ODE的演变

从数学上(通过Fokker-Planck方程)也可以证明存在不同方差的前向过程,它们产生的边缘分布 \(p({x_t})\) 是一样的,从而可以推出对于所有扩散过程,都存在一个相应的确定性过程,其轨迹与SDE具有相同的边缘概率密度,\(\sigma_{t}=0\) 此时SDE退化为常微分方程ODE:

\[\mathrm{d}\mathbf{x}=\Big[\mathbf{f}(\mathbf{x},t)-\frac{1}{2}g(t)^{2}\nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})\Big]\mathrm{d}t\]

这个ODE称为概率流ODE(Probability flow ODE),此时扩散过程是确定过程的变换,因而可以基于这个式子做精确的似然计算、在隐空间编码并编辑等任务。同时,对于反向传播 \(x_T\)\(x_0\) 的过程也可以采用一些ODE算法去加速。

伊藤公式与Fokker-Planck方程

要通过 伊藤方程 (Itô formula) 求得 \(p(x_t)\),利用SDE对应的 Fokker-Planck方程 来描述概率密度\(p(x_t)\) 的时间演化过程。以下是详细的推导过程:

给定的SDE:假设随机过程 \(x_t\) 满足以下SDE:

\[dx_t = f(x_t, t) dt + g(x_t, t) dw_t\]

其中:

  • \(f(x_t, t)\) 是漂移项 (drift term),描述了系统的确定性趋势。
  • \(g(x_t, t)\) 是扩散系数 (diffusion coefficient),描述了随机性。
  • \(w_t\) 是标准维纳过程 (Wiener process)。

目标是求得 \(p(x_t)\) ,即随机变量 \(x_t\) 的概率密度函数。

利用伊藤公式,可以推导出概率密度 \(p(x_t)\) 的演化方程,即 Fokker-Planck方程。这个方程描述了 \(p(x_t)\) 随时间 \(t\) 的变化。

伊藤公式

对于任意光滑函数 \(\phi(x_t, t)\),伊藤公式为:

\[d\phi(x_t, t) = \frac{\partial \phi}{\partial t} dt + \frac{\partial \phi}{\partial x} dx_t + \frac{1}{2} \frac{\partial^2 \phi}{\partial x^2} (g(x_t, t))^2 dt\]

\(dx_t\) 替换为SDE的形式 \(dx_t = f(x_t, t)dt + g(x_t, t)dw_t\),并取期望值,可以得到:

\[\frac{d}{dt} \mathbb{E}[\phi(x_t, t)] = \mathbb{E}\left[ \frac{\partial \phi}{\partial t} + f(x_t, t) \frac{\partial \phi}{\partial x} + \frac{1}{2} (g(x_t, t))^2 \frac{\partial^2 \phi}{\partial x^2} \right]\]

转化为概率密度形式

根据概率密度的定义:

\[\mathbb{E}[\phi(x_t, t)] = \int \phi(x, t) p(x, t) dx\]

将期望值形式代入上式,并通过分部积分,将导数作用从 \(\phi(x, t)\) 转移到 \(p(x, t)\) 。最终得到Fokker-Planck方程:

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

Fokker-Planck方程的物理意义

Fokker-Planck方程包含两个部分:

  1. 漂移项\(\frac{\partial}{\partial x} [f(x, t)p(x, t)]\) 这一项描述了系统的确定性变化趋势,源于SDE中的漂移项 \(f(x, t)\)
  2. 扩散项\(\frac{1}{2} \frac{\partial^2}{\partial x^2} [(g(x, t))^2 p(x, t)]\) 这一项描述了随机扰动对概率密度的影响,源于扩散系数 \(g(x, t)\)

证明存在不同方差的前向过程(SDE和ODE),它们产生的边缘分布 \(p({x_t})\)是一样的

  1. 假设我们有两个不同的SDE过程,它们具有相同的边缘分布\(p(x,t)\)
    • SDE1: \(dx = f₁dt + σ₁dw\)
    • SDE2: \(dx = f₂dt + σ₂dw\)
  1. 这两个过程对应的Fokker-Planck方程必须相同(因为描述同一个概率密度演化):
\[-\frac{\partial}{\partial x}(f_1p) + \frac{1}{2}\frac{\partial^2}{\partial x^2}(\sigma_1^2p) = -\frac{\partial}{\partial x}(f_2p) + \frac{1}{2}\frac{\partial^2}{\partial x^2}(\sigma_2^2p)\]
  1. 特别地,当\(σ₂ = 0\)时(即ODE情况),方程变为:
\[-\frac{\partial}{\partial x}(f_1p) + \frac{1}{2}\frac{\partial^2}{\partial x^2}(\sigma_1^2p) = -\frac{\partial}{\partial x}(f_2p)\]
  1. 从这个等式可以解出\(f₂\)
\[f_2 = f_1 - \frac{1}{2p}\frac{\partial}{\partial x}(\sigma_1^2p)\]
  1. \(\frac{1}{p}\frac{\partial p}{\partial x}\)可以转换为\(\nabla_{\mathbf{x}}\log p\)\(\sigma_1^2\)又对应sde中的\(g(t)^2\),变换:
\[f_2 = f_1 - \frac{1}{2}g(t)^2\nabla_{\mathbf{x}}\log p\]

得证

求解 \(p(x_t)\)

通过Fokker-Planck方程,可以求得 \(p(x_t)\) 的具体形式。以下是几种常见的情况:

简单情况:常系数漂移和扩散

如果 \(f(x, t) = f_0\)\(g(x, t) = g_0\) 是常数,则Fokker-Planck方程简化为:

\[\frac{\partial p(x, t)}{\partial t} = -f_0 \frac{\partial p(x, t)}{\partial x} + \frac{g_0^2}{2} \frac{\partial^2 p(x, t)}{\partial x^2}\]

这是一个线性偏微分方程,其解是高斯分布:

\[p(x, t) = \frac{1}{\sqrt{2\pi \sigma_t^2}} \exp\left(-\frac{(x - \mu_t)^2}{2\sigma_t^2}\right)\]

其中:

  • 均值: \(\mu_t = x_0 + f_0 t\)
  • 方差:\(\sigma_t^2 = g_0^2 t\)

更复杂情况

对于一般的 \(f(x, t)\)\(g(x, t)\),Fokker-Planck方程通常需要通过数值方法求解,例如:

  • 有限差分法
  • 蒙特卡洛模拟

通过伊藤方程,我们可以将随机微分方程 (SDE) 转化为Fokker-Planck方程,从而描述概率密度 \(p(x_t)\) 的演化过程。在扩散模型中,利用这一理论可以明确地求解出前向过程的概率分布 \(p(x_t)\),并为后向生成过程提供理论基础。

总结

扩散模型通过定义一个不断加噪声的前向过程来将图片逐步变为高斯噪声,再通过定义了一个逆向过程将高斯噪声逐步去噪变为清晰图片以得到采样。

在采样过程中,根据是否添加额外的噪声,可以将扩散模型分为两类:

  • 一类是扩散随机微分方程模型(Diffusion SDE) ,其反向过程可以建模为 \(d{x}=[{f}({x},t)-g^2(t)s_{\theta}({x}(t),t)]dt+g(t)d\bar{\mathbf{w}}\)
  • 另一类是扩散常微分方程(Diffusion ODE),其反向过程可以建模为 \({d}{x}=[{f}({x},t)-\frac{1}{2}g(t)^{2}s_{\theta}({x}(t),t)]\mathrm{d}t\)

Reference

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

Generative Modeling by Estimating Gradients of the Data Distribution

扩散模型Diffusion model | 基于分数的生成模型