DDPM

The Annotated Diffusion Model

在这篇博文中,我们将更深入地研究去噪扩散概率模型(也称为 DDPM、扩散模型、基于分数的生成模型或简称为自编码器),因为研究人员已经能够使用它们进行高质量的条件(or 非条件)图像/音频/视频生成。

我们将重温 (Ho et al., 2020) DDPM的原始论文,基于 Phil Wang 的实现在 PyTorch 中逐步实现它——(原论文为TensorFlow 实现)。

请注意,生成建模的扩散思想实际上已经在 (Sohl-Dickstein et al., 2015) 中引入。 然而,直到 (Song et al., 2019)(斯坦福大学),和之后的(Ho et al., 2020)(谷歌大脑)才独立改进了该方法。

请注意,扩散模型有多种视角解读。 在这里,我们采用离散时间(潜在变量模型)视角,但有必要理解其他视角。

好吧,让我们开始吧!

 

一.什么是扩散模型

如果将(去噪)扩散模型与其他生成模型(例如归一化流、GAN 或 VAE)进行比较,它并没有那么复杂:它们都将噪声从一些简单的分布转换为数据样本。 这也是神经网络学习从纯噪声开始逐渐对数据进行去噪的情况。

具体来说包括 2 个过程:

  1. 我们选择的固定(或预定义)前向扩散过程 q,逐渐将高斯噪声添加到图像中,直到最终得到纯噪声。
  2. 学习到的反向去噪扩散过程 pθ,其中训练神经网络从纯噪声开始逐渐对图像进行去噪,直到最终得到实际图像。

image-20230224162710062

t 索引的前向和反向过程都发生在一定数量的有限时间步长 T内(DDPM 作者使用 T=1000)。 你从t=0开始,在那里你采样了一个真实的图像x0,来自您的数据分布(假设是来自 ImageNet 的图像),前向过程在每个时间步 t从高斯分布中采样一些噪声,这些噪声被添加到前一个时间步的图像中。 给定足够大的T和在每个时间步适宜地添加噪声,您最终会通过在t=T处得到所谓的各向同性高斯分布。

二.数学形式

因为需要损失函数来定义优化目标,所以我们以更公式化的方式解读扩散过程。

1.前向扩散过程:q(xtxt1)

基本上,每个新(稍微嘈杂的)图像xt都是从具有 μt=1βtxt1σt2=βt条件高斯分布中提取的。可以通过xt=1βtxt1+βtϵ,其中ϵN(0,I)实现,两种形式是等价的。

这是添加的信号相关噪声。

请注意,βt 在每个时间步长 t 都不是常数(因此也有下标)——事实上,它定义了一个所谓的“方差表(variance schedule)”,它可以是线性的、二次的、余弦的,等等,我们将进一步看到 (有点像学习率表)。

感觉和GDN(广义分歧归一化)有点像。

如果方差表设置得合适,xT可以是完全的高斯噪声。

如果我们知道条件分布p(xt1xt),然后我们可以反向运行这个过程:通过采样一些随机的高斯xT,然后不断去噪,最终得到一个真实分布的样本x0

然而,我们不知道 p(xt1xt)。这很棘手,因为它需要知道所有可能图像的分布才能计算这个条件概率。因此,我们将利用神经网络

2.反向扩散过程:pθ(xt1xt)

我们将利用神经网络来近似(学习)反向的条件概率分布,我们称之为 pθ(xt1|xt), 其中 θ 是神经网络的参数,通过梯度下降更新。

如果我们假设这个逆向过程也是高斯分布,那么回想一下任何高斯分布都由 2 个参数定义:均值μθ和方差Σθ

我们可以将该过程参数化为:

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

其中均值和方差也以噪声水平 t为条件。

因此,我们的神经网络需要学习/表示均值和方差。 然而,DDPM 作者决定保持方差固定,让神经网络只学习(表示)均值 μθ

后续有论文也学习方差Σθ

三.定义目标函数

(通过重参数化均值)

为了导出目标函数来学习反向过程的均值μθ,作者观察到 qpθ 的组合可以看作是变分自动编码器 (VAE)。 因此,变分下界(也称为 ELBO)可用于最小化关于GT数据样本 x0 的负对数似然(关于 ELBO 的详细信息,我们参考 VAE 论文)。 事实证明,这个过程的 ELBO 是每个时间步 t 的损失总和,L=L0+L1+...+LT .

通过前向q和后向过程pθ的构建,损失的每一项(L0除外)实际上是2个高斯分布之间的KL散度,可以明确地写成关于均值的l2损失!

如 Sohl-Dickstein 等人所示,构造正向过程 q 的直接结果是我们可以以 x0 为条件在任意噪声水平下对 xt 进行采样(因为高斯的和 也是高斯分布)。 这非常方便:我们不需要重复应用 q 来对 xt 进行采样。我们有:

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

也就是T步的q实际一步到位。其中 αt:=1βtα¯t:=Πs=1tαs

我们认为这个等式有着“nice property”。

  1. 这意味着我们可以对高斯噪声进行采样并对其进行适当缩放,并将其添加到 x0 中以直接得到 xt。 请注意,α¯t 是已知 方差表βt 的函数,因此也是已知的并可以预先计算。 然后,这允许我们在训练期间优化损失函数 L 的随机项(或者换句话说,在训练期间随机采样 t 并优化 Lt)。

  2. 如 Ho 等人所示,此属性的另一个优点 是可以(经过一些数学运算后)重新参数化均值以使神经网络学习(预测)在噪声水平 t 添加的噪声(通过网络 ϵθ(xt,t) )。

    这意味着我们的神经网络成为噪声预测器,而不是均值预测器。 均值可以计算如下:

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

如此:

μ~t(xt,x0):=α¯t1βt1α¯tx0+αt(1α¯t1)1α¯txt

β~t:=1α¯t11α¯tβt

最终的目标函数 Lt 就会看起来如下:

ϵϵθ(xt,t)2=ϵϵθ(α¯tx0+(1α¯t)ϵ,t)2.

这里, x0是初始(真实的、未损坏的)图像,我们看到固定前向过程给出的直接噪声水平 t 样本。 ϵ 是在时间步长 t 采样的纯噪声,而 ϵθ(xt,t) 是我们的神经网络。 神经网络使用真实和预测高斯噪声之间的简单均方误差 (MSE) 进行优化。

训练算法现在可以总结如下:

image-20230301141221490

换句话说:

  1. 我们从真实未知且可能复杂的数据分布 q(x0) 中随机抽取样本 x0
  2. 我们在 1T 之间均匀采样噪声水平 t (即随机时间步长)
  3. 我们从高斯分布中采样一些噪声,并在 t 级通过该噪声破坏输入(使用上面提到的“nice property”)
  4. 神经网络被训练为根据损坏的图像 xt(即根据已知时间表 βt 应用于x0的噪声)预测此噪声

实际上,所有这些都是在批量数据上完成的,因为人们使用随机梯度下降来优化神经网络。

四.神经网络

神经网络需要在特定时间步长处接收噪声图像并返回预测噪声。 请注意,预测噪声是与输入图像具有相同大小/分辨率的张量。 所以从技术上讲,网络接收和输出相同形状的张量。 我们可以为此使用哪种类型的神经网络?

这里通常使用的与自动编码器非常相似,您可能还记得典型的“深度学习入门”教程。 自动编码器在编码器和解码器之间有一个所谓的“瓶颈”层。 编码器首先将图像编码成称为“瓶颈”的较小隐藏表示,然后解码器将该隐藏表示解码回实际图像。 这迫使网络只在瓶颈层保留最重要的信息。

在架构方面,DDPM 作者选择了 U-Net。 这个网络,像任何自动编码器一样,中间有一个瓶颈,确保网络只学习最重要的信息。 重要的是,它在编码器和解码器之间引入了残差连接,极大地改善了梯度流(受 ResNet [He et al., 2015] 启发)。

五.网络模块

0.位置嵌入

由于神经网络的参数跨时间共享(噪声水平),作者受 Transformer 的启发,采用正弦位置向量(position embeddings)t 进行编码。 这使得神经网络“知道”它在哪个特定时间步长(噪声水平)运行。

SinusoidalPositionEmbeddings 模块将形状为 (batch_size, 1) 的张量作为输入(即一批中几个噪声图像的噪声水平),并将其转换为形状为 (batch_size, dim) 的张量,其中 dim 是图像的维数 位置嵌入。 然后将其添加到每个残差块中,我们将进一步介绍。

1.残差块

接下来,我们定义 U-Net 模型的核心构建块。 DDPM 作者采用了 Wide ResNet 块(Zagoruyko 等人,2016 年),但 Phil Wang 已将标准卷积层替换为“权重标准化”版本,该版本与组归一化结合使用效果更好(参见(Kolesnikov 等人, 2019)了解详情)。

2.注意力模块

接下来,我们定义注意力模块,DDPM 作者将其添加到卷积块之间。这里使用了两种注意力变体:一种是常规的多头自注意力(在 Transformer 中使用),另一种是线性注意力变体(Shen 等人,2018),是non-local模块的直接替代品。

3.组归一化

DDPM 作者将 U-Net 的卷积层/注意力层与组归一化(GN)交织在一起。 下面,我们定义了一个 PreNorm 类,它将用于在注意力层之前应用 groupnorm。请注意,关于在 Transformers 中该在注意力之前还是之后应用归一化一直存在争议。

六.条件U-Net

现在我们已经定义了所有构建块(位置嵌入、残差块、注意力和组归一化),是时候定义整个神经网络了。 回想一下网络 ϵθ(xt,t) 的工作是接收一批噪声图像及其各自的噪声水平,并输出添加到输入中的噪声。 更正式地说:

网络将一批形状为 (batch_size, num_channels, height, width) 的噪声图像和一批形状为 (batch_size, 1) 的噪声水平作为输入,并返回一个形状为 (batch_size, num_channels, height, width) 的张量

网络构建如下:

最终,神经网络层层叠叠,就好像它们是乐高积木一样(但了解它们的工作原理很重要)。

七.定义前向扩散过程

前向扩散过程在多个时间步长 T 中逐渐向真实分布的图像添加噪声。 这是根据方差表发生的。 最初的 DDPM 作者采用了线性表:

We set the forward process variances to constants increasing linearly from β1=104 to βT=0.02.

然而,在 (Nichol et al., 2021) 中表明,使用余弦方差表可以获得更好的结果。

下面,我们为 T 时间步定义了各种scheduler(稍后我们将选择一个)。

首先,让我们使用 T=300 时间步长的线性方差表,并定义我们需要的 βt相关的各种变量,例如方差的累积乘积 α¯t 。 下面的每个变量都是一维张量,存储从 tT 的值。 重要的是,我们还定义了一个提取函数,它将允许我们为一个batch的索引提取适当的 t 的索引。

βt (betas), α¯t (alphas),

α¯t (sqrt_alphas_cumprod), 1α¯t(sqrt_one_minus_alphas_cumprod) ,

α¯t1βt1α¯t(posterior_variance),

这些全是T个元素的一维向量。

torch.cumprod 和 torch.cumsum很熟悉,一个累乘一个累加。

我们将用一张猫的图像来说明如何在扩散过程的每个时间步添加噪声

八.定义数据集及DataLoader

这里我们定义了一个常规的 PyTorch 数据集。 数据集仅包含来自真实数据集的图像,如 Fashion-MNIST、CIFAR-10 或 ImageNet,线性缩放至 [−1, 1]]。

每个图像都resize为相同的大小。 有趣的是,图像也会随机水平翻转。

来自论文:我们在 CIFAR10 的训练过程中使用了随机水平翻转; 我们尝试了使用翻转和不使用翻转的训练,发现翻转可以稍微提高样本质量。

在这里,我们使用 🤗 Datasets 库从中心轻松加载 Fashion MNIST 数据集。 该数据集由已经具有相同分辨率(即 28x28)的图像组成。

 

应用:https://cloud.tencent.com/developer/article/2090159

扩散模型综述

参考曾面试的候选人的一篇调研。

横向比较:

纵向比较:

DDPM

前向过程:image-20230301131220228

DDPM 将给定图像𝑥0逐步编码为高斯噪声𝑥𝑇,这被称为前向过程。

反向过程:image-20230301131242203

⚠️这里p(xt1|xt)误写为q(xt1|xt)

解码阶段,也是一个反向过程,逆向进行。在这个过程中,𝜇和Σ通过神经网络建模。

反向过程从 N(0,I) 开始。 在每一步中,根据当前噪声图像 𝑥𝑡 推导高斯分布。 (𝑡 − 1) 步图像是从新分布中采样的。 因此,图像合成过程实际上转化为去噪过程,正如名称:去噪扩散模型。

在反向过程中,𝜇 和 Σ 应该通过神经网络推断。 然而,在前向过程中,细心的读者应该注意到,𝑞(𝑥T|𝑥0)实际上是一系列已知高斯分布N(1βtxt1,βtI)的组合。因此,𝑞(𝑥T|𝑥0)也是已知的高斯分布:

αt=1βt ,

α¯t=Πs=1tαs ,

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

训练:训练损失很简单:

ϵϵθ(xt,t)2=ϵϵθ(α¯tx0+(1α¯t)ϵ,t)2

𝐸 是一个噪声估计器,通常是一个由 𝑡 调节的 U-Net。 𝐸旨在预测高斯噪声𝜖。