写在前面
网上关于扩散模型的教程很多,但常见讲法要么从离散贝叶斯推导切入,要么从 score matching、能量模型、Langevin dynamics 切入,读起来往往枝节很多。
笔者当初也是从贝叶斯概率开始推起,过程可以说十分痛苦,推完了也是云里雾里。
后面在做diffusion强化学习的过程中,重新从高观点SDE的角度审视了扩散模型的发展历程,发现可以完全扔掉复杂的概率论公式,只专注于随机微分方程,用简洁的定理,把整个扩散模型的主线串起来,从一开始的DDPM到DDIM再到DPM-solver,甚至最近的Flow matching 和对应的flowgrpo里面的SDE采样,这些都可以使用一条主线全部串联起来。
因此有了这篇文章,希望可以帮助后来的朋友们有新的理解。本文的诸多观点都是来源于宋飏大佬,可惜宋飏大佬的许多工作数学比较难,读来颇为费力,导致许多朋友很难理解宋飏大佬从SDE观点理解diffusion的精彩之处,本文试着将其背后的物理直觉和强大的数学结合起来,读者将会看到,我们将如何绕过原本一页又一页的复杂推导与贝叶斯变换,凭借简单的思路,重新推导出所有主流扩散模型的公式,直击本质。
我敢说,无论你是第一次接触扩散模型,还是已经对它十分熟悉,看完本文后,相信你一定会产生一种醍醐灌顶的感觉: 原来这些核心概念是如此优雅地完美闭环的。
扩散模型SDE背后的物理直觉
背景入门
我们知道,生成式模型就是吃一个高斯噪声,然后恢复出来一个高清4K图片或者视频。
从数学的角度来看,这个过程其实就是从一个分布到另一个分布,就是一堆像素点本来分布在这里,突然跑到了另一块区域。
这个过程很像物理学的现象,就是一堆粒子本来在这个位置,然后飘移到另一个位置,这个过程可以用一个工具来描述,就是随机微分方程。
随机微分方程和扩散模型
随机微分方程的一般形式如下:
dzt=f(t,zt)dt+g(t)dWt
- f(t,zt) 是漂移系数(Drift)。
- g(t) 是扩散系数(Diffusion coefficient)
- 其中 dWt∼N(0,dt)
方程的第一项是漂移项,决定了粒子**“应该”去哪里**,物理直觉是水流的方向、风向、或者重力。如果除去所有随机干扰,粒子将完全沿着漂移项规定的轨迹运动(这就是 ODE)
方程的第二项是扩散项,物理直觉是随机的“热运动”或“能量,水分子的无规则撞击(布朗运动)等等。
扩散模型分成两个过程,前向加噪过程,把一张图片加噪变成高斯噪声,逆向解噪过程,从一个高斯噪声中生成一个图片。
对应的,在前向加噪过程中,漂移项的作用是最后把图片拉回均值为0,方差为1的高斯噪声。(严谨来说是对于VP-DIffusion或者VP-SDE来说,具体这两个概念后面会解释)。设想一下,如果没有漂移项,那么我每次加的高斯噪声都是期望为0,方差为1,加了无数次之后,最终的分布的期望是初始值,方差爆炸无穷大,所以这就是为什么我们需要漂移项往回拉。在逆向解噪过程中,漂移项的作用更重要,是为了把杂乱无章的高斯分布,强行推向看起来像猫、像狗等真实图片的数据流形(Manifold)上。
而在前向加噪过程中,扩散项的含义就很明显了,通过高斯扰动打乱图像信息,变成高斯噪声。在逆向解噪过程中,扩散项是为了保持多样性,就是在生成的过程中注入随机噪声,来增加多样性。不过现在很多的著名sampler(比如DDIM)都是ODE确定性采样,没有多样性了,因此在他们的逆向方程中扩散系数就是0。
以上这些,大概有个概念就可以了,下面我会细说。
线性随机微分方程
当 f(t,zt) 可以拆成 h(t)zt 的时候(即 z 的线性函数),我们叫线性SDE,如下:
dzt=h(t)zt dt+g(t)dwt
线性SDE有一些非常好的性质,可以直接用伊藤积分算出精确解,这个我们待会说。
根据伊藤积分我们可以得到两个推论(具体推导略):
dtdE[zt2]=2h(t)E[zt2]+g2(t)
dtdVar(zt)=2h(t)Var(zt)+g2(t)
这两个性质揭示了 zt 的一些分布特征,后面可能会用到。
SDE反演
我们还是从物理学的比喻讲起:
假如我们把一滴墨水滴到一杯清水里,墨水粒子受到布朗运动的随机撞击(扩散项),慢慢散开,最后变成一杯均匀的淡墨水。这个过程很容易用我们上一节说的 SDE(漂移项 + 扩散项)来描述。
现在,给你这杯均匀的淡墨水,让你写出一个方程,让这些墨水粒子“自动”聚拢回原本那一滴墨水的状态。
能不能做,能做。Anderson 在1982 年提出了一个非常牛逼的方程可以做这个事情,叫逆向时间 SDE 定理(Reverse-Time SDE Theorem)
如下:
dzt=[f(t,zt)−g(t)2∇zlogpt(zt)]dt+g(t)dwt
注意两个事情:
相比于前向,唯一多出来的项是:∇zlogpt(zt)
这个是什么东西?看起来是 zt 的概率密度分布函数取对数然后求导,这个东西叫 Score Function
怎么理解这个东西,我们还是举墨水扩散的例子,p(zt) 代表墨水粒子的密度分布情况,如果某个区域墨水粒子分布比较集中,那么这个区域的 p(zt) 大,而求导则是使其密度聚集的方向。
试想一下,扩散墨水想要恢复成聚拢回原本那一滴的状态,其实就是不断地集中的过程,那也意味着对于每个状态来说,我其实希望下一个状态是更聚集的,所以要往导数的方向去走,因为有这样一个 ∇zlogpt(zt) 项,同时注意到我之前扩散过程被前向的drift 项 f 牵引着,所以我现在要把这块数据的缩放牵引拉回去,因为前面再加了一个 drift 项 (注意dt是负的,因此实际上真实的方向是 g(t)2∇zlogpt(zt)−f(t,zt) )
如果用一句话来总结,逆向的过程就是,熵减收缩变得更聚集(score function)和补偿之前的拉伸(drift项)
问题是,这玩意是怎么推出来的?怎么用呢?
Fokker-Planck 方程
刚刚我们描述的 zt 本质上都是单个墨水粒子,我们可以用 FP 方程来描述这些粒子的整体概率云方案(绝大多数情况下的SDE都满足FP方程)。
∂t∂p(z,t)=−∇⋅(fp)+21g2Δp
这个方程是一个PDE(偏微分方程),我们希望可以通过求解这个PDE,来算出概率密度p(z) 函数(已知初始 条件的情况下)。
一般PDE的特点是,当维数比较少的时候,我们可以通过数值计算等方式求解出来。当然,对于也有些著名的PDE 方程,例如 Navier-Stokes (纳维-斯托克斯) 方程,即使是3维也很难解。
上面的 anderson 定理就是把 FP 方程结合贝叶斯公式推导出来的,同时如果我们想使用 anderson 定理得出逆向过程,就必须求解出 score function,而score function的求解可以通过上面的FP这个PDE方程的求解来得到。
在上个世纪的时候,假设你是某个军事基地的科学家,你现在得到了一个任务,雷达上我们观测到一架飞机运行,由于雷达精度、气流的影响,我们的雷达的监控轨迹中存在着大量的噪声,飞机的点震荡不稳定,现在我们想恢复出精确这个飞机的行驶路线,怎么办?
就是解这个FP方程,FP方程求解出来以后,我们拿着概率密度 p 就可以找到这个飞机最大概率的运行轨迹和路线了。
当然对于我们的扩散模型,z的维度可以说非常高,所以正常的解PDE是搞不定了,因此用 神经网络 拟合这个 score function 从而完成逆向。
Probability FLOW ODE
利用FP方程,我们还可以完成的一个事情是,我们可以使得上述的SDE过程等价于一个ODE过程:
dzt=[f(t,zt)−21g(t)2∇zlogpt(zt)]dt
注意这里的等价指的是宏观上 p(z) 在每个时刻的分布是等价一样的,联合分布是等价 的,但是微观轨迹是不同的。
即对于这部分的SDE和ODE,我们可以找到同样一条演化路线使得 p(zt)=p(xt)
这里说明设 zt 是SDE的演化轨迹,xt 是ODE的演化轨迹,仅在这里是这个含义。
但是并不代表 p(x1,x2,...) 联合分布等于 p(z1,z2,....) ,不过对于anderson逆向方程来说,前向过程的演化和逆向过程的演化是联合分布等价的。
然后你可以注意到,因为ODE的过程是前向和逆向等价的,所以你可以观察到,这个ODE对于前向SDE和逆向SDE 都是等价的。
比如 anderson 逆向公式:
dzt=[f(t,zt)−g(t)2∇zlogpt(zt)]dt+g(t)dwt
你把这里的dt前面的漂移项反过来(dt变成正的)然后减去 21g(t)2∇zlogpt(zt) 是不是还是这个ODE。
事实上,如果我们不考虑联合分布等价的话,就可以通过这个结论把 逆向的anderson 定理推出来。(把ODE反过来然后加上21g(t)2∇zlogpt(zt) 再反过去)。PF-ODE 结论 不依赖于 anderson 定理,只需要 FP 方程就可以推导出来。
Probability Flow ODE非常重要,他可以把一个SDE采样变成ODE采样(这就是DDIM等ODE sampler的来历),也可以把一个 ODE 采样改成 SDE 采样(flowgrpo等算法在flow matching 上做强化学习需要这么做)。
整体的背景知识已经介绍完毕了,接下来就是具体讲一下各大diffusion算法为例子。
Tweedie’s 公式
前面说过,逆向 SDE 需要用到 Score Function ∇ztlogpt(zt)。为了求这个东西,我们先要弄明白前向过程 p(zt∣x) 到底长什么样。
现在我们直接求解线性SDE dzt=f(t)ztdt+g(t)dwt
类似于常微分方程中的积分因子法,我们利用**伊藤引理(Itô’s Lemma)**对它进行积分,可以得到 zt 的闭式解:记x=z0
zt=λtexp(∫0tf(s)ds)⋅x+∫0tH(t,s)exp(∫stf(τ)dτ)g(s)dws
其中,x前面的系数记为 λt ,而第二项 ∫0texp(∫stf(τ)dτ)g(s)dws 这一坨叫维纳积分,可以证明服从均值为0的高斯分布,方差是:∫0texp(2∫stf(τ)dτ)g2(s)ds, 记为 σ2(t)
因此根据高斯分布的重参数化特性,我们可以把前向过程写为:
zt=λtx+σtη, 其中 η∼N(0,1)
这意味着,虽然加噪是一个连续的微分过程,但在任意宏观时刻 t,给定原图 x,我们都能直接写出从 x 跃迁到 zt 的精确转移概率(条件分布),且它必然是一个高斯分布:
pt(zt∣x)=N(zt∣λtx,σt2I)
对于不同的模型来说,其 λt 和 σt 是不同的,这就决定了数据在向噪声演化时的“几何轨迹”。大部分模型我们不会去定义SDE的系数,而是直接去定义重参数化后的系数。
比如对于 Flow Matching 来说,就是最简单的线性插值(Linear Interpolation),满足 λt+σt=1(即 λt=1−t,σt=t),数据是沿着直线走向噪声的;而对于 DDPM 来说,则是高维空间中的球面插值(Spherical Interpolation),满足 λt2+σt2=1(即 λt=αˉt,σt=1−αˉt)。
很多 flow matching 论文定义 t=0 时为噪声 , t = 1时为原始图片,本文延续diffusion的传统,t=0时定义为原始图片 ,t=1时为噪声,望读者注意。
由于原图 x 和纯噪声 η 的方差都近似为 1,根据独立随机变量方差相加的性质,任意时刻 zt 的方差 Var(zt)=λt2Var(x)+σt2Var(η)=λt2+σt2=1。这恰恰就是为什么 DDPM 对应的 SDE 会被称为 Variance Preserving (VP-SDE,保方差) 的原因。
这里的 zt 的分布为:
pt(zt∣x)=N(zt∣λtx,σt2I)
因此我们得到:
∇ztlogpt(zt∣x)=−σt2(zt−λtx)=−σtη
对数求导就是 对高斯函数的概率密度函数对数求导
然而,anderson定理下的score function 是 ∇ztlogpt(zt) 边缘分布,而上面的这个假定x已知,是条件分布,怎么办呢?
Tweedie’s 定理说:边缘分布的分数,等于条件分布的分数在后验分布下的期望。
∇ztlogpt(zt)=Ex∼p(x∣zt)[∇ztlogpt(zt∣x)]
因此 score function为:
∇ztlogpt(zt)=−σtE[η∣zt]≈−σtηθ(zt,t)
这里的 ηθ 可以理解为神经网络预测的我们加噪的噪声,DDPM训练过程中就是对整体加噪的期望进行估计。其他的形式:估计 x 或者 估计 v ,也都只差一个线性变换而已。So,我们直接绕开复杂的ELBO变分法的各种变换,从直观的角度理解了diffusion loss的设计原因,本质就是为了估计 score function。
各类diffusion变种
VP-SDE/DDPM
传统的DDPM论文里面是定义前向概率转移过程,然后贝叶斯公式先验分布和后验分布推上好几页把解噪公式推出来。这里我们抛弃掉概率论推导, 从纯粹的微分方程出发,使用微积分技巧,把DDPM公式推出来。
定义前向加噪过程为如下SDE形式:
dzt=−21β(t)ztdt+β(t)dwt
其中这里的 β(t) 表示的是噪声强度,漂移项系数和扩散项系数均是 β(t) 的函数。
注意这个 噪声强度 代表的是从一个时间步状态到另一个时间步状态转移的过程中,这一步加噪的强度,时间步越大,噪声强度越大。但是噪声强度不是时间步,千万不要弄混了,时间步是0-1000分布的,噪声强度不是,一般是0.0001到0.02之间均匀分布。
noise level 的调度(噪声调度)和 时间步调度 (time scheduler) 是两个完全不同的事情。我发现我的实习生经常弄混这个,这主要是因为不同论文里面符号不一致,如果不搞清楚本质的话,确实是很容易着迷。
一般我们说采样器是二次采样高阶采样,这里值得都是时间步调度算法,大部分的噪声强度调度算法,要么是常数,要么是线性调度。VP-SDE的调度算法就是下面线性调度,均匀分布在0.1在20之间:
β(t)=βminSDE+t⋅(βmaxSDE−βminSDE)β(t)=0.1+19.9⋅t(t∈[0,1])
下面的一个问题是,为什么要把drift项的系数和diffusion项的系数设计成这样一个关系?
从前面的伊藤积分推论我们知道:
dtdVar(zt)=2h(t)Var(zt)+g2(t)=−β(t)Var(zt)+β(t)=−β(t)(Var(zt)−1)=0
可以得出,整个扩散过程中,变量 zt 的方差是不变的,这就是 variance perseving VP-SDE的由来。
btw,我在面试的时候问过不少同学为什么diffusion模型的latent空间在训练之前要把方差变成1(事实上VAE并不能保证latent方差近似1,这也是为什么早期社区要乘以0.18215这种奇怪的系数),不过很少有人能答的很好,前面tweedie部分已经给了一个解答方式(根据高斯分布随机变量相加得出可以使得方差恒为1),这里再给出一个更高观点下的解答方式,即根据伊藤积分的推论,变量variance会一直被往1这里拉,当初始的variance为1的时候,整个过程的variance都将保持不变。
前向过程我们离散化后就是DDPM,即:
zi=1−βizi−1+βiη,η∼N(0,I)zt=αˉtx+1−αˉtη,η∼N(0,I)
离散化需要一点数学技巧,注意到:1−βi≈1−21βi 和 βi=β(t)Δt 即可
这里面规定符号: αt=1−βt 和 αˉt=∏αi , 其中 x 是原始图片 z0
直接从 anderson 定理逆向
dzt=[−21β(t)zt−β(t)∇zlogpt(zt)]dt+β(t)dWt
推导出DDPM的解噪公式:
zt−1≈zt−(−21βtzt−βt∇zlogpt(zt))+βtη=系数 A(1+21βt)zt+βtScore∇zlogpt(zt)+βtη=(1+21βt)zt−1−αˉtβtηθ(zt,t)+βtη=αt1(zt−1−αˉt1−αtηθ(zt,t))+βtη,η∼N(0,I)
推导过程只注意到两个事情,第一个是刚刚说的tweedie’s公式,第二个是泰勒展开。
∇logpt(zt)=−1−αˉtηθ1−βt1=(1−βt)−21≈1+(−21)(−βt)=1+21βt
所以从SDE高观点的角度来看,DDPM 本质上就是用 欧拉-丸山 (Euler-Maruyama) 方法去求解上面的逆向 SDE(离散化),但在参数化上做了一些特定的重整化(预测噪声使用tweedie公式变成score function)
注:DDPM的原文推导最后的 η 前面的系数是 β~t=1−αˉt1−αˉt−1βt ,但DDPM的原作者也发现直接用 βt 效果也还行,因为时间步划分的很细的时候,二者比例接近于1。
这其中的差别在于,我使用的是 EM (欧拉-丸山) 方法,它的近似粒度较粗,假定在极短时间内漂移项是完全冻结的常数;但其实我们可以把粒度做得更细:我们不做 EM 近似,只做局部线性近似。即假定神经网络预测的终点 x^0 是常数,此时逆向方程就变成了一个关于 zt 的线性 SDE(保留了系统随 zt 移动的连续收缩动态)。对于线性 SDE,我们就可以直接用伊藤积分硬解出 β~t=1−αˉt1−αˉt−1βt 这个精确方差。
逆向过程天生是非线性 SDE,我们无法求出全局解析解,所以要么做粗粒度近似(EM方法得 βt),要么做细粒度局部线性化(求精确解得 β~t)。
DDIM
DDIM 的思路是一样的,不过我们需要上的微积分技术会稍微更硬核一些。思路简单,计算技巧就多一些。传统的DDIM推导虽然计算技巧上简单,但是思路太复杂了。
VP-SDE 根据 PF-ODE的过程可以等价写为:
dz=[−21β(t)zt−21(β(t))2∇logp(z)]dt
代入 tweedie’s公式有:
dzt=[−21β(t)zt+21β(t)1−αˉtηθ(zt)]dt
使用简单的微积分技巧,我们可以观察到 αtˉ 和 β(t)的关系:
β(t)=−αˉt1dtdαˉt
已知简单的微积分除法公式:
dtd(αˉtzt)=αˉtdtdztαˉt−ztdtdαˉt
除法公式代入上面两个式子可以得到:
dtd(αˉtzt)=−2αtαtˉ1−αtˉ1dtdαtˉηθ=dtd(αtˉ1−αtˉ)ηθ
第二个等号是因为我们注意到:
dxd(x1−1)=−2xx−x21
所以我们拿到的ODE是:
d(αˉtzt)=d(αˉt1−αˉt)⋅ηθ(zt)
我们这里做一个欧拉法离散化,假定足够短的时间内,ηθ 是常数:
αˉszs−αˉtzt=ηθ(αˉs1−αˉs−αˉt1−αˉt)
移项,解得DDIM的解噪公式:
zs=αˉtαˉs(zt−1−αˉtη)+1−αˉsηzs=αˉs预测的 x0(αˉtzt−1−αˉtη)+1−αˉsη
所以,DDIM如果我们从高观点下理解,就是对VP-SDE转换成ODE使用欧拉法求解。
ODE有很多的高阶解法,因此针对ODE的数值分析的高阶解法就可以衍生出各种 Sampler 求解器,比如 DPM-solver,UNIPC等等。这背后使用的数学trick要更多,推导起来更复杂,读者有兴趣不妨自己读论文自行推导。
Flow Matching
flow matching 有意思的地方在于,前向加噪和逆向解噪过程我们不是把其建立为一个SDE,而是直接写成ODE:
dtdzt=v(t,zt)
定义重参数化系数形式为:
zt=(1−t)x0+tη
很多 flow matching 论文定义 t=0 时为噪声 , t = 1时为原始图片,本文延续diffusion的传统,t=0时定义为原始图片 ,t=1时为噪声,望读者注意。
根据PF-ODE的公式,该ODE可以等价为如下的SDE方程:
dzt=[v(t,zt)+21g(t)2∇zlogpt(zt)]dt+g(t)dwt
根据anderson定理,我们可以推出,逆向的SDE方程为:
dzt=[v(t,zt)−21g(t)2∇zlogpt(zt)]dt+g(t)dwˉt
分数怎么求解呢?注意到 Tweedie’s 公式有:
∇ztlogpt(zt)≈−σtE[η∣zt]=−tE[η∣zt]
然而 Flow matching 来说,预测的并不是 ηθ ,而是速度场 vθ,因此我们需要改成速度场的形式。
注意到 Flow matching 满足如下两个方程:
vθ≈E[η−x0∣zt]zt=(1−t)x0+tη
其中,vθ 一般我们在Flow matching 中就是使用神经网络近似:
vθ≈E[η−x0∣zt,t]
因此我们可以解出:
∇zlogpt(zt)≈−tzt+(1−t)vθ
代回去有:
dzt=[(1+2tg(t)2(1−t))vθ(t,zt)+2tg(t)2zt]dt+g(t)dwˉt
我们做一阶离散化,使用欧拉丸山解法,得到:
zt−Δt=zt−[(1+2tg(t)2(1−t))vθ(t,zt)+2tg(t)2zt]Δt+g(t)Δtξ
其中 ξ∼N(0,I)
很多论文里面会把扩散系数定义为 g(t)=2ϵ(t) (这就是为什么我一直不把噪声定义为 ϵ 的原因,因为很多人弄混了)
因此我们可以有:
zt−Δt=zt−[(1+tϵ(t)(1−t))vθ(t,zt)+tϵ(t)zt]Δt+2ϵ(t)Δtξ
一般我们训练的时候 ϵ(t) 设置成 0-1之间的一个常数,当然也有一些论文是t的函数,比如线性调度或非线性调度函数,需要再次强调的是,这是 noise scheduler(噪声调度),而非 time scheduler。
最后
大概就这些吧,像VE-SDE, NCSN,DPM-Solver 等有意思的工作都没有细讲,读者有兴趣可以下来按照这种思路自行推导了解。
严格来说抛开贝叶斯概率论是不可能的,我只是尽可能地把那些复杂繁琐的概率论推导包装成已有的一个个数学定理,我们直接从数学的角度出发,最快最直观地帮助我们了解复杂的扩散模型背后的原理。
水平有限,内容表述有诸多不严谨之处,还望指正。