论文链接:Denoising Diffusion Probabilistic Models

这篇论文的贡献主要有两个

  1. 证明扩散模型确实有能力生成高质量样本
  2. 证明扩散模型的一种特定参数化方法与降噪分数匹配(denoising score matching)等价

一些主要前置知识:

变分推理(Variational Inference)

扩散模型:Deep Unsupervised Learning using Nonequilibrium Thermodynamics

Score Matching:Estimation of Non-Normalized Statistical Models by Score Matching

Score-based Generative Model:Generative Modeling by Estimating Gradients of the Data Distribution

另一篇非常好的参考博客:https://aman.ai/primers/ai/diffusion-models/

Diffusion Model

设$\mathbf{x}{0}q(\mathbf{x}{0})$

x1,,xT是与x0维度相同的隐变量,它们通过一个固定的马尔可夫链采样得到

该马尔可夫链定义为根据方差表β1,,βT逐步向x0添加随机高斯噪声,称为前向过程(forward process)扩散过程(diffusion process)

q(xt|xt1):=N(xt;1βtxt1,βtI)

前向过程的一个重要性质是可以解析地在任意timestep进行对x采样,令$\alphat=1-\beta_t\overline{\alpha}_t=\prod{s=1}^t\alpha_s$,则有

q(xt|x0)=N(xt;α¯tx0,(1α¯t)I)

证明:

使用重参数化(reparameterize)技巧,设ϵN(0,I),则

xt=αtxt1+1αtϵt1=αt(αt1xt2+1αt1ϵt2)+1αtϵt1=αtαt1xt2+αtαtαt1ϵt2+1αtϵt1=αtαt1xt2+(αtαtαt1)+(1αt)ϵt2=αtαt1xt2+1αtαt1ϵt2=...=α¯tx0+1α¯tϵ0

采样得到隐变量序列后,可以得到后验q(x1,T|x0)

q(x1:T|x0):=t=1Tq(xt|xt1)

现在再定义另一个逆向马尔可夫链,其初始分布为pθ(xT)=N(xT;0,I),转移概率是通过学习得到的高斯分布,称为逆过程(reverse process)

pθ(xt1|xt):=N(xt1;μθ(xt,t),Σθ(xt,t))

进一步可得联合分布$p{\theta}\left(\mathbf{x}{0: T}\right)$

pθ(x0:T):=p(xT)t=1Tpθ(xt1xt)

扩散模型(Diffusion Models, DM)的思路是用$p{\theta}(\mathbf x{t-1}|\mathbf xt)q(\mathbf x{t-1}|\mathbf x{t})p{\theta}(\mathbf{x}{0})=\int p{\theta}\left(\mathbf{x}{0: T}\right)d \mathbf{x}{1: T}q(\mathbf{x}_{0})$的近似

通俗的说,DM对样本不断加噪直到几乎只有高斯噪声,然后再学习加噪的逆过程,即从高斯噪声不断降噪直到获得一个样本

diffusion

$p{\theta}(\mathbf{x}{0})$的训练是通过优化变分下界进行的,即

Eq[logpθ(x0)]Eq[logpθ(x0:T)q(x1:T|x0)]=Eq[logp(xT)t1logpθ(xt1|xt)q(xt|xt1)]=:L

直接计算这个式子需要MCMC采样,方差会比较大,因此将L进一步改写为

L=Eq[DKL(q(xTx0)p(xT))LT+t>1DKL(q(xt1xt,x0)pθ(xt1xt))Lt1logpθ(x0x1)L0]

证明:

L=Eq[logp(xT)t1logpθ(xt1|xt)q(xt|xt1)]=Eq[logp(xT)logpθ(x0|x1)q(x1|x0)t>1logpθ(xt1|xt)q(xt|xt1)]=Eq[logp(xT)logpθ(x0|x1)q(x1|x0)t>1logpθ(xt1|xt)q(xt1|xt,x0)q(xt1|x0)q(xt|x0)]=Eq[logp(xT)logpθ(x0|x1)q(x1|x0)logq(x1|x0)q(xT|x0)t>1logpθ(xt1|xt)q(xt1|xt,x0)]=Eq[logp(xT)q(xT|x0)logpθ(x0|x1)t>1logpθ(xt1|xt)q(xt1|xt,x0)]=Eq[DKL(q(xT|x0)p(xT))+t>1DKL(q(xt1|xt,x0)pθ(xt1|xt))logpθ(x0|x1)]

改写后的L含义更加明显了,最小化L其实就是最小化$D{KL}(q(\mathbf x{t-1}|\mathbf x{t})|p{\theta}(\mathbf x{t-1}|\mathbf x_t))p{\theta}(\mathbf x{t-1}|\mathbf x_t)q(\mathbf x{t-1}|\mathbf x_{t})$尽可能相似

$L{t-1}q(\mathbf{x}{t-1}|\mathbf{x}{t}, \mathbf{x}{0})$根据贝叶斯公式可得

q(xt1|xt,x0)=N(xt1;μ~t(xt,x0),β~tI)whereμ~t(xt,x0):=α¯t1βt1α¯tx0+αt(1α¯t1)1α¯txtandβ~t:=1α¯t11α¯tβt

正态分布间的KL散度是可以直接计算的,这样就避免了MCMC

Score-based Model

Score Matching最初的目的是计算非归一化概率模型的归一化常量

其论文中将Score Function定义为对数概率密度的梯度(the gradient of the log-density)

ψ(ξ;θ)=(logp(ξ;θ)ξ1logp(ξ;θ)ξn)=ξlogp(ξ;θ)

Score-based Model的目标是训练是一个分数网络sθ(x)估计数据的Score Function,即

12Epdata[sθ(x)xlogpdata(x)22]

由于数据的Score Function是未知的,该式需要进一步推导,可以证明,该式在常数差距内等价于

Epdata(x)[tr(xsθ(x))+12sθ(x)22]

此时$\operatorname{tr}(\nabla{\mathbf{x}}\mathbf{s}{\theta}(\mathbf{x}))$仍是难以计算的,不同的Score-based Model目标就是解决这个问题

一种比较常用的方法是降噪分数匹配(denoising score matching)

其首先使用一个噪声分布qσ(x~|x)对数据点x进行搅动,然后使用分数网络估计搅动后的数据分布

12Eqσ(x~x)pdata (x)[sθ(x~)x~logqσ(x~x)22]

可以证明该式最小化时有$\mathbf{s}^*{\boldsymbol{\theta}}(\tilde{\mathbf{x}})=\nabla{\mathbf{x}} \log q{\sigma}(\mathbf{x})q{\sigma}(\mathbf{x})\approx p_{data}(\mathbf x)$

获得了分数网络$\mathbf s{\theta}(\mathbf x)$后,我们可以使用Lagevin Dynamics(或称Lagevin Sampling)从Score Function $\nabla{\mathbf x} p_{\mathrm{data}}(\mathbf{x})$中进行采样,它是一种MCMC采样方法

给定固定步长ϵ>0、先验分布π和初始值x~0π(x),Lagevin Sampling循环计算

x~t=x~t1+ϵ2xlogp(x~t1)+ϵzt

其中$\mathbf{z}t\sim \mathcal N(0, I)\epsilon \to 0, T\to \infty\tilde{\mathbf{x}}{t}p(\mathbf x)$的准确采样

Diffusion Models and Denoising Autoencoders

通过指定不同的$\betatp{\theta}\left(\mathbf{x}{t-1}|\mathbf{x}{t}\right)$的参数化形式,可以得到扩散模型的很多不同实现

下面要讨论的一种实现将使扩散模型产生与降噪分数匹配等价的效果

对于βt,我们将其固定为常数,此时后验q无可学习参数,进而LT也是常数,因此可以忽略

对于逆向过程 $p{\theta}\left(\mathbf{x}{t-1}|\mathbf{x}{t}\right):=\mathcal{N}\left(\mathbf{x}{t-1} ; \boldsymbol{\mu}{\theta}\left(\mathbf{x}{t}, t\right), \boldsymbol{\Sigma}{\theta}\left(\mathbf{x}{t}, t\right)\right)$

首先,令$\boldsymbol{\Sigma}{\theta}\left(\mathbf{x}{t}, t\right)=\sigma_t^2 \mathbf I\sigma^2_t$是常量且相互独立,不参与训练

从实验结果来看,$\sigma^2t=\beta_t\sigma^2_t=\tilde{\beta}{t}$具有相似效果(但是论文中的解释暂时没看懂)

The first choice is optimal for x0N(0,I), and the second is optimal for x0 deterministically set to one point. These are the two extreme choices corresponding to upper and lower bounds on reverse process entropy for data with coordinatewise unit variance.

此时我们可以根据正态分布间的KL散度公式进一步推导得

Lt1=Eq[12σt2μ~t(xt,x0)μθ(xt,t)2]+C

其中C是与θ无关的常量

前面提到前向过程可以在任意时间步采样,即q(xt|x0)=N(xt;α¯tx0,(1α¯t)I)

我们对其使用重参数化技巧,令ϵN(0,I),则 $\mathbf xt(\mathbf x_0, \epsilon)=\sqrt{\bar\alpha_t}\mathbf x_0+\sqrt{1-\bar\alpha_t}\epsilonL{t}$得

Lt1C=Ex0,ϵ[12σt2μ~t(xt(x0,ϵ),1α¯t(xt(x0,ϵ)1α¯tϵ))μθ(xt(x0,ϵ),t)2]=Ex0,ϵ[12σt21αt(xt(x0,ϵ)βt1α¯tϵ)μθ(xt(x0,ϵ),t)2]

从这个推导结果容易得到,对于均值μθ(xt,t),我们应将其参数化为

μθ(xt,t)=μ~t(xt,1α¯t(xt1α¯tϵθ(xt)))=1αt(xtβt1α¯tϵθ(xt,t))

其中$\mathbf xt\epsilon{\theta}\mathbf x_t\epsilon$的近似函数

确定均值的参数化方法后,Lt可化简为

Ex0,ϵ[βt22σt2αt(1α¯t)ϵϵθ(α¯tx0+1α¯tϵ,t)2]

此时可以明显看出该式与降噪分数匹配非常类似,实际上是在多个尺度的噪声下进行降噪分数匹配,其中的ϵθ就相当于学习到的数据梯度

下面给出该参数化方法下扩散模型的算法伪码,如Algorithm2所示,采样$\mathbf x{t-1}\sim p{\theta}(\mathbf x{t-1}|\mathbf x_t)\mathbf{x}{t-1}=\frac{1}{\sqrt{\alphat}}\left(\mathbf{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\theta(\mathbf{x}_t,t)\right)+\sigma_t\mathbf{z}z \sim \mathcal N(\mathbf 0,\mathbf I)$,这个过程与Langevin dynamics类似

DDPM_algorithm

Discrete Decoder of Reverse Process

由RGB表示的图像中每个像素点离散地在0,1,,255中取值,但现在逆过程中每一步都是连续的高斯分布

为了获得离散似然,我们需要将逆过程的最后一步改写为一个独立的离散解码器(independent discrete decoder)

以下我们假设所有图像的像素值都从[0,255]均匀的缩放到了[1,1]

由于上述推导中假定逆过程的协方差矩阵是对角阵,所以显然有

pθ(x0|x1)=i=1Dpθ(x0i|x1i)=i=1DN(x;μθi(x1,1),σ12)

其中D是数据的维度

接下来使用分箱进行离散化,即像素点取值在区间[x1255,x+1255]时视为取离散值xx=1,2,,255,于是有

pθ(x0x1)=i=1Dδ(x0i)δ+(x0i)N(x;μθi(x1,1),σ12)dxδ+(x)={ if x=1x+1255 if x<1δ(x)={ if x=1x1255 if x>1

Simplified training object

作者在实验中发现,将前面推导的得到的训练目标进一步化简可以得到更好的生成效果

Lsimple (θ):=Et,x0,ϵ[ϵϵθ(α¯tx0+1α¯tϵ,t)2]

其中tU(1,T),即均匀分布

Experiment

IS和FID是评价(图像)生成模型的常用量化方法

IS(Inception Score)使用Inception Net输出图像x1000维分类向量y,并从两方面量化生成图像的质量

  • 清晰度:对单一的生成图像,其类别分布的熵应尽可能小,即对于清晰的图片,其属于某一个类别的概率应趋近于1,其余则趋近于0。因此最小化

    Expg[H(p(yx))]
  • 多样性:对成批的生成图像,其类别分布的熵应尽量大,即生成模型生成的图像类别因尽可能丰富,因此最大化

    H(Expg[p(yx))])

将前者取负,两者相加后得IS表达式,IS越大,生成图像质量越高

IS=exp(Expg[DKL(p(y|x)p(y))])

FID(Frechet Inception Distance)使用Inception Net-V3并删除最后的分类层,得到图像的2048维特征向量

FID的思想是直接计算真实图像分布和生成图像分布的距离,但图像分布维度过大不易计算,因此使用Inception输出的2048维特征向量计算Frechet Distance,显然FID越小越好

FID=μxμg2+tr(Σx+Σg2(ΣxΣg)12)

下图Table 1、Table2展示了DDPM的IS、FID和NLL比较

由Tabel 1可知DPPM的IS仅次于StyleGAN2+ADA,而FID则最优,对于负对数似然NLL,论文称使用未化简的变分目标得到的NLL更好,但图片质量不如化简的目标

Table 2对比的是不同的L和不同的逆过程参数化方法

DDPM_IS_FID