『扩散模型』一篇文章入门随机微分方程SDE

随机微分方程SDE

笔者建议，学完DDPM再来看SDE的作用和推导过程

标准布朗运动

在学习随机微分方程之前，我们先来看一下什么是标准布朗运动
假设有一个一维的直线，有个小人从原点出发，每次随机地选择向左走1格或者向右走1格，且向左走和向右走的两个选项，被选择的概率相等$\to$用$S_t$代表小人离原点的距离，$t$代表代表选择的次数，如果选择的次数越多，那么$S_t$将会逐渐服从一个均值为0、方差为$t$的正态分布
布朗运动$W(t)$是期望为0、方差为$t$的正态分布$\iff$$W_t\sim \mathcal{N}(0,t)$$\Rightarrow$ $W_{t+\Delta t}-W_t\sim \mathcal{N}(0,\Delta t)$，当$\Delta t\to 0$时，$dw=\sqrt{dt}\epsilon$（重参数技巧）

SDE加噪

在DDPM中，扩散过程被划分为固定的T步$\Rightarrow$DDPM=拆楼+建楼$\Rightarrow$“拆楼”和“建楼”都被事先划分为了T步，这个划分有着相当大的人为性。事实上，真实的“拆”、“建”过程应该是没有刻意划分的步骤$\Rightarrow$可以将它们理解为一个在时间上连续的变换过程，可以用随机微分方程（Stochastic Differential Equation，SDE）来描述，即$d\mathbf{x}={\mathbf{f}}_t(\mathbf{x})dt+g_{t}d{\mathbf{w}}_{\mathbf{t}}$，其中$f_t(x_t)$是漂移项，描述数据的确定性演化；$g_t$是扩散项，描述的是噪声的扩散程度；$d{w}_t$是维纳运动（布朗运动）的微小增量，表示随机波动
随机微分方程：$dx=$确定的变化$+$随机的变化，其中随机的变化代表着随机性
随机微分方程描述了系统从$t$时刻到$t+\Delta t$时刻的变化
我们可以将随机微分方程看成是${\mathbf{x}}_{t+\Delta t}-{\mathbf{x}}_t={\mathbf{f}}_t({\mathbf{x}}_t)\Delta t+g_t\sqrt{\Delta t}\mathbf{\epsilon },\mathbf{\epsilon }\sim \mathcal{N}(0,\mathbf{I})$在$\Delta t\to 0$时的极限$\Rightarrow$如果建楼要1天，那么拆楼就是$x$从$t=0$到$t=1$时刻的变化
越小的步数$\Delta t$意味着对原始噪声越好的近似，如果$\Delta t=0.001$，对应着$T=1000$；如果$\Delta t=0.01$，则对应$T=100$（总时间步数$T$是模拟的总时间跨度被步长$\Delta t$划分的次数$T=\frac{t_{\max }}{\Delta t}$）$\Rightarrow$引入SDE的本质好处是“将理论分析和代码实现分离开来”
DDPM的加噪过程本质上是一个SDE，而SDE本质上描述的是微小时间变化下系统状态的变化

• DDPM的加噪：$x_{t+1}=\sqrt{1-{\beta }_t}{x}_t+\sqrt{{\beta }_t}ϵ$
• SDE的加噪：$d\mathbf{x}={\mathbf{f}}_t(\mathbf{x})dt+g_{t}d{\mathbf{w}}_{\mathbf{t}}$

在这里，笔者介绍一下将DDPM加噪公式映射到SDE加噪公式的推导过程：

1. 重写DDPM加噪公式：$x_{t+1}-x_t=(\sqrt{1-{\beta }_t}-1)x_t+\sqrt{{\beta }_t}ϵ$$\Rightarrow$$\sqrt{1-{\beta }_t}\approx 1-\frac{{\beta }_t}{2}$$\Rightarrow$将DDPM加噪公式重新表示为一个确定项和随机噪声项的和：$x_{t+1}-x_t\approx -\frac{{\beta }_t}{2}{x}_t+\sqrt{{\beta }_t}ϵ$
   在这里，使用泰勒展开得到$\sqrt{1-{\beta }_t}\approx 1-\frac{{\beta }_t}{2}$
   先来介绍一下泰勒展开：如果$f(x)$在$x=a$处是可微的，则它的泰勒展开可以写为$f(x)\approx f(a)+f^{\prime }(a)(x-a)+\frac{f^{\prime \prime }(a)}{2!}(x-a{)}^2+\dots$，其中$f^{\prime }(a)$和$f^{\prime \prime }(a)$分别是$f(x)$在$a$处的一阶导数和二阶导数；在泰勒展开中，若函数依赖多个变量，需要对每个变量分别进行展开
   $f({\beta }_t)=\sqrt{1-{\beta }_t}$在${\beta }_t=0$处展开$\Rightarrow$零阶项：$f(0)=\sqrt{1-0}=1$；一阶导数：$f^{\prime }({\beta }_t)=\frac{d}{d{\beta }_t}\sqrt{1-{\beta }_t}=\frac{-1}{2\sqrt{1-{\beta }_t}}$，在${\beta }_t=0$处$f^{\prime }(0)=\frac{-1}{2\sqrt{1-0}}=-\frac{1}{2}$$\Rightarrow$$f({\beta }_t)\approx f(0)+f^{\prime }(0){\beta }_t$$\Rightarrow$$\sqrt{1-{\beta }_t}\approx 1-\frac{1}{2}{\beta }_t$
2. 引入$\Delta t$：$\Delta t$只是在数学上引入的时间增量，而${\beta }_t$在离散模型中的定义是独立于$\Delta t$的，将DDPM离散的加噪过程转换为连续时间的随机微分方程描述：$x_{t+\Delta t}-x_t\approx -\frac{{\beta }_t}{2}{x}_t\Delta t+\sqrt{{\beta }_t\Delta t}ϵ$$\Rightarrow$$dx=-\frac{1}{2}{\beta }_t{x}_{t}dt+\sqrt{{\beta }_t}dw$
3. SDE的形式：漂移项$f_t(x_t)=-\frac{{\beta }_t}{2}{x}_t$，扩散系数$g_t=\sqrt{{\beta }_t}$
   左侧是数据分布，右侧是正态分布，t是连续时间

SDE去噪

SDE去噪的目标是求$p(x_t|x_{t+\Delta t})$
已知：$x_{t+\Delta t}$和前向SDE过程$p(x_{t+\Delta t}|x_t)$$\Rightarrow$贝叶斯公式：$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|x_{t+\Delta t})$的分布满足正态分布

1. $x_{t+\Delta t}=x_t+f_t(x_t)\Delta t+g_t\sqrt{\Delta t}ϵ$$\Rightarrow$根据重参数可得：$x_{t+\Delta t}\sim \mathcal{N}(x_t+f_t(x_t)\Delta t,g_t^2\Delta t)$
2. 正态分布的概率密度函数：$f(x)=\frac{1}{\sigma \sqrt{2\pi }}exp(-\frac{(x-\mu {)}^2}{2{\sigma }^2})$
3. $\begin{array}{rlrl}p(x_t|x_{t+\Delta t})& =exp(-\frac{(x_{t+\Delta t}-x_t-f_t(x_t)\Delta t{)}^2}{2g_t^2\Delta t}+logp(x_t)-logp(x_{t+\Delta t}))& & (1)\end{array}$
4. 在$x_t$处泰勒展开$logp(x_{t+\Delta t})$：$\log p(x_{t+\Delta t},t+\Delta t)\approx \log p(x_t,t)+\underset{\text{状杰变化的影响}}{\underbrace{(x_{t+\Delta t}-x_t){\nabla }_x\log p(x_t,t)}}+\underset{\text{时间推移的影响}}{\underbrace{\Delta t\cdot {\nabla }_t\log p(x_t,t)}}$，其中状态梯度${\nabla }_x\log p(x_t)$描述概率密度在空间中的“漂移”趋势（例如粒子倾向于从高密度区向低密度区移动），时间梯度${\nabla }_t\log p(x_t)$描述概率密度随时间的整体衰减或增长
   在这里笔者介绍一下为什么会多出一项${\nabla }_t\log p(x_t)$：$p(x_t)$实际上是“$t$时刻随机变量等于$x_t$的概率密度”，$p(x_{t+\Delta t})$实际上是“$t+\Delta t$时刻随机变量等于$x_{t+\Delta t}$的概率密度”，即$p(x_t)$实际上同时是时间$t$和状态变量$x_t$的函数
   时间梯度项${\nabla }_t\log p(x_t)$的必要性：${\nabla }_t\log p(x_t)$描述了概率密度的动态演化，即使状态$x_t$不变，概率密度$p(x_t)$也会随时间$t$变化
5. 当$\Delta t\to 0$时，${\Delta }^{2}t=0$
6. $\begin{array}{rlrl}(1)=p(x_t|x_{t+\Delta t})& =exp(-\frac{(x_{t+\Delta t}-x_t-(f_t(x_t)-g_t^2\nabla xlogp(x_t))\Delta t{)}^2}{2g_t^2\Delta t})& & \\ & =exp(-\frac{(x_t-(x_{t+\Delta t}-(f_t(x_t)-g_t^2\nabla xlogp(x_t))\Delta t){)}^2}{2g_t^2\Delta t})& & (2)\end{array}$
   将$x_{t+\Delta t}-x_t$和$(f_t(x_t)-g_t^2\nabla xlogp(x_t))\Delta t$分别看作一个整体，$\frac{(x_{t+\Delta t}-x_t-(f_t(x_t)-g_t^2\nabla xlogp(x_t))\Delta t{)}^2}{2g_t^2\Delta t}$开方后可以得到和$\frac{(x_{t+\Delta t}-x_t-f_t(x_t)\Delta t{)}^2}{2g_t^2\Delta t}+logp(x_t)-logp(x_{t+\Delta t})$相同的结果
7. 当$\Delta t\to 0$时，$t+\Delta t\to t$
8. $\begin{array}{r}(2)=exp(-\frac{(x_t-(x_{t+\Delta t}-(f_{t+\Delta t}(x_{t+\Delta t})-g_{t+\Delta t}^2\nabla xlogp(x_{t+\Delta t}))\Delta t){)}^2}{2g_{t+\Delta t}^2\Delta t})\end{array}$
9. 均值$\mu =x_{t+\Delta t}-(f_{t+\Delta t}(x_{t+\Delta t})-g_{t+\Delta t}^2\nabla xlogp(x_{t+\Delta t}))\Delta t$、方差$g_{t+\Delta t}^2\Delta t$
10. $x_t=x_{t+\Delta t}-(f_{t+\Delta t}(x_{t+\Delta t})-g_{t+\Delta t}^2\nabla xlogp(x_{t+\Delta t}))\Delta t+g_{t+\Delta t}\sqrt{\Delta t}ϵ$
11. 当$\Delta t\to 0$时，$\Delta t\to dt$
12. $x_{t+\Delta t}-x_t=(f_{t+\Delta t}(x_{t+\Delta t})-g_{t+\Delta t}^2\nabla xlogp(x_{t+\Delta t}))\Delta t+g_{t+\Delta t}\sqrt{\Delta t}ϵ$
13. $\begin{array}{r}d{x}_t=(f_t(x_t)-g_t^2\nabla xlogp(x_t))dt+g_t\sqrt{dt}ϵ=(f_t(x_t)-g_t^2\nabla xlogp(x_t))dt+g_{t}d\bar{w}\end{array}$

Score Matching

Score Matching 是一种用于估计概率密度函数$p(x)$ 的方法，核心方法是直接估计概率密度的梯度（称为score）而非密度本身
通过最小化模型估计的score与真实score之间的差异，可以简介学习数据的分布
我们先来看一下${\nabla }_{x_t}\log p({\mathbf{x}}_t)$的含义：对数概率密度函数$\log p(x_t)$关于$x_t$的梯度方向指向概率密度增加最快的方向、梯度大小反映密度变化的速率$\Rightarrow$使用一个$\theta$参数化的概率分布$p_{\theta }$模拟$p$，通过学习参数$\theta$使$p_{\theta }$接近$p$
我们可以将$p_{\theta }$看成是由两部分组成的，分别是表示密度的函数$\widetilde{p_{\theta }}$、归一化因子$Z_{\theta }$$\Rightarrow$$p_{\theta }(x)=\frac{{\tilde{p}}_{\theta }(x)}{Z_{\theta }}=\frac{{\tilde{p}}_{\theta }(x)}{{\int }_{x\in X}{\tilde{p}}_{\theta }(x)dx}$，其中未归一化的概率密度函数$\widetilde{p_{\theta }}$给出某个数据点$x$相对于其他数据点的可能性大小，但并不能给出直接用于表示$x$发生的真实概率
目前，使用极大似然估计求解$\theta$的问题：不知道归一化因子$Z_{\theta }$的值
解决方法：

1. 引入得分函数（score function）：概率密度函数的梯度${\nabla }_x\log p_{\theta }(x)$
2. 将$p_{\theta }(x)$通过$log$拆分成两项${\nabla }_x\log {\tilde{p}}_{\theta }(x)-{\nabla }_x\log Z_{\theta }$$\Rightarrow$由于求解的是$x$的梯度，所以可以直接消掉${\nabla }_x\log Z_{\theta }$，因为${\nabla }_x\log Z_{\theta }$与$x$无关；同时$\widetilde{p_{\theta }}$不受“概率分布”的约束，可以使用神经网络作为$\widetilde{p_{\theta }}$，因为$\widetilde{p_{\theta }}$本身就不是概率密度函数，$\widetilde{p_{\theta }}$只是密度函数
3. 目标：选择一个loss让${\nabla }_x\log {\tilde{p}}_{\theta }(x)$尽可能接近${\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t)$

新的问题：不知道数据分布的 score function${\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t)$
为了简化公式，下面公式中的${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$等同于${\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t)$
解决方法：Score Matching

1. Score Matching：用于估计概率密度函数的梯度（得分函数$\mathrm{score}(x)={\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t)$），而无需知道密度函数的归一化常数
2. Score Matching的目标：学习一个模型$q(x;\theta )$，使得模型得分函数${\nabla }_x\log q(x;\theta )$与真实分布$p(x)$的得分函数尽可能接近
3. Score Matching的损失函数：$L(\theta )={\mathbb{E}}_{x\sim p(x)}\left[\frac{1}{2}{\Vert {\nabla }_x\log q(x;\theta )-{\nabla }_x\log p(x)\Vert }^2\right]$，其中的期望差异可以帮助模型更全面地学习到真实分布的特征

接下来，对Score Matching的损失函数$L(\theta )={\mathbb{E}}_{x\sim p(x)}\left[\frac{1}{2}{\Vert {\nabla }_x\log q(x;\theta )-{\nabla }_x\log p(x)\Vert }^2\right]$进行推导：

1. 展开欧几里得范数的平方项：$\Vert {\nabla }_x\log q(x;\theta )-{\nabla }_x\log p(x){\Vert }^2=\Vert {\nabla }_x\log q(x;\theta ){\Vert }^2-2{\nabla }_x\log q(x;\theta )\cdot {\nabla }_x\log p(x)+\Vert {\nabla }_x\log p(x){\Vert }^2$
2. 将上式代入原始损失函数中可得$L(\theta )={\mathbb{E}}_{x\sim p(x)}\left[\frac{1}{2}\left(\Vert {\nabla }_x\log q(x;\theta ){\Vert }^2-2{\nabla }_x\log q(x;\theta )\cdot {\nabla }_x\log p(x)+\Vert {\nabla }_x\log p(x){\Vert }^2\right)\right]$
3. 消除不可计算的项：由于不知道真实分布的${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$，我们无法直接计算$\Vert {\nabla }_x\log p(x){\Vert }^2$和${\nabla }_x\log q(x;\theta )\cdot {\nabla }_x\log p(x)$

接下来，笔者给出如何消除不可计算项的过程：

1. 由于$\Vert {\nabla }_x\log p(x){\Vert }^2$和$\theta$无关，它仅仅依赖于真实数据分布$p(x)$，所以可以直接消掉$\Vert {\nabla }_x\log p(x){\Vert }^2$
2. 对损失函数中的项${\nabla }_x\log q(x;\theta )\cdot {\nabla }_x\log p(x)$进行分部积分$\Rightarrow$$\int p(x){\nabla }_{\mathbf{x}}\log q(\mathbf{x};\theta )\cdot {\nabla }_{\mathbf{x}}\log p(\mathbf{x})d\mathbf{x}=-\int p(\mathbf{x}){\nabla }_x^2\log q(\mathbf{x};\theta )d\mathbf{x}$

$\int p(x){\nabla }_{\mathbf{x}}\log q(\mathbf{x};\theta )\cdot {\nabla }_{\mathbf{x}}\log p(\mathbf{x})d\mathbf{x}=-\int p(\mathbf{x}){\nabla }_x^2\log q(\mathbf{x};\theta )d\mathbf{x}$的推导过程：

1. $\int p(\mathbf{x}){\nabla }_x\log q(\mathbf{x};\theta )\cdot \frac{{\nabla }_{\mathbf{x}}p(\mathbf{x})}{p(\mathbf{x})}d\mathbf{x}=\int {\nabla }_x\log q(\mathbf{x};\theta )\cdot {\nabla }_{x}p(\mathbf{x})d\mathbf{x}$，其中${\nabla }_x\log p(\mathbf{x})=\frac{{\nabla }_{x}p(\mathbf{x})}{p(\mathbf{x})}$
2. 分部积分：$\int udv=uv{|}_a^b-\int vdu$、$dv={\nabla }_{x}p(x)dx$，$u={\nabla }_x\log q(\mathbf{x};\theta )$$\Rightarrow$$du={\nabla }_x^{2}logq(x;\theta )dx、$$v=p(x)$
3. $\int {\nabla }_x\log q(\mathbf{x};\theta )\cdot {\nabla }_{x}p(\mathbf{x})d\mathbf{x}={\nabla }_x\log q(\mathbf{x};\theta )\cdot p(x){|}_{x=-\infty }^{x=\infty }-\int p(\mathbf{x}){\nabla }_x^2\log q(\mathbf{x};\theta )d\mathbf{x}$
   
   梯度是一个向量，表示在函数在某一点处变化最快的方向和速率；散度是一个标量，表示向量场在某一点处的“扩散”程度。散度为正，表示向量场从该点扩散；散度为负，表示向量场向该点聚集
4. 当x趋于无穷大时，概率密度通常趋于0，,所以$\int {\nabla }_x\log q(\mathbf{x};\theta )\cdot {\nabla }_{x}p(\mathbf{x})d\mathbf{x}=-\int p(\mathbf{x}){\nabla }_x^2\log q(\mathbf{x};\theta )d\mathbf{x}$
5. $\begin{array}{rl}L(\theta )& ={\mathbb{E}}_{x\sim p(x)}\left[\frac{1}{2}{\Vert {\nabla }_x\log q(x;\theta )-{\nabla }_x\log p(x)\Vert }^2\right]\\ & ={\mathbb{E}}_{\mathbf{x}\sim p(\mathbf{x})}\left[\frac{1}{2}\Vert {\nabla }_{\mathbf{x}}\log q(\mathbf{x};\theta ){\Vert }^2+{\nabla }_{\mathbf{x}}^2\log q(\mathbf{x};\theta )\right]\end{array}$

至此，我们可以通过损失函数$L(\theta )$使${\nabla }_x\log q(x;\theta )$接近${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$，进而求出SDE去噪过程中的$dx$
笔者也是刚刚接触SDE，如果文中出现错误，请各位读者指正

参考文献

1、生成扩散模型漫谈（五）：一般框架之SDE篇
2、SDE公式推导
3、SDE的底层原理
4、AIGC: SGM (Score-based Generative Model) 笔记
5、Score Matching（得分匹配）
6、从零开始的扩散模型 | 基于分数的生成模型解释