Diffusion扩散模型

关于扩散模型（Diffusion Models）有很多种理解，本文的介绍是基于denoising diffusion probabilistic model （DDPM），DDPM已经在（无）条件图像/音频/视频生成领域取得了较多显著的成果，现有的比较受欢迎的的例子包括由OpenAI主导的GLIDE和DALL-E 2、由海德堡大学主导的潜在扩散和由Google Brain主导的图像生成。

实际上生成模型的扩散概念已经在（Sohl-Dickstein et al., 2015）中介绍过。然而，直到（Song et al., 2019）（斯坦福大学）和（Ho et al., 2020）（在Google Brain）才各自独立地改进了这种方法。

本文是在Phil Wang基于PyTorch框架的复现的基础上（而它本身又是基于TensorFlow实现），迁移到MindSpore AI框架上实现的。

模型简介

什么是Diffusion Model？

如果将Diffusion与其他生成模型（如Normalizing Flows、GAN或VAE）进行比较，它并没有那么复杂，它们都将噪声从一些简单分布转换为数据样本，Diffusion也是从纯噪声开始通过一个神经网络学习逐步去噪，最终得到一个实际图像。 Diffusion对于图像的处理包括以下两个过程：

• 我们选择的固定（或预定义）正向扩散过程 𝑞𝑞 ：它逐渐将高斯噪声添加到图像中，直到最终得到纯噪声

• 一个学习的反向去噪的扩散过程 𝑝𝜃𝑝𝜃 ：通过训练神经网络从纯噪声开始逐渐对图像去噪，直到最终得到一个实际的图像

由 𝑡𝑡 索引的正向和反向过程都发生在某些有限时间步长 𝑇𝑇（DDPM作者使用 𝑇=1000𝑇=1000）内。从𝑡=0𝑡=0开始，在数据分布中采样真实图像 𝐱0𝑥0（本文使用一张来自ImageNet的猫图像形象的展示了diffusion正向添加噪声的过程），正向过程在每个时间步长 𝑡𝑡 都从高斯分布中采样一些噪声，再添加到上一个时刻的图像中。假定给定一个足够大的 𝑇𝑇 和一个在每个时间步长添加噪声的良好时间表，您最终会在 𝑡=𝑇𝑡=𝑇 通过渐进的过程得到所谓的各向同性的高斯分布。

扩散模型实现原理

Diffusion 前向过程

所谓前向过程，即向图片上加噪声的过程。虽然这个步骤无法做到图片生成，但这是理解diffusion model以及构建训练样本至关重要的一步。 首先我们需要一个可控的损失函数，并运用神经网络对其进行优化。

设 𝑞(𝑥0)𝑞(𝑥0) 是真实数据分布，由于 𝑥0∼𝑞(𝑥0)𝑥0∼𝑞(𝑥0) ，所以我们可以从这个分布中采样以获得图像 𝑥0𝑥0 。接下来我们定义前向扩散过程 𝑞(𝑥𝑡|𝑥𝑡−1)𝑞(𝑥𝑡|𝑥𝑡−1) ，在前向过程中我们会根据已知的方差 0<𝛽1<𝛽2<...<𝛽𝑇<10<𝛽1<𝛽2<...<𝛽𝑇<1 在每个时间步长 t 添加高斯噪声，由于前向过程的每个时刻 t 只与时刻 t-1 有关，所以也可以看做马尔科夫过程：

回想一下，正态分布（也称为高斯分布）由两个参数定义：平均值 𝜇𝜇 和方差 𝜎2≥0𝜎2≥0 。基本上，在每个时间步长 𝑡𝑡 处的产生的每个新的（轻微噪声）图像都是从条件高斯分布中绘制的，其中

我们可以通过采样然后设置

请注意， 𝛽𝑡𝛽𝑡 在每个时间步长 𝑡𝑡 （因此是下标）不是恒定的：事实上，我们定义了一个所谓的“动态方差”的方法，使得每个时间步长的 𝛽𝑡𝛽𝑡 可以是线性的、二次的、余弦的等（有点像动态学习率方法）。

因此，如果我们适当设置时间表，从 𝐱0𝑥0 开始，我们最终得到 𝐱1,...,𝐱𝑡,...,𝐱𝑇𝑥1,...,𝑥𝑡,...,𝑥𝑇，即随着 𝑡𝑡 的增大 𝐱𝑡𝑥𝑡 会越来越接近纯噪声，而 𝐱𝑇𝑥𝑇 就是纯高斯噪声。

那么，如果我们知道条件概率分布 𝑝(𝐱𝑡−1|𝐱𝑡)𝑝(𝑥𝑡−1|𝑥𝑡) ，我们就可以反向运行这个过程：通过采样一些随机高斯噪声 𝐱𝑇𝑥𝑇，然后逐渐去噪它，最终得到真实分布 𝐱0𝑥0 中的样本。但是，我们不知道条件概率分布 𝑝(𝐱𝑡−1|𝐱𝑡)𝑝(𝑥𝑡−1|𝑥𝑡) 。这很棘手，因为需要知道所有可能图像的分布，才能计算这个条件概率。

Diffusion 逆向过程

为了解决上述问题，我们将利用神经网络来近似（学习）这个条件概率分布 𝑝𝜃(𝐱𝑡−1|𝐱𝑡)𝑝𝜃(𝑥𝑡−1|𝑥𝑡) , 其中 𝜃𝜃 是神经网络的参数。如果说前向过程(forward)是加噪的过程，那么逆向过程(reverse)就是diffusion的去噪推断过程，而通过神经网络学习并表示 𝑝𝜃(𝐱𝑡−1|𝐱𝑡)𝑝𝜃(𝑥𝑡−1|𝑥𝑡) 的过程就是Diffusion 逆向去噪的核心。

现在，我们知道了需要一个神经网络来学习逆向过程的（条件）概率分布。我们假设这个反向过程也是高斯的，任何高斯分布都由2个参数定义：

• 由 𝜇𝜃𝜇𝜃 参数化的平均值

• 由 𝜇𝜃𝜇𝜃 参数化的方差

综上，我们可以将逆向过程公式化为

其中平均值和方差也取决于噪声水平 𝑡𝑡 ，神经网络需要通过学习来表示这些均值和方差。

• 注意，DDPM的作者决定保持方差固定，让神经网络只学习（表示）这个条件概率分布的平均值 𝜇𝜃𝜇𝜃 。

• 本文我们同样假设神经网络只需要学习（表示）这个条件概率分布的平均值 𝜇𝜃𝜇𝜃 。

为了导出一个目标函数来学习反向过程的平均值，作者观察到 𝑞𝑞 和 𝑝𝜃𝑝𝜃 的组合可以被视为变分自动编码器(VAE)。因此，变分下界（也称为ELBO）可用于最小化真值数据样本 𝐱0𝑥0 的似然负对数（有关ELBO的详细信息，请参阅VAE论文（Kingma等人，2013年）），该过程的ELBO是每个时间步长的损失之和 𝐿=𝐿0+𝐿1+...+𝐿𝑇𝐿=𝐿0+𝐿1+...+𝐿𝑇 ，其中，每项的损失 𝐿𝑡𝐿𝑡 （除了 𝐿0𝐿0 ）实际上是2个高斯分布之间的KL发散，可以明确地写为相对于均值的L2-loss!

如Sohl-Dickstein等人所示，构建Diffusion正向过程的直接结果是我们可以在条件是 𝐱0𝑥0 （因为高斯和也是高斯）的情况下，在任意噪声水平上采样 𝐱𝑡𝑥𝑡 ，而不需要重复应用 𝑞𝑞 去采样 𝐱𝑡𝑥𝑡 ，这非常方便。使用

我们就有

这意味着我们可以采样高斯噪声并适当地缩放它，然后将其添加到 𝐱0𝑥0 中，直接获得 𝐱𝑡𝑥𝑡 。

请注意，𝛼¯𝑡𝛼¯𝑡 是已知 𝛽𝑡𝛽𝑡 方差计划的函数，因此也是已知的，可以预先计算。这允许我们在训练期间优化损失函数 𝐿𝐿 的随机项。或者换句话说，在训练期间随机采样 𝑡𝑡 并优化 𝐿𝑡𝐿𝑡 。

正如Ho等人所展示的那样，这种性质的另一个优点是可以重新参数化平均值，使神经网络学习（预测）构成损失的KL项中噪声的附加噪声。这意味着我们的神经网络变成了噪声预测器，而不是（直接）均值预测器。其中，平均值可以按如下方式计算：

最终的目标函数 𝐿𝑡𝐿𝑡 如下 （随机步长 t 由 (𝜖∼𝑁(0,𝐈))(𝜖∼𝑁(0,𝐼)) 给定）：

在这里， 𝐱0𝑥0 是初始（真实，未损坏）图像， 𝜖𝜖 是在时间步长 𝑡𝑡 采样的纯噪声，𝜖𝜃(𝐱𝑡,𝑡)𝜖𝜃(𝑥𝑡,𝑡)是我们的神经网络。神经网络是基于真实噪声和预测高斯噪声之间的简单均方误差（MSE）进行优化的。

训练算法现在如下所示：

换句话说：

• 我们从真实未知和可能复杂的数据分布中随机抽取一个样本 𝑞(𝐱0)𝑞(𝑥0)

• 我们均匀地采样11和𝑇𝑇之间的噪声水平𝑡𝑡（即，随机时间步长）

• 我们从高斯分布中采样一些噪声，并使用上面定义的属性在 𝑡𝑡 时间步上破坏输入

• 神经网络被训练以基于损坏的图像 𝐱𝑡𝑥𝑡 来预测这种噪声，即基于已知的时间表 𝐱𝑡𝑥𝑡 上施加的噪声

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

U-Net神经网络预测噪声

神经网络需要在特定时间步长接收带噪声的图像，并返回预测的噪声。请注意，预测噪声是与输入图像具有相同大小/分辨率的张量。因此，从技术上讲，网络接受并输出相同形状的张量。那么我们可以用什么类型的神经网络来实现呢？

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

在模型结构方面，DDPM的作者选择了U-Net，出自（Ronneberger et al.，2015）（当时，它在医学图像分割方面取得了最先进的结果）。这个网络就像任何自动编码器一样，在中间由一个bottleneck组成，确保网络只学习最重要的信息。重要的是，它在编码器和解码器之间引入了残差连接，极大地改善了梯度流。

可以看出，U-Net模型首先对输入进行下采样（即，在空间分辨率方面使输入更小），之后执行上采样。

构建Diffusion模型

DDPM论文表明扩散模型是（非）条件图像有希望生成的方向。自那以后，diffusion得到了（极大的）改进，最明显的是文本条件图像生成。下面，我们列出了一些重要的（但远非详尽无遗的）后续工作：

• 改进的去噪扩散概率模型(Nichol et al., 2021)：发现学习条件分布的方差（除平均值外）有助于提高性能

• 用于高保真图像生成的级联扩散模型([Ho et al., 2021)：引入级联扩散，它包括多个扩散模型的流水线，这些模型生成分辨率提高的图像，用于高保真图像合成

• 扩散模型在图像合成上击败了GANs(Dhariwal et al., 2021)：表明扩散模型通过改进U-Net体系结构以及引入分类器指导，可以获得优于当前最先进的生成模型的图像样本质量

• 无分类器扩散指南([Ho et al., 2021)：表明通过使用单个神经网络联合训练条件和无条件扩散模型，不需要分类器来指导扩散模型

• 具有CLIP Latents (DALL-E 2) 的分层文本条件图像生成 (Ramesh et al., 2022)：在将文本标题转换为CLIP图像嵌入之前使用，然后扩散模型将其解码为图像

• 具有深度语言理解的真实文本到图像扩散模型（ImageGen）(Saharia et al., 2022)：表明将大型预训练语言模型（例如T5）与级联扩散结合起来，对于文本到图像的合成很有效

请注意，此列表仅包括在撰写本文，即2022年6月7日之前的重要作品。

目前，扩散模型的主要（也许唯一）缺点是它们需要多次正向传递来生成图像（对于像GAN这样的生成模型来说，情况并非如此）。然而，有正在进行中的研究表明只需要10个去噪步骤就能实现高保真生成。