扩散模型最近火得不行，你是不是也想搞懂它背后的原理？这篇文章会从理论到实践，一步步带你玩转扩散模型。先别急着翻代码，咱们先把数学基础捋清楚——毕竟，知其然更要知其所以然。接下来，我们会先深入扩散模型的理论框架，再演示如何在 PyTorch 中把它用起来。

介绍

扩散模型属于生成模型，核心任务就是生成与训练数据相似的新样本。它的工作原理其实很直观：先通过不断添加高斯噪声把训练数据“破坏”掉，再学习一个逆向过程把噪声“复原”成数据。训练完成后，你只需要扔进去一个随机噪声，模型就能通过反向去噪一步步变出图像来。

更准确地说，扩散模型是一种潜变量模型，它利用一条固定的马尔可夫链将数据映射到潜在空间。这条链会逐步向数据中添加噪声，得到近似后验分布 \(q(x_t | x_{t-1})\)，其中 \(x_t\) 与原始数据 \(x_0\) 维度相同。下图展示了这样一个马尔可夫链：

最终，图像会逐渐变成纯高斯噪声。训练扩散模型的目标就是学习逆向过程，即训练 \(p_\theta(x_{t-1}|x_t)\)。沿着这条链往回走，就能生成新的数据了。

扩散模型的优点

近年来，扩散模型的研究热度可以用“爆炸式增长”来形容。受非平衡热力学启发，当前扩散模型已经能生成顶尖质量的图像，甚至在某些方面超越了 GAN。

除了图像质量一流，扩散模型还有其他好处：不需要对抗训练。对抗训练有多难搞，大家都知道。在训练效率方面，扩散模型还具备可伸缩性和并行性，这是它的一大优势。

别看扩散模型的结果像是凭空蹦出来的，背后其实有一整套严谨的数学选择和细节在支撑，而且最佳实践还在不断演进。下面我们就来仔细看看这些数学理论。

扩散模型——深入

如前所述，扩散模型由两个过程组成：正向过程（扩散过程）和反向过程（逆向扩散过程）。正向过程对数据（通常是图像）逐步加噪，反向过程则把噪声一步步转化回目标样本。

当噪声水平足够低时，正向过程中的采样链转换可以设为条件高斯。结合马尔可夫假设，就能得到正向过程的简单参数化：

其中 \(\beta_t\) 是一个方差策略（可学习也可固定），只要设计得当，就能保证当 \(T\) 足够大时，\(x_T\) 几乎是一个各向同性的高斯噪声。

在马尔可夫假设下，潜变量的联合分布是高斯条件链变换的乘积。

扩散模型的“魔力”恰恰来自反向过程。训练时，模型学习这个扩散过程的逆过程，从而生成新数据。从纯高斯噪声出发，模型学习联合分布：

其中高斯变换的随时间变化的参数（均值 \(\mu_\theta\) 和方差 \(\Sigma_\theta\)）被学习到。注意，马尔可夫性质意味着给定的反向扩散变换分布只依赖于前一个时间步：

训练

扩散模型通过寻找反向马尔可夫变换来训练，目标是最大化训练数据的似然。实际训练等价于最小化负对数似然的变分上界。

我们尝试用 Kullback-Leibler (KL) 散度来重写这个目标。KL 散度是一种不对称的统计距离度量，衡量一个概率分布 P 与参考分布 Q 的差异程度。之所以想用 KL 散度来重写，是因为我们的马尔可夫链中的过渡分布都是高斯分布，而高斯分布之间的 KL 散度有封闭形式。

什么是 KL 散度？

连续分布的 KL 散度数学形式如下：

双杠表示该函数关于其参数不对称。

下面这张图展示了分布 P（蓝色）与参考分布 Q（红色）的 KL 散度变化。绿色曲线是上述积分内的函数，曲线下的总面积就是任意给定时刻 P 与 Q 的 KL 散度值。

将目标转换为 KL 散度的形式

如上所述，我们可以把变分下界重写成 KL 散度的形式：

其中

对后验 \(q(x_{t-1}|x_t,x_0)\) 中的前向过程进行条件化，会得到一个易于处理的形式，使得所有 KL 散度都是高斯分布之间的比较。这意味着我们可以用封闭表达式精确计算，而不必用蒙特卡洛估计。

模型选择

建立了目标函数的数学基础后，接下来需要为扩散模型的具体实现做几个关键选择。对于前向过程，唯一需要定义的是方差策略 \(\beta_t\)，其值通常随时间递增。

对于逆向过程，我们需要选择高斯分布的参数化方式以及模型架构。注意扩散模型的一大优势：它对架构的唯一要求是输入和输出维度相同。下面我们详细探讨这些选择。

前向过程和 \(\beta_t\)

关于前向过程，我们必须定义方差策略。通常，我们把它设为依赖时间的常数，虽然理论上也可以学习。例如，\(\beta_1\) 到 \(\beta_T\) 可以用线性策略，也可以用几何级数。

无论选择什么具体值，方差策略固定后，整个 \(L_T\) 相对于可学习参数就变成了常数，训练时可以忽略它。

反向过程和 \(\Sigma_\theta\)

现在讨论反向过程需要的定义。回想一下，逆马尔可夫变换被定义为高斯：

我们需要定义 \(\Sigma_\theta\) 或 \(\mu_\theta\) 的函数形式。虽然有更复杂的参数化方式，但简单做法是设置：

也就是说，假设多元高斯分布是独立高斯分布的乘积，方差相同，且随时间变化。我们将这些方差设为前向过程中的方差策略中的值。

给定新的 \(\Sigma_\theta\) 形式，我们有：

这允许我们进行变换，将目标函数中与 \(\mu_\theta\) 相关的项转化为：

变成：

其中第一项是 \(x_0\) 和 \(x_t\) 的线性组合，取决于方差策略。具体形式这里不展开。

上述比例的意义在于：最直接的参数化方式是预测扩散的后验均值 \(\mu_\theta\)。但重要的是，有学者发现，训练模型预测噪声 \(\epsilon\) 在任何给定时间步长下效果更好。具体地，令：

这里：

这会导出另一个损失函数，有学者发现它能带来更稳定的训练和更好的结果：

该学者还注意到，这种扩散模型公式与基于 Langevin 动力学的得分匹配生成模型存在联系。事实上，扩散模型和基于分数的模型似乎是同一枚硬币的两面——就像基于波的量子力学和基于矩阵的量子力学，独立发展却揭示了同一现象的两个等价公式。

网络结构

虽然简化后的损失函数旨在训练模型 \(\epsilon_\theta\)，但我们还没定义模型架构。唯一的要求是输入和输出维度相同。

鉴于这个限制，图像扩散模型通常采用类似 U-Net 的架构。

反向过程解码和 \(L_0\)

反向过程的路径由连续条件高斯分布下的多个变换组成。在反向过程结束时，我们要生成一张图像，它由整数像素值组成。因此，必须设计一种方法获得所有像素中每个可能像素值的离散（对数）似然。

做法是将反向扩散链的最后一步设为独立的离散解码器。为了确定给定生成图像的可能性，首先在数据维度之间施加独立性：

其中 \(D\) 为数据维数，上标 \(i\) 表示特定坐标。现在的目标是，在 \(t=1\) 时，一个给定像素的概率分布与轻微噪声图中对应像素的相似程度：

其中 \(t=1\) 的像素分布来自多元高斯分布，其对角协方差矩阵允许我们将分布拆分为单变量高斯分布的乘积，每个高斯分布对应数据的一个维度：

假设图像由 0,1,…,255（标准 RGB 图像）组成，这些整数已线性缩放到 [-1,1]。对于给定像素值 \(x\)，该像素值的连续变化范围是 \([x-1/255, x+1/255]\)。给定 \(x_1\) 中对应像素的单变量高斯分布，像素值 \(x\) 的概率就是以 \(x\) 为中心的 \([x-1/255, x+1/255]\) 范围内的单变量高斯分布下的面积。

下面这张图展示了每个范围内的面积及其均值为 0 的高斯概率——这里对应平均像素值为 255/2（半亮度）的分布。

对于每个像素，给定 \(t=0\) 时的像素值，可以简单相乘，过程用下式表示：

其中

并且

给定 \(p_\theta(x_0|x_1)\) 的等式，可以计算出最终 \(L_0\) 的形式，它与 KL 散度的形式不同：

最终目标

如上一节所述，作者发现预测给定时间步长的噪声效果最好，最终使用以下目标：

扩散模型的训练和采样算法如下图所示：

扩散模型总结

这一节我们详细探讨了扩散模型的理论，容易陷入数学细节。下面把最重要的要点整理出来，方便从全局角度把握：

• 扩散模型被参数化为马尔可夫链，潜变量只依赖于前一个（或后一个）时间步。
• 马尔可夫链中的变换分布是高斯的；正向过程需要方差策略，逆向过程的参数是学习的。
• 扩散过程确保当 \(T\) 足够大时，渐近分布为各向同性高斯分布。
• 方差策略可以固定（通常随时间递增，几何级数往往比线性级数效果更好），也可以学习。
• 扩散模型高度灵活，允许使用任何输入输出维度相同的架构——常见实现用 U-Net 类结构。
• 训练目标是最大化训练数据的似然，表现为调整模型参数以最小化负对数似然的变分上界。
• 由于马尔可夫假设，目标函数中几乎所有项都可转化为 KL 散度，而高斯分布使这些值可计算，无需蒙特卡洛近似。
• 最终，使用简化的训练目标（预测噪声）能得到最佳、最稳定的结果。
• 反向扩散过程的最后一步，用离散解码器获取像素值的对数似然。

有了扩散模型的高阶概述，接下来看看如何在 PyTorch 中实际使用它。

PyTorch 中的扩散模型

虽然扩散模型还没有像 CNN 或 Transformer 那样有海量的实现，但仍有可用的开源库。在 PyTorch 中使用扩散模型最简单的方式是借助 denoising-diffusion-pytorch 包，它实现了本文讨论的图像扩散模型。安装只需在终端执行：

pip install denoising_diffusion_pytorch

最小示例

要训练模型生成图像，首先导入必要的包：

import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

然后定义网络结构，这里用 U-Net。参数 dim 表示第一次下采样前的特征图数量，dim_mults 表示每次下采样时通道数的乘数：

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
)

网络定义好后，需要定义扩散模型本身。将 U-Net 模型作为参数传入，还需指定生成图像的尺寸、扩散过程的步数、选择 L1 还是 L2 损失：

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,   # number of steps
    loss_type = 'l1'    # L1 or L2
)

现在扩散模型定义好了，生成随机数据来训练，使用标准的训练流程：

training_images = torch.randn(8, 3, 128, 128)
loss = diffusion(training_images)
loss.backward()

训练完成后，使用 diffusion.sample() 生成图像。这里生成 4 张图像，由于训练数据是随机的，得到的也只会是噪声：

sampled_images = diffusion.sample(batch_size = 4)

在自定义数据集上训练

denoising-diffusion-pytorch 包也支持在特定数据集上训练。只需将下面的 Trainer 对象中的 'path/to/your/images' 替换为数据集目录路径，并将 image_size 改为适当的值。之后运行代码即可训练，然后像之前一样采样。注意，PyTorch 必须在启用 CUDA 的情况下编译才能使用 Trainer 类：

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
).cuda()

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,   # number of steps
    loss_type = 'l1'    # L1 or L2
).cuda()

trainer = Trainer(
    diffusion,
    'path/to/your/images',
    train_batch_size = 32,
    train_lr = 2e-5,
    train_num_steps = 700000,         # total training steps
    gradient_accumulate_every = 2,    # gradient accumulation steps
    ema_decay = 0.995,                # exponential moving a verage decay
    amp = True                        # turn on mixed precision
)
trainer.train()

下面这张图展示了从多元高斯噪声到 MNIST 数字的渐进去噪过程，类似于反向扩散：