深度学习（生成式模型）——score-based generative modeling through stochastic differential equations

前言

yang song博士在《Score-Based Generative Modeling Through Stochastic Differential Equations》一文中提出可以使用SDE（随机微分方程）来刻画Diffusion model的前向过程，并且用SDE统一了Score-based Model (NCSN)和DDPM的前向过程与反向过程。此外，SDE对应了多个前向过程，即从一张图到某个噪声点的加噪方式有多种，但其中存在一个ODE（常微分方程）形式的前向过程，即不存在随机变量的确定性的前向过程。

本文将总结SDE与DDPM的关系，并给出相应推导

SDE是什么

SDE具体的数学形式如下：
$$dx=f(x,t)dt+g(t)dw\text{\hspace{1em}(1.0)}$$

$f(x,t)$表示自变量$x$随着时间$t$确定性的变化（又被称为drift coefficients），$g(t)$是一项与时间$t$相关的函数（又被称为diffusion coefficients），$dw$为布朗运动的增量，是一个随机项（可以理解为噪声）

SDE与DDPM前向过程的关系

我们将上述部分微分项展开并移位可得

$$\begin{array}{rl}{x}_{t+\Delta t}-x_t& =f(x,t)dt+g(t)dw\\ x_{t+\Delta t}& =x_t+f(x,t)dt+g(t)dw\end{array}\text{\hspace{1em}(2.0)}$$

我们将$x_t$看成是前向过程$t$时刻的图像，则下一时刻$t+\Delta t$的图像$x_{t+\Delta t}$可通过式2.0加噪得到。

接下来，我们将简单推导式2.0与DDPM前向过程的关系，已知DDPM的前向过程为
$$x_{t+\Delta t}=\sqrt{1-{\beta }_{t+\Delta t}}{x}_t+\sqrt{{\beta }_{t+\Delta t}}{ϵ}_t\text{\hspace{1em}(2.1)}$$

设${\overline{\beta }}_{t+\Delta t}=T{\beta }_{t+\Delta t}$，$\Delta t=\frac{1}{T}$，则式2.1为

$$\begin{array}{rl}{x}_{t+\Delta t}=& \sqrt{1-{\beta }_{t+\Delta t}}{x}_t+\sqrt{{\beta }_{t+\Delta t}}{ϵ}_t\\ =& \sqrt{1-\frac{{\overline{\beta }}_{t+\Delta t}}{T}}{x}_t+\sqrt{{\beta }_{t+\Delta t}}{ϵ}_t\\ =& \sqrt{1-{\overline{\beta }}_{t+\Delta t}\Delta t}{x}_t+\sqrt{{\beta }_{t+\Delta t}}{ϵ}_t\end{array}\text{\hspace{1em}(2.2)}$$

当$\Delta t$趋近于0，依据等价无穷小代换，式2.2有
$$\begin{array}{rl}{x}_{t+\Delta t}=& \sqrt{1-{\beta }_{t+\Delta t}}{x}_t+\sqrt{{\beta }_{t+\Delta t}}{ϵ}_t\\ =& \sqrt{1-{\overline{\beta }}_{t+\Delta t}\Delta t}{x}_t+\sqrt{{\beta }_{t+\Delta t}}{ϵ}_t\\ \approx & (1-\frac{1}{2}{\overline{\beta }}_{t+\Delta t}\Delta t)x_t+\sqrt{{\beta }_{t+\Delta t}}{ϵ}_t\\ =& x_t-\frac{1}{2}{\overline{\beta }}_{t+\Delta t}{x}_{t}dt+\sqrt{{\beta }_{t+\Delta t}}{ϵ}_t\end{array}\text{\hspace{1em}(2.3)}$$

比对式2.3与2.0，则有

$$\begin{array}{rl}f(x,t)& =-\frac{1}{2}{\overline{\beta }}_{t+\Delta t}{x}_t\\ g(t)& =\sqrt{{\beta }_{t+\Delta t}}\\ dw& ={ϵ}_t\end{array}$$

逆向过程的SDE

前文我们已经介绍了Diffusion model的前向过程可以用SDE描述，本节将推导出逆向过程的SDE形式。

令$dw=\sqrt{\Delta t}ϵ$，由式2.0，可得
$$p(x_{t+\Delta t}|x_t)=\mathcal{N}(x_{t+\Delta t};x_t+f(x_t,\Delta t)\Delta t,g^2(t)\Delta t)\text{\hspace{1em}(3.0)}$$

利用贝叶斯公式，则逆向过程为
q ( x t ∣ x t + Δ t ) = q ( x t + Δ t ∣ x t ) q ( x t ) q ( x t + Δ t ) = q ( x t + Δ t ∣ x t ) exp ⁡ { log ⁡ p ( x t ) − log ⁡ p ( x t + Δ t ) } (3.1) \begin{aligned} q(x_{t}|x_{t+\Delta t})&=\frac{q(x_{t+\Delta t}|x_{t})q(x_{t})}{q(x_{t+\Delta t})}\\ &=q(x_{t+\Delta t}|x_{t})\exp\{\log p(x_t)-\log p(x_{t+\Delta t})\}\tag{3.1} \end{aligned} q(xt​∣xt+Δt​)​=q(xt+Δt​)q(xt+Δt​∣xt​)q(xt​)​=q(xt+Δt​∣xt​)exp{logp(xt​)−logp(xt+Δt​)}​(3.1)

利用泰勒展开，则有

$$\log p(x_{t+\Delta t})\approx \log p(x_t)+(x_{t+\Delta t}-x_t){\nabla }_x\log p(x_t)\text{\hspace{1em}(3.2)}$$

代入式3.1，并且结合式3.0，则有
q ( x t ∣ x t + Δ t ) = q ( x t + Δ t ∣ x t ) exp ⁡ { − ( x t + Δ t − x t ) ∇ x log ⁡ p ( x t ) } ≈ exp ⁡ { − ( x t + Δ t − x t − f ( x t , t ) Δ t ) 2 + 2 g 2 ( t ) Δ t ( x t + Δ t − x t ) ∇ x log ⁡ p ( x t ) 2 g 2 ( t ) Δ t } (3.3) \begin{aligned} q(x_{t}|x_{t+\Delta t})&=q(x_{t+\Delta t}|x_{t})\exp \{-(x_{t+\Delta t}-x_t)\nabla_{x}\log p(x_t)\}\\ &\approx\exp\{-\frac{(x_{t+\Delta t}-x_t-f(x_t,t)\Delta t)^2+2g^2(t)\Delta t (x_{t+\Delta t}-x_t)\nabla_{x}\log p(x_t)}{2g^2(t)\Delta t }\}\tag{3.3} \end{aligned} q(xt​∣xt+Δt​)​=q(xt+Δt​∣xt​)exp{−(xt+Δt​−xt​)∇x​logp(xt​)}≈exp{−2g2(t)Δt(xt+Δt​−xt​−f(xt​,t)Δt)2+2g2(t)Δt(xt+Δt​−xt​)∇x​logp(xt​)​}​(3.3)

为了后续书写方便，令
$$\begin{array}{rl}a& =f(x_t,t)\Delta t\\ b& =g^2(t)\Delta t\end{array}$$

则有

$$\begin{array}{rl}& (x_{t+\Delta t}-x_t-f(x_t,t)\Delta t{)}^2+2g^2(t)\Delta t(x_{t+\Delta t}-x_t){\nabla }_x\log p(x_t)\\ & =（x_{t+\Delta t}-x_t-a{）}^2+2b(x_{t+\Delta t}-x_t){\nabla }_x\log p(x_t)\\ & =(x_{t+\Delta t}-x_t{)}^2-2a(x_{t+\Delta t}-x_t)+a^2+2b(x_{t+\Delta t}-x_t){\nabla }_x\log p(x_t)\\ & =(x_{t+\Delta t}-x_t{)}^2-2(a-b{\nabla }_x\log p(x_t))(x_{t+\Delta t}-x_t)+(a-b{)}^2+a^2-(a-b{)}^2\\ & =(x_{t+\Delta t}-x_t-(a-b{\nabla }_x\log p(x_t)){)}^2+a^2-(a-b{\nabla }_x\log p(x_t){)}^2\end{array}$$

当$\Delta t$趋近0时，则有

$$\begin{array}{rl}\frac{a^2}{2b}& =\frac{f(x_t,t{)}^2\Delta t}{2g^2(t)}\to 0\\ \frac{(a-b{\nabla }_x\log p(x_t){)}^2}{2b}& =\frac{(f(x_t,t)-g^2(t){\nabla }_x\log p(x_t))\Delta t}{g^2(t)}\to 0\end{array}$$

则当$\Delta t$趋近0，式3.3为

q ( x t ∣ x t + Δ t ) = q ( x t + Δ t ∣ x t ) exp ⁡ { − ( x t + Δ t − x t ) ∇ x log ⁡ p ( x t ) } ≈ exp ⁡ { − ( x t + Δ t − x t − f ( x t , t ) Δ t ) 2 + 2 g 2 ( t ) Δ t ( x t + Δ t − x t ) ∇ x log ⁡ p ( x t ) 2 g 2 ( t ) Δ t } = exp ⁡ { − ( x t + Δ t − x t − ( f ( x t , t ) Δ t − g 2 ( t ) Δ t ∇ x log ⁡ p ( x t ) ) ) 2 2 g 2 ( t ) Δ t } (3.4) \begin{aligned} q(x_{t}|x_{t+\Delta t})&=q(x_{t+\Delta t}|x_{t})\exp \{-(x_{t+\Delta t}-x_t)\nabla_{x}\log p(x_t)\}\\ &\approx\exp\{-\frac{(x_{t+\Delta t}-x_t-f(x_t,t)\Delta t)^2+2g^2(t)\Delta t (x_{t+\Delta t}-x_t)\nabla_{x}\log p(x_t)}{2g^2(t)\Delta t }\}\\ &=\exp\{-\frac{(x_{t+\Delta t}-x_t-(f(x_t,t)\Delta t-g^2(t)\Delta t\nabla_{x}\log p(x_t)))^2}{2g^2(t)\Delta t}\}\tag{3.4} \end{aligned} q(xt​∣xt+Δt​)​=q(xt+Δt​∣xt​)exp{−(xt+Δt​−xt​)∇x​logp(xt​)}≈exp{−2g2(t)Δt(xt+Δt​−xt​−f(xt​,t)Δt)2+2g2(t)Δt(xt+Δt​−xt​)∇x​logp(xt​)​}=exp{−2g2(t)Δt(xt+Δt​−xt​−(f(xt​,t)Δt−g2(t)Δt∇x​logp(xt​)))2​}​(3.4)

则有

$$q(x_t|x_{t+\Delta t})=\mathcal{N}(x_t|x_{t+\Delta t}-f(x_t,t)\Delta t+g^2(t)\Delta t{\nabla }_x\log p(x_t),2g^2(t)\Delta t)\text{\hspace{1em}(3.5)}$$

设噪声$z$服从标准正态分布，则式3.5写成SDE的形式为

$$\begin{array}{rl}{x}_t& =x_{t+\Delta t}-f(x_t,t)\Delta t+g^2(t)\Delta t{\nabla }_x\log p(x_t)+g(t)\sqrt{2\Delta t}z\\ dx& =(f(x_t,t)-g^2(t){\nabla }_x\log p(x_t))dt-g(t)\sqrt{2\Delta t}z\\ & =(f(x_t,t)-g^2(t){\nabla }_x\log p(x_t))dt+g(t)d\overline{w}\end{array}$$

${\nabla }_{x_t}\log p(x_t)$与DDPM预测的噪声$ϵ$的关系

score base model一般会用神经网络拟合${\nabla }_{x_t}\log p(x_t)$，DDPM其实是一种特殊的score base model，已知DDPM的前向过程为

$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t\text{\hspace{1em}(4.0)}$$

依据Tweedie方法，我们有

$$\begin{array}{r}\sqrt{{\bar{\alpha }}_t}{x}_0=x_t+(1-{\bar{\alpha }}_t){\nabla }_x\log p(x_t)\end{array}$$
进而有
$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0-(1-{\bar{\alpha }}_t){\nabla }_x\log p(x_t)\text{\hspace{1em}(4.2)}$$

结合式4.0与4.2，则有

$${\nabla }_{x_t}\log p(x_t)=-\frac{1}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t\text{\hspace{1em}(4.3)}$$

逆向过程SDE与DDPM逆向过程的关系

在进行正式的推导前，我们先对式3.1做个简单的变化，利用泰勒展开，则有
$$\log p(x_t)\approx \log p(x_{t+\Delta t})+(x_t-x_{t+\Delta t}){\nabla }_x\log p(x_{t+\Delta t})\text{\hspace{1em}(5.0)}$$

代入式3.1，并结合式3.0，则有
q ( x t ∣ x t + Δ t ) = = q ( x t + Δ t ∣ x t ) exp ⁡ { log ⁡ p ( x t ) − log ⁡ p ( x t + Δ t ) = q ( x t + Δ t ∣ x t ) exp ⁡ { − ( x t + Δ t − x t ) ∇ x log ⁡ p ( x t + Δ t ) } ≈ exp ⁡ { − ( x t + Δ t − x t − f ( x t , t ) Δ t ) 2 + 2 g 2 ( t ) Δ t ( x t + Δ t − x t ) ∇ x log ⁡ p ( x t + Δ t ) 2 g 2 ( t ) Δ t } (5.1) \begin{aligned} q(x_{t}|x_{t+\Delta t})&==q(x_{t+\Delta t}|x_{t})\exp\{\log p(x_t)-\log p(x_{t+\Delta t})\\ &=q(x_{t+\Delta t}|x_{t})\exp \{-(x_{t+\Delta t}-x_t)\nabla_{x}\log p(x_{t+\Delta t})\}\\ &\approx\exp\{-\frac{(x_{t+\Delta t}-x_t-f(x_t,t)\Delta t)^2+2g^2(t)\Delta t (x_{t+\Delta t}-x_t)\nabla_{x}\log p(x_{t+\Delta t})}{2g^2(t)\Delta t }\}\tag{5.1} \end{aligned} q(xt​∣xt+Δt​)​==q(xt+Δt​∣xt​)exp{logp(xt​)−logp(xt+Δt​)=q(xt+Δt​∣xt​)exp{−(xt+Δt​−xt​)∇x​logp(xt+Δt​)}≈exp{−2g2(t)Δt(xt+Δt​−xt​−f(xt​,t)Δt)2+2g2(t)Δt(xt+Δt​−xt​)∇x​logp(xt+Δt​)​}​(5.1)

因此我们可将式3.5重写为
$$q(x_t|x_{t+\Delta t})=\mathcal{N}(x_t|x_{t+\Delta t}-f(x_t,t)\Delta t+g^2(t)\Delta t{\nabla }_x\log p(x_{t+\Delta t}),2g^2(t)\Delta t)\text{\hspace{1em}(5.2)}$$

已知用SDE表示DDPM的前向过程时，有

$$\begin{array}{rl}f(x,t)& =-\frac{1}{2}{\overline{\beta }}_{t+\Delta t}{x}_t\\ g(t)& =\sqrt{{\beta }_{t+\Delta t}}\end{array}$$

其中
$${\overline{\beta }}_{t+\Delta t}=T{\beta }_{t+\Delta t}=\frac{{\beta }_{t+\Delta t}}{\Delta t}$$

$T=\frac{1}{\Delta t}$，设$z$服从标准正态分布，代入式5.2并结合式4.1，当$\Delta t$趋近于0时，有

$$\begin{array}{rl}{x}_t& =x_{t+\Delta t}-f(x_t,t)\Delta t+g^2(t)\Delta t{\nabla }_x\log p(x_{t+\Delta t})+g(t)\sqrt{2\Delta t}z\\ x_t& =x_{t+\Delta t}+\frac{1}{2}{\overline{\beta }}_{t+\Delta t}{x}_t\Delta t+{\beta }_{t+\Delta t}\Delta t{\nabla }_x\log p(x_{t+\Delta t})+\sqrt{2{\beta }_{t+\Delta t}\Delta t}z\\ (1-\frac{1}{2}{\overline{\beta }}_{t+\Delta t}\Delta t)x_t& =x_{t+\Delta t}+{\beta }_{t+\Delta t}\Delta t{\nabla }_x\log p(x_{t+\Delta t})+\sqrt{2{\beta }_{t+\Delta t}\Delta t}z\\ \sqrt{1-{\beta }_{t+\Delta t}}{x}_t& \approx x_{t+\Delta t}+{\beta }_{t+\Delta t}\Delta t{\nabla }_x\log p(x_{t+\Delta t})+\sqrt{2{\beta }_{t+\Delta t}\Delta t}z\\ x_t& \approx \frac{1}{\sqrt{1-{\beta }_{t+\Delta t}}}(x_{t+\Delta t}+{\beta }_{t+\Delta t}\Delta t{\nabla }_x\log p(x_{t+\Delta t}))+\sqrt{\frac{2{\beta }_{t+\Delta t}\Delta t}{1-{\beta }_{t+\Delta t}}}z\\ x_t& \approx \frac{1}{\sqrt{1-{\beta }_{t+\Delta t}}}(x_{t+\Delta t}+{\beta }_{t+\Delta t}\Delta t{\nabla }_x\log p(x_{t+\Delta t}))+\sqrt{\frac{2{\beta }_{t+\Delta t}\Delta t}{1-{\beta }_{t+\Delta t}}}z\\ x_t& \approx \frac{1}{\sqrt{1-{\beta }_{t+\Delta t}}}(x_{t+\Delta t}-\frac{{\beta }_{t+\Delta t}\Delta t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{t+\Delta t}))+\sqrt{\frac{2{\beta }_{t+\Delta t}\Delta t}{1-{\beta }_{t+\Delta t}}}z\end{array}$$

当$\Delta t=1$时，则有

$$x_t\approx \frac{1}{\sqrt{1-{\beta }_{t+1}}}(x_{t+1}-\frac{{\beta }_{t+1}}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{t+1}))+\sqrt{\frac{2{\beta }_{t+1}}{1-{\beta }_{t+1}}}z$$

不过这种约等于号是真的很膈应，是不能做完全等价的。

Probability Flow (PF) ODE

有多种SDE，可以将一张图像变为某个噪声点，其中也包括一个ODE（即去除掉SDE中布朗运动增量）

对于前向过程

$$dx=f(x,t)dt+g(t)dw$$

由Fokker-Planck方程可得

$$\begin{array}{rl}\frac{\partial p(x,t)}{\partial t}& =-\sum _i\frac{\partial }{\partial x_i}[f_i(x,t)p(x,t)]+\frac{1}{2}\sum _{i,j}\frac{{\partial }^2}{\partial x_i{x}_j}\left\{[g^2(t)I{]}_{ij}p(x,t)\right\}\\ & =-\sum _i\frac{\partial }{\partial x_i}[f_i(x,t)p(x,t)]+\frac{1}{2}\sum _i\frac{{\partial }^2}{\partial x_i^2}[g^2(t)p(x,t)]\end{array}$$

对上述式子做个等价变换，则有
$$\begin{array}{rl}\frac{\partial p(x,t)}{\partial t}& =-\sum _i\frac{\partial }{\partial x_i}[f_i(x,t)p(x,t)]+\frac{1}{2}\sum _i\frac{{\partial }^2}{\partial x_i^2}[g^2(t)p(x,t)]\\ & =-\sum _i\frac{\partial }{\partial x_i}[f_i(x,t)p(x,t)]+\frac{1}{2}\sum _i\frac{{\partial }^2}{\partial x_i^2}[(g^2(t)-{\sigma }^2(t))p(x,t)]+\frac{1}{2}\sum _i\frac{{\partial }^2}{\partial x_i^2}[{\sigma }^2(t)p(x,t)]\\ & =-\sum _i\frac{\partial }{\partial x_i}[f_i(x,t)p(x,t)]+\frac{1}{2}\sum _i\frac{\partial }{\partial x_i}(g^2(t)-{\sigma }^2(t))\frac{\partial }{\partial x_i}p(x,t)+\frac{1}{2}\sum _i\frac{{\partial }^2}{\partial x_i^2}[{\sigma }^2(t)p(x,t)]\\ & =-\sum _i\frac{\partial }{\partial x_i}[f_i(x,t)p(x,t)]+\frac{1}{2}\sum _i\frac{\partial }{\partial x_i}(g^2(t)-{\sigma }^2(t))p(x,t)\frac{\partial }{\partial x_i}\log p(x,t)+\frac{1}{2}\sum _i\frac{{\partial }^2}{\partial x_i^2}[{\sigma }^2(t)p(x,t)]\\ & =-\sum _i\frac{\partial }{\partial x_i}\left[\left(f_i(x,t)-\frac{1}{2}(g^2(t)-{\sigma }^2(t))\frac{\partial }{\partial x_i}\log p(x,t)\right)p(x,t)\right]+\frac{1}{2}\sum _i\frac{{\partial }^2}{\partial x_i^2}[{\sigma }^2(t)p(x,t)]\end{array}$$

利用Fokker-Planck方程的对应关系，则有
$$dx=\left(f(x,t)-\frac{1}{2}(g^2(t)-{\sigma }^2(t)){\nabla }_x\log p(x,t)\right)dt+\sigma (t)dw$$

特别的，当$\delta (t)=0$时，则有
$$dx=\left(f(x,t)-\frac{1}{2}{g}^2(t){\nabla }_x\log p(x,t)\right)dt$$
上述式子又被称为Probability Flow (PF) ODE