Diffusion学习笔记
Published:
扩散模型学习过程中一些文章的阅读笔记以及模型的相关内容的简单整理
一般框架SDE
扩散模型的前向过程是由随机微分方程(SDE)描述的线性扩散
dx=Ftxdt+Gtdω其中,M、N为输入数据的维度,Ft∈RM×N、Gt∈RM×N分别为漂移系数和扩散系数,ω为标准的为维纳过程
目前对SDE的研究较为成熟,前向过程(1)对应的反向过程可由式(2)所示的SDE函数族表示:
dx=[Ftx−1+λ22GtG⊤t∇logpt(x)]dt+λGtdω其中,λ≥0,当λ=1时,即为式(3)对应的SDE反向过程
dx=[Ftxdt−GtG⊤t∇logpt(x)]+Gtdω而当λ=0时,反向过程中的方差为0,SDE退化为概率流ODE。
在实际应用中,通过对网络训练得到∇logpt(x)的近似值sθ(xt,t),借助离散化的方式对式(2)实现数值求解,完成扩散模型的反向过程
拟合误差
概率流ODE
概率流ODE为确定性的微分方程,
dx=[Ftx−12GtG⊤t∇logpt(x)]dtEuler解法
ˆxt−Δt=ˆxt−[Ftxt+12GtG⊤tL−⊤tϵθ(xt,t)]Δt指数积分(EI)解法
xt−Δt=e∫t−ΔttFτdτxt+∫t−Δtt12e∫t−ΔtτFrdrGτG⊤τL−⊤tϵθ(xτ,τ)dτDDIM
在DDIM中,作者提出,p(x1:t)分解为马尔可夫过程不是必须的,只要保证p(xt|x0)和p(xt−1|xt,x0)与DDPM是相同的就可以得到DDPM的等价模型。文中,作者将p(xt−1|xt,x0)的分布定义为(7):
p(xt−1|xt,x0)∼N(xt−1;√ˉαt−1 x0+√ˉβt−1−σ2t⋅xt−√ˉαt x0√1−ˉαt,σ2tI)其中,
σt=η√ˉβt−1βtˉβtη为可调节的参数,使用网络估计得到的
$\boldsymbol{\epsilon}{\theta}\left(\boldsymbol{x}{t}, t\right)$
替换
x0
得到
xt−1=1√αt(xt−(√ˉβt−√αt√ˉβt−1−σ2t)ϵθ(xt,t))+σtϵDPM-Solver
半线性公式
DPM-Solver提出,ODE方程是一个半线性的方程,线性项ftxt是可以准确的计算的,之前的采样算法忽略了这一点,从而导致对ODE方程的数值解法会产生较大的拟合误差,算法的加速性能不好。因此,在DPM-Solver中,作者将线性项与非线性项分开,对非线性项采用数值解法,从而减少了拟合误差。
引入参数$\lambda_{t} = log({\bar{\alpha}{t}} / {\sigma{t}})将g^2(t)表示为\lambda_t$的函数:
g2(t)=ˉα2tddt(σ2tˉα2t)=−2σ2dλtdt对式(6)进行参数替换,同时代入VP-SDE前向过程对应的参数,得到:
xt=ˉαtˉαsxs−ˉαt∫ts(dλτdτ)στˉατϵθ(xτ,τ)dτ考虑到λt为前向过程中信噪比的一半,是严格单调递减的,因此存在一个函数tλ(⋅)使得t=tλ(λ(t)),因此,对式(11)进行变量替代后,得到:
xt=ˉαtˉαsxs−ˉαt∫λtλse−λϵθ(xλ,λ)dλ式(12)给出了ODE解法的新视角——只需要对指数积分项进行估计,从而避免了估计线性项带来的误差。
数值估计
在对非线性项的估计中,DPM-Solver对$\boldsymbol{\epsilon}{\theta}\left(\boldsymbol{x}{\lambda}, \lambda\right)进行泰勒展开(式(13))得到式(14)$:
ϵθ(xλ,λ)=k−1∑n=0(λ−λti)nn!ϵ(n)θ(xλti,λti)+O((λ−λti)k) xti→ti−1=αti−1αtixti−αti−1k−1∑n=0ϵ(n)θ(xλti,λti)∫λti−1λtie−λ(λ−λti)nn!dλ+O((λ−λti)k+1)其中,$\boldsymbol{\epsilon}{\theta}^{(n)}\left(\boldsymbol{x}{\lambda_{t_{i}}}, \lambda_{t_{i}}\right)表示\boldsymbol{\epsilon}{\theta}\left(\boldsymbol{x}{\lambda_{t_{i}}}, \lambda_{t_{i}}\right)的n阶导数。可以发现,对于任意的泰勒展开阶数k,\int e^{-\lambda} \frac{(\lambda - \lambda_{t_{i}})^{n}}{n!} \mathrm{d}\lambda都是可以准确计算出结果的,因此,只要估计\boldsymbol{\epsilon}{\theta}\left(\boldsymbol{x}{\lambda_{t_{i}}}, \lambda_{t_{i}}\right)$的导数,即可实现对非线性部分的估计,而对于其导数的估计已经在现有的文章中有较好的研究。
考虑到k比较大时需要引入过多的中间点来进行导数的估计,因此作者只使用了k=1,2,3,三种不同阶数的DPM-Solver。
与DDIM联系
DDIM是较早提出的确定性采样算法,但是一直没有较好的理论将其与ODE联系起来,将λti带入到一阶DPM-Solver(式(15))中可以得到DDIM对应的微分表达式(待引用DDIM)
xti→ti−1=αti−1αtixti−αti−1ϵθ(xλti,λti)(e−λti−e−λti−1)因此,DDIM可以看作是DPM-Solver的一种特殊情况,由于充分利用了半线性的特点,因此DDIM相比于传统的Euler数值解法,具有更好的性能。
DEIS
Diffusion Exponential Integrator Sampler (DEIS)同样利用了ODE方程的半线性的性质,该方法与DPM-Solver最本质的区别在于对非线性项的估计中使用端点处的ϵθ(xti,ti)代替积分区间内的ϵθ(xt,t),同时构建r阶$\boldsymbol{P}{r}(t)项式(16)减少\epsilon{\theta}(x_{t}, t)$估计的估计误差:
Pr(t)=r∑j=0[∏k≠jt−ti+jti+j−ti+k]ϵθ(xti+j,ti+j)因此,DEIS的采样过程为:
xti→ti−1=Ψ(ti−1,ti)xti+r∑j=0∫ti−1ti12Ψ(ti−1,τ)GτG⊤τL−⊤τ[∏k≠jτ−ti+jti+j−ti+k]ϵθ(xti+j,ti+j)dτ其中,
Ψ(ti−1,ti)=e∫ti−1tiFτdτ同时,DEIS还提出利用$\boldsymbol{y}{t} = \Psi(0, t)\boldsymbol{x}{t}$进行参数替换,消除ODE方程的非线性,从而使现有成熟的ODE数值解法具有更好的表现。
拟合误差
模型训练
方差估计
Analytic-DPM
Σ(xt)=Ex0∼p(x0|xt)[(x0−ˉμ(xt))(x0−ˉμ(xt))⊤]=Ex0∼p(x0|xt)[((x0−xt√ˉαt)+√ˉβtˉαtϵθ(xt,t))((x0−xt√ˉαt)+√ˉβtˉαtϵθ(xt,t))⊤]=Ex0∼p(x0|xt)[(x0−xt√ˉαt)(x0−xt√ˉαt)⊤]−ˉβtˉαtϵθ(xt,t)ϵθ(xt,t)⊤=1ˉαtEx0∼p(x0|xt)[(xt−√ˉαtx0)(xt−√ˉαtx0)⊤]−ˉβtˉαtϵθ(xt,t)ϵθ(xt,t)⊤ Ext∼p(xt)Ex0∼p(x0|xt)[(xt−√ˉαtx0)(xt−√ˉαtx0)⊤]=Ex0∼p(x0)Ext∼p(xt|x0)[(xt−√ˉαtx0)(xt−√ˉαtx0)⊤] ˉσ2t=ˉβtˉαt(1−1dExt∼p(xt)[‖ϵθ(xt,t)‖2])SN-DPM
Σ(xt)=Ex0∼p(x0|xt)[(x0−ˉμ(xt))(x0−ˉμ(xt))⊤]=1ˉαtEx0∼p(x0|xt)[(xt−√ˉαtx0)(xt−√ˉαtx0)⊤]−ˉβtˉαtϵθ(xt,t)ϵθ(xt,t)⊤=ˉβtˉαtEx0∼p(x0|xt)[ϵθ(xt,t)ϵθ(xt,t)⊤]−ˉβtˉαtϵθ(xt,t)ϵθ(xt,t)⊤NPR-DPM
Σ(xt)=Ex0∼p(x0|xt)[(x0−ˉμ(xt))(x0−ˉμ(xt))⊤]=ˉβtˉαtEx0∼p(x0|xt)[(ϵt−ϵθ(xt,t))(ϵt−ϵθ(xt,t))⊤]乔列斯基(Cholesky)分解
##
Latent Diffusion Model(LDM)
LDM在原本的DDPM的基础上使用预训练的VAE将输入压缩到潜空间,模型被训练用来生成图像在潜空间的表示。
VAE编码
潜空间训练
Blurring Diffusion Model(BDM)
BDM利用DCT将模型定义在了频率空间,对图像在频率空间的表征进行扩散模型的训练,令 $\boldsymbol{u}{t} = \boldsymbol{V}^{\top}\boldsymbol{x}{t},\boldsymbol{u}{\boldsymbol{\epsilon},t} = \boldsymbol{V}^{\top}\boldsymbol{\epsilon}{t},其中\boldsymbol{V}^{\top}表示DCT变换矩阵,则前向过程重构为式(24)$所示:
ut=αtut+σtuϵ,t同时,由于对噪声的估计在标准的像素空间表现更好,因此在使用网络去噪时,使用逆变换将频率空间内的图像表示转换到像素空间输入网络进行噪声预测,如式(25)所示:
L:=‖ϵθ(zt,t)−ϵt‖2其中,
$\boldsymbol{z}{t} = \boldsymbol{V}(\boldsymbol{\alpha}{t}\boldsymbol{u}{t} + \boldsymbol{\sigma}{t} \boldsymbol{u}_{\boldsymbol{\epsilon},t})$
,
V
表示DCT逆变化,在频率空间的采样过程与原DDPM保持相同。