生成扩散模型漫谈(五):一般框架之SDE篇
By 苏剑林 | 2022-08-03 | 186155位读者 |在写生成扩散模型的第一篇文章时,就有读者在评论区推荐了宋飏博士的论文《Score-Based Generative Modeling through Stochastic Differential Equations》,可以说该论文构建了一个相当一般化的生成扩散模型理论框架,将DDPM、SDE、ODE等诸多结果联系了起来。诚然,这是一篇好论文,但并不是一篇适合初学者的论文,里边直接用到了随机微分方程(SDE)、Fokker-Planck方程、得分匹配等大量结果,上手难度还是颇大的。
不过,在经过了前四篇文章的积累后,现在我们可以尝试去学习一下这篇论文了。在接下来的文章中,笔者将尝试从尽可能少的理论基础出发,尽量复现原论文中的推导结果。
随机微分 #
在DDPM中,扩散过程被划分为了固定的$T$步,还是用《生成扩散模型漫谈(一):DDPM = 拆楼 + 建楼》的类比来说,就是“拆楼”和“建楼”都被事先划分为了$T$步,这个划分有着相当大的人为性。事实上,真实的“拆”、“建”过程应该是没有刻意划分的步骤的,我们可以将它们理解为一个在时间上连续的变换过程,可以用随机微分方程(Stochastic Differential Equation,SDE)来描述。
为此,我们用下述SDE描述前向过程(“拆楼”):
\begin{equation}d\boldsymbol{x} = \boldsymbol{f}_t(\boldsymbol{x}) dt + g_t d\boldsymbol{w}\label{eq:sde-forward}\end{equation}
相信很多读者都对SDE很陌生,笔者也只是在硕士阶段刚好接触过一段时间,略懂皮毛。不过不懂不要紧,我们只需要将它看成是下述离散形式在$\Delta t\to 0$时的极限:
\begin{equation}\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})\label{eq:sde-discrete}\end{equation}
再直白一点,如果假设拆楼需要$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进行数值计算。
对于式$\eqref{eq:sde-discrete}$,读者可能比较有疑惑的是为什么右端第一项是$\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})$才能在长期结果中体现出随机效应的作用。
逆向方程 #
用概率的语言,式$\eqref{eq:sde-discrete}$意味着条件概率为
\begin{equation}\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}\label{eq:sde-proba}\end{equation}
简单起见,这里没有写出无关紧要的归一化因子。按照DDPM的思想,我们最终是想要从“拆楼”的过程中学会“建楼”,即得到$p(\boldsymbol{x}_t|\boldsymbol{x}_{t+\Delta t})$,为此,我们像《生成扩散模型漫谈(三):DDPM = 贝叶斯 + 去噪》一样,用贝叶斯定理:
\begin{equation}\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}\label{eq:bayes-dt}\end{equation}
不难发现,当$\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$足够接近时的情形做近似分析,为此,我们可以用泰勒展开:
\begin{equation}\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)\end{equation}
注意不要忽略了$\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$的偏导数。代入到式$\eqref{eq:bayes-dt}$后,配方得到
\begin{equation}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)\end{equation}
当$\Delta t\to 0$时,$\mathcal{O}(\Delta t)\to 0$不起作用,因此
\begin{equation}\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}\end{equation}
即$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:
\begin{equation}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}\label{eq:reverse-sde}\end{equation}
这就是反向过程对应的SDE,最早出现在《Reverse-Time Diffusion Equation Models》中。这里我们特意在$p$处标注了下标$t$,以突出这是$t$时刻的分布。
得分匹配 #
现在我们已经得到了逆向的SDE为$\eqref{eq:reverse-sde}$,如果进一步知道$\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})$,那么就可以通过离散化格式
\begin{equation}\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}\label{eq:reverse-sde-discrete}\end{equation}
来逐步完成“建楼”的生成过程【其中$\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)$解析解的模型,这意味着
\begin{equation}\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}\end{equation}
是可以直接求出的,比如当$\boldsymbol{f}_t(\boldsymbol{x})$是关于$\boldsymbol{x}$的线性函数时,$p(\boldsymbol{x}_t|\boldsymbol{x}_0)$就可以解析求解。在此前提下,有
\begin{equation}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]\end{equation}
于是
\begin{equation}\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]}\end{equation}
可以看到最后的式子具有“$\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)$。
很多读者应该对如下结果并不陌生(或者推导一遍也不困难):
\begin{equation}\mathbb{E}[\boldsymbol{x}] = \mathop{\text{argmin}}_{\boldsymbol{\mu}}\mathbb{E}_{\boldsymbol{x}}\left[\Vert \boldsymbol{\mu} - \boldsymbol{x}\Vert^2\right]\end{equation}
即要让$\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$的加权平均,即
\begin{equation} \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]}\end{equation}
分母的$\mathbb{E}_{\boldsymbol{x}_0}\left[p(\boldsymbol{x}_t|\boldsymbol{x}_0)\right]$只是起到调节Loss权重的作用,简单起见我们可以直接去掉它,这不会影响最优解的结果。最后我们再对$\boldsymbol{x}_t$积分(相当于对于每一个$\boldsymbol{x}_t$都要最小化上述损失),得到最终的损失函数
\begin{equation}\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}\label{eq:score-match}\end{equation}
这就是“(条件)得分匹配”的损失函数,之前我们在《从去噪自编码器到生成模型》推导的去噪自编码器的解析解,也是它的一个特例。得分匹配的最早出处可以追溯到2005年的论文《Estimation of Non-Normalized Statistical Models by Score Matching》,至于条件得分匹配的最早出处,笔者追溯到的是2011年的论文《A Connection Between Score Matching and Denoising Autoencoders》。
不过,虽然该结果跟得分匹配是一样的,但其实在这一节的推导中,我们已经抛开了“得分”的概念了,纯粹是由目标自然地引导出来的答案,笔者认为这样的处理过程更有启发性,希望这一推导能降低大家对得分匹配的理解难度。
结果倒推 #
至此,我们构建了生成扩散模型的一般流程:
1、通过随机微分方程$\eqref{eq:sde-forward}$定义“拆楼”(前向过程);
2、求$p(\boldsymbol{x}_t|\boldsymbol{x}_0)$的表达式;
3、通过损失函数$\eqref{eq:score-match}$训练$\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$(得分匹配);
4、用$\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$替换式$\eqref{eq:reverse-sde}$的$\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})$,完成“建楼”(反向过程)。
可能大家看到SDE、微分方程等字眼,天然就觉得“恐慌”,但本质上来说,SDE只是个“幌子”,实际上将对SDE的理解转换到式$\eqref{eq:sde-discrete}$和式$\eqref{eq:sde-proba}$上后,完全就可以抛开SDE的概念了,因此概念上其实是没有太大难度的。
不难发现,定义一个随机微分方程$\eqref{eq:sde-forward}$是很容易的,但是从$\eqref{eq:sde-forward}$求解$p(\boldsymbol{x}_t|\boldsymbol{x}_0)$却是不容易的。原论文的剩余篇幅,主要是对两个有实用性的例子推导和实验。然而,既然求解$p(\boldsymbol{x}_t|\boldsymbol{x}_0)$不容易,那么按照笔者的看法,与其先定义$\eqref{eq:sde-forward}$再求解$p(\boldsymbol{x}_t|\boldsymbol{x}_0)$,倒不如像DDIM一样,先定义$p(\boldsymbol{x}_t|\boldsymbol{x}_0)$,然后再来反推对应的SDE?
例如,我们先定义
\begin{equation} p(\boldsymbol{x}_t|\boldsymbol{x}_0) = \mathcal{N}(\boldsymbol{x}_t; \bar{\alpha}_t \boldsymbol{x}_0,\bar{\beta}_t^2 \boldsymbol{I})\end{equation}
并且不失一般性假设起点是$t=0$,终点是$t=1$,那么$\bar{\alpha}_t,\bar{\beta}_t$要满足的边界就是
\begin{equation} \bar{\alpha}_0 = 1,\quad \bar{\alpha}_1 = 0,\quad \bar{\beta}_0 = 0,\quad \bar{\beta}_1 = 1\end{equation}
当然,上述边界条件理论上足够近似就行,也不一定非要精确相等,比如上一篇文章我们分析过DDPM相当于选择了$\bar{\alpha}_t = e^{-5t^2}$,当$t=1$时结果为$e^{-5}\approx 0$。
有了$p(\boldsymbol{x}_t|\boldsymbol{x}_0)$,我们去反推$\eqref{eq:sde-forward}$,本质上就是要求解$p(\boldsymbol{x}_{t+\Delta t}|\boldsymbol{x}_t)$,它要满足
\begin{equation} 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\end{equation}
我们考虑线性的解,即
\begin{equation}d\boldsymbol{x} = f_t\boldsymbol{x} dt + g_t d\boldsymbol{w}\end{equation}
跟《生成扩散模型漫谈(四):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{equation}\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}\end{equation}
令$\Delta t\to 0$,分别解得
\begin{equation}
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)\end{equation}
取$\bar{\alpha}_t\equiv 1$时,结果就是论文中的VE-SDE(Variance Exploding SDE);而如果取$\bar{\alpha}_t^2 + \bar{\beta}_t^2=1$时,结果就是原论文中的VP-SDE(Variance Preserving SDE)。
至于损失函数,此时我们可以算得
\begin{equation}\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}\end{equation}
第二个等号是因为$\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}$,此时式$\eqref{eq:score-match}$为
\begin{equation}\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]\end{equation}
忽略系数后就是DDPM的损失函数,而用$-\frac{\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\boldsymbol{x}_{t+\Delta t}, t+\Delta t)}{\bar{\beta}_{t+\Delta t}}$替换掉式$\eqref{eq:reverse-sde-discrete}$的$\nabla_{\boldsymbol{x}_{t+\Delta t}}\log p(\boldsymbol{x}_{t+\Delta t})$后,结果与DDPM的采样过程具有相同的一阶近似(意味着$\Delta t\to 0$时两者等价)。
文章小结 #
本文主要介绍了宋飏博士建立的利用SDE理解扩散模型的一般框架,其中包括以尽可能直观的语言推导了反向SDE、得分匹配等结果,并对方程的求解给出了自己的想法。
转载到请包括本文地址:https://kexue.fm/archives/9209
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Aug. 03, 2022). 《生成扩散模型漫谈(五):一般框架之SDE篇 》[Blog post]. Retrieved from https://kexue.fm/archives/9209
@online{kexuefm-9209,
title={生成扩散模型漫谈(五):一般框架之SDE篇},
author={苏剑林},
year={2022},
month={Aug},
url={\url{https://kexue.fm/archives/9209}},
}
August 8th, 2023
但我还是不明白为什么在公式(5)中多了t 的偏导数这一项。 一般的二元函数p(x,y)中的x和y是相互独立的,这种我知道如何泰利展开,但是文章中的p(x(t))里的x和t是嵌套的,能用二元泰利展开吗?我感觉这种带嵌套的情况好像不能用二元泰利展开。因为按照多元泰利公式证明不出这种带嵌套的情况是合理的(可以使用多元泰利展开),有没有更具体的相关参考资料可以查阅?多元泰利展开的公式证明见网址:https://max.book118.com/html/2018/0415/161594799.shtm
$(5)$式下面其实已经简单注释过了,$p(\boldsymbol{x}_{t+\Delta t}),p(\boldsymbol{x}_t)$只是习惯的简单记号,它并不是同一个函数在不同点的取值,而是“不同”函数在“不同”点的取值,更严谨地它们应该记做$p_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t}),p_t(\boldsymbol{x}_t)$,其中$t,\boldsymbol{x}_t$确实是相互独立的,在此基础上用正常的二元泰勒展开。
谢谢您的回复,我明白了,$p(x_t)$实际上写成$p_t(x)$更合适,表示时刻$t$的条件下粒子处于位置$x$的概率,在计算概率$p_t(x)$时,$t$与$x$的取值是独立的。
我理解你的困惑,事实上有时候我也有困惑,但有时候这又是比较简单而习惯的记号。比如贝叶斯公式$p(A|B)=\frac{p(B|A)p(A)}{p(B)}$,这里的$p(A)$和$p(B)$其实就是两个不同随机变量的不同分布,但这种情况下我们又不会混淆。
August 8th, 2023
因为按照多元泰利公式的证明过程证明不出这种带嵌套的情况是合理的。
October 12th, 2023
苏老师好,请问在第三个式子中的$\propto$正比符号的后面的exp里 的式子是怎么得过来的呢?需要什么定理吗?而且分母上还有个2也不太明白是为什么?请老师赐教
大致明白了,因为∝后面的式子要满足积分为1,也就是按照标准正态分布那个形式来写,所以分母有个2,分子是(x-u)平方,对面苏老师
就是正态分布的概率密度表达式...
October 31st, 2023
请问苏神,最后算ft,gt的时候,$ \frac{d\bar \alpha_t}{dt} $这些该怎么算呢?是不是$\bar \alpha_t=exp(\int ln(\alpha_t))$?这个积分好像有点难求解析形式,是要把$\alpha_t=(1-\frac{0.002t}{T})$带入再求吗?
另外,VP SDE的形式是$-\frac{1}{2}\beta_t x_t dt+\sqrt{\beta_t}$,我演算了很久,都和DDPM,和本文的结论对不上啊?请问这三个具体该怎么打通呢?
1、一般给定$\bar{\alpha}_t,\bar{\beta}_t$求$f_t,g_t$,如果你要反过来,那就只能算积分了,对于实际训练来说,有个数值积分的结果就行,不一定要解析解。
2、@windSpeak|comment-22965:当$\bar{\alpha}_t^2 + \bar{\beta}_t^2=1$时,
$$g_t^2 = \bar{\alpha}_t^2 \frac{d}{dt}\left(\frac{\bar{\beta}_t^2}{\bar{\alpha}_t^2}\right) = \bar{\alpha}_t^2 \frac{d}{dt}\left(\frac{1-\bar{\alpha}_t^2}{\bar{\alpha}_t^2}\right)=-\frac{2}{\bar{\alpha}_t}\frac{d\bar{\alpha}_t}{dt}=-2f_t$$
所以$g_t = \sqrt{-2f_t}$。
苏神这里第二个式子少了一个负号。
这里如果用DDPM的p(xt|x0)再反推一遍f_t,然后结合这里推导出来的f_t表达式是不是可以直接求解出a_t的表达式呀?还是两者是一样的?a_t的表达式就是解不出来,只能得到g_t与f_t之间的关系?
谢谢,已更正。
这里没有$\alpha_t$,(如果有必要)用的是DDPM的$\bar{\alpha}_t$。
November 8th, 2023
https://www.zhihu.com/question/616179189
按照这里直接将$x_{t+1}$与$x_{t}$之间关系连续化,可以得到另一个关于$f_t$与$\beta_t$的表达式,结合这里的表达式,可以得到$\beta_t$的解。但是问题在于理论上没有两个推导都没有新条件的引入啊,两个$f_t$表达式应该等价才对,为啥可解呢?应该是哪里弄错了?
这个链接用的是标准的DDPM记号,我这里用的是我自己认为更直观的记号,即两者的$\bar{\alpha},\bar{\beta}$含义是有差别的。请参考第一篇文章:https://kexue.fm/archives/9119
December 7th, 2023
[...]在生成扩散模型的发展史上,DDIM和同期Song Yang的扩散SDE都称得上是里程碑式的工作,因为它们建立起了扩散模型与随机微分方程(SDE)、常微分方程(ODE)这两个数学领域的紧密联系,从而允许我们可以利用SDE、ODE已有的各种数学工具来对分析、求解和拓展扩散模型,比如后续大量的加速采样工作都以此为基础,可以说这打开了生成扩散模型的一个全新视角。[...]
December 27th, 2023
在第一节随机微分中,离散状态下的扩散方程第二项为$\sqrt{\Delta t}\epsilon$,关于$\sqrt{\Delta t}$文中给出了直观上的解释,我想补充一个客观上的解释供讨论。
一般而言,随机微分方程第二项中的$\pmb w$为随机维纳过程,或者说为布朗运动,其定义中描述了$\pmb w_{t+s} - \pmb w_{s}$是期望为0、方差为$t$的正态随机变量。而离散定义中的$\sqrt{\Delta t}\epsilon$ 显然也是期望为0、方差为$\Delta t$的正态随机变量。
因此,可以认为,离散定义中的$\sqrt{\Delta t}$就是将布朗运动的方差项提取到了分布外面。
感谢补充。但有一说一,这好像就是换了一个新名词的感觉?本质上就是要回答为什么布朗运动(的方差)要这样定义呀?
January 4th, 2024
请问苏老师,(7)式的约等于是怎么证明的,为什么可以把$f_t(x_t),g_t,p(x_t)$换成$f_{t+\Delta_t},g_{t+\Delta_t},p(x_{t+\Delta_t})$呢?
因为中括号部分乘了$\Delta t$,也就是说整一项(连同$\Delta t$)就是一个一阶无穷小项,中括号内$t\to t+\Delta t$给整一项带来的扰动是二阶无穷小,在当前情况下可以忽略。
January 23rd, 2024
苏老师,请问(5)式代入(4)式配平得到(6)式的过程能否再更详尽地指教一下?
要不然反过来,你详尽地写一下你的尝试过程,然后指出哪里遇到了瓶颈?我好针对性地回答。
February 16th, 2024
苏老师,请问一下,式(18)不应该是$p(x_{t+\Delta t}|x_0)=\int p(x_{t+\Delta t}|x_t,x_0)p(x_t|x_0)dx_t$吗?为什么$p(x_{t+\Delta t}|x_t,x_0)$可以换成$p(x_{t+\Delta t}|x_t)$?这里是用到了markov chain的假定吗?
本文讨论的前向过程是SDE,即式$\eqref{eq:sde-forward}$,离散化后是式$\eqref{eq:sde-discrete}$,显然$\boldsymbol{x}_{t+\Delta t}$只跟$\boldsymbol{x}_t$有关。
十分感谢苏老师的回答,我是不是可以理解成式(1)的SDE限制了$p(x_{t+1}|x_t,x_0)$的可选形式,即只能把$p(x_{t+1}|x_t,x_0)$取为与$x_0$无关的$p(x_{t+1}|x_t)$的形式?
是的
谢谢苏老师的解答~