昇思25天学习打卡营第十二天|Diffusion扩散模型

背景

提供免费算力支持，有交流群有值班教师答疑的华为昇思训练营进入第十二天了。
今天是第十二天，从第十天开始，进入了应用实战阶段，前九天都是基础入门阶段，具体的学习内容可以看链接
基础学习部分
昇思25天学习打卡营第一天|快速入门
昇思25天学习打卡营第二天|张量 Tensor
昇思25天学习打卡营第三天|数据集Dataset
昇思25天学习打卡营第四天|数据变换Transforms
昇思25天学习打卡营第五天|网络构建
昇思25天学习打卡营第六天|函数式自动微分
昇思25天学习打卡营第七天|模型训练
昇思25天学习打卡营第八天|保存与加载
昇思25天学习打卡营第九天|使用静态图加速
应用实践部分
昇思25天学习打卡营第十天|CycleGAN图像风格迁移互换
昇思25天学习打卡营第十一天|DCGAN生成漫画头像

学习内容

扩散模型是目前最流行的文生图模型，这篇学习笔记就不写代码了，只记录一下实现原理。

什么是Diffusion Model？

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

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

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

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

扩散模型实现原理

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 }_T<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,\mathbf{I})$ 然后设置

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

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

因此，如果我们适当设置时间表，从 ${\mathbf{x}}_0$ 开始，我们最终得到 ${\mathbf{x}}_1,...,{\mathbf{x}}_t,...,{\mathbf{x}}_T$，即随着 $t$ 的增大 ${\mathbf{x}}_t$ 会越来越接近纯噪声，而 ${\mathbf{x}}_T$ 就是纯高斯噪声。

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

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 }$ 。

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

训练算法现在如下所示：

换句话说：

• 我们从真实未知和可能复杂的数据分布中随机抽取一个样本 $q({\mathbf{x}}_0)$

• 我们均匀地采样$1$和$T$之间的噪声水平$t$（即，随机时间步长）

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

• 神经网络被训练以基于损坏的图像 ${\mathbf{x}}_t$ 来预测这种噪声，即基于已知的时间表 ${\mathbf{x}}_t$ 上施加的噪声

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

U-Net神经网络预测噪声

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

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

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

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

总结

扩散模型最核心的还是加噪去噪过程，重点理解UNet非常重要。