SDE in diffusion models

SDE in diffusion models

参考：https://www.bilibili.com/video/BV19M411z7hS/

论文：Score-Based Generative Modeling through Stochastic Differential Equations

本文被认为是 diffusion models 方向中最重要的一篇的论文。作者通过将之前的两类 diffusion models（DDPM、NCSN）通过 SDE 在理论上统一在了一起。注意虽然统一了两种形式，但是本质上 SDE 还是通过估计 score。

为什么要用 SDE 来描述扩散模型？

对于扩散模型，我们的研究对象就是噪声图 $x_t$ 。对于 $x_t$ ，实际就是一个关于 $x$ 和 $t$ 的函数：

• 当 $t$ 固定时，$x_t$ 可看做是一个随机变量，$x_t\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,\sqrt{1-{\bar{\alpha }}_t}I)$
• 当 $x$ 固定时，我们就得到一条采样轨迹：$x_T,x_{T-1},\dots ,x_0$

基于上述观察，可以发现 $x_t$ 实际上就是一个随机过程，而描述随机过程的数学工具，自然就是 SDE。

连续的扩散过程

之前的工作，无论是 DDPM 还是 NCSN（score-based model），都是将 diffusion 过程看作是一个离散的过程：

离散 diffusion 过程：$x_0,x_1,...x_t,...,x_T$，reverse 过程：$x_T,x_{T-1},...x_t,...,x_0$ 。

实际当然我们也可以将它考虑为一个连续的过程：

连续 diffusion 过程：$t\in [0,1]$ ，扩散过程是 $x_t\to x_{t+\Delta t},\Delta t\to 0$ ，reverse 过程则是 $x_{t+\Delta t}\to x_t$

在 SDE 的框架内，我们使用更一般的连续形式来研究扩散过程。

SDE 框架描述扩散模型的公式

论文中给出了扩散过程的 SDE 公式：
$$dx=f(x,t)dt+g(t)dw$$
式中，$x$ 就是我们的研究对象噪声图，$f(x,t)$ 称为漂移系数（drift coefficient），$g(t)$ 称为扩散系数（diffusion coefficient），$w$ 是布朗运动（ brownian motion）。

SDE 公式共有两项，其中前一项 $f(x,t)dt$ 是一个确定性的变化过程，后一项 $g(t)dw$ 则是一个不确定性的过程。

下面，我们基于 SDE 形式的扩散公式，来推到其重建公式。

我们先把 $dt$ 写成离散形式 $\Delta t$ ，有：
$$\begin{gathered} x_{t+\Delta t}-x_t=f(x_t,t)\Delta t+g(t)\sqrt{\Delta t}ϵ,ϵ\sim \mathcal{N}(0,I) \\ {x}_{t+\Delta t}=x_t+f(x_t,t)\Delta t+g(t)\sqrt{\Delta t}ϵ \end{gathered}$$
注意这里 $dw$ 的系数从 $g(t)$ 变成 $\sqrt{\Delta t}$ ，从一阶的两变成了一个 $1/2$ 阶的量，这是一个结论性的东西，先记住就行了。

然后写出 $P(x_{t+\Delta t}|x_t)$ ：
$$P(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}$ ，而我们推导重建公式的目标是由 $x_{t+\Delta t}$ 推导 $x_t$ ，即 $P(x_t|x_{t+\Delta t})$ ，有：
$$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)$ 已经写出来了，现在要做的就是处理 $P(x_t)$ 和 $P(x_{t+\Delta t})$ ，我们将他们构造成 log 的形式：
P ( x t ∣ x t + Δ t ) = P ( x t + Δ t ∣ x t ) exp ⁡ { log ⁡ P ( x t ) − log ⁡ P ( x t + Δ t ) } P(x_t|x_{t+\Delta t})=P(x_{t+\Delta t}|x_t)\exp \{\log P(x_t)-\log P(x_{t+\Delta t})\} P(xt​∣xt+Δt​)=P(xt+Δt​∣xt​)exp{logP(xt​)−logP(xt+Δt​)}
对式中的 $\log P(x_{t+\Delta t})$ ，对其进行一阶泰勒展开的近似：
$$\log (x_{t+\Delta t})\approx \log P(x_t)+(x_{t+\Delta t}-x_t){\nabla }_{x_t}\log P(x_t)+\Delta t\frac{\partial }{\partial t}\log P(x_t)$$
代回去，有：
P ( x t ∣ x t + Δ t ) = P ( x t + Δ t ∣ x t ) exp ⁡ { − ( x t + Δ t − x t ) ∇ x t log ⁡ P ( x t ) − Δ t ∂ ∂ t log ⁡ P ( x t ) } ∝ exp ⁡ { ∣ ∣ x t + Δ t − x t − f ( x t , t ) Δ t ∣ ∣ 2 2 2 g 2 ( t ) Δ t − ( x t + Δ t − x t ) ∇ x t log ⁡ p ( x t ) − Δ t ∂ ∂ t log ⁡ P ( x t ) } \begin{align} P(x_t|x_{t+\Delta t})&=P(x_{t+\Delta t}|x_t)\exp \{-(x_{t+\Delta t}-x_t)\nabla_{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)\nabla_{x_t}\log p(x_t)-\Delta t\frac{\partial}{\partial t}\log P(x_t)\} \end{align} P(xt​∣xt+Δt​)​=P(xt+Δt​∣xt​)exp{−(xt+Δt​−xt​)∇xt​​logP(xt​)−Δt∂t∂​logP(xt​)}∝exp{2g2(t)Δt∣∣xt+Δt​−xt​−f(xt​,t)Δt∣∣22​​−(xt+Δt​−xt​)∇xt​​logp(xt​)−Δt∂t∂​logP(xt​)}​​
接下来，我们对 $x_{t+\Delta t}$ 进行配方：
P ( x t ∣ x t + Δ t ) ∝ exp ⁡ { − 1 2 g 2 ( t ) Δ t [ ( x t + Δ t − x t ) 2 − ( 2 f ( x t , t ) Δ t − 2 g 2 ( t ) Δ t ∇ x t log ⁡ p ( x t ) ) ( x t + Δ t − x t ) ] } − Δ t ∂ ∂ t log ⁡ P ( x t ) − f 2 ( x t , t ) Δ t 2 g 2 ( t ) = exp ⁡ { − 1 2 g 2 ( t ) Δ t ∣ ∣ ( x t + Δ t ) − x t − ( f ( x t , t ) − g 2 ( t ) ∇ x t log ⁡ P ( x t ) ) Δ t ∣ ∣ 2 2 − Δ ∂ ∂ t log ⁡ P ( x t ) − f 2 ( x t , t ) Δ t 2 g 2 ( t ) + f ( x t , t ) − g 2 ( t ) ∇ x t log ⁡ P ( x t ) 2 g 2 ( t ) Δ t } \begin{align} P(x_t|x_{t+\Delta t})&\propto\exp\{-\frac{1}{2g^2(t)\Delta t}[(x_{t+\Delta t}-x_t)^2-(2f(x_t,t)\Delta t-2g^2(t)\Delta t\nabla_{x_t}\log p(x_t))(x_{t+\Delta t}-x_t)]\}-\Delta t\frac{\partial}{\partial t}\log P(x_t)-\frac{f^2(x_t,t)\Delta t}{2g^2(t)} \\ &=\exp \{-\frac{1}{2g^2(t)\Delta t}||(x_{t+\Delta t)}-x_t-(f(x_t,t)-g^2(t)\nabla_{x_t}\log P(x_t))\Delta t||_2^2-\Delta \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)\nabla_{x_t}\log P(x_t)}{2g^2(t)\Delta t}\} \end{align} P(xt​∣xt+Δt​)​∝exp{−2g2(t)Δt1​[(xt+Δt​−xt​)2−(2f(xt​,t)Δt−2g2(t)Δt∇xt​​logp(xt​))(xt+Δt​−xt​)]}−Δt∂t∂​logP(xt​)−2g2(t)f2(xt​,t)Δt​=exp{−2g2(t)Δt1​∣∣(xt+Δt)​−xt​−(f(xt​,t)−g2(t)∇xt​​logP(xt​))Δt∣∣22​−Δ∂t∂​logP(xt​)−2g2(t)f2(xt​,t)Δt​+2g2(t)Δtf(xt​,t)−g2(t)∇xt​​logP(xt​)​}​​

接下来做一下近似，我们的 $\Delta t\to 0$，可以将 $\Delta t$ 的一次项丢掉，而 $t=t+\Delta t$ ，则：
P ( x t ∣ x t + Δ t ) ∝ exp ⁡ { − 1 2 g 2 ( t + Δ t ) ∣ ∣ ( x t + Δ t − x t ) − f ( x t + Δ t , t + Δ t ) − g 2 ( t + Δ t ) ∇ x t + Δ t log ⁡ P ( x t + Δ t ) Δ t ∣ ∣ 2 2 } P(x_t|x_{t+\Delta t})\propto\exp \{-\frac{1}{2g^2(t+\Delta t)}||(x_{t+\Delta t}-x_t)-f(x_{t+\Delta t},t+\Delta t)-g^2(t+\Delta t)\nabla_{x_{t+\Delta t}}\log P(x_{t+\Delta t})\Delta t||^2_2\} P(xt​∣xt+Δt​)∝exp{−2g2(t+Δt)1​∣∣(xt+Δt​−xt​)−f(xt+Δt​,t+Δt)−g2(t+Δt)∇xt+Δt​​logP(xt+Δt​)Δt∣∣22​}
从而我们就可以写出 $P(x_t|x_{t+\Delta t})$ 的均值：
$$x_{t+\Delta t}-(f(x_{t+\Delta t},t+\Delta t)-g^2(t+\Delta t){\nabla }_{x_{t+\Delta t}}\log P(x_{t+\Delta t}))\Delta t$$
方差：
$$g^2(t+\Delta t)\Delta t$$
从而：
$$dx=[f(x,t)-g^2(t){\nabla }_{x_t}\log P(x_t)]dt+g(t)dw$$
即得到了连续的扩散模型采样过程的 SDE 公式。

当然，代码实现的时候我们还是按照离散的形式来实现，写出其离散形式为：
$$x_{t+\Delta t}-x_t=[f(x_{t+\Delta t},t+\Delta t)-g^2(t+\Delta t){\nabla }_{x_{t+\Delta t}}\log P(x_{t+\Delta t})]\Delta t+g(t+\Delta t)\sqrt{\Delta t}ϵ$$
结果中，只有 score 的部分 ${\nabla }_{x_{t+\Delta t}}\log P(x_{t+\Delta t})$ 是我们不知道的，需要模型来预测，别的部分都是已知的。

如果写成 DDPM 那样的离散形式，即 $\Delta t=1$，有：
$$x_{t-1}=x_t-[f(x_t,t)-g^2(t){\nabla }_{x_t}\log P(x_t)]\times 1+g(t)ϵ$$
至此，我们已经有了 SDE 框架下扩散模型的扩散公式和采样公式：

扩散公式：
$$dx=f(x,t)dt+g(t)dw$$
采样公式：
$$dx=[f(x,t)-g^2(t){\nabla }_x\log P(x)]dt+g(t)dw$$

SDE框架下的NCSN(VE)和DDPM(VP)

我们最开始的时候提到，SDE 这篇工作统一了扩散模型中的 DDPM 和 NCSN 两类模型。本节将具体介绍 SDE 框架是如何统一这两类模型的。

在文章中，NCSN 和 DDPM 分别对应于 VE-SDE 和 VP-SDE，其意义分别为 Variance Exploding 和 Variance Preseving。这里我们先写出两类模型的扩散公式：

方法	NCSN	DDPM
SDE	VE	VP
一步到位扩散公式	$x_t=x_0+{\sigma }_tϵ$	$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$
单步扩散公式	$x_{t+1}=\sqrt{{\sigma }_{t+1}^2-{\sigma }_t^2}ϵ$	$x_{t+1}=\sqrt{1-{\beta }_{t+1}}{x}_t+\sqrt{{\beta }_{t+1}}ϵ$

这两个名字是从何得来呢？我们知道，在扩散模型中，加到最大的噪声强度时，噪声图 $x_T$ 需要几乎完全是一个高斯分布。

• 在 NCSN 中，$x_T=x_0+{\sigma }_Tϵ$ 要是一个完全的噪声，其中方差 ${\sigma }_T$ 就要非常大，故称为 Variance Exploding，噪声爆炸；

• 在 DDPM 中，$x_T=\sqrt{{\bar{\alpha }}_T}{x}_0+\sqrt{1-{\bar{\alpha }}_T}ϵ$ 要是一个完全的噪声，其中 ${\bar{\alpha }}_T$ 就要非常小，所以噪声 $\sqrt{1-{\bar{\alpha }}_t}$ 最大也只有 1，故称为 Variance Preserving，方差收紧。

接下来我们就正式介绍如何使用 SDE 框架统一这两种扩散模型。我们先再写出 SDE 的扩散公式：
$$\begin{gathered} dx=f(x,t)dt+g(t)dw \\ {x}_{t+\Delta t}=x_t+f(x_t,t)\Delta t+g(t)\sqrt{\Delta t}ϵ \end{gathered}$$
我们先看 DDPM（VE）的情况，有：
$$\begin{array}{rl}{x}_{t+\Delta t}& =x_t+\sqrt{{\sigma }_{t+\Delta t}^2-{\sigma }_t^2}ϵ\\ & =x_t+\sqrt{\frac{{\sigma }_{t+\Delta t}^2-{\sigma }_t^2}{\Delta t}}\sqrt{\Delta t}ϵ\\ & =x_t+\sqrt{\frac{\Delta {\sigma }_t^2}{\Delta t}}\sqrt{\Delta t}ϵ\end{array}$$
与上式比对，则有：
$$\begin{gathered} f(x,t)=0 \\ g(t)=\frac{d}{dt}{\sigma }_t^2 \end{gathered}$$

VE 的情况还是比较简单的。

再来看 DDPM（VP） 的情况，其单步扩散公式为：
$$x_{t+1}=\sqrt{1-{\beta }_{t+1}}{x}_t+\sqrt{{\beta }_t}ϵ$$
其中 { β i } i = 1 T \{\beta_i\}_{i=1}^T {βi​}i=1T​ 是我们人为定义的一组超参数。为了将离散函数连续化，现在我们再构造 { β ˉ i = T β i } i = 1 T \{\bar\beta_i=T\beta_i\}_{i=1}^T {βˉ​i​=Tβi​}i=1T​ 。当 $T\to \infty$ 时，$\beta$ 就趋近于一个函数：
{ β ˉ i } i = 1 T → β ( t ) ,    t ∈ [ 0 , 1 ] \{\bar\beta_i\}_{i=1}^T\rightarrow \beta(t),\ \ t\in [0,1] {βˉ​i​}i=1T​→β(t),  t∈[0,1]
则有：$\beta (\frac{i}{T})={\bar{\beta }}_i$

再看扩散公式：
$$x_{t+1}=\sqrt{1-\frac{{\bar{\beta }}_{t+1}}{T}}{x}_t+\sqrt{\frac{{\bar{\beta }}_{t+1}}{T}}ϵ$$
当 $\Delta t=\frac{1}{T}$ 时：
$$x_{t+\Delta t}=\sqrt{1-\beta (t+\Delta t)}{x}_t+\sqrt{\beta (t+\Delta t)\Delta t}ϵ$$
这里做一下近似。我们知道，当 $x\to 0$ 时，有 $(1-x{)}^a\approx 1-ax$，则有：
$$x_{t+\Delta t}\approx (1-\frac{1}{2}\beta (t+\Delta t)\Delta t)x_t+\sqrt{\beta (t+\Delta t)}\sqrt{\Delta t}ϵ$$
又由于 $\Delta t\to 0$，我们再做一下近似，把不想要的 $\Delta t$ 拿掉，有：
$$x_{t+\Delta t}\approx (1-\frac{1}{2}\beta (t)\Delta t)x_t+\sqrt{\beta (t)}\sqrt{\Delta t}ϵ$$
再对比 SDE 的扩散公式，有：
$$\begin{gathered} f(x,t)=-\frac{1}{2}\beta (t)x_t \\ g(t)=\sqrt{\beta (t)} \end{gathered}$$
到这里，我们已经把 NCSN 和 DDPM 在 SDE 框架中的形式推导出来了。通过指定不同的 $f(x,t)$ 和 $g(t)$ ，扩散模型的 SDE 框架就分别变成了 NCSN 和 DDPM 两类模型。

关于NCSN中的score estimator和DDPM中的denoiser之间关系的分析

在 NCSN 中，我们使用一个模型 $s_{\theta }(x_t,t)$ 来预测模型的分数，而在 DDPM 中，则是使用一个模型 ${ϵ}_{\theta }(x_t,t)$ 来估计模型的噪声。既然 SDE 框架下，NCSN 能与 DDPM 统一起来，那么我们实际训练的这两个模型，是否也存在一定的关系呢？

$s_{\theta }(x_t,t)$ 是对分数的估计，即：
$$s_{\theta }(x_t,t)\approx {\nabla }_{x_t}\log P(x_t)$$
而我们知道，$x_t$ 是服从一个正态分布的，即：
$$x_t\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)$$

则有：
P ( x t ) ∝ exp ⁡ { − ∣ ∣ x t − α ˉ t x 0 ∣ ∣ 2 2 2 ( 1 − α ˉ t ) } P(x_t)\propto\exp \{-\frac{||x_t-\sqrt{\bar\alpha_t}x_0||_2^2}{2(1-\bar\alpha_t)}\} P(xt​)∝exp{−2(1−αˉt​)∣∣xt​−αˉt​ ​x0​∣∣22​​}
然后我们可以直接把 score 求出来：
$$s_{\theta }(x_t,t)\approx {\nabla }_{x_t}\log P(x_t)=-\frac{x_t-\sqrt{{\bar{\alpha }}_t}{x}_0}{1-{\bar{\alpha }}_t}$$
我们还知道噪声预测器：
$${ϵ}_{\theta }(x_t,t)=\frac{x_t-\sqrt{{\bar{\alpha }}_t}{x}_0}{\sqrt{1-{\bar{\alpha }}_t}}$$
比较以上两式，有：
$$s_{\theta }(x_t,t)=-\frac{1}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)$$
可以看到，NCSN 中的分数估计模型和 DDPM 中的噪声预测模型实际就只差了一个系数而已。如果训好了一个 NSCN 的分数估计模型，甚至可以直接乘个系数直接拿到 DDPM 中用。当然，在介绍 score-based model 的时候我们也提到了，估计分数，实际上就是估计噪声，在这里也再一次得到了验证。

关于VE和VP等价性的简单证明

本节对 VP 和 VE 的等价性做简单的证明。

我们从 VE 开始：
$$x_t=x_0+{\sigma }_tϵ$$
同除 $\sqrt{1+{\sigma }_t^2}$，有：
$$\frac{x_t}{\sqrt{1+{\sigma }_t^2}}=\frac{x_0}{\sqrt{1+{\sigma }_t^2}}+\frac{{\sigma }_tϵ}{\sqrt{1+{\sigma }_t^2}}$$
这里我们记
$$\begin{gathered} x_t^{\prime }=\frac{x_t}{\sqrt{1+{\sigma }_t^2}} \\ {\bar{\alpha }}_t=\frac{1}{\sqrt{1+{\sigma }_t^2}} \end{gathered}$$
从而
$$\sqrt{1-{\bar{\alpha }}_t}=\frac{{\sigma }_t}{\sqrt{1+{\sigma }_t^2}}$$
则有：
$$x_t^{\prime }=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$$
到这里发现就已经推到了 VP 的形式，从而就完成了 VE 和 VP 等价性的简单证明。

如果在 VE 上做了些理论的分析，那是可以直接照搬到 VP 上去的，因为它们是等价的。