Skip to content

Denoising Diffusion Probabilistic Models

Ho J , Jain A , Abbeel P .Denoising Diffusion Probabilistic Models[J]. 2020.DOI:10.48550/arXiv.2006.11239.

Denoising Diffusion Probabilistic Models


前向过程

前向扩散过程就是向图像不断加高斯噪声

给定初始数据点 x0q(x0)q(x0) 是未被噪声破坏的原始数据分布

t 时刻的噪化状态与上一时刻 t1 的关系为:

q(xtxt1)=N(xt;1βtxt1,βtI)

在给定前一个状态 xt1 的条件下,当前状态 xt 的分布是一个高斯分布。这个高斯分布的均值是 1βtxt1,方差是 βtI

其中 βt 是一个预先定义好的小于 1 的正数,通常随着时间步 t 的增加而增加(一开始,加一点点噪声就能比较明显的看出和原图的区别,越到后面,图像退化的越厉害,轻微的扰动已经看不出明显的变化,所以 βt 的值需要更大)。I 是单位矩阵。

这个过程中,我们逐渐将噪声加入到数据中,因为均值是 1βt 乘以 xt1,而 1βt 小于1,所以相当于在上一时刻的数据上缩放了一下,然后加上一个方差为 βt 的噪声。这样,随着 t 的增加,数据中原始的信息逐渐减少,噪声逐渐增加。

输入 x0 的条件下,x1,x2,...,xT 的联合分布可以表示为:

q(x1:Tx0)=t=1Tq(xtxt1)

重参数化技巧 (reparameterization trick)

对于采样操作,采样的输出不是输入参数(如 μσ)的一个确定函数,而是一个随机值。即使我们固定了分布参数,每次采样得到的值都是不同的。因此,采样操作不具备输入与输出之间的确定性关系,无法计算导数,进而导致采样操作不可微。

在深度学习中,我们通常使用梯度下降来优化模型参数,梯度计算依赖于链式法则。由于采样结果随机,这种不确定性使得梯度无法通过采样的结果反向传播到模型参数。

重参数化的核心思想是将随机过程与可微过程分离,这样我们可以对可微部分进行梯度计算,同时保留随机性。具体而言,重参数化技巧将从分布 N(μ,σ2) 采样的过程转换为一个确定的线性变换加上随机噪声的过程

对于一个高斯分布:

N(x;μ,σ2)

它的采样过程可以写成:

x=μ+σϵ,ϵN(0,1)

其中,ϵN(0,1) 是从标准正态分布中采样的噪声,μσ 为模型参数。

通过这种方式,将随机性分离到了 ϵ 中,其与模型参数 μσ 无关,即在梯度计算时,ϵ 不会影响 μσ 的梯度计算。

对于:

q(xtxt1)=N(xt;1βtxt1,βtI)

采样过程就是:

xt=1βtxt1+βtϵt,ϵtN(0,I)

为了简化计算,令:αt=1βt,代入:

有:

xt=α¯tx0+1α¯tϵt,ϵtN(0,I)α¯t=αtαt1αt2α2α1=i=1tαi

即:

xtN(α¯tx0,(1α¯t)I)q(xtx0)=N(xt;α¯tx0,(1α¯t)I)

反向过程

定义反向过程为一个以 θ 为参数的马尔科夫链,它试图近似真实但未知的逆向分布:

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

其中:

  • 起点:p(xT)=N(xT;0,I) ,即纯噪声
  • 每一步:pθ(xt1xt)=N(xt1;μθ(xt,t),Σθ(xt,t))

反向过程的每一步也是一个高斯分布,其均值和方差由神经网路 θ 预测,我们需要一个目标函数来训练 θ,最自然的目标是就是最大化模型生成真实数据的似然 pθ(x0)

直接计算很困难,所以引入了变分推断的思想:我们有一个由前向过程定义的、已知的真实数据分布 q(x0),以及从数据到噪声的前向过程 q(x1:Tx0),我们可以通过最小化模型分布 $$p_{\theta}$ 和真实后验分布 q 的 KL散度 来训练模型。

将损失函数定义为目标函数的负值,目标函数越大,损失越小:

pθ(x0)=pθ(x0:T)dx1:TL=logpθ(x0)

前向过程 q(x1:Tx0) 是我们已知的加噪过程,如果我们的反向模型 pθ(x1:Tx0) 是完美的,那么它应该和“从真实数据加噪再完美去噪”的过程一致。

数学上,我们希望两个条件分布接近:

pθ(x1:Tx0)q(x1:Tx0)

KL散度 DKL(PQ) 衡量分布 PQ 的差异,对于任意 x0

DKL(q(x1:Tx0)pθ(x1:Tx0))0

展开KL散度:

Eq[logq(x1:Tx0)pθ(x1:Tx0)]0

由于:

pθ(x1:Tx0)=pθ(x0:T)pθ(x0)

有:

Eq[logq(x1:Tx0)logpθ(x0:T)pθ(x0)]0Eq[logq(x1:Tx0)]Eq[logpθ(x0:T)]+logpθ(x0)0

即:

logpθ(x0)Eq[logpθ(x0:T)logq(x1:Tx0)]L=logpθ(x0)Eq[log(pθ(x0:T))q(x1:T|x0)]

根据我们之前的定义:

  • 反向过程联合分布:
pθ(x0:T)=p(xT)t=1Tpθ(xt1xt)
  • 前向过程联合分布:
q(x1:Tx0)=t=1Tq(xtxt1)

代入:

pθ(x0:T)q(x1:Tx0)=p(xT)t=1Tpθ(xt1xt)t=1Tq(xtxt1)logpθ(x0:T)q(x1:Tx0)=logp(xT)t=1Tlogpθ(xt1xt)+t=1Tlogq(xtxt1)L=Eq[logp(xT)t=1Tlogpθ(xt1xt)q(xtxt1)]

要计算这个期望,

  • 我们需要从 q(x0) 中采样一个数据点 x0
  • 然后从 q(x1:Tx0) 中采样整个加噪轨迹 x1,x2,,xT
  • 对每个时间步,计算比值
  • 最后取期望

我们希望的是对每个时间步单独计算损失,而不是对整个轨迹进行采样后再计算。

另外,公式中混合了前向和反向转移,没有清晰地分离出“去噪误差”,即缺乏明确的优化目标。

根据贝叶斯定理:

q(xtxt1)=q(xt1xt,x0)q(xtx0)q(xt1x0)

有:

t=1Tlogpθ(xt1xt)q(xtxt1)=t=1T[logpθ(xt1xt)q(xt1xt,x0)+logq(xt1x0)q(xtx0)]t=1Tlogq(xt1x0)q(xtx0)=logq(x0x0)logq(xTx0)=logq(xTx0)

于是:

L=Eq[logp(xT)t=1Tlogpθ(xt1xt)q(xt1xt,x0)+logq(xTx0)]

整理:

L=Eq[logp(xT)q(xTx0)t=1Tlogpθ(xt1xt)q(xt1xt,x0)]

第一项是 q(xTx0)p(xT) 的 KL散度:

Eq[logq(xTx0)p(xT)]=DKL(q(xTx0)p(xT))

第二项同理,但注意 t=1 的情况。

最后。我们得到了负对数似然的变分下界(ELBO)的分解形式:

L=Eq[DKL(q(xTx0)p(xT))+t=2TDKL(q(xt1xt,x0)pθ(xt1xt))logpθ(x0x1)]

这样,损失函数被分解为三项:

  1. 第一项是最终噪声分布的 KL散度。由于前向过程确保 q(xTx0) 接近标准高斯,且 p(xT) 也是标准高斯,这一项近似为 0,可以忽略。

  2. 第二项是一系列 KL散度,要求每个时间步上模型的反向分布 pθ(xt1xt) 尽可能接近真实的后验分布 q(xt1xt,x0)

  3. 第三项是最后一步的对数似然,即从 x1 生成 x0 的概率。

TIP

真实后验 q(xt1xt,x0) 是一个高斯分布,下面计算其均值 μ~t

已知:

q(xtxt1)=N(xt;αtxt1,βtI),αt=1βtq(xtx0)=N(xt;α¯tx0,(1α¯t)I),α¯t=i=1tαiq(xt1x0)=N(xt1;α¯t1x0,(1α¯t1)I)

根据贝叶斯定理,真实后验分布为:

q(xt1xt,x0)=q(xtxt1)q(xt1x0)q(xtx0)

代入三个高斯分布的概率密度函数,关注指数部分(忽略常数项):

12[(xtαtxt1)2βt+(xt1α¯t1x0)21α¯t1(xtα¯tx0)21α¯t]

展开前两项中与 xt1 相关的部分:

第一项展开:

1βt(xt22αtxtxt1+αtxt12)

第二项展开:

11α¯t1(xt122α¯t1x0xt1+α¯t1x02)

合并同类项:

  • 二次项系数(xt12 的系数)
αtβt+11α¯t1
  • 一次项系数(xt1 的系数)
2αtβtxt2α¯t11α¯t1x0

高斯分布的精度(方差的倒数)为二次项系数:

β~t1=αtβt+11α¯t1

化简:

β~t1=αt(1α¯t1)+βtβt(1α¯t1)=(1βt)(1α¯t1)+βtβt(1α¯t1)=1α¯t1βt(1α¯t1)+βtβt(1α¯t1)=1α¯t1+βtα¯t1βt(1α¯t1)=1α¯t1(1βt)βt(1α¯t1)

注意到 α¯t=αtα¯t1=(1βt)α¯t1,所以:

1α¯t=1(1βt)α¯t1=1α¯t1+βtα¯t1

因此:

β~t1=1α¯tβt(1α¯t1)

方差为:

β~t=βt(1α¯t1)1α¯t

高斯分布的均值满足:

β~t1μ~t=αtβtxt+α¯t11α¯t1x0

解得:

μ~t=β~t(αtβtxt+α¯t11α¯t1x0)

代入 β~t

μ~t=βt(1α¯t1)1α¯t(αtβtxt+α¯t11α¯t1x0)=αt(1α¯t1)1α¯txt+βtα¯t11α¯tx0

由:

xt=α¯tx0+1α¯tϵ

解出 x0

x0=1α¯t(xt1α¯tϵ)

有:

μ~t=αt(1α¯t1)1α¯txt+βtα¯t11α¯t1α¯t(xt1α¯tϵ)

由于 α¯t=αtα¯t1

μ~t=αt(1α¯t1)1α¯txt+βtαt(1α¯t)(xt1α¯tϵ)

利用 αt(1α¯t1)+βt=1α¯t,可以简化为:

μ~t=1αtxtβtαt(1α¯t)1α¯tϵ=1αt(xtβt1α¯tϵ)

真实后验分布为:

q(xt1xt,x0)=N(xt1;μ~t,β~tI)μ~t=1αt(xtβt1α¯tϵ),β~t=βt(1α¯t1)1α¯t

这个结果表明在已知原始数据 x0 和当前噪声数据 xt 时,真实的去噪均值可以通过从 xt 中减去一定比例的噪声 ϵ 并缩放得到。

我们也将反向分布 pθ(xt1xt) 设为高斯分布,并固定其方差为 β~tI(或 βtI,实验表明两者效果相近)。当两个高斯分布的方差相同时,它们的 KL 散度正比于均值之差的平方:

DKL(qpθ)μ~tμθ(xt,t)2.

为了匹配 μ~t 的形式,我们令模型的均值参数化为:

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

其中 ϵθ(xt,t) 是神经网络预测的噪声。代入后,KL 散度项简化为:

μ~tμθ(xt,t)2=βt2(1α¯t)αtϵϵθ(xt,t)2

在论文中,作者发现忽略权重系数 βt2(1α¯t)αt 能使训练更稳定,因此最终损失函数采用简化的均方误差形式:

Lsimple =Et,x0,ϵ[ϵϵθ(xt,t)2]

其中:

  • x0q(x0)
  • ϵN(0,I)
  • xt=α¯tx0+1α¯tϵ

训练算法

重复直到收敛:

  • 从训练集中采样 x0
  • 均匀采样时间步 t1,,T
  • 采样噪声 ϵN(0,I)
  • 计算加噪样本 xt=α¯tx0+1α¯tϵ
  • 计算损失 L=|ϵϵθ(xt,t)|2
  • 反向传播更新参数 θ

采样算法

  • 从标准高斯分布采样初始噪声:xTN(0,I)
  • t=Tt=1 循环:
    • 预测噪声:ϵθ=ϵθ(xt,t)
    • 计算均值:μθ=1αt(xtβt1α¯tϵθ)
    • 采样 zN(0,I)(若 t=1z=0
    • 更新:xt1=μθ+βtz
  • 返回生成样本 x0