昇思Mindspore25天学习打卡Day21：Diffusion扩散模型

本文基于Link: Hugging Face：The Annotated Diffusion Model一文翻译迁移而来，同时参考了Link: 由浅入深了解Diffusion Model一文。

• 本教程在Jupyter Notebook上成功运行。如您下载本文档为Python文件，执行Python文件时，请确保执行环境安装了GUI界面。

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

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

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

实验中我们采用离散时间（潜在变量模型）的观点，另外，读者也可以查看有关于扩散模型的其他Link:几个观点！

实验开始之前请确保安装并导入所需的库（假设您已经安装了Link:MindSpore、download、dataset、matplotlib以及tqdm）。

1 模型简介

1.1 什么是Diffusion

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

• 我们选择的固定（或预定义）正向扩散过程 $q$ :它逐渐将高斯噪声添加到图像中，直到最终得到纯噪声
• 一个学习的反向去噪的扩散过程 $p_{\theta }$ ：通过训练神经网络从纯噪声开始逐渐对图像去噪，直到最终得到一个实际的图像-
• 由 $t$ 索引的正向和反向过程都发生在某些有限时间步长 $T$（DDPM作者使用$T$=1000）内。从$t$=0开始，在数据分布中采样真实图像$X_0$（本文使用一张来自ImageNet的猫图像形象的展示了diffusion正向添加噪声的过程），正向过程在每个时间步长$t$都从高斯分布中采样一些噪声，再添加到上一个时刻的图像中。假定给定一个足够大的$T$ 和一个在每个时间步长添加噪声的良好时间表，您最终会在$t=T$通过渐进的过程得到所谓的各向同性的Link:高斯分布。

1.2 扩散模型实现原理

1.2.1 Diffusion前向过程

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

设 $q(x_0)$ 是真实数据分布，由于 $x_0\sim q(x_0)$ ，所以我们可以从这个分布中采样以获得图像 $x_0$ 。接下来我们定义前向扩散过程 $q(x_t|x_{t-1})$ ，在前向过程中我们会根据已知的方差 $0<{\beta }_1<{\beta }_2<...<{\beta }_r<1$ 在每个时间步长$t$ 添加高斯噪声，由于前向过程的每个时刻 $t$ 只与时刻 $t-1$ 有关，所以也可以看做马尔科夫过程：

$$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$$

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

$$q({\mu }_t)=\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}$$

我们可以通过采样 $ϵ\sim \mathcal{N}(0,I)$ 然后设置

$$q({\mathbf{x}}_t)=\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}ϵ$$

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

因此，如果我们适当设置时间表，从 $X_0$ 开始，我们最终得到 $x_1,...,x_t,...,x_T$，即随着 $t$ 的增大 $x_t$ 会越来越接近纯噪声，而 $x_T$就是纯高斯噪声。

那么，如果我们知道条件概率分布$p(x_t-i|x_t)$ ，我们就可以反向运行这个过程：通过采样一些随机高斯噪声$x_T$，然后逐渐去噪它，最终得到真实分布 ${\mathbf{x}}_0$中的样本。但是，我们不知道条件概率分布$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$。这很棘手，因为需要知道所有可能图像的分布，才能计算这个条件概率。

1.2.2 Diffusion逆向过程

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

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

• 由 ${\mu }_{\theta }$ 参数化的平均值

• 由 ${\mu }_{\theta }$ 参数化的方差

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

$$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t),{\Sigma }_{\theta }({\mathbf{x}}_t,t))$$

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

• 注意，DDPM的作者决定保持方差固定，让神经网络只学习（表示）这个条件概率分布的平均值${\mu }_{\theta }$ 。
• 本文我们同样假设神经网络只需要学习（表示）这个条件概率分布的平均值 ${\mu }_{\theta }$。

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

如Sohl-Dickstein等人所示，构建Diffusion正向过程的直接结果是我们可以在条件是 x0 （因为高斯和也是高斯）的情况下，在任意噪声水平上采样 xt ，而不需要重复应用 q 去采样 xt ，这非常方便。使用
$${\alpha }_t:=1-{\beta }_t$$$$\bar{\alpha }t:=\Pi {{}^t}_{s=1}{\alpha }_s$$
我们就有
$$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)$$
这意味着我们可以采样高斯噪声并适当地缩放它，然后将其添加到$x_0$ 中，直接获得 $x_t$ 。

请注意，${\bar{\alpha }}_t$ 是已知 ${\beta }_t$ 方差计划的函数，因此也是已知的，可以预先计算。这允许我们在训练期间优化损失函数 $L$ 的随机项。或者换句话说，在训练期间随机采样 $t$ 并优化 $L_t$ 。

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

$${\mu }_{\theta }({\mathbf{x}}_t,t)=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }({\mathbf{x}}_t,t)\right)$$

最终的目标函数 $L_t$ 如下 （随机步长 $t$ 由 $(ϵ\sim \mathcal{N}(0,I))$ 给定）：

$$||ϵ-{ϵ}_{\theta }({\mathbf{x}}_t,t)|{|}^2=||ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{(1-{\bar{\alpha }}_t)}ϵ,t)|{|}^2$$

在这里， ${\mathbf{x}}_0$是初始（真实，未损坏）图像，$ϵ$是在时间步长$t$采样的纯噪声，${ϵ}_{\theta }({\mathbf{x}}_t,t)$是我们的神经网络。神经网络是基于真实噪声和预测高斯噪声之间的简单均方误差（MSE）进行优化的。

训练算法现在如下所示：

换句话说：

• 我们从真实未知和可能复杂的数据分布中随机抽取一个样本$q({\mathbf{x}}_0)$
• 我们均匀地采样$1$和$T$之间的噪声水平$t$（即，随机时间步长）
• 我们从高斯分布中采样一些噪声，并使用上面定义的属性在 $t$ 时间步上破坏输入
• 神经网络被训练以基于损坏的图像 ${\mathbf{x}}_t$ 来预测这种噪声，即基于已知的时间表 ${\mathbf{x}}_t$ 上施加的噪声

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

1.2.3 U-Net神经网络预测噪声

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

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

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

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

2 构建Diffusion模型

下面，我们逐步构建Diffusion模型。
首先，我们定义了一些帮助函数和类，这些函数和类将在实现神经网络时使用。

我们还定义了上采样和下采样操作的别名。

2.1 位置向量

由于神经网络的参数在时间（噪声水平）上共享，作者使用正弦位置嵌入来编码$t$，灵感来自Transformer（Link:Vaswani et al., 2017）。对于批处理中的每一张图像，神经网络"知道"它在哪个特定时间步长（噪声水平）上运行。

SinusoidalPositionEmbeddings模块采用(batch_size, 1)形状的张量作为输入（即批处理中几个有噪声图像的噪声水平），并将其转换为(batch_size, dim)形状的张量，其中dim是位置嵌入的尺寸。然后，我们将其添加到每个剩余块中。

2.2 ResNet/ConvNeXT块

接下来，我们定义U-Net模型的核心构建块。DDPM作者使用了一个Wide ResNet块Link:（Zagoruyko et al., 2016），但Phil Wang决定添加Link：ConvNeXT（Liu et al., 2022）替换ResNet，因为后者在图像领域取得了巨大成功。

在最终的U-Net架构中，可以选择其中一个或另一个，本文选择ConvNeXT块构建U-Net模型。

2.3 Attention模块

接下来，我们定义Attention模块，DDPM作者将其添加到卷积块之间。Attention是著名的ransformer架构(Link:Vaswani et al., 2017)，在人工智能的各个领域都取得了巨大的成功，从NLP到Link:蛋白质折叠。Phil Wang使用了两种注意力变体：一种是常规的multi-head self-attention（如Transformer中使用的），另一种是Link:LinearAttentionLink:(Shen et al., 2018)，其时间和内存要求在序列长度上线性缩放，而不是在常规注意力中缩放。 要想对Attention机制进行深入的了解，请参照Jay Allamar的Link:精彩的博文。

2.4 组归一化

DDPM作者将U-Net的卷积/注意层与群归一化Link:（Wu et al., 2018）。下面，我们定义一个PreNorm类，将用于在注意层之前应用groupnorm。

2.5 条件U-Net

我们已经定义了所有的构建块（位置嵌入、ResNet/ConvNeXT块、Attention和组归一化），现在需要定义整个神经网络了。请记住，网络${ϵ}_{\theta }({\mathbf{x}}_t,t)$ 的工作是接收一批噪声图像+噪声水平，并输出添加到输入中的噪声。

更具体的： 网络获取了一批(batch_size, num_channels, height, width)形状的噪声图像和一批(batch_size, 1)形状的噪音水平作为输入，并返回(batch_size, num_channels, height, width)形状的张量。

网络构建过程如下：

• 首先，将卷积层应用于噪声图像批上，并计算噪声水平的位置

• 接下来，应用一系列下采样级。每个下采样阶段由2个ResNet/ConvNeXT块 + groupnorm + attention + 残差连接 + 一个下采样操作组成

• 在网络的中间，再次应用ResNet或ConvNeXT块，并与attention交织

• 接下来，应用一系列上采样级。每个上采样级由2个ResNet/ConvNeXT块+ groupnorm + attention + 残差连接 + 一个上采样操作组成

• 最后，应用ResNet/ConvNeXT块，然后应用卷积层

最终，神经网络将层堆叠起来，就像它们是乐高积木一样Link:（但重要的是了解它们是如何工作的）。

2.6 正向扩散

我们已经知道正向扩散过程在多个时间步长$T$中，从实际分布逐渐向图像添加噪声，根据差异计划进行正向扩散。最初的DDPM作者采用了线性时间表：

• 我们将正向过程方差设置为常数，从${\beta }_t=10^{-4}$- 线性增加到${\beta }_T=0.02$。

• 但是，它在（Link:Nichol et al., 2021）中表明，当使用余弦调度时，可以获得更好的结果。

下面，我们定义了$T$时间步的时间表。

首先，让我们使用 $T$=200 时间步长的线性计划，并定义我们需要的 ${\beta }_t$ 中的各种变量，例如方差$\bar{\alpha }t$
的累积乘积。下面的每个变量都只是一维张量，存储从 $t$ 到 $T$ 的值。重要的是，我们还定义了extract函数，它将允许我们提取一批适当的$t$ 索引。

我们将用猫图像说明如何在扩散过程的每个时间步骤中添加噪音。

噪声被添加到mindspore张量中，而不是Pillow图像。我们将首先定义图像转换，允许我们从PIL图像转换到mindspore张量（我们可以在其上添加噪声），反之亦然。

这些转换相当简单：我们首先通过除以255来标准化图像（使它们在 [0,1] 范围内），然后确保它们在 [−1,1] 范围内。PPM论文中有介绍到：

• 假设图像数据由 {0,1,…,255} 中的整数组成，线性缩放为 [−1,1] ， 这确保了神经网络反向过程在从标准正常先验 $p({\mathbf{x}}_T)$开始的一致缩放输入上运行。
  
  我们还定义了反向变换，它接收一个包含 [−1,1] 中的张量，并将它们转回 PIL 图像：
  
  让我们验证一下：
  
  我们现在可以定义前向扩散过程，如本文所示：

让我们在特定的时间步长上测试它：

让我们为不同的时间步骤可视化此情况：

这意味着我们现在可以定义给定模型的损失函数，如下所示：

denoise_model将是我们上面定义的U-Net。我们将在真实噪声和预测噪声之间使用Huber损失。

3 数据准备与处理

在这里我们定义一个正则数据集。数据集可以来自简单的真实数据集的图像组成，如Fashion-MNIST、CIFAR-10或ImageNet，其中线性缩放为 [−1,1]
。

每个图像的大小都会调整为相同的大小。有趣的是，图像也是随机水平翻转的。根据论文内容：我们在CIFAR10的训练中使用了随机水平翻转；我们尝试了有翻转和没有翻转的训练，并发现翻转可以稍微提高样本质量。

本实验我们选用Fashion_MNIST数据集，我们使用download下载并解压Fashion_MNIST数据集到指定路径。此数据集由已经具有相同分辨率的图像组成，即28x28。

接下来，我们定义一个transform操作，将在整个数据集上动态应用该操作。该操作应用一些基本的图像预处理：随机水平翻转、重新调整，最后使它们的值在 [−1,1] 范围内。

3.1 采样

由于我们将在训练期间从模型中采样（以便跟踪进度），我们定义了下面的代码。采样在本文中总结为算法2：

从扩散模型生成新图像是通过反转扩散过程来实现的：我们从$T$开始，我们从高斯分布中采样纯噪声，然后使用我们的神经网络逐渐去噪（使用它所学习的条件概率），直到我们最终在时间步$t$=0结束。如上图所示，我们可以通过使用我们的噪声预测器插入平均值的重新参数化，导出一个降噪程度较低的图像 KaTeX parse error: Undefined control sequence: \segmat at position 1: \̲s̲e̲g̲m̲a̲t̲_t−1。请注意，方差是提前知道的。

理想情况下，我们最终会得到一个看起来像是来自真实数据分布的图像。

下面的代码实现了这一点。

请注意，上面的代码是原始实现的简化版本。

4 训练过程

下面，我们开始训练吧！

5 推理过程（从模型中采样）

要从模型中采样，我们可以只使用上面定义的采样函数：

可以看到这个模型能产生一件衣服！

请注意，我们训练的数据集分辨率相当低（28x28）。

我们还可以创建去噪过程的gif：

6 总结

请注意，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个去噪步骤就能实现高保真生成。

参考

1. Link:The Annotated Diffusion Model

2. Link:由浅入深了解Diffusion Model