DDPM

预备知识:

  1. 贝叶斯公式

最基本的贝叶斯公式是:

P(AB)=P(BA)P(A)P(B)P(A|B)=\frac{P(B|A)P(A)}{P(B)}

这里并不要求 AABB 独立,只需要 P(B)>0P(B)>0。如果 AABB 独立,那么 P(AB)=P(A)P(A|B)=P(A),反而就没有“用 BB 更新对 AA 的认识”的意义了。

追加一个条件 CC 后,贝叶斯公式写成:

P(AB,C)=P(BA,C)P(AC)P(BC)P(A|B,C)=\frac{P(B|A,C)P(A|C)}{P(B|C)}

它的推导可以直接从条件概率定义出发。先把“在 CC 已经发生的前提下”看作新的样本空间,此时所有概率都变成条件在 CC 上的概率。于是:

P(AB,C)=P(A,BC)P(BC)P(A|B,C)=\frac{P(A,B|C)}{P(B|C)}

这里的意思是:在 CC 已知的情况下,BB 中同时满足 AA 的比例是多少。

同理,另一个条件概率也可以写成:

P(BA,C)=P(A,BC)P(AC)P(B|A,C)=\frac{P(A,B|C)}{P(A|C)}

把这个式子移项,可以得到:

P(A,BC)=P(BA,C)P(AC)P(A,B|C)=P(B|A,C)P(A|C)

再把它代回第一个式子:

P(AB,C)=P(A,BC)P(BC)=P(BA,C)P(AC)P(BC)P(A|B,C) =\frac{P(A,B|C)}{P(B|C)} =\frac{P(B|A,C)P(A|C)}{P(B|C)}

这就是带条件 CC 的贝叶斯公式。它和普通贝叶斯公式的结构完全一样,只是每一项都额外限定在 CC 这个条件下。

也可以用联合概率形式再看一遍。由条件概率定义:

P(AB,C)=P(A,B,C)P(B,C)P(A|B,C)=\frac{P(A,B,C)}{P(B,C)}

而:

P(A,B,C)=P(BA,C)P(A,C)=P(BA,C)P(AC)P(C)P(A,B,C)=P(B|A,C)P(A,C)=P(B|A,C)P(A|C)P(C)

P(B,C)=P(BC)P(C)P(B,C)=P(B|C)P(C)

所以:

P(AB,C)=P(BA,C)P(AC)P(C)P(BC)P(C)=P(BA,C)P(AC)P(BC)P(A|B,C) =\frac{P(B|A,C)P(A|C)P(C)}{P(B|C)P(C)} =\frac{P(B|A,C)P(A|C)}{P(B|C)}

这个推导里需要 P(C)>0P(C)>0P(BC)>0P(B|C)>0,否则条件概率没有定义。

在 DDPM 里,我们会反复使用这个形式,例如根据已知的 x0x_0xtx_t 去推导中间变量 xt1x_{t-1} 的分布。

  1. 两个高斯分布相加

如果 XN(μ1,σ12I)X\sim \mathcal N(\mu_1,\sigma_1^2 I)YN(μ2,σ22I)Y\sim \mathcal N(\mu_2,\sigma_2^2 I),并且 X,YX,Y 相互独立,那么:

X+YN(μ1+μ2,(σ12+σ22)I)X+Y\sim \mathcal N(\mu_1+\mu_2,(\sigma_1^2+\sigma_2^2)I)

更一般地,如果 Z=aX+bYZ=aX+bY,那么:

ZN(aμ1+bμ2,(a2σ12+b2σ22)I)Z\sim \mathcal N(a\mu_1+b\mu_2,(a^2\sigma_1^2+b^2\sigma_2^2)I)

这个性质会让加噪过程的多步递推变得非常干净。

简介

DDPM 是 diffusion 模型的开山之作。它的主要思想是在一张图像上进行多步加噪,让模型学习反向去噪。

DDPM图示

这是模型反向去噪的过程,从一张纯噪声图,逐步去噪还原成模型“认为”的原始图像。

数学推导

加噪过程

DDPM 有很多巧妙的假设,它假设模型加的都是高斯噪声。高斯噪声的意思就是噪声图像中的每一个像素点都符合给定的高斯分布。这个假设的好处是,模型去噪的时候也只需要预测高斯噪声,让预测的分布简单了很多。另外,它假设加噪过程是一个马尔科夫链,当前步的加噪结果只依赖于上一步加噪得到的图像(在上一步的图像上叠加高斯噪声)。同时,模型不只加噪声,还会弱化图像本身的信号,也就是说最终图像一定会变成纯噪声。这个设计也很巧妙,等下公式推导时会体现出来。

公式:

q(xtxt1)=N(αtxt1,(1αt)I)q(x_t|x_{t-1})=\mathcal N(\sqrt{\alpha_t}x_{t-1},(1-\alpha_t)I)

这个公式代表从 xt1xtx_{t-1}\to x_t 的加噪分布。公式里的 αt\alpha_t 是人为设置的参数。
通常也会记 βt=1αt\beta_t=1-\alpha_t,因此也可以写成:

q(xtxt1)=N(1βtxt1,βtI)q(x_t|x_{t-1})=\mathcal N(\sqrt{1-\beta_t}x_{t-1},\beta_t I)

用重参数化的方式表达,就是:

xt=αtxt1+1αtϵt,ϵtN(0,I)x_t=\sqrt{\alpha_t}x_{t-1}+\sqrt{1-\alpha_t}\epsilon_t,\quad \epsilon_t\sim\mathcal N(0,I)

这里的 αt\sqrt{\alpha_t} 用来保留上一张图像的信号,1αt\sqrt{1-\alpha_t} 用来加入新的高斯噪声。

当我们考虑 xt2xtx_{t-2}\to x_t 的加噪分布时,可以做一步替换:

xt=αtxt1+1αtϵt=αt(αt1xt2+1αt1ϵt1)+1αtϵt=αtαt1xt2+αt(1αt1)ϵt1+1αtϵt\begin{aligned} x_t &=\sqrt{\alpha_t}x_{t-1}+\sqrt{1-\alpha_t}\epsilon_t\\ &=\sqrt{\alpha_t}(\sqrt{\alpha_{t-1}}x_{t-2}+\sqrt{1-\alpha_{t-1}}\epsilon_{t-1})+\sqrt{1-\alpha_t}\epsilon_t\\ &=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2} +\sqrt{\alpha_t(1-\alpha_{t-1})}\epsilon_{t-1} +\sqrt{1-\alpha_t}\epsilon_t \end{aligned}

后面两项都是相互独立的高斯噪声的线性组合,所以仍然是高斯噪声。它的方差为:

αt(1αt1)+(1αt)=1αtαt1\alpha_t(1-\alpha_{t-1})+(1-\alpha_t) =1-\alpha_t\alpha_{t-1}

因此可以把两项合并成一个新的标准高斯噪声 ϵ\epsilon

xt=αtαt1xt2+1αtαt1ϵx_t=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2}+\sqrt{1-\alpha_t\alpha_{t-1}}\epsilon

也就是:

q(xtxt2)=N(αtαt1xt2,(1αtαt1)I)q(x_t|x_{t-2})=\mathcal N(\sqrt{\alpha_t\alpha_{t-1}}x_{t-2},(1-\alpha_t\alpha_{t-1})I)

我们惊奇地发现,化简出的式子,是一个类似先前分布的分布,除了参数不一样。参数变成了乘积。如果以此类推,可以一直推到 x0x_0 的情况,此时的参数就变成了连乘:

αˉt=i=1tαi\bar\alpha_t=\prod_{i=1}^t\alpha_i

于是有:

q(xtx0)=N(αˉtx0,(1αˉt)I)q(x_t|x_0)=\mathcal N(\sqrt{\bar\alpha_t}x_0,(1-\bar\alpha_t)I)

对应的重参数化形式为:

xt=αˉtx0+1αˉtϵ,ϵN(0,I)x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon,\quad \epsilon\sim\mathcal N(0,I)

这启示我们,任意一步加噪得到的结果,都符合原图像有关的正态分布!

这个结论非常重要:训练时我们不需要真的从 x0x_0 一步步加噪到 xtx_t,而是可以直接随机采样一个 tt,再用上面的闭式公式一次性得到 xtx_t

参数设置

参数设置的核心是决定每一步加多少噪声,也就是选择 βt\beta_tαt\alpha_t。早期 DDPM 常用线性调度:让 βt\beta_t 从一个很小的值线性增加到较大的值。这样实现简单,但问题是图像信号的保留量 αˉt\bar\alpha_t 不一定按“感知上均匀”的速度下降:有些阶段噪声增长太快,图像很早就接近纯噪声,后面的时间步训练价值变低。

更常见的改进是余弦调度。它不直接线性设置 βt\beta_t,而是先设计 αˉt\bar\alpha_t 的变化:

f(t)=cos2(t/T+s1+sπ2)f(t)=\cos^2\left(\frac{t/T+s}{1+s}\cdot\frac{\pi}{2}\right)

αˉt=f(t)f(0)\bar\alpha_t=\frac{f(t)}{f(0)}

其中 TT 是总步数,ss 是一个很小的偏移量,常用 s=0.008s=0.008。然后再由相邻两步的 αˉ\bar\alpha 反推出:

βt=1αˉtαˉt1\beta_t=1-\frac{\bar\alpha_t}{\bar\alpha_{t-1}}

实际实现里通常会对 βt\beta_t 做上限裁剪,避免某一步噪声过大导致数值不稳定。

余弦调度的直觉是:前期慢慢加噪,让模型看到更多仍然保留结构的样本;中间阶段平滑过渡;后期再逐渐接近纯噪声。这样每个时间步都更有训练价值。

调度方式 做法 优点 缺点
线性调度 βt\beta_t 线性增大 简单,容易实现 αˉt\bar\alpha_t 下降不够平滑,后期可能浪费时间步
二次/平方调度 按非线性方式增大 βt\beta_t 比线性更灵活 仍然需要手动调形状
余弦调度 直接让 αˉt\bar\alpha_t 按余弦曲线下降 信号衰减更平滑,训练通常更稳定 公式稍复杂,需要从 αˉt\bar\alpha_t 反推 βt\beta_t

去噪过程

考虑到以上加噪的过程,去噪过程该如何表示呢?DDPM 的做法是假设去噪过程中对应减去的噪声,也服从高斯分布。当我们从 xtx_t 去噪到 xt1x_{t-1} 时,对应的分布应当是:

pθ(xt1xt)=N(μθ(xt,t),σt2I)p_\theta(x_{t-1}|x_t)=\mathcal N(\mu_\theta(x_t,t),\sigma_t^2I)

其中 μθ\mu_\theta 由神经网络预测,σt2\sigma_t^2 可以固定为某个与 βt\beta_t 有关的值,也可以让模型学习。为了推导这个均值应该长什么样,我们先看真实后验分布:

q(xt1xt,x0)q(x_{t-1}|x_t,x_0)

为什么均值、方差和 x0x_0xtx_t 有关呢?我们首先来考虑下面这个式子:

q(xt1xt,x0)=q(xtxt1,x0)q(xt1x0)q(xtx0)q(x_{t-1}|x_t,x_0)=\frac{q(x_t|x_{t-1},x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)}

由于加噪过程是马尔科夫链,xtx_t 在给定 xt1x_{t-1} 后不再依赖 x0x_0,所以上式可以写成:

q(xt1xt,x0)=q(xtxt1)q(xt1x0)q(xtx0)q(x_{t-1}|x_t,x_0)=\frac{q(x_t|x_{t-1})q(x_{t-1}|x_0)}{q(x_t|x_0)}

这三个分布我们都已经知道:

q(xtxt1)=N(αtxt1,βtI)q(x_t|x_{t-1})=\mathcal N(\sqrt{\alpha_t}x_{t-1},\beta_tI)

q(xt1x0)=N(αˉt1x0,(1αˉt1)I)q(x_{t-1}|x_0)=\mathcal N(\sqrt{\bar\alpha_{t-1}}x_0,(1-\bar\alpha_{t-1})I)

q(xtx0)=N(αˉtx0,(1αˉt)I)q(x_t|x_0)=\mathcal N(\sqrt{\bar\alpha_t}x_0,(1-\bar\alpha_t)I)

把高斯分布的指数项展开,只保留和 xt1x_{t-1} 有关的部分:

logq(xt1xt,x0)=12((xtαtxt1)2βt+(xt1αˉt1x0)21αˉt1)+C\begin{aligned} \log q(x_{t-1}|x_t,x_0) &= -\frac{1}{2}\left( \frac{(x_t-\sqrt{\alpha_t}x_{t-1})^2}{\beta_t} +\frac{(x_{t-1}-\sqrt{\bar\alpha_{t-1}}x_0)^2}{1-\bar\alpha_{t-1}} \right)+C \end{aligned}

这是一个关于 xt1x_{t-1} 的二次型,因此后验仍然是高斯分布:

q(xt1xt,x0)=N(μ~t(xt,x0),β~tI)q(x_{t-1}|x_t,x_0)=\mathcal N(\tilde\mu_t(x_t,x_0),\tilde\beta_tI)

其中:

β~t=1αˉt11αˉtβt\tilde\beta_t=\frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\beta_t

μ~t(xt,x0)=αˉt1βt1αˉtx0+αt(1αˉt1)1αˉtxt\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

可以发现最终的式子中的变量包含 x0x_0xtx_t。去噪时,我们会从 xtx_t 转化到 xt1x_{t-1}xtx_t 是已知的。但是显然 x0x_0 是未知的,这下该怎么办?这里就是去噪模型大放异彩的地方,我们可以根据当前去噪的信息,通过模型估计一个 x0x_0,再带回到式子中,这样就能保证所有的量都是确定的,进而采样 xt1x_{t-1}

Tips: 在反向去噪前期,对 x0x_0 的预测是粗略的,往往会丢失掉很多细节。我们并不能一步到位,通过早期预测的 x0x_0 直接生成图像,这里的 x0x_0 只是指明了去噪的大致方向。

但是 DDPM 实际操作中并没有选择这么做。 x0x_0 的分布是复杂的、难拟合的。上面的 Tips 也说明直接拟合 x0x_0 是一种反直觉的操作。实际上,DDPM 选择拟合高斯噪声。考虑在公式中把 x0x_0 替换成噪声 ϵ\epsilon

由前面的闭式加噪公式:

xt=αˉtx0+1αˉtϵx_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon

可以反解出:

x0=1αˉt(xt1αˉtϵ)x_0=\frac{1}{\sqrt{\bar\alpha_t}}\left(x_t-\sqrt{1-\bar\alpha_t}\epsilon\right)

把它代入 μ~t(xt,x0)\tilde\mu_t(x_t,x_0)

μ~t=αˉt1βt1αˉtxt1αˉtϵαˉt+αt(1αˉt1)1αˉtxt\begin{aligned} \tilde\mu_t &= \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1-\bar\alpha_t} \cdot \frac{x_t-\sqrt{1-\bar\alpha_t}\epsilon}{\sqrt{\bar\alpha_t}} +\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}x_t \end{aligned}

利用 αˉt=αtαˉt1\bar\alpha_t=\alpha_t\bar\alpha_{t-1} 继续化简,可以得到更常见的形式:

μ~t(xt,ϵ)=1αt(xtβt1αˉtϵ)\tilde\mu_t(x_t,\epsilon)= \frac{1}{\sqrt{\alpha_t}} \left( x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon \right)

既然真实噪声 ϵ\epsilon 在采样时不知道,就让模型预测它:

ϵθ=ϵθ(xt,t)\epsilon_\theta=\epsilon_\theta(x_t,t)

于是反向过程的均值写成:

μθ(xt,t)=1αt(xtβt1αˉtϵθ(xt,t))\mu_\theta(x_t,t)= \frac{1}{\sqrt{\alpha_t}} \left( x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t,t) \right)

训练目标也可以写得非常简单:

Lsimple=Et,x0,ϵ[ϵϵθ(xt,t)2]L_{\text{simple}}= \mathbb E_{t,x_0,\epsilon} \left[ \left\|\epsilon-\epsilon_\theta(x_t,t)\right\|^2 \right]

也就是说,随机抽一张原图 x0x_0、随机抽一个时间步 tt、随机抽一个标准高斯噪声 ϵ\epsilon,构造出 xtx_t,然后让模型预测当初加进去的那份噪声。

最终效果是一样的,而且每个时间步噪声的分布相比原始图像的分布是更简单、更容易拟合的。

Score-based 视角下的 DDPM

从 score-based model 的角度看,DDPM 做的事情也可以理解为学习每个噪声等级下的数据分布梯度,也就是 score:

xtlogq(xt)\nabla_{x_t}\log q(x_t)

具体可以分成几步理解:

  1. score 表示“往高概率区域走”的方向

    对一个概率分布 p(x)p(x) 来说,xlogp(x)\nabla_x\log p(x) 指向概率密度上升最快的方向。生成时如果我们知道这个方向,就可以从噪声样本逐步移动到更像真实数据的位置。

  2. DDPM 的噪声预测和 score 等价

    对条件分布 q(xtx0)=N(αˉtx0,(1αˉt)I)q(x_t|x_0)=\mathcal N(\sqrt{\bar\alpha_t}x_0,(1-\bar\alpha_t)I),它的 score 是:

    xtlogq(xtx0)=xtαˉtx01αˉt\nabla_{x_t}\log q(x_t|x_0) =-\frac{x_t-\sqrt{\bar\alpha_t}x_0}{1-\bar\alpha_t}

    又因为:

    xtαˉtx0=1αˉtϵx_t-\sqrt{\bar\alpha_t}x_0=\sqrt{1-\bar\alpha_t}\epsilon

    所以:

    xtlogq(xtx0)=ϵ1αˉt\nabla_{x_t}\log q(x_t|x_0) =-\frac{\epsilon}{\sqrt{1-\bar\alpha_t}}

    因此预测噪声 ϵθ(xt,t)\epsilon_\theta(x_t,t),本质上就等价于预测 score:

    sθ(xt,t)=ϵθ(xt,t)1αˉts_\theta(x_t,t)=-\frac{\epsilon_\theta(x_t,t)}{\sqrt{1-\bar\alpha_t}}

  3. 时间步对应不同噪声等级

    tt 越大,噪声越强,分布越接近标准高斯;tt 越小,分布越接近真实图像。模型需要在所有噪声等级上都学会“应该往哪里去”,所以训练时必须把 tt 输入给网络。

  4. 采样就是沿着 score 逐步回到数据流形

    从纯噪声开始,模型在每个时间步预测当前样本中的噪声,也就给出了一个往高概率图像区域移动的方向。DDPM 的反向采样公式可以看成是在这个方向上走一步,并保留适量随机性。

这个视角的好处是,它把 DDPM 和后来的 score matching、SDE diffusion、Langevin dynamics 等方法统一了起来:它们都在学习不同噪声强度下的 score,只是参数化方式和采样器不同。

模型结构

DDPM 原始论文的实现采用了 U-Net 作为去噪的模型。当然不是传统的 U-Net,相比传统的 U-Net 有以下的修改:

  1. 输入不是普通图像,而是带噪图像 xtx_t
  2. 网络必须额外接收时间步 tt,因为不同时间步的噪声强度不同。
  3. 主体使用残差块,而不是最朴素的卷积块。
  4. 在低分辨率层加入 attention,让模型能建模远距离像素关系。
  5. 输出通常不是去噪后的图像,而是噪声预测 ϵθ(xt,t)\epsilon_\theta(x_t,t)

展开来看:

  1. 整体 U-Net 结构

    U-Net 由下采样路径、中间瓶颈层、上采样路径组成。下采样路径不断降低分辨率、增加通道数,用来提取高级语义;上采样路径逐步恢复分辨率,用来输出和输入同尺寸的噪声图。对应层之间会有 skip connection,把浅层的空间细节直接传到上采样阶段。

  2. 残差块

    DDPM U-Net 没有采用简单的 Conv -> ReLU -> Conv 结构,而是使用了 ResNet block。残差连接可以让深层网络更容易训练,也能让模型在不同噪声等级下稳定学习。一个典型块大致是:

    h=Conv(Act(Norm(x)))h=\text{Conv}(\text{Act}(\text{Norm}(x)))

    然后把时间嵌入加进去,再经过下一组 Norm -> Act -> Conv,最后和输入做残差相加:

    ResBlock(x,t)=x+Conv2(Act(Norm(h+TimeProj(t))))\text{ResBlock}(x,t)=x+\text{Conv}_2(\text{Act}(\text{Norm}(h+\text{TimeProj}(t))))

    但是纯粹的残差连接产生了一个问题,输入通道数和输出通道数不同,残差分支不能直接相加。解决方法是额外用一个 1×11\times1 卷积把 xx 投影到相同通道数:

    ResBlock(x,t)=Conv1×1(x)+F(x,t)\text{ResBlock}(x,t)=\text{Conv}_{1\times1}(x)+F(x,t)

    这里的归一化和激活函数不是随便选的,它们对扩散模型很重要。

    归一化:通常使用 GroupNorm

    DDPM U-Net 里常见的归一化是 GroupNorm,而不是 BatchNorm。GroupNorm 会把通道分成若干组,在每个样本内部、每个组内计算均值和方差。假设输入是 xRB×C×H×Wx\in\mathbb R^{B\times C\times H\times W},把 CC 个通道分成 GG 组后,对每个样本的每一组做:

    x^=xμgroupσgroup2+ϵ\hat x=\frac{x-\mu_{\text{group}}}{\sqrt{\sigma_{\text{group}}^2+\epsilon}}

    再通过可学习参数恢复表达能力:

    y=γx^+βy=\gamma \hat x+\beta

    选择 GroupNorm 的原因主要有三个:

    1. 扩散模型训练很吃显存,batch size 经常比较小。BatchNorm 依赖 batch 维度统计量,小 batch 下均值和方差会不稳定;GroupNorm 不依赖 batch 维度,所以更稳定。
    2. 每个样本的时间步 tt 可能不同,噪声强度也不同。BatchNorm 会把不同噪声等级的样本混在一起统计,可能削弱时间条件;GroupNorm 在样本内部归一化,更适合这种“同一个 batch 里噪声等级不同”的训练方式。
    3. 采样时通常是单张或少量图片生成。BatchNorm 的训练/推理统计切换会带来额外问题,而 GroupNorm 在训练和推理时行为一致。

    激活函数:通常使用 SiLU / Swish

    DDPM U-Net 里常见的激活函数是 SiLU,也常被称作 Swish:

    SiLU(x)=xσ(x)\text{SiLU}(x)=x\cdot\sigma(x)

    其中 σ(x)\sigma(x) 是 sigmoid 函数。它和 ReLU 最大的区别是:SiLU 是平滑的,并且在负数区域不会直接变成 0,而是保留一个较小的负值响应。

    这对扩散模型有几个好处:

    1. 去噪网络本质上在拟合连续的噪声预测函数 ϵθ(xt,t)\epsilon_\theta(x_t,t)。SiLU 比 ReLU 更平滑,输出变化不会在 0 点突然截断,更适合拟合这种连续映射。
    2. ReLU 会把所有负值的梯度直接置零,会丢掉一部分细微信号。扩散模型的中间特征包含噪声、纹理和边缘信息,弱响应也可能有用,SiLU 对这些信息更温和。

    所以一个更贴近实际的残差块可以理解成:

    h=Conv1(SiLU(GroupNorm(x)))h=\text{Conv}_1(\text{SiLU}(\text{GroupNorm}(x)))

    h=h+Linear(TimeEmbedding(t))h=h+\text{Linear}(\text{TimeEmbedding}(t))

    F(x,t)=Conv2(SiLU(GroupNorm(h)))F(x,t)=\text{Conv}_2(\text{SiLU}(\text{GroupNorm}(h)))

    out=x+F(x,t)\text{out}=x+F(x,t)

    简单总结:GroupNorm 负责让特征尺度稳定,SiLU 负责提供平滑且保留弱信号的非线性,残差连接负责让深层 U-Net 容易训练。这三者组合起来,才构成了 DDPM 里常见的 ResNet block。

  3. 时间步注入

    时间步 tt 是 DDPM U-Net 最关键的额外条件。因为同一个 xtx_t 在不同时间步代表的噪声强度不同,模型如果不知道 tt,就不知道应该去掉多少噪声。

    常见做法是先把整数时间步 tt 转成 sinusoidal position embedding,类似 Transformer 的位置编码:

    emb(t)=[sin(ω1t),cos(ω1t),,sin(ωkt),cos(ωkt)]\text{emb}(t)=[\sin(\omega_1t),\cos(\omega_1t),\ldots,\sin(\omega_kt),\cos(\omega_kt)]

    然后通过 MLP 得到时间向量。这个时间向量会被投影到各个残差块所需的通道维度,再加到 feature map 上。形式上可以理解为:

    h=h+Linear(TimeEmbedding(t))h = h + \text{Linear}(\text{TimeEmbedding}(t))

    实现时会把时间向量 reshape 成 [batch, channels, 1, 1],再广播到整张特征图。这样每个空间位置都知道当前处于哪个噪声阶段。

  4. Attention 模块

    卷积擅长局部建模,但图像里经常存在远距离依赖,比如左右眼、物体边缘、整体构图。如果只靠卷积,信息需要经过很多层才能从图像的一端传到另一端;self-attention 则可以让任意两个空间位置直接交互。

    假设某一层的特征图形状是:

    hRB×C×H×Wh\in\mathbb R^{B\times C\times H\times W}

    其中 BB 是 batch size,CC 是通道数,H,WH,W 是特征图尺寸。做 attention 前,通常先把空间维度展平:

    hhRB×N×C,N=HWh\to h'\in\mathbb R^{B\times N\times C},\quad N=H\cdot W

    这样每一个空间位置都可以看作一个 token。接着通过 1×11\times1 卷积生成 Query、Key、Value:

    Q=hWQ,K=hWK,V=hWVQ=h'W_Q,\quad K=h'W_K,\quad V=h'W_V

    self-attention 的核心计算是:

    Attention(Q,K,V)=softmax(QKd)V\text{Attention}(Q,K,V)= \text{softmax}\left(\frac{QK^\top}{\sqrt d}\right)V

    这里 QKQK^\top 会得到一个 N×NN\times N 的矩阵,表示每个位置和其它所有位置的相关程度。softmax 后,每个位置都会根据相关性从所有位置聚合信息。比如生成一只眼睛时,模型可以直接参考另一只眼睛的位置;生成物体边缘时,也可以参考远处的整体轮廓。

    在 U-Net 中,attention 通常不会放在所有分辨率层。原因是计算量和显存开销与 N2=(HW)2N^2=(H\cdot W)^2 成正比。如果特征图是 64×6464\times64,那么 token 数是 4096,attention 矩阵就有一千多万个元素;如果特征图降到 16×1616\times16,token 数只有 256,计算量会小很多。因此 DDPM 常在低分辨率层,例如 16×1616\times168×88\times8 的位置加入 attention。

    attention 模块同样也有归一化和残差连接残差连接:

    hout=h+Attention(Norm(h))h_{\text{out}}=h+\text{Attention}(\text{Norm}(h))

  5. 输出层

    输出张量和输入图像 xtx_t 形状相同。如果输入是 RGB 图像,那么输出也是 3 个通道;如果输入是 latent diffusion 里的 latent,那么输出就是 latent 的通道数。输出值表示模型认为当前 xtx_t 中包含的噪声 ϵθ(xt,t)\epsilon_\theta(x_t,t)

后续也有多种改进或替换,例如:

  1. Improved DDPM:改进噪声调度、方差学习和训练目标,让采样质量和效率更好。
  2. DDIM:把采样过程改成确定性或半确定性形式,可以用更少步数生成图像。
  3. Latent Diffusion / Stable Diffusion:不在像素空间扩散,而是在 VAE 的 latent 空间扩散,大幅降低计算成本。
  4. DiT:用 Transformer 替代 U-Net 主体结构,把图像 patch 或 latent patch 当成 token 来建模。

代码

huggingface 提供了有关 DDPM U-Net 实现的代码,非常详细、清晰。


电波交流