【理论推导】随机微分方程(SDE)视角下的Diffusion Model与Score-based Model
SDE与DDPM
以 DDPM 为例，DDPM的通项公式为
x t ∼ N ( α ‾ t x 0 , ( 1 − α ‾ t ) I ) x_t \sim \mathcal N(\sqrt{\overline{\alpha}_t}x_0, (1-\overline{\alpha}_t) I) xt∼N(αtx0,(1−αt)I)
当我们固定 t t t 的取值时， x t x_t xt 是定义在样本空间上的函数，即为一随机变量，当我们固定 x x x的随机性时，即为关于变量 t t t 的一个函数，因此 x t x_t xt 是一随机过程，而对于一组确定的 { x t } t = 0 T \{x_t\}_{t=0}^T {xt}t=0T，称为随机过程的一个实现，或是一条轨迹/轨道，而随机过程可以使用随机微分方程(Stochastic Differential Equation)进行描述
SDE定义为
d X t = b ( t , X t ) d t + σ ( t , X t ) d B t dX_t = b(t,X_t)dt + \sigma(t,X_t)dB_t dXt=b(t,Xt)dt+σ(t,Xt)dBt
其中 B t B_t Bt 表示布朗运动，此方程的解称为伊藤(Itô)过程或是扩散过程。如果我们将 DDPM的迭代式从离散扩展到连续区间，即 x t → x t + Δ t ( Δ t → 0 ) x_t\rightarrow x_{t+\Delta t}\;\;(\Delta t\rightarrow 0) xt→xt+Δt(Δt→0)，即可得到SDE形式的扩散过程，论文中表示为
d x = f ( x , t ) d t + g ( t ) d w dx = f(x,t)dt + g(t)dw dx=f(x,t)dt+g(t)dw
其中 w w w 表示一个标准布朗运动， f ( ⋅ , t ) f(\cdot,t) f(⋅,t) 称作是漂移系数 (drift coefficient)，描述了确定性的变化过程， g ( ⋅ ) g(\cdot) g(⋅) 称作是扩散系数 (diffusion coefficient)，描述了不确定的变化过程
SDE视角下的生成模型
SDE的离散形式表示如下所示
x t + Δ t − x t = f ( x , t ) Δ t + g ( t ) Δ t ε x_{t+\Delta_t}-x_t = f(x,t)\Delta t + g(t)\sqrt{\Delta t} \varepsilon xt+Δt−xt=f(x,t)Δt+g(t)Δtε
其中 ε ∼ N ( 0 , I ) \varepsilon \sim \mathcal N(0,I) ε∼N(0,I)，因此有条件概率分布
x t + Δ t ∣ x t ∼ N ( x t + f ( x t , t ) Δ t , g 2 ( t ) Δ t I ) x_{t+\Delta_t}|x_t \sim \mathcal N(x_t+f(x_t,t)\Delta t, g^2(t)\Delta t I) xt+Δt∣xt∼N(xt+f(xt,t)Δt,g2(t)ΔtI)
考虑逆过程 x t ∣ x t + Δ t x_t|x_{t+\Delta t} xt∣xt+Δt，有
p ( x t ∣ x t + Δ t ) = p ( x t + Δ t ∣ x t ) p ( x t ) p ( x t + Δ t ) = p ( x t + Δ t ∣ x t ) exp ( log p ( x t ) − log p ( 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 ) } = 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 ∂ ∂ 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 Δ t 2 g 2 ( t ) } = Δ t → 0 exp { − 1 2 g 2 ( t + Δ 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 } \begin{align} p(x_t|x_{t+\Delta t}) &= \frac{p(x_{t+\Delta t}|x_t)p(x_t)}{p(x_{t+\Delta t})} \nonumber \\&= p(x_{t+\Delta t}|x_t)\exp(\log p(x_t)-\log p(x_{t+\Delta t})) \nonumber \\&\approx p(x_{t+\Delta t}|x_t)\exp \{ - (x_{t+\Delta t}-x_t)\triangledown_{x_t} \log p(x_t)-\Delta t \frac{\partial }{\partial t}\log p(x_{t})\} \nonumber \\&\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)\triangledown_{x_t}\log p(x_t)- \Delta t \frac{\partial }{\partial t}\log p(x_t)\} \nonumber \\&= \exp \left\{ -\frac{1}{2g^2(t)\Delta t}||(x_{t+\Delta t}-x_t)-(f(x_t,t)-g^2(t)\triangledown_{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)\triangledown_{x_t}\log p(x_t))^2\Delta t}{2g^2(t)} \right\} \nonumber \\&\overset{\Delta t\rightarrow 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)\triangledown_{x_{t+\Delta t}}\log p(x_{t+\Delta t}))\Delta t||_2^2 \right \} \nonumber \end{align} p(xt∣xt+Δt)=p(xt+Δt)p(xt+Δt∣xt)p(xt)=p(xt+Δt∣xt)exp(logp(xt)−logp(xt+Δt))≈p(xt+Δt∣xt)exp{−(xt+Δt−xt)▽xtlogp(xt)−Δt∂t∂logp(xt)}∝exp{−2g2(t)Δt∣∣xt+Δt−xt−f(xt,t)Δt∣∣22−(xt+Δt−xt)▽xtlogp(xt)−Δt∂t∂logp(xt)}=exp{−2g2(t)Δt1∣∣(xt+Δt−xt)−(f(xt,t)−g2(t)▽xtlogp(xt))Δt∣∣22−Δt∂t∂logp(xt)−2g2(t)f2(xt,t)Δt+2g2(t)(f(xt,t)−g2(t)▽xtlogp(xt))2Δt}=Δt→0exp{−2g2(t+Δt)Δt1∣∣(xt+Δt−xt)−(f(xt+Δt,t+Δt)−g2(t+Δt)▽xt+Δtlogp(xt+Δt))Δt∣∣22}
因此， x t ∣ x t + Δ t x_t|x_{t+\Delta t} xt∣xt+Δt 服从均值方差如下的高斯分布
μ = x t + Δ t − ( f ( x t + Δ t , t + Δ t ) − g 2 ( t + Δ t ) ▽ x t + Δ t log p ( x t + Δ t ) ) Δ t σ 2 = g 2 ( t + Δ t ) Δ t \mu = x_{t+\Delta t}-(f(x_{t+\Delta t},t+\Delta t)-g^2(t+\Delta t)\triangledown_{x_{t+\Delta t}}\log p(x_{t+\Delta t}))\Delta t \\\sigma^2 = g^2(t+\Delta t)\Delta t μ=xt+Δt−(f(xt+Δt,t+Δt)−g2(t+Δt)▽xt+Δtlogp(xt+Δt))Δtσ2=g2(t+Δt)Δt
因此，有SDE表示逆过程离散形式与连续形式如下所示
x t + Δ t − x t = ( f ( x t + Δ t , t + Δ t ) − g 2 ( t + Δ t ) ▽ x t + Δ t log p ( x t + Δ t ) ) Δ t + g ( t + Δ t ) Δ t ε d x = [ f ( x , t ) − g 2 ( t ) ▽ x t log p ( x t ) ] + g ( t ) d w \begin{align} x_{t+\Delta t}-x_t &= (f(x_{t+\Delta t},t+\Delta t)-g^2(t+\Delta t)\triangledown_{x_{t+\Delta t}}\log p(x_{t+\Delta t}))\Delta t + g(t+\Delta t)\sqrt{\Delta t}\varepsilon \\dx &= [f(x,t)-g^2(t)\triangledown_{x_t}\log p(x_t)]+g(t)dw \end{align} xt+Δt−xtdx=(f(xt+Δt,t+Δt)−g2(t+Δt)▽xt+Δtlogp(xt+Δt))Δt+g(t+Δt)Δtε=[f(x,t)−g2(t)▽xtlogp(xt)]+g(t)dw
考虑 NCSN 与 DDPM，两种形式均可以统一到SDE的理论表示形式下面，分别称为 VE-SDE (Variance Exploding) 与 VP-SDE (Variance Preserving)，对应 NCSN 与 DDPM
对于NCSN，正向过程（加噪声）如下所示
x t = x 0 + σ t ε x t + 1 = x t + σ t + 1 2 − σ t 2 ε x_t = x_0 + \sigma_t\varepsilon \\ x_{t+1} = x_t+\sqrt{\sigma_{t+1}^2 - \sigma_t^2}\varepsilon xt=x0+σtεxt+1=xt+σt+12−σt2ε
因此，对应的SDE表示形式中
f ( x t , t ) = 0 g ( t ) = d d t σ t 2 f(x_{t},t)=0 \\g(t)= \frac{d}{dt}\sigma_t^2 f(xt,t)=0g(t)=dtdσt2
对于 DDPM，正向过程（加噪声）如下所示
x t = α ‾ t x 0 + 1 − α ‾ t ε x t + 1 = 1 − β t + 1 x t + β t + 1 ε x_t = \sqrt{\overline{\alpha}_t}x_0 + \sqrt{1-\overline{\alpha}_t}\varepsilon \\ x_{t+1} = \sqrt{1-\beta_{t+1}}x_t + \sqrt{\beta_{t+1}}\varepsilon xt=αtx0+1−αtεxt+1=1−βt+1xt+βt+1ε
我们令 β : [ 0 , 1 ] → R \beta:[0,1]\rightarrow \R β:[0,1]→R 代替 β t \beta_{t} βt，满足 β ( i T ) = T β i \beta(\frac{i}{T}) = T\beta_i β(Ti)=Tβi， Δ t = 1 T \Delta t = \frac{1}{T} Δt=T1，则有
x t + 1 = 1 − β ( t + Δ t ) Δ t x t + β ( t + Δ t ) Δ t ε = Δ t → 0 ( 1 − 1 2 β ( t ) Δ t ) x t + β ( t ) Δ t ε \begin{align} x_{t+1} &= \sqrt{1-\beta(t+\Delta t)\Delta t}x_t +\sqrt{\beta(t+\Delta t)\Delta t}\varepsilon \\&\overset{\Delta t\rightarrow 0}{=}(1-\frac{1}{2}\beta(t)\Delta t)x_t+\sqrt{\beta(t)}\sqrt{\Delta t}\varepsilon \end{align} xt+1=1−β(t+Δt)Δtxt+β(t+Δt)Δtε=Δt→0(1−21β(t)Δt)xt+β(t)Δtε
因此，对应的SDE表现形式中，有
f ( x t , t ) = − 1 2 β ( t ) x t g ( t ) = β ( t ) f(x_{t},t)=-\frac{1}{2}\beta(t)x_t \\g(t)= \sqrt{\beta(t)} f(xt,t)=−21β(t)xtg(t)=β(t)
当我们希望 t → T t\rightarrow T t→T 时，图像为纯粹的噪声图像，那么 σ t → ∞ \sigma_t\rightarrow \infty σt→∞，但 α ‾ t → 0 \overline{\alpha}_t \rightarrow 0 αt→0 即可，因此分别称作是 VE-SDE 和 VP-SDE
DDPM Denoiser ϵ θ ( x t , t ) \epsilon_\theta(x_t,t) ϵθ(xt,t) 与 NCSN Estimator s θ ( x t , t ) s_\theta(x_t,t) sθ(xt,t) ：在 DDPM 正向过程中，有 x t ∼ N ( α ‾ t x 0 , ( 1 − α ‾ t ) I ) x_t \sim \mathcal N(\sqrt{\overline \alpha}_t x_0, (1-\overline \alpha_t)I) xt∼N(αtx0,(1−αt)I)，代入 s θ ( x t , t ) = ▽ x t log p ( x t ) s_\theta(x_t,t) = \triangledown_{x_t} \log p(x_t) sθ(xt,t)=▽xtlogp(xt)，可以得到
s θ ( x t , t ) = − x t − α ‾ t x 0 1 − α ‾ t = − 1 1 − α ‾ t ϵ θ ( x t , 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}} \epsilon_\theta(x_t,t) sθ(xt,t)=−1−αtxt−αtx0=−1−αt1ϵθ(xt,t)
参考资料
Score-based Generative Modeling Through Stochastic Differential Equations
扩散模型 Diffusion Model SDE