注解扩散模型

在这篇博客文章中，我们将深入探讨**去噪扩散概率模型**（也称为 DDPM、扩散模型、基于分数的生成模型或简称自编码器），因为研究人员已通过它们在（有）条件图像/音频/视频生成方面取得了卓越的成果。目前（撰写本文时）流行的例子包括 OpenAI 的GLIDE和DALL-E 2，海德堡大学的Latent Diffusion和 Google Brain 的ImageGen。

我们将回顾 (Ho 等人，2020) 的原始 DDPM 论文，并根据 Phil Wang 的实现（该实现本身基于原始 TensorFlow 实现）在 PyTorch 中逐步实现它。请注意，用于生成建模的扩散思想实际上在 (Sohl-Dickstein 等人，2015) 中已经提出。然而，直到 (Song 等人，2019)（斯坦福大学）和随后的 (Ho 等人，2020)（Google Brain）独立改进了该方法。

请注意，扩散模型有多种视角。这里，我们采用离散时间（潜在变量模型）的视角，但请务必也查看其他视角。

好的，让我们开始吧！

from IPython.display import Image
Image(filename='assets/78_annotated-diffusion/ddpm_paper.png')

我们首先安装并导入所需的库（假设您已安装 PyTorch）。

!pip install -q -U einops datasets matplotlib tqdm

import math
from inspect import isfunction
from functools import partial

%matplotlib inline
import matplotlib.pyplot as plt
from tqdm.auto import tqdm
from einops import rearrange, reduce
from einops.layers.torch import Rearrange

import torch
from torch import nn, einsum
import torch.nn.functional as F

什么是扩散模型？

去噪扩散模型与其他生成模型（如归一化流、GAN 或 VAE）相比并不复杂：它们都将来自某个简单分布的噪声转换为数据样本。这里也是如此，**神经网络学习从纯噪声开始逐渐去噪数据**。

更具体地针对图像，该设置包含 2 个过程：

• 我们选择的固定（或预定义）正向扩散过程 $q$，它逐渐向图像添加高斯噪声，直到最终变为纯噪声
• 一个学习到的反向去噪扩散过程 $p_\theta$，其中训练神经网络从纯噪声开始逐渐去噪图像，直到最终得到真实图像。

正向和反向过程都由 $t$ 索引，在有限的时间步长 $T$ 内进行（DDPM 作者使用 $T=1000$）。您从 $t=0$ 开始，从数据分布中采样真实图像 $\mathbf{x}_0$（例如 ImageNet 中的猫图像），并且正向过程在每个时间步长 $t$ 从高斯分布中采样一些噪声，并将其添加到上一个时间步长的图像中。给定足够大的 $T$ 和在每个时间步长添加噪声的良好调度，您最终将在 $t=T$ 处通过渐进过程得到所谓的各向同性高斯分布。

更数学化的形式

让我们更正式地写下来，因为最终我们需要一个可处理的损失函数，我们的神经网络需要对其进行优化。

令 $q(\mathbf{x}_0)$ 为真实数据分布，例如“真实图像”。我们可以从这个分布中采样得到图像，$\mathbf{x}_0 \sim q(\mathbf{x}_0)$。我们定义正向扩散过程 $q(\mathbf{x}_t | \mathbf{x}_{t-1})$，它在每个时间步长 $t$ 根据已知方差调度 $0 < \beta_1 < \beta_2 < ... < \beta_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 \geq 0$。基本上，每个在时间步长 $t$ 处的新（略微嘈杂的）图像都从**条件高斯分布**中抽取，其均值为 $\mathbf{\mu}_t = \sqrt{1 - \beta_t} \mathbf{x}_{t-1}$，方差为 $\sigma^2_t = \beta_t$，我们可以通过采样 $\mathbf{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$，然后设置 $\mathbf{x}_t = \sqrt{1 - \beta_t} \mathbf{x}_{t-1} + \sqrt{\beta_t} \mathbf{\epsilon}$。

请注意，$\beta_t$ 在每个时间步长 $t$ 并不恒定（因此带有下标）——实际上，我们定义了一个所谓的**“方差调度”**，它可以是线性的、二次的、余弦的等等，我们将在后面看到（有点像学习率调度）。

所以从 $\mathbf{x}_0$ 开始，我们最终得到 $\mathbf{x}_1, ..., \mathbf{x}_t, ..., \mathbf{x}_T$