概述

去噪扩散模型全称为 Denoising Diffusion Probabilistic Models，简称 DDPM 或 Diffusion 模型（下简称 Diffusion 模型）。相比于 GAN, VAE 等模型，其核心优势在于:通过学习训练数据的真实分布，来生成与训练数据风格、质量一致的逼真样本。

Diffusion 通过学习训练数据（比如大量真实图片）的概率分布，最终能从纯噪声出发，生成出分布与训练数据高度相似的逼真图片。在生成图片的过程中，Diffusion 模型主要有以下两个流程

  1. 从高随机性的噪声分布中采样
  2. 将这张纯噪声图进行处理，这一步取决于去噪模型(Denoise Model)

Diffusion 模型通常被描述为一个马尔可夫链，这个马尔可夫链包括前向过程和反向过程。前向过程是有具体的表达式可以计算的，后向过程是利用神经网络来学习的。和 GAN, VAE 类似，Diffusion 模型想要通过某种方式估计给定数据样本 {x(i)}\{x^{(i)}\} 服从的概率分布，因此我们将花费更多笔墨介绍这一部分。

基础模型

反向过程

用 x0\boldsymbol{x}_0 表示生成的目标样本，我们定义一系列与 x0\boldsymbol{x}_0 维数相同的随机变量 x1,⋯ ,xT\boldsymbol{x}_1,\cdots,\boldsymbol{x}_T，其表示从 x0\boldsymbol{x}_0 开始不断添加随机高斯噪声，最终得到 xT∼p(xT)=N(xT;0,I)\boldsymbol{x}_T\sim p(\boldsymbol{x}_T)=N(\boldsymbol{x}_T;\boldsymbol{0},\boldsymbol{I})。这一过程称为前向过程或扩散过程，我们在下一小节再详细介绍。

设 x0\boldsymbol{x}_0 的先验分布为 q(x0)q(\boldsymbol{x}_{0})，我们希望用参数为 θ\theta 的神经网络估计 x0\boldsymbol{x}_0 的概率密度函数，记为 pθ(x0)p_{\theta}(\boldsymbol{x}_0)。这可以看做对联合概率密度 pθ(x0:T)p_{\theta}(\boldsymbol{x}_{0:T}) 求变量 x0\boldsymbol{x}_0 的边缘概率密度。由概率论知识有

pθ(x0)=∫pθ(x0:T)dx1:Tp_{\theta}(\boldsymbol{x}_0)=\int p_{\theta}(\boldsymbol{x}_{0:T})\mathrm{d}\boldsymbol{x}_{1:T}

其中的 dx1:T\mathrm{d}\boldsymbol{x}_{1:T} 可以被视为 dx1⋯dxT\mathrm{d}\boldsymbol{x}_{1}\cdots\mathrm{d}\boldsymbol{x}_{T}。接下来，我们将 pθ(x0:T)p_{\theta}(\boldsymbol{x}_{0:T}) 定义为一个马尔科夫链，其转移方程被描述为高斯过程。具体形式可以由条件概率乘积链表示

pθ(x0:T)=p(xT)∏t=1Tpθ(xt−1∣xt)p_{\theta}(\boldsymbol{x}_{0:T})=p(\boldsymbol{x}_T)\prod_{t=1}^Tp_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)

这一过程称为反向过程，因为其描述了从噪音 xT\boldsymbol{x}_T 恢复为 x0\boldsymbol{x}_0 的过程。在恢复的每一步中，我们可以证明条件概率仍然服从高斯分布，并且其均值，协方差来自上一步由神经网络 θ\theta 给出的结果

pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)=N\big(\boldsymbol{x}_{t-1};\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t,t),\boldsymbol{\Sigma}_{\theta}(\boldsymbol{x}_t,t)\big)

最终，整个 Diffusion 模型的前向过程和反向过程可以被表示为下图。

[{"url":"https://pic3.zhimg.com/v2-8f48241c47d13649d11a471ae53881fc_1440w.jpg","alt":"Figure 2","title":""}]

前向过程

现在来看前向过程的概率分布，用记号 qq 表示人为规定的高斯噪声添加过程。整个前向过程 q(x1:T∣x0)q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0) 仍然被定义为一个马尔科夫链，表示从已知样本 x0\boldsymbol{x}_0 得到纯高斯噪声 xT\boldsymbol{x}_T，因此有

q(x1:T∣x0)=∏t=1Tq(xt∣xt−1)q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)=\prod_{t=1}^Tq(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1})

指定标量序列 β1,⋯ ,βT\beta_1,\cdots,\beta_T，令前向过程的单步转移满足

q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)q(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1})=N(\boldsymbol{x}_t;\sqrt{1-\beta_t}\boldsymbol{x}_{t-1},\beta_t\boldsymbol{I})

前向过程具有一个特殊的性质：我们可以借助 x0\boldsymbol{x}_0 对任意 xt\boldsymbol{x}_t 进行采样。定义 αt=1−βt\alpha_t=1-\beta_t，他们的所有前步乘积表示为

αˉt=∏s=1tαs\bar{\alpha}_t=\prod_{s=1}^t\alpha_s

利用高斯分布的可加性，可以证明前向传播的任意采样满足

q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)q(\boldsymbol{x}_t|\boldsymbol{x}_0)=N\big(\boldsymbol{x}_t;\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0,(1-\bar{\alpha}_t)\boldsymbol{I}\big)

原文指出了序列 β1,⋯ ,βT\beta_1,\cdots,\beta_T 的意义：将 βt\beta_t 从 β1=10−4\beta_1=10^{-4} 线性增长到 βT=0.02\beta_T=0.02，当 βt\beta_t 越大时 αt\alpha_t 和 αˉt\bar{\alpha}_t 越小，此时代入重参数化采样得到

xt=αˉtx0+1−αˉtϵ∼approxN(xT;0,I)=p(xT)\boldsymbol{x}_t=\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\epsilon\overset{\mathrm{approx}}{\sim}N(\boldsymbol{x}_T;\boldsymbol{0},\boldsymbol{I})=p(\boldsymbol{x}_T)

因此噪声公式能够从慢到快地改变原图像，让噪声 xT\boldsymbol{x}_T 最终服从标准高斯分布。

KL散度优化

训练目标在于对 pθ(x0)p_{\theta}(\boldsymbol{x}_0) 进行最大似然估计，这可以转化为负对数优化问题

min⁡θ−log⁡pθ(x0)\min_{\theta}-\log p_{\theta}(\boldsymbol{x}_0)

对于前向过程分布 qq，可以将上式的估计内容改写为以 qq 为基准分布的期望

−log⁡pθ(x0)=−log⁡∑x1:Tq(x1:T∣x0)pθ(x0)q(x1:T∣x0)=−log⁡Eq[pθ(x0)q(x1:T∣x0)]-\log p_{\theta}(\boldsymbol{x}_0)=-\log \sum_{\boldsymbol{x}_{1:T}}q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)\frac{p_{\theta}(\boldsymbol{x}_0)}{q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}=-\log\mathbb{E}_q\left[\frac{p_{\theta}(\boldsymbol{x}_0)}{q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}\right]

利用对数函数的凸性，由Jensen不等式得到

−log⁡Eq[pθ(x0)q(x1:T∣x0)]⩽Eq[−log⁡pθ(x0)q(x1:T∣x0)]-\log\mathbb{E}_q\left[\frac{p_{\theta}(\boldsymbol{x}_0)}{q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}\right]\leqslant \mathrm{E}_q\left[-\log\frac{p_{\theta}(\boldsymbol{x}_0)}{q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}\right]

将分子分母拆分为乘积链的形式，利用对数函数的性质有

Eq[−log⁡pθ(x0)q(x1:T∣x0)]=Eq[−log⁡p(xT)−∑t=1Tlog⁡pθ(xt−1∣xt)q(xt∣xt−1)]=defL\mathbb{E}_q\left[-\log\frac{p_{\theta}(\boldsymbol{x}_0)}{q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}\right]=\mathbb{E}_q\left[-\log p(\boldsymbol{x}_T)-\sum_{t=1}^T\log\frac{p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)}{q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})}\right]\overset{\mathrm{def}}{=}\mathcal{L}

利用贝叶斯定理，对于 2⩽t⩽T2\leqslant t\leqslant T 可以把 qq 反向化书写

q(xt∣xt−1)=q(xt−1∣xt,x0)q(xt∣x0)q(xt−1∣x0)q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})=q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\frac{q(\boldsymbol{x}_t|\boldsymbol{x}_0)}{q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0)}

此时，L\mathcal{L} 内求和部分的每一项可以被拆分为

pθ(xt−1∣xt)q(xt∣xt−1)=pθ(xt−1∣xt)q(xt−1∣xt,x0)q(xt−1∣x0)q(xt∣x0)\frac{p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)}{q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1})}=\frac{p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)}{q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)}\frac{q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0)}{q(\boldsymbol{x}_t|\boldsymbol{x}_{0})}

注意到下式可以通过裂项相消的方式得到

∏t=2Tq(xt−1∣x0)q(xt∣x0)=q(x1∣x0)q(xT∣x0)\prod_{t=2}^T\frac{q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0)}{q(\boldsymbol{x}_t|\boldsymbol{x}_{0})}=\frac{q(\boldsymbol{x}_{1}|\boldsymbol{x}_0)}{q(\boldsymbol{x}_T|\boldsymbol{x}_{0})}

因此不难得到

pθ(x0)q(x1:T∣x0)=p(xT)q(xT∣x0)⋅∏t=2Tpθ(xt−1∣xt)q(xt−1∣xt,x0)⋅pθ(x0∣x1)\frac{p_{\theta}(\boldsymbol{x}_0)}{q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}=\frac{p(\boldsymbol{x}_T)}{q(\boldsymbol{x}_T|\boldsymbol{x}_0)}\cdot\prod_{t=2}^T\frac{p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)}{q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)}\cdot p_{\theta}(\boldsymbol{x}_0|\boldsymbol{x}_1)

对上式左端取负对数并对 q(x1:T∣x0)q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0) 求期望得到 L\mathcal{L}，右端化简为

L=Eq[DKL[q(xT∣x0)∥p(xT)]+∑t=2TDKL[q(xt−1∣xt,x0)∥pθ(xt−1∣xt)]−log⁡pθ(x0∣x1)]\mathcal{L}=\mathbb{E}_q\left[\mathbb{D}_{\mathrm{KL}}\Big[q(\boldsymbol{x}_T|\boldsymbol{x}_0)\|p(\boldsymbol{x}_T)\Big]+\sum_{t=2}^T\mathbb{D}_{\mathrm{KL}}\Big[q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\|p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) \Big]-\log p_{\theta}(\boldsymbol{x}_0|\boldsymbol{x}_1)\right]

用符号 LL 简记上式几项，可以记作

L=Eq[LT+∑t=2TLt−1−L0]\mathcal{L}=\mathbb{E}_q[L_{T}+\sum_{t=2}^TL_{t-1}-L_0]

对于条件概率 q(xt−1∣xt,x0)q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)，可以证明其形式为

q(xt−1∣xT,x0)=N(xt−1;αˉt−1βt1−αˉtx0+αt(1−αˉt−1)1−αˉtxt,(1−αˉt−1)βt1−αˉtI)q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_T,\boldsymbol{x}_0)=N\left(\boldsymbol{x}_{t-1};\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\boldsymbol{x}_0+\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\boldsymbol{x}_t,\frac{(1-\bar{\alpha}_{t-1})\beta_t}{1-\bar{\alpha}_t}\boldsymbol{I}\right)

将上述高斯分布的参数记为 μ~t(xt,x0),β~tI\tilde{\boldsymbol{\mu}}_t(\boldsymbol{x}_t,\boldsymbol{x}_0),\tilde{\beta}_t\boldsymbol{I}

扩散与去噪

ELBO分析

前面我们得到了 Diffusion 模型的变分界形式为

L=Eq[LT+∑t=2TLt−1−L0]\mathcal{L}=\mathbb{E}_q[L_{T}+\sum_{t=2}^TL_{t-1}-L_0]

这若干项可以看做三个部分

  • LTL_T：前向过程结果 xT\boldsymbol{x}_T 与标准高斯先验 N(0,I)N(0,\boldsymbol{I}) 的区别
  • Lt−1L_{t-1}：后验分布 q(xt−1∣xt,x0)q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0) 与神经网络 pθ(xt−1∣xt)p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) 的区别
  • L0L_0：从 x1\boldsymbol{x}_1 恢复到 x0\boldsymbol{x}_0 的重建项

目前，我们先假定 βi\beta_i 是固定常数，因此 LTL_T 可以视为常数不参与训练。现在重点关注 Lt−1L_{t-1} 项，它是 Diffusion 模型的理论核心。考虑训练目标神经网络

pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t)),1<t⩽Tp_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)=N\big(\boldsymbol{x}_{t-1};\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t,t),\boldsymbol{\Sigma}_{\theta}(\boldsymbol{x}_t,t)\big),\quad 1<t\leqslant T

这个反向过程中，我们假定方差的形式为 Σθ(xt,t)=σt2I\boldsymbol{\Sigma}_{\theta}(\boldsymbol{x}_t,t)=\sigma_t^2\boldsymbol{I}，其中的 σt2\sigma^2_t 有两种取法

  1. σt2=βt\sigma_t^2=\beta_t：此时若 x0∼N(x0;0,I)\boldsymbol{x}_0\sim N(\boldsymbol{x}_0;\boldsymbol{0},\boldsymbol{I})，则模型表现更佳
  2. σt2=βt~\sigma_t^2=\tilde{\beta_t}：此时若 x0\boldsymbol{x}_0 服从单点分布，则模型表现更佳

预测新噪声

上述两种采样情况都属于极端情况，划定了数据损失熵的上界和下界。接着，我们需要对均值 μθ(xt,t)\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t,t) 进行参数化。利用 KLKL 散度的性质，我们可以把 Lt−1L_{t-1} 项转化为如下 L2L2 范数形式

Lt−1=Eq[12σt2∥μ~t(xt,x0)−μθ(xt,t)∥2]+CL_{t-1}=\mathbb{E}_q\left[\frac{1}{2\sigma_t^2}\big\|\tilde{\boldsymbol{\mu}}_t(\boldsymbol{x}_t,\boldsymbol{x}_0)-\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t,t)\big\|^2\right]+C

直观上说，由于 μ~t(xt,x0)\tilde{\boldsymbol{\mu}}_t(\boldsymbol{x}_t,\boldsymbol{x}_0) 表示后验 q(xt−1∣xt,x0)q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0) 的均值，因此上式试图使用神经网络 θ\theta 的输出 μθ(xt,t)\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t,t) 以逼近真实后验均值。参考VAE中的重参数化技巧，我们考虑随机 ϵ∼N(ϵ;0,I)\epsilon\sim N(\epsilon;\boldsymbol{0},\boldsymbol{I}) 并对 xt\boldsymbol{x}_t 采样

xt=xt(x0,ϵ)=αˉtx0+1−αˉtϵ\boldsymbol{x}_t=\boldsymbol{x}_t(\boldsymbol{x}_0,\epsilon)=\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\epsilon

该采样方法的依据为

q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)q(\boldsymbol{x}_t|\boldsymbol{x}_0)=N\big(\boldsymbol{x}_t;\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0,(1-\bar{\alpha}_t)\boldsymbol{I}\big)

将该采样结果代入 Lt−1L_{t-1} 可以得到如下形式

Lt−1−C=Eq[12σt2∥μ~t(xt(x0,ϵ),x0)−μθ(xt(x0,ϵ),t)∥2]=Ex0,ϵ[12σt2∥μ~t(xt(x0,ϵ),xt(x0,ϵ)−1−αˉtϵαˉt)−μθ(xt(x0,ϵ),t)∥2]=Ex0,ϵ[12σt2∥1αˉt(xt(x0,ϵ)−βt1−αˉtϵ)−μθ(xt(x0,ϵ),t)∥2]\begin{aligned}L_{t-1}-C&=\mathbb{E}_{q}\left[\frac{1}{2\sigma_t^2}\big\|\tilde{\boldsymbol{\mu}}_t(\boldsymbol{x}_t(\boldsymbol{x}_0,\epsilon),\boldsymbol{x}_0)-\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t(\boldsymbol{x}_0,\epsilon),t)\big\|^2\right]\\&=\mathbb{E}_{\boldsymbol{x}_0,\epsilon}\left[\frac{1}{2\sigma_t^2}\Bigg\|\tilde{\boldsymbol{\mu}}_t\left(\boldsymbol{x}_t(\boldsymbol{x}_0,\epsilon),\frac{\boldsymbol{x}_t(\boldsymbol{x}_0,\epsilon)-\sqrt{1-\bar{\alpha}_t}\epsilon}{\sqrt{\bar{\alpha}_t}}\right)-\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t(\boldsymbol{x}_0,\epsilon),t)\Bigg\|^2\right]\\&=\mathbb{E}_{\boldsymbol{x}_0,\epsilon}\left[\frac{1}{2\sigma_t^2}\Bigg\|\frac{1}{\sqrt{\bar{\alpha}_t}}\left(\boldsymbol{x}_t(\boldsymbol{x}_0,\epsilon)-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon\right)-\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t(\boldsymbol{x}_0,\epsilon),t)\Bigg\|^2\right]\end{aligned}

上式从代数角度证明了，对前向过程 qq 的期望可以转化为对 x0,ϵ\boldsymbol{x}_0,\epsilon 的期望。对于给定的 xt\boldsymbol{x}_t，神经网络 θ\theta 的目标是用 μθ\boldsymbol{\mu}_{\theta} 去预测真实后验均值，而真实后验均值现在可以被视为 ϵ\epsilon 的函数。我们记 ϵ=ϵθ(xt,t)\epsilon=\epsilon_{\theta}(\boldsymbol{x}_t,t)，那么神经网络 θ\theta 的输出将尝试逼近下式。

μθ(xt,t)=1αˉt(xt−βt1−αˉtϵθ(xt,t))\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t,t)=\frac{1}{\sqrt{\bar{\alpha}_t}}\left(\boldsymbol{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_{\theta}(\boldsymbol{x}_t,t)\right)

这一过程巧妙的把对 μ~t\tilde{\boldsymbol{\mu}}_t 的预测转化为对简单噪声 ϵ\epsilon 的预测，我们就只需要对标准高斯噪声 ϵθ\epsilon_{\theta} 进行预测，模型就能根据上式计算出对应的 μθ\boldsymbol{\mu}_{\theta}，从而完成输出预测。换言之，我们现在的神经网络只需要输出 ϵθ\epsilon_{\theta} 即可，对应的 Lt−1L_{t-1} 项被记为

Lt−1−C=Ex0,ϵ[βt22σt2αt(1−αˉt)∥ϵ−ϵθ(αˉtx0+1−αˉtϵ,t)∥2]L_{t-1}-C=\mathbb{E}_{\boldsymbol{x}_0,\epsilon}\left[\frac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar{\alpha}_t)}\big\|\epsilon-\epsilon_{\theta}(\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\epsilon,t)\big\|^2\right]

图像解码近似

一个常见的图像像素值可以被定义为 i∈{0,1,⋯ ,255}i\in\{0,1,\cdots,255\}，在 Diffusion 处理图像时，常将这些像素值线性映射到 [−1,1][-1,1] 中，以保证参数量级与先验分布 p(xT)p(\boldsymbol{x}_T) 对应。现在考虑 L0=−log⁡pθ(x0∣x1)L_0=-\log p_{\theta}(\boldsymbol{x}_0|\boldsymbol{x}_1)，为了处理离散形式的 x0\boldsymbol{x}_0，设其像素维度数为 DD，线性处理后的像素表示为 x0(i)x_0^{(i)}，则有

pθ(x0∣x1)=∏i=1D∫δ−(x0(i))δ+(x0(i))N(x;μθ(i)(x1,1),σ12)dxp_{\theta}(\boldsymbol{x}_0|\boldsymbol{x}_1)=\prod_{i=1}^D\int_{\delta_{-}(x_0^{(i)})}^{\delta_{+}(x_0^{(i)})}N\big(x;\mu_{\theta}^{(i)}(\boldsymbol{x}_1,1),\sigma_1^2\big)\mathrm{d}x

其中的积分界函数被定义为

δ+(x)={+∞,x=1x+1255,x<1,δ−(x)={−∞,x=−1x−1255,x>−1\delta_{+}(x)=\begin{cases}+\infty,&x=1\\x+\dfrac{1}{255},&x<1\end{cases},\quad \delta_{-}(x)=\begin{cases}-\infty,&x=-1\\x-\dfrac{1}{255},&x>-1\end{cases}

这里通过在离散点周围进行积分，以求取近似的高斯分布值。实际中可近似为

∫δ−(x0(i))δ+(x0(i))N(x;μθ(i)(x1,1),σ12)dx≈2255N(x0(i);μθ(i),σ12)\int_{\delta_{-}(x_0^{(i)})}^{\delta_{+}(x_0^{(i)})}N\big(x;\mu_{\theta}^{(i)}(\boldsymbol{x}_1,1),\sigma_1^2\big)\mathrm{d}x\approx\frac{2}{255}N(x_0^{(i)};\mu_{\theta}^{(i)},\sigma_1^2)

训练与预测

在实际训练中，下述目标函数近似通常表现的更佳

Lsimple(θ)=Et,x0,ϵ[∥ϵ−ϵθ(αˉtx0+1−αˉtϵ,t)∥2]\mathcal{L}_{\mathrm{simple}}(\theta)=\mathbb{E}_{t,\boldsymbol{x}_0,\epsilon}\left[\big\|\epsilon-\epsilon_{\theta}(\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\epsilon,t)\big\|^2\right]

接下来我们来简单讨论 Diffusion 模型的训练步骤。

  1. 从数据集 q(x0)q(\boldsymbol{x}_0) 采样 x0\boldsymbol{x}_0
  2. 采样时间步 t∼U({1,⋯ ,T})t\sim U(\{1,\cdots,T\})
  3. 采样噪音 ϵ∼N(0,I)\epsilon\sim N(\boldsymbol{0},\boldsymbol{I})
  4. 一步合成含噪音样本 xt=αˉtx0+1−αˉtϵ\boldsymbol{x}_t=\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\epsilon
  5. 计算梯度 ∇θ∥ϵ−ϵθ(xt,t)∥2\nabla_{\theta}\|\epsilon-\epsilon_{\theta}(\boldsymbol{x}_t,t)\|^2 并更新参数
  6. 重复上述步骤直到梯度收敛

当神经网络 θ\theta 训练完毕后，我们使用 ϵθ\epsilon_{\theta} 生成新数据。需要注意的是，我们在反向过程中的采样 xt−1∼pθ(xt−1∣x)\boldsymbol{x}_{t-1}\sim p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}) 需要同样引入一个全新的标准高斯噪音 z∼N(0,I)\boldsymbol{z}\sim N(\boldsymbol{0},\boldsymbol{I})，具体而言有

xt−1=μθ(xt,t)+σtz=1αˉt(xt−βt1−αˉtϵθ(xt,t))+σtz\boldsymbol{x}_{t-1}=\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t,t)+\sigma_t\boldsymbol{z}=\frac{1}{\sqrt{\bar{\alpha}_t}}\left(\boldsymbol{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_{\theta}(\boldsymbol{x}_t,t)\right)+\sigma_t\boldsymbol{z}

当 T⩾t>1T\geqslant t>1 时，我们反复执行上述转移。当 t=1t=1 时不需要添加噪声，即有最终输出

x0=μθ(x1,1)\boldsymbol{x}_0=\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_1,1)

U-Net

最后一节，我们来简要介绍神经网络 θ\theta 的构建。在 Diffusion 论文中，作者主要使用了 U-Net 架构。U-Net常用于CV领域的图像分割任务，可以被视为编码器-解码器架构。

简单来说，U-Net 通过一系列的卷积、反卷积、池化操作提取图像信息，其架构如下图所示。

[{"url":"https://pic3.zhimg.com/v2-9fff7b43ed249e5241af058ff410a904_1440w.jpg","alt":"U-Net","title":""}]

不同图例的含义如下
若已经理解前面U-net的基础知识，将能更好地看懂整个U-net的结构。首先看一下图例的含义：

  • 深蓝色箭头：利用 3×33\times 3的卷积核对图片进行卷积后，通过ReLU\mathrm{ReLU} 激活函数输出
  • 灰色箭头：对左边下采样过程中的图片进行裁剪复制
  • 红色箭头：通过最大池化对图片进行下采样，池化核大小为 2×22\times 2
  • 绿色箭头：反卷积，对图像进行上采样，卷积核大小为 2×22\times 2
  • 青色箭头：使用 1×11\times 1 的卷积核对图片进行卷积

DiTs

融合Transformer

众所周知，人工智能领域中最有影响力的模型之一便是完全基于注意力机制的 Transformer 架构，因此人们开始尝试把这种注意力架构范式融入生成式模型中，这便得到了 Diffusion in Transformer 架构。下面我们简称其为 DiT 模型。

具体而言，DiT 模型可以被简单拆解为以下几点

  1. DiT 完全保留了 DDPM 的扩散生成原理
  2. DiT 不直接在像素级别扩散，而是先使用预训练的 VAE 把高维原图编码为低维的潜在表示 z\boldsymbol{z}，然后在 z\boldsymbol{z} 上执行扩散，再解码为高维图片
  3. 使用 (Vision in) Transformer 替换 U-Net，这是本模型的核心贡献

输入处理

DiT 的扩散过程只发生在低维潜空间，这使得计算更加容易。以 ImageNet 的 256×256×3256\times 256\times 3 数据集为例，我们抽取图像 x∼q(x)\boldsymbol{x}\sim q(\boldsymbol{x}) 和其图像标签 cc，进行编码操作

z0=E(x)∈R32×32×4\boldsymbol{z}_0=\mathcal{E}(\boldsymbol{x})\in\mathbb{R}^{32\times 32\times 4}

其中 E\mathcal{E} 采用的是 Stable Diffusion 中的 KL-f8 VAE 编码器，下采样因子为 88。

需要注意的是，DiT 的训练和 VAE 的训练是完全解耦的。因此在扩散模型的训练过程中，VAE 的参数是完全冻结的，这使得我们可以使用前人的预训练 VAE。

使用 DDPM 的采样方法，可以一步合成得到含噪潜变量

zt=αˉtz0+1−αˉtϵ,t∼U({1,2,⋯ ,T})\boldsymbol{z}_t=\sqrt{\bar{\alpha}_t}\boldsymbol{z}_0+\sqrt{1-\bar{\alpha}_t}\epsilon,\quad t\sim U(\{1,2,\cdots,T\})

类似于 ViT，我们把 zt\boldsymbol{z}_t 拆分为若干个 patches\mathrm{patches}，其大小表示为 p×p×4p\times p\times 4。这总计 TT 个 patches\mathrm{patches} 被拉平为并线性投影，然后植入正弦-余弦编码以表示位置信息。为避免混淆，我们称这部分输入为图像词元。

原文设置了以下四种 Transformer Block，用不同方式添加条件信息（t,c,t,c,自然语言等）

  1. 上下文条件生成：简单的把 tt 和 cc 通过向量嵌入添加到输入序列中。这意味着上下文 t,ct,c 只能在自注意力层中影响图像词元。在输出头把这两个特殊词元拿走，以保证生成效果。
  2. 交叉注意力：在每个 Transformer Block 中添加一个交叉注意力层，其查询为图像词元，键值为向量嵌入后的 t,ct,c。
  3. adaLN： 我们考虑一个标准的层归一化 x^=γ∘x−μσ+β\hat{x}=\gamma\circ \dfrac{x-\mu}{\sigma}+\beta，其中的参数 γ,β\gamma,\beta 可以使用一个 MLP 得到。这个 MLP 接收向量嵌入后的 t,ct,c 作为输出，并将结果作用于 Transformer Block 的两个层归一化中。这意味着每个模块将多出四个参数 (γ1,2,β1,2)(\gamma_{1,2},\beta_{1,2})
  4. adaLN-Zero：另外考虑两个初始为零的门控参数 α1,2\alpha_{1,2}，其同样由 MLP 输出。该门控参数作用于每次残差连接之前，保证噪声信号的量级稳定。

下面的图片展示了上述的模块。

[{"url":"/img/world/DiT.png","alt":"DiT模块架构","title":""}]

训练与生成

DiT的训练目标函数基本遵循 DDPM 的形式

L=∥ϵ−ϵθ(zt,t,c)∥\mathcal{L}=\|\epsilon-\epsilon_{\theta}(\boldsymbol{z}_t,t,c)\|

如果我们不采用分类标签 cc，可以将其设为空词元 <null> 或 ∅\varnothing。

训练过程便是对扩散模型的参数进行梯度参数优化。需要注意的是，DiT 同时对 μθ\boldsymbol{\mu}_{\theta} 和 Σθ\boldsymbol{\Sigma}_{\theta} 进行了参数优化，具体原理不在此赘述。

在生成过程中，我们首先进行先验分布采样

zT∼N(zT;0,I)\boldsymbol{z}_T\sim N(\boldsymbol{z}_T;\boldsymbol{0},\boldsymbol{I})

在逐步去噪中，我们考虑一种双条件前向过程

ϵ~θ(zt,t,c)=ϵθ(zt,t,∅)+s[ϵθ(zt,t,c)−ϵθ(zt,t,∅)]\tilde{\epsilon}_{\theta}(\boldsymbol{z}_t,t,c)=\epsilon_{\theta}(\boldsymbol{z}_t,t,\varnothing)+s\big[\epsilon_{\theta}(\boldsymbol{z}_t,t,c)-\epsilon_{\theta}(\boldsymbol{z}_t,t,\varnothing)\big]

这使得我们不需要 Dropout 就可以把类标签 cc 看做可学习的嵌入类型，其中的 ss 是超参数，在原文中取值范围为 [1,4][1,4] 以进行实验评估。

最后，我们把得到的 z0\boldsymbol{z}_0 经过冻结的 VAE 解码器便得到了生成图像。