DDPM:从噪声到图像的冷蓝灰色主题插画

DDPM

Denoising Diffusion Probabilistic Models——从逐步去噪到一步加噪

GO UP! Diffusion + RL学习 01

理解扩散模型时,可以从得分匹配、能量模型和朗之万采样的角度出发。但这篇笔记里,我更想沿着 DDPM 的概率建模过程,看看它为什么能一步步去噪。

DDPM 给我的感觉更像 VAE:它引入了多个中间变量,也就是不断加噪、不断去噪的过程。这样每一次只拟合一点,有点像把曲线分段处理。这个类比并不准确,但能帮助我理解为什么要把生成过程拆成许多小步。

PAPER · HO, JAIN & ABBEELDenoising Diffusion Probabilistic ModelsNeurIPS 2020 · arXiv:2006.11239 ↗

#整体思路

我们想从一张高斯噪声图开始,通过不断去噪,得到一张图像。但如何训练?如何引导去噪过程?

这也是图像生成的核心问题。Diffusion 和 GAN 给出了不同的方案:这里的扩散训练可以自己构造监督信号,把已知的噪声作为 noise prediction 模块的标签,再让网络学会去噪。GAN 则通过生成器与判别器之间的对抗来训练,并不采用这里的逐步加噪、去噪过程。

原文提到,也可以加入文本来引导生成。这属于条件生成的扩展;下面先只讨论原始 DDPM 的无条件图像生成。

原文中的前向加噪示意:图像经过多个步骤,逐渐被随机噪声覆盖
原文中的前向过程示意图:逐步加噪,逐渐失去图像结构。

但是要注意:图中的噪声是逐步产生的,而在 DDPM 的实际训练中,任意一个中间图像都可以从 x0x_0 出发,一步得到。这是我读这篇论文时特别想弄明白的一点。

一方面,这来自高斯噪声的性质;另一方面,这样不用为每次训练都模拟完整的加噪链,训练会更高效。更让我感兴趣的是:从 x0x_0 到 xtx_t 的这份噪声,恰好又出现在反向过程的均值里。下面就推一遍公式,看看前后两部分是怎么接上的。

#前向过程

前向过程就是加噪声。来自训练集的图像 x0x_0 经过 TT 次扰动,最终接近标准正态噪声。每一步只依赖前一步,因此它是一个马尔可夫过程:

q(x1:T∣x0)=∏t=1Tq(xt∣xt−1).q(x_{1:T}\mid x_0)=\prod_{t=1}^{T}q(x_t\mid x_{t-1}).
q(xt∣xt−1)=N ⁣(xt;1−βt xt−1,βtI).q(x_t\mid x_{t-1})=\mathcal N\!\left(x_t;\sqrt{1-\beta_t}\,x_{t-1},\beta_t I\right).

这里 0<βt<10<\beta_t<1 是第 tt 步预先设置的噪声方差。通过一个标准正态样本,就能写出 xtx_t 与 xt−1x_{t-1} 的关系:

xt=1−βt xt−1+βt zt,zt∼N(0,I).x_t=\sqrt{1-\beta_t}\,x_{t-1}+\sqrt{\beta_t}\,z_t,\qquad z_t\sim\mathcal N(0,I).

上面两个公式是同一件事的两种写法:既可以理解为从高斯条件分布中采样,也可以理解为缩放上一张图像,再加上一份噪声。不同时间步的 ztz_t 相互独立,图片大小始终不变。

为什么能够一步加噪?

先记:

αt=1−βt,αˉt=∏s=1tαs,αˉ0=1.\alpha_t=1-\beta_t,\qquad \bar\alpha_t=\prod_{s=1}^{t}\alpha_s,\qquad\bar\alpha_0=1.

把 xt−1x_{t-1} 展开一次:

xt=αt xt−1+βt zt=αtαt−1 xt−2+αtβt−1 zt−1+βt zt.\begin{aligned}x_t&=\sqrt{\alpha_t}\,x_{t-1}+\sqrt{\beta_t}\,z_t\\&=\sqrt{\alpha_t\alpha_{t-1}}\,x_{t-2}+\sqrt{\alpha_t\beta_{t-1}}\,z_{t-1}+\sqrt{\beta_t}\,z_t.\end{aligned}

两个独立的零均值高斯噪声相加,协方差相加。因此,后面两项可以合并:

αtβt−1 zt−1+βt zt=dαtβt−1+βt ϵ=1−αtαt−1 ϵ,ϵ∼N(0,I).\begin{aligned}\sqrt{\alpha_t\beta_{t-1}}\,z_{t-1}+\sqrt{\beta_t}\,z_t &\overset{d}{=}\sqrt{\alpha_t\beta_{t-1}+\beta_t}\,\epsilon\\ &=\sqrt{1-\alpha_t\alpha_{t-1}}\,\epsilon,\qquad\epsilon\sim\mathcal N(0,I).\end{aligned}

其中 =d\overset{d}{=} 表示同分布。原来笔记里写的是“均值相同,方差可以相加”,更准确的条件是这些高斯噪声相互独立。

继续展开到 x0x_0,就得到原文手写推导中的乘积形式:

xt=d∏s=1t(1−βs) x0+1−∏s=1t(1−βs) ϵ.x_t\overset{d}{=}\sqrt{\prod_{s=1}^{t}(1-\beta_s)}\,x_0+\sqrt{1-\prod_{s=1}^{t}(1-\beta_s)}\,\epsilon.

于是,任意时间步的直接采样公式为:

xt=αˉt x0+1−αˉt ϵ,ϵ∼N(0,I).\boxed{x_t=\sqrt{\bar\alpha_t}\,x_0+\sqrt{1-\bar\alpha_t}\,\epsilon,\qquad\epsilon\sim\mathcal N(0,I).}
q(xt∣x0)=N ⁣(xt;αˉt x0,(1−αˉt)I).q(x_t\mid x_0)=\mathcal N\!\left(x_t;\sqrt{\bar\alpha_t}\,x_0,(1-\bar\alpha_t)I\right).

这里的 ϵ\epsilon 是从 x0x_0 直接构造 xtx_t 时使用的合成噪声,与前面某一个单步的 ztz_t 不是同一个变量。

展开推导:所有噪声项为什么恰好得到这个方差?
xt=αˉt x0+∑j=1tβj∏s=j+1tαs zj.x_t=\sqrt{\bar\alpha_t}\,x_0+\sum_{j=1}^{t}\sqrt{\beta_j\prod_{s=j+1}^{t}\alpha_s}\,z_j.

空乘积取 1。利用独立性,噪声部分的协方差系数为:

∑j=1tβj∏s=j+1tαs=∑j=1t(∏s=j+1tαs−∏s=jtαs)=1−∏s=1tαs=1−αˉt.\begin{aligned}\sum_{j=1}^{t}\beta_j\prod_{s=j+1}^{t}\alpha_s &=\sum_{j=1}^{t}\left(\prod_{s=j+1}^{t}\alpha_s-\prod_{s=j}^{t}\alpha_s\right)\\ &=1-\prod_{s=1}^{t}\alpha_s=1-\bar\alpha_t.\end{aligned}

所以这一整串噪声,在分布上可以替换成 1−αˉt ϵ\sqrt{1-\bar\alpha_t}\,\epsilon。

最后为什么接近纯噪声?

在 DDPM 论文的设置中,T=1000T=1000,βt\beta_t 从 β1=10−4\beta_1=10^{-4} 到 βT=0.02\beta_T=0.02 线性增长。因此 αt\alpha_t 逐渐减小,连乘得到的 αˉt\bar\alpha_t 也越来越小:

αˉT≈0⟹q(xT∣x0)≈N(0,I).\bar\alpha_T\approx0 \quad\Longrightarrow\quad q(x_T\mid x_0)\approx\mathcal N(0,I).

前向过程逐步削弱原图、增加噪声,但对训练来说,重要的不只是终点,而是中间任意时间步都可以构造出“带噪图像—真实噪声”这一对监督信号。

#反向过程

前向过程从 xt−1x_{t-1} 得到 xtx_t,反向过程则希望从 xtx_t 得到 xt−1x_{t-1}。我们可以先从贝叶斯公式入手。这里原文参考了苏神的讲解,非常详细透彻,respect。

原文参考 · 科学空间苏剑林的扩散模型推导原文所引用的贝叶斯与高斯配方过程 ↗
q(xt−1∣xt)=q(xt∣xt−1)q(xt−1)q(xt).q(x_{t-1}\mid x_t)=\frac{q(x_t\mid x_{t-1})q(x_{t-1})}{q(x_t)}.

但这里的边缘分布并不好求。训练时有原图 x0x_0,因此先给各项加上这个条件:

q(xt−1∣xt,x0)=q(xt∣xt−1,x0) q(xt−1∣x0)q(xt∣x0)=q(xt∣xt−1) q(xt−1∣x0)q(xt∣x0).\begin{aligned}q(x_{t-1}\mid x_t,x_0) &=\frac{q(x_t\mid x_{t-1},x_0)\,q(x_{t-1}\mid x_0)}{q(x_t\mid x_0)}\\ &=\frac{q(x_t\mid x_{t-1})\,q(x_{t-1}\mid x_0)}{q(x_t\mid x_0)}.\end{aligned}

第二行使用了前向过程的马尔可夫性:知道 xt−1x_{t-1} 之后,下一步不再额外依赖 x0x_0。右边这三个分布,都已经能写出来:

q(xt∣xt−1)=N ⁣(xt;αt xt−1,βtI),q(xt−1∣x0)=N ⁣(xt−1;αˉt−1 x0,(1−αˉt−1)I),q(xt∣x0)=N ⁣(xt;αˉt x0,(1−αˉt)I).\begin{aligned}q(x_t\mid x_{t-1})&=\mathcal N\!\left(x_t;\sqrt{\alpha_t}\,x_{t-1},\beta_t I\right),\\ q(x_{t-1}\mid x_0)&=\mathcal N\!\left(x_{t-1};\sqrt{\bar\alpha_{t-1}}\,x_0,(1-\bar\alpha_{t-1})I\right),\\ q(x_t\mid x_0)&=\mathcal N\!\left(x_t;\sqrt{\bar\alpha_t}\,x_0,(1-\bar\alpha_t)I\right).\end{aligned}

下面的推导先取 t≥2t\ge2。我们的目标是把这个后验整理成高斯分布,求出均值和方差:

q(xt−1∣xt,x0)=N ⁣(xt−1;μ~t(xt,x0),β~tI).q(x_{t-1}\mid x_t,x_0) =\mathcal N\!\left(x_{t-1};\tilde\mu_t(x_t,x_0),\tilde\beta_t I\right).
xt−1=μ~t(xt,x0)+β~t z,z∼N(0,I).x_{t-1}=\tilde\mu_t(x_t,x_0)+\sqrt{\tilde\beta_t}\,z,\qquad z\sim\mathcal N(0,I).

把高斯分布代入,再对 xt−1x_{t-1} 配方

原笔记是把三个高斯密度逐项代入,然后合并系数和指数。这里保留这条推导主线,把长式子放在下面:

展开推导:代入密度、合并二次项与一次项

记样本维度为 dd,则:

q(xt−1∣xt,x0)=[1−αˉt2πβt(1−αˉt−1)]d/2exp⁡ ⁣[−12R(xt−1)].q(x_{t-1}\mid x_t,x_0) =\left[\frac{1-\bar\alpha_t}{2\pi\beta_t(1-\bar\alpha_{t-1})}\right]^{d/2} \exp\!\left[-\frac12R(x_{t-1})\right].

为了看得清楚,令 y=xt−1y=x_{t-1},并把指数中的各项写成:

R(y)=∥xt−αty∥22βt+∥y−αˉt−1x0∥221−αˉt−1−∥xt−αˉtx0∥221−αˉt.\begin{aligned}R(y) &=\frac{\|x_t-\sqrt{\alpha_t}y\|_2^2}{\beta_t}\\ &\quad+\frac{\|y-\sqrt{\bar\alpha_{t-1}}x_0\|_2^2}{1-\bar\alpha_{t-1}}\\ &\quad-\frac{\|x_t-\sqrt{\bar\alpha_t}x_0\|_2^2}{1-\bar\alpha_t}.\end{aligned}

按 yy 展开:

R(y)=At∥y∥22−2bt⊤y+Ct.R(y)=A_t\|y\|_2^2-2b_t^\top y+C_t.

CtC_t 与 yy 无关。二次项的系数为:

At=αtβt+11−αˉt−1=αt(1−αˉt−1)+βtβt(1−αˉt−1)=1−αˉtβt(1−αˉt−1).\begin{aligned}A_t &=\frac{\alpha_t}{\beta_t}+\frac1{1-\bar\alpha_{t-1}}\\ &=\frac{\alpha_t(1-\bar\alpha_{t-1})+\beta_t}{\beta_t(1-\bar\alpha_{t-1})}\\ &=\frac{1-\bar\alpha_t}{\beta_t(1-\bar\alpha_{t-1})}.\end{aligned}

一次项对应的向量为:

bt=αtβtxt+αˉt−11−αˉt−1x0.b_t=\frac{\sqrt{\alpha_t}}{\beta_t}x_t+\frac{\sqrt{\bar\alpha_{t-1}}}{1-\bar\alpha_{t-1}}x_0.

完成平方:

R(y)=At∥y−btAt∥22+Ct′.R(y)=A_t\left\|y-\frac{b_t}{A_t}\right\|_2^2+C'_t.

因此,高斯后验的方差系数是 At−1A_t^{-1},均值是 bt/Atb_t/A_t。这也补全了原手写图里从“系数合并”直接读出方差的中间步骤。

最后得到:

β~t=1−αˉt−11−αˉt βt.\boxed{\tilde\beta_t=\frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\,\beta_t.}
μ~t(xt,x0)=αˉt−1βt1−αˉtx0+αt(1−αˉt−1)1−αˉtxt.\boxed{\tilde\mu_t(x_t,x_0)= \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1-\bar\alpha_t}x_0+ \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}x_t.}

方差只和预先设置的 α\alpha、β\beta 有关;但均值中既有 xtx_t,也有 x0x_0。所以接下来就是要去掉显式的 x0x_0:生成时,我们并不知道原图是什么样子。

#为什么预测噪声?

再拿出前向的一步加噪公式:

xt=αˉt x0+1−αˉt ϵ.x_t=\sqrt{\bar\alpha_t}\,x_0+\sqrt{1-\bar\alpha_t}\,\epsilon.

把它移项,就能用 xtx_t 和噪声表示 x0x_0:

x0=xt−1−αˉt ϵαˉt.x_0=\frac{x_t-\sqrt{1-\bar\alpha_t}\,\epsilon}{\sqrt{\bar\alpha_t}}.

代回刚才求出的后验均值,化简得到:

μ~t(xt,x0)=1αt(xt−βt1−αˉt ϵ).\boxed{\tilde\mu_t(x_t,x_0)= \frac1{\sqrt{\alpha_t}}\left(x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\,\epsilon\right).}
展开推导:代入后,两个 xtx_t 的系数如何合并?
μ~t=βtαt(1−αˉt)(xt−1−αˉt ϵ)+αt(1−αˉt−1)1−αˉtxt=βt+αt(1−αˉt−1)αt(1−αˉt)xt−βtαt1−αˉt ϵ=1αt(xt−βt1−αˉt ϵ).\begin{aligned}\tilde\mu_t &=\frac{\beta_t}{\sqrt{\alpha_t}(1-\bar\alpha_t)} \left(x_t-\sqrt{1-\bar\alpha_t}\,\epsilon\right)\\ &\quad+\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}x_t\\ &=\frac{\beta_t+\alpha_t(1-\bar\alpha_{t-1})}{\sqrt{\alpha_t}(1-\bar\alpha_t)}x_t\\ &\quad-\frac{\beta_t}{\sqrt{\alpha_t}\sqrt{1-\bar\alpha_t}}\,\epsilon\\ &=\frac1{\sqrt{\alpha_t}}\left(x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\,\epsilon\right).\end{aligned}

最后一步使用了 βt+αt(1−αˉt−1)=1−αˉt\beta_t+\alpha_t(1-\bar\alpha_{t-1})=1-\bar\alpha_t。

这个式子告诉我们:如果知道从 x0x_0 到 xtx_t 的合成噪声,就能确定这个后验均值。这就是为什么 DDPM 可以加入一个神经网络,去拟合这份噪声。

生成时没有真实的 ϵ\epsilon,我们使用网络的预测 ϵθ(xt,t)\epsilon_\theta(x_t,t),定义反向模型的均值:

μθ(xt,t)=1αt(xt−1−αt1−αˉt ϵθ(xt,t)).\mu_\theta(x_t,t)= \frac1{\sqrt{\alpha_t}}\left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\,\epsilon_\theta(x_t,t)\right).
pθ(xt−1∣xt)=N ⁣(xt−1;μθ(xt,t),σt2I).p_\theta(x_{t-1}\mid x_t)=\mathcal N\!\left(x_{t-1};\mu_\theta(x_t,t),\sigma_t^2I\right).

至此,我最开始的疑问就接上了:前向过程为网络构造噪声标签,反向过程则利用预测出的噪声来计算每一步的均值。两边用到的是同一种参数化,体现了前后过程设计的严密。

#训练与采样

现在再看原文中的两段伪代码,就能理解为什么训练时一步加噪,生成时却要逐步去噪。

Algorithm 1 · Training

  1. 从训练集采样一张原图:x0∼qdatax_0\sim q_{\mathrm{data}}。
  2. 随机选择时间步:t∼Uniform({1,…,T})t\sim\mathrm{Uniform}(\{1,\ldots,T\})。
  3. 采样已知的高斯噪声:ϵ∼N(0,I)\epsilon\sim\mathcal N(0,I)。
  4. 直接构造 xt=αˉtx0+1−αˉtϵx_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon,对下面的损失做一次梯度下降:
    ∇θ∥ϵ−ϵθ ⁣(αˉtx0+1−αˉtϵ,t)∥22.\nabla_\theta\left\|\epsilon-\epsilon_\theta\!\left(\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon,t\right)\right\|_2^2.
  5. 重复以上步骤,直到训练完成。

对应的简化训练目标就是预测噪声的均方误差:

Lsimple(θ)=Ex0,t,ϵ[∥ϵ−ϵθ ⁣(αˉtx0+1−αˉtϵ,t)∥22].\mathcal L_{\mathrm{simple}}(\theta)= \mathbb E_{x_0,t,\epsilon}\left[\left\|\epsilon-\epsilon_\theta\!\left(\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon,t\right)\right\|_2^2\right].

不需要把整条前向链都走一遍;只采样一个时间步,构造当前的带噪图像和标签就可以训练。

Algorithm 2 · Sampling

  1. 从标准正态分布采样初始图像:xT∼N(0,I)x_T\sim\mathcal N(0,I)。
  2. 依次遍历 t=T,T−1,…,1t=T,T-1,\ldots,1,每一步执行下面的更新:
    z∼N(0,I) if t>1,z=0 if t=1.z\sim\mathcal N(0,I)\ \text{if }t>1,\qquad z=0\ \text{if }t=1.
    xt−1=1αt(xt−1−αt1−αˉtϵθ(xt,t))+σtz.x_{t-1}=\frac1{\sqrt{\alpha_t}} \left(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t,t)\right)+\sigma_tz.
  3. 返回最后得到的 x0x_0。

这里 σt\sigma_t 是反向采样的标准差。原论文讨论了 σt2=βt\sigma_t^2=\beta_t 和 σt2=β~t\sigma_t^2=\tilde\beta_t 两种固定设置;它与预测的噪声 ϵθ\epsilon_\theta 不是一回事。最后一步按伪代码令 z=0z=0,不再加入新的随机噪声。

训练时,噪声来自我们自己的采样,所以标签已知;生成时,先从随机噪声出发,再由网络预测噪声、计算均值、逐步更新图像。这样再回看 DDPM 的公式,前向和反向过程就串起来了。

#参考与延伸阅读

  1. Jonathan Ho, Ajay Jain & Pieter Abbeel. Denoising Diffusion Probabilistic Models, NeurIPS 2020. 前向闭式分布、反向后验及训练与采样算法。
  2. 苏剑林,原文引用的扩散模型推导,科学空间。本文沿用原笔记的贝叶斯与高斯配方思路,并统一为 DDPM 论文的记号。
  3. 系列学习 03:Score-based Generative Models——从梯度角度理解加噪去噪。