【AI理论学习】深入理解扩散模型：Diffusion Models（DDPM）（理论篇）

本文综合最近阅读的关于扩散模型的一些基础博客和文章整理而成。主要参考的内容来自 Calvin Luo 的论文，针对的对象主要是对扩散模型已经有一些基础了解的读者。

引言

继OpenAI在2021提出的文本转图像模型DALLE之后，越来越多的大公司卷入这个方向，例如谷歌相继推出了Imagen和Parti。一些主流的文本转图像模型，例如DALL·E 2，stable-diffusion和Imagen采用了扩散模型（Diffusion Model）作为图像生成模型，这也引发了对扩散模型的研究热潮。

与GAN相比，扩散模型训练更稳定，而且能够生成更多样的样本，OpenAI的论文Diffusion Models Beat GANs on Image Synthesis也证明了<font color='red>扩散模型能够超越GAN。

论文地址：https://arxiv.org/abs/2105.05233

简单来说，扩散模型包含两个过程：前向扩散过程和反向生成过程。前向扩散过程是对一张图像逐渐添加高斯噪音直至变成随机噪音，而反向生成过程是去噪音过程，我们将从一个随机噪音开始逐渐去噪音直至生成一张图像，这也是我们要求解或者训练的部分。扩散模型与其它主流生成模型的对比如下所示：

目前所采用的扩散模型大都是来自于2020年的工作DDPM: Denoising Diffusion Probabilistic Models，DDPM对之前的扩散模型（具体见Deep Unsupervised Learning using Nonequilibrium Thermodynamics）进行了简化，并通过变分推断（variational inference）来进行建模，这主要是因为扩散模型也是一个隐变量模型（latent variable model），相比VAE这样的隐变量模型，扩散模型的隐变量是和原始数据是同维度的，而且推理过程（即扩散过程）往往是固定的。本文将基于DDPM详细介绍扩散模型的原理，并给出具体的代码实现和分析。

论文地址：https://arxiv.org/abs/2006.11239

扩散模型的原理

扩散模型包括两个过程：前向过程（forward process）和反向过程（reverse process），其中前向过程又称为为扩散过程（diffusion process），如下图所示。无论是前向过程还是反向过程都是一个参数化的马尔可夫链（Markov chain ），其中反向过程可以用来生成数据，这里我们将通过变分推断来进行建模和求解。

扩散过程

扩散过程是指的对数据逐渐增加高斯噪声直至数据完全变成随机噪声的过程。对于原始数据$x_0\sim q(x_0)$，总共包含$T$步的扩散过程的每一步都是对上一步得到的数据$x_{t-1}$按如下方式增加高斯噪声：
$$q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$$
这里 { β t } t = 1 T \{\beta_t\}_{t=1}^T { βt​}t=1T​为每一步所采用的方差，它介于$01$之间。对于扩散模型，往往称不同step的方差设定为variance schedule或者noise schedule，通常情况下，越后面的step会采用更大的方差，即满足${\beta }_1<{\beta }_2<...<{\beta }_T$。在一个设计好的variance schedule下，如果扩散步数$T$足够大，那么最终得到的$x_T$就完全丢失了原始数据而变成了一个随机噪音。 扩散过程的每一步都生成一个带噪音的数据$x_t$，整个扩散过程也就是一个马尔卡夫链：
$$q(x_{1:T}|x_0)={\prod }_{t=1}^{T}q(x_t|x_{t-1})$$

另外，要指出的是，扩散过程往往是固定的，即采用一个预先定义好的variance schedule，比如DDPM就是一个线性的variance schedule。扩散过程的一个重要特性是我们可以直接基于原始数据$x_0$来对任意$t$步的$x_t$进行采样：$x_{t}q(x_t|x_0)$。这里定义${\alpha }_t=1-{\beta }_t$和${\bar{\alpha }}_t={\prod }_{i=1}^t{\alpha }_i$，通过重参数技巧（和VAE类似），那么有：

上述推导过程利用了两个方差不同的高斯分布$\mathcal{N}(0,{\sigma }_1^2\mathtt{I})$和$\mathcal{N}(0,{\sigma }_2^2\mathtt{I})$相加等于一个新的高斯分布$\mathcal{N}(0,({\sigma }_1^2+{\sigma }_2^2)\mathtt{I})$。反重参数化后，我们得到：
$$q({\mathrm{x}}_t|{\mathrm{x}}_0)=\mathcal{N}({\mathrm{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathrm{x}}_0,(1-{\bar{\alpha }}_t)\mathtt{I})$$

扩散过程的这个特性很重要。首先，可以看到${\mathrm{x}}_t$其实可以看成是原始数据${\mathrm{x}}_0$和随机噪音$ϵ$的线性组合，其中$\sqrt{{\bar{\alpha }}_t}$和$\sqrt{1-{\bar{\alpha }}_t}$为组合系数，它们的平均和等于1，可以称两者分别为signal_rate和noise_rate

参考：https://keras.io/examples/generative/ddim/#diffusion-schedule和Variational Diffusion Models

更进一步，可以基于${\bar{\alpha }}_t$而不是${\beta }_t$来定义noise schedule(见Improved Denoising Diffusion Probabilistic Models所设计的cosine schedule)，因为这样处理更直接，比如我们直接将${\bar{\alpha }}_T$设定为一个接近0的值，那么就可以保证最终得到的${\mathrm{x}}_T$近似为一个随机噪声。其次，后面的建模和分析过程将使用这个特性。

反向过程

扩散过程是将数据噪声化，那么反向过程就是一个去噪的过程。如果我们知道反向过程的每一步的真实分布$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)$，那么从一个随机噪音${\mathrm{x}}_t\sim \mathcal{N}(0,\mathtt{I})$开始，逐渐去噪，就能生成一个真实的样本，所以反向过程也就是生成数据的过程。

估计分布$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)$需要用到整个训练样本，可以用神经网络来估计这些分布。这里，将反向过程也定义为一个马尔可夫链，只不过它是由一系列用神经网络参数化的高斯分布来组成：$$p_{\theta }({\mathrm{x}}_{0;T})=p({\mathrm{x}}_T)\prod _{t=1}^{T}p({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)$$$$p_{\theta }({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)=\mathcal{N}({\mathrm{x}}_{t-1};{\mu }_{\theta }({\mathrm{x}}_t,t),{\sum }_{\theta }({\mathrm{x}}_t,t))$$
这里$p({\mathrm{x}}_T)=\mathcal{N}({\mathrm{x}}_T;0,\mathtt{I})$，而$p_{\theta }({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)$为参数化的高斯分布，它们的均值和方差由训练的网络${\mu }_{\theta }({\mathrm{x}}_t,t)$和${\sum }_{\theta }({\mathrm{x}}_t,t)$给出。实际上，扩散模型就是要得到这些训练好的网络，因为它们构成了最终的生成模型。

虽然分布$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)$是不可直接处理的，但是加上条件${\mathrm{x}}_0$的后验分布$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t,{\mathrm{x}}_0)$却是可处理的，这里有：$$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t,{\mathrm{x}}_0)=\mathcal{N}({\mathrm{x}}_{t-1};\tilde{\mu }({\mathrm{x}}_t,{\mathrm{x}}_0,{\tilde{\beta }}_t\mathtt{I}))$$
下面来具体推导这个分布。首先根据贝叶斯公式，有：$$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t,{\mathrm{x}}_0)=q({\mathrm{x}}_t|{\mathrm{x}}_{t-1},{\mathrm{x}}_0)\frac{q({\mathrm{x}}_{t-1}|{\mathrm{x}}_0)}{q({\mathrm{x}}_t|{\mathrm{x}}_0)}$$
由于扩散过程的马尔可夫链特性，知道分布$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t,{\mathrm{x}}_0)=q({\mathrm{x}}_t|{\mathrm{x}}_{t-1})=\mathcal{N}({\mathrm{x}}_t;\sqrt{1-{\beta }_t}{\mathrm{x}}_{t-1},{\beta }_t\mathtt{I})$（这里条件${\mathrm{x}}_0$是多余的），而由前面得到的扩散过程特性可知：$$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_0)=\mathcal{N}({\mathrm{x}}_{t-1};\sqrt{{\bar{\alpha }}_{t-1}}{\mathrm{x}}_0,(1-{\bar{\alpha }}_{t-1})\mathtt{I})$$$$q({\mathrm{x}}_t|{\mathrm{x}}_0)=\mathcal{N}({\mathrm{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathrm{x}}_0,(1-{\bar{\alpha }}_t)\mathtt{I})$$
所以，有:

这里的$C({\mathrm{x}}_t,{\mathrm{x}}_0)$是一个和${\mathrm{x}}_{t-1}$无关的部分，所以省略。根据高斯分布的概率密度函数定义和上述结果（配平方）,可以得到分布$q({\mathrm{x}}_t|{\mathrm{x}}_{t-1},{\mathrm{x}}_0)$的均值和方差：

可以看到方差是一个定量（扩散过程参数固定），而均值是一个依赖${\mathrm{x}}_0$和${\mathrm{x}}_t$的函数。这个分布将会被用于推导扩散模型的优化目标。

优化目标

上面介绍了扩散模型的扩散过程和反向过程，现在我们来从另外一个角度来看扩散模型：如果把中间产生的变量看成隐变量的话，那么扩散模型其实是包含T个隐变量的隐变量模型（latent variable model），它可以看成是一个特殊的Hierarchical VAEs（见Understanding Diffusion Models: A Unified Perspective）


相比VAE来说，扩散模型的隐变量是和原始数据同维度的，而且encoder（即扩散过程）是固定的。既然扩散模型是隐变量模型，那么我们就可以基于变分推断来得到variational lower bound（VLB，又称ELBO）作为最大化优化目标，这里有：

这里最后一步是利用了Jensen's inequality（不采用这个不等式的推导见博客What are Diffusion Models?），对于网络训练来说，其训练目标为VLB取负：

进一步对训练目标进行分解可得：

可以看到最终的优化目标共包含$T+1$项，其中$L_0$可以看成是原始数据重建，优化的是负对数似然，$L_0$可以用估计的$\mathcal{N}({\mathrm{x}}_0;{\mu }_{\theta }({\mathrm{x}}_1,1),{\sum }_{\theta }({\mathrm{x}}_1,1)$来构建一个离散化的decoder来计算（见DDPM论文3.3部分）；
而$L_T$计算的是最后得到的噪音的分布和先验分布的KL散度，这个KL散度没有训练参数，近似为0，因为先验$p({\mathrm{x}}_T)=\mathcal{N}(0,\mathtt{I})$，而扩散过程最后得到的随机噪音$p({\mathrm{x}}_T|{\mathrm{x}}_0)$也近似为$\mathcal{N}(0,\mathtt{I})$；
而$L_{t-1}$则是计算的是估计分布$p_{\theta }({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)$和真实后验分布$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t,{\mathrm{x}}_0)$的KL散度，这里希望我们估计的去噪过程和依赖真实数据的去噪过程近似一致：

之所以在前面将$p_{\theta }({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)$定义为一个用网络参数化的高斯分布$\mathcal{N}({\mathrm{x}}_{t-1};{\mu }_{\theta }({\mathrm{x}}_t,t),{\sum }_{\theta }({\mathrm{x}}_t,t))$，是因为要匹配的后验分布$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t,{\mathrm{x}}_0)$也是一个高斯分布。对于训练目标$L_0$和$L_{t-1}$来说，都是希望得到训练好的网络${\mu }_{\theta }({\mathrm{x}}_t,t)$和${\sum }_{\theta }({\mathrm{x}}_t,t)$（对于$L_0$，$t=1$）。
DDPM对$p_{\theta }({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)$做了进一步简化，采用固定的方差：${\sum }_{\theta }({\mathrm{x}}_t,t)={\sigma }_t^2\mathtt{I}$，这里的${\sigma }_t^2$可以设定为${\beta }_t$或者${\tilde{\beta }}_t$（这其实是两个极端，分别是上限和下限，也可以采用可训练的方差，见论文Improved Denoising Diffusion Probabilistic Models和Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models）。


这里假定${\sigma }_t^2={\tilde{\beta }}_t$，那么：
$$q({\mathrm{x}}_{t-1}|{\mathrm{x}}_t,{\mathrm{x}}_0)=\mathcal{N}({\mathrm{x}}_{t-1};\tilde{\mu }({\mathrm{x}}_t,{\mathrm{x}}_0),{\sigma }_t^2\mathtt{I})p_{\theta }({\mathrm{x}}_{t-1}|{\mathrm{x}}_t)=\mathcal{N}({\mathrm{x}}_{t-1};{\mu }_{\theta }({\mathrm{x}}_t,t),{\sigma }_t^2\mathtt{I})$$
对于两个高斯分布的KL散度，其计算公式为：

那么就有：

那么优化目标$L_{t-1}$即为：$$L_{t-1}={\mathbb{E}}_{q({\mathrm{x}}_t|{\mathrm{x}}_0)}[\frac{1}{2{\sigma }_t^2}||{\tilde{\mu }}_t({\mathrm{x}}_t,{\mathrm{x}}_0)-{\mu }_{\theta }({\mathrm{x}}_t,t)|{|}^2]$$
从上述公式来看，我们希望网络学习到的均值${\mu }_{\theta }({\mathrm{x}}_t,t)$和后验分布的均值${\tilde{\mu }}_t({\mathrm{x}}_t,{\mathrm{x}}_0)$一致。不过DDPM发现预测均值并不是最好的选择。根据前面得到的扩散过程的特性，有：$${\mathrm{x}}_t({\mathrm{x}}_0,ϵ)=\sqrt{{\bar{\alpha }}_t}{\mathrm{x}}_0+\sqrt{1-{\bar{\alpha }}_tϵ}\text{ where }ϵ\sim \mathcal{N}(0,\mathtt{I})$$
将这个公式带入上述优化目标，可以得到：

进一步地，对${\mu }_{\theta }({\mathrm{x}}_t({\mathrm{x}}_0,ϵ),t)$也进行重参数化，变成：$${\mu }_{\theta }({\mathrm{x}}_t({\mathrm{x}}_0,ϵ),t)=\frac{1}{{\sqrt{\alpha }}_t}({\mathrm{x}}_t({\mathrm{x}}_0,ϵ)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }({\mathrm{x}}_t({\mathrm{x}}_0,ϵ),t))$$
这里的${ϵ}_{\theta }$是一个基于神经网络的拟合函数，这意味着由原来的预测均值而换成预测噪音$ϵ$。我们将上述等式代入优化目标，可以得到：

DDPM进一步对上述目标进行了简化，即去掉了权重参数，变成了$L_{t-1}^{simple}={\mathbb{E}}_{{\mathrm{x}}_0,ϵ\sim \mathcal{N}(0,\mathtt{I})}[||ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathrm{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)|{|}^2]$。这里的$t$在$[1,T]$范围内取值（如前所述，其中取1时对应$L_0$）。由于去掉了不同$t$的权重系数，所以这个简化的目标其实是VLB优化目标进行了reweight。从DDPM的对比实验结果来看，预测噪音比预测均值效果要好，采用简化版本的优化目标比VLB目标效果要好：

虽然扩散模型背后的推导比较复杂，但是最终得到的优化目标非常简单，就是让网络预测的噪音和真实的噪音一致。DDPM的训练过程也非常简单，如下图所示：随机选择一个训练样本->从1-T随机抽样一个t->随机产生噪音-计算当前所产生的带噪音数据（红色框所示）->输入网络预测噪音->计算产生的噪音和预测的噪音的L2损失->计算梯度并更新网络。

一旦训练完成，其采样过程也非常简单，如上所示：我们从一个随机噪音开始，并用训练好的网络预测噪音，然后计算条件分布的均值（红色框部分），然后用均值加标准差乘以一个随机噪音，直至t=0完成新样本的生成（最后一步不加噪音）。

不过实际的代码实现和上述过程略有区别（见https://github.com/hojonathanho/diffusion/issues/5：先基于预测的噪音生成${\mathrm{x}}_0$，并进行了clip处理（范围[-1, 1]，原始数据归一化到这个范围），然后再计算均值。这应该算是一种约束，既然模型预测的是噪音，那么我们也希望用预测噪音重构处理的原始数据也应该满足范围要求。

模型设计

前面我们介绍了扩散模型的原理以及优化目标，那么扩散模型的核心就在于训练噪音预测模型，由于噪音和原始数据是同维度的，所以我们可以选择采用AutoEncoder架构来作为噪音预测模型。DDPM所采用的模型是一个基于residual block和attention block的U-Net模型。如下所示：

U-Net属于encoder-decoder架构，其中encoder分成不同的stages，每个stage都包含下采样模块来降低特征的空间大小（H和W），然后decoder和encoder相反，是将encoder压缩的特征逐渐恢复。U-Net在decoder模块中还引入了skip connection，即concat了encoder中间得到的同维度特征，这有利于网络优化。DDPM所采用的U-Net每个stage包含2个residual block，而且部分stage还加入了self-attention模块增加网络的全局建模能力。

另外，扩散模型其实需要的是个噪音预测模型，实际处理时，我们可以增加一个time embedding（类似transformer中的position embedding）来将timestep编码到网络中，从而只需要训练一个共享的U-Net模型。具体地，DDPM在各个residual block都引入了time embedding，如上图所示。

代码实现

最后，我们基于PyTorch框架给出DDPM的具体实现，这里主要参考了三套代码实现：

• GitHub - hojonathanho/diffusion: Denoising Diffusion Probabilistic Models（官方TensorFlow实现）
• GitHub - openai/improved-diffusion: Release for Improved Denoising Diffusion Probabilistic Models （OpenAI基于PyTorch实现的DDPM+）
• GitHub - lucidrains/denoising-diffusion-pytorch: Implementation of Denoising Diffusion Probabilistic Model in Pytorch

Stable Diffusion、DALL-E、Imagen背后共同的套路

Framework

Component组件：

1. Text Encoder: 使用Text Encoder将文本生成一个embedding，
2. Generation Model：然后和噪声一起生成一个Generation Model，得到一个中间产物
3. Decoder： 使用Decoder都得到带有噪声图片

Stable Diffusion

High-Resolution Image Synthesis with Latent Diffusion Models
https://arxiv.org/abs/2112.10752

DALL-E series

Imagen

Text encoder

Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding
https://arxiv.org/abs/2205.11487

Decoder

Decoder can be trained without labelled data.

什么是FID（Frechet Inception Distance）

什么是CLIP（Contrastive Language-Image Pre-Training）

Learning Transferable Visual Models From Natural Language Supervision
https://arxiv.org/abs/2103.00020

Diffusion Model: Summary

Diffusion Model in PyTorch

The easier way to use a Diffusion Model in PyTorch is to use the denoising-diffusion-pytorch package.

pip install denoising_diffusion_pytorch

一个示例代码：

import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion
model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
)
model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
)
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()
sampled_images = diffusion.sample(batch_size = 4)

参考资料

1. 从大一统视角理解扩散模型（Diffusion Models）
2. What are Diffusion Models?
3. Understanding Diffusion Models: A Unified Perspective
4. Generative Modeling by Estimating Gradients of the Data Distribution
5. Stable Diffusion、DALL-E、Imagen 背後共同的套路
6. 带你深入理解扩散模型DDPM
7. https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_diffusion.ipynb#scrollTo=AAVZStIokTVv