【理论推导】随机微分方程(SDE)视角下的Diffusion Model与Score-based Model

SDE与DDPM

以 DDPM 为例，DDPM的通项公式为
$$x_t\sim \mathcal{N}(\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)I)$$
当我们固定 $t$ 的取值时，$x_t$ 是定义在样本空间上的函数，即为一随机变量，当我们固定$x$的随机性时，即为关于变量 $t$ 的一个函数，因此 $x_t$ 是一随机过程，而对于一组确定的 $\{x_t{\}}_{t=0}^T$，称为随机过程的一个实现，或是一条轨迹/轨道，而随机过程可以使用随机微分方程(Stochastic Differential Equation)进行描述

SDE定义为
$$d{X}_t=b(t,X_t)dt+\sigma (t,X_t)d{B}_t$$
其中 $B_t$ 表示布朗运动，此方程的解称为伊藤(Itô)过程或是扩散过程。如果我们将 DDPM的迭代式从离散扩展到连续区间，即 $x_t\to x_{t+\Delta t}(\Delta t\to 0)$，即可得到SDE形式的扩散过程，论文中表示为
$$dx=f(x,t)dt+g(t)dw$$
其中 $w$ 表示一个标准布朗运动，$f(\cdot ,t)$ 称作是漂移系数 (drift coefficient)，描述了确定性的变化过程，$g(\cdot )$ 称作是扩散系数 (diffusion coefficient)，描述了不确定的变化过程

SDE视角下的生成模型

SDE的离散形式表示如下所示
$$x_{t+{\Delta }_t}-x_t=f(x,t)\Delta t+g(t)\sqrt{\Delta t}\epsilon$$
其中 $\epsilon \sim \mathcal{N}(0,I)$，因此有条件概率分布
$$x_{t+{\Delta }_t}|x_t\sim \mathcal{N}(x_t+f(x_t,t)\Delta t,g^2(t)\Delta tI)$$
考虑逆过程 $x_t|x_{t+\Delta t}$，有
$$\begin{array}{rl}p(x_t|x_{t+\Delta t})& =\frac{p(x_{t+\Delta t}|x_t)p(x_t)}{p(x_{t+\Delta t})}\\ & =p(x_{t+\Delta t}|x_t)\exp (\log p(x_t)-\log p(x_{t+\Delta t}))\\ & \approx p(x_{t+\Delta t}|x_t)\exp \{-(x_{t+\Delta t}-x_t){▽}_{x_t}\log p(x_t)-\Delta t\frac{\partial }{\partial t}\log p(x_t)\}\\ & \propto \exp \{-\frac{||x_{t+\Delta t}-x_t-f(x_t,t)\Delta t|{|}_2^2}{2g^2(t)\Delta t}-(x_{t+\Delta t}-x_t){▽}_{x_t}\log p(x_t)-\Delta t\frac{\partial }{\partial t}\log p(x_t)\}\\ & =\exp \left\{-\frac{1}{2g^2(t)\Delta t}||(x_{t+\Delta t}-x_t)-(f(x_t,t)-g^2(t){▽}_{x_t}\log p(x_t))\Delta t|{|}_2^2-\Delta t\frac{\partial }{\partial t}\log p(x_t)-\frac{f^2(x_t,t)\Delta t}{2g^2(t)}+\frac{(f(x_t,t)-g^2(t){▽}_{x_t}\log p(x_t){)}^2\Delta t}{2g^2(t)}\right\}\\ & \stackrel{\Delta t\to 0}{=}\exp \left\{-\frac{1}{2g^2(t+\Delta t)\Delta t}||(x_{t+\Delta t}-x_t)-(f(x_{t+\Delta t},t+\Delta t)-g^2(t+\Delta t){▽}_{x_{t+\Delta t}}\log p(x_{t+\Delta t}))\Delta t|{|}_2^2\right\}\end{array}$$

因此，$x_t|x_{t+\Delta t}$ 服从均值方差如下的高斯分布
$$\begin{gathered} \mu =x_{t+\Delta t}-(f(x_{t+\Delta t},t+\Delta t)-g^2(t+\Delta t){▽}_{x_{t+\Delta t}}\log p(x_{t+\Delta t}))\Delta t \\ {\sigma }^2=g^2(t+\Delta t)\Delta t \end{gathered}$$
因此，有SDE表示逆过程离散形式与连续形式如下所示
$$\begin{array}{rl}{x}_{t+\Delta t}-x_t& =(f(x_{t+\Delta t},t+\Delta t)-g^2(t+\Delta t){▽}_{x_{t+\Delta t}}\log p(x_{t+\Delta t}))\Delta t+g(t+\Delta t)\sqrt{\Delta t}\epsilon \\ dx& =[f(x,t)-g^2(t){▽}_{x_t}\log p(x_t)]+g(t)dw\end{array}$$

考虑 NCSN 与 DDPM，两种形式均可以统一到SDE的理论表示形式下面，分别称为 VE-SDE (Variance Exploding) 与 VP-SDE (Variance Preserving)，对应 NCSN 与 DDPM

对于NCSN，正向过程（加噪声）如下所示
$$\begin{gathered} x_t=x_0+{\sigma }_t\epsilon \\ {x}_{t+1}=x_t+\sqrt{{\sigma }_{t+1}^2-{\sigma }_t^2}\epsilon \end{gathered}$$
因此，对应的SDE表示形式中
$$\begin{gathered} f(x_t,t)=0 \\ g(t)=\frac{d}{dt}{\sigma }_t^2 \end{gathered}$$

对于 DDPM，正向过程（加噪声）如下所示
$$\begin{gathered} x_t=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}\epsilon \\ {x}_{t+1}=\sqrt{1-{\beta }_{t+1}}{x}_t+\sqrt{{\beta }_{t+1}}\epsilon \end{gathered}$$
我们令 $\beta :[0,1]\to \mathbb{R}$ 代替 ${\beta }_t$，满足 $\beta (\frac{i}{T})=T{\beta }_i$，$\Delta t=\frac{1}{T}$，则有
$$\begin{array}{rl}{x}_{t+1}& =\sqrt{1-\beta (t+\Delta t)\Delta t}{x}_t+\sqrt{\beta (t+\Delta t)\Delta t}\epsilon \\ & \stackrel{\Delta t\to 0}{=}(1-\frac{1}{2}\beta (t)\Delta t)x_t+\sqrt{\beta (t)}\sqrt{\Delta t}\epsilon \end{array}$$
因此，对应的SDE表现形式中，有
$$\begin{gathered} f(x_t,t)=-\frac{1}{2}\beta (t)x_t \\ g(t)=\sqrt{\beta (t)} \end{gathered}$$

当我们希望 $t\to T$ 时，图像为纯粹的噪声图像，那么 ${\sigma }_t\to \infty$，但 ${\overline{\alpha }}_t\to 0$ 即可，因此分别称作是 VE-SDE 和 VP-SDE

DDPM Denoiser ${ϵ}_{\theta }(x_t,t)$ 与 NCSN Estimator $s_{\theta }(x_t,t)$ ：在 DDPM 正向过程中，有 $x_t\sim \mathcal{N}({\sqrt{\overline{\alpha }}}_t{x}_0,(1-{\overline{\alpha }}_t)I)$，代入 $s_{\theta }(x_t,t)={▽}_{x_t}\log p(x_t)$，可以得到
$$s_{\theta }(x_t,t)=-\frac{x_t-{\sqrt{\overline{\alpha }}}_t{x}_0}{1-{\overline{\alpha }}_t}=-\frac{1}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)$$

参考资料

Score-based Generative Modeling Through Stochastic Differential Equations

扩散模型 Diffusion Model SDE