『扩散模型』一篇文章入门DDPM

扩散模型DDPM

2022年，Stable Diffusion模型的问世标志着AI绘画行业从传统深度学习时代迈向AIGC时代。在这一转变中，扩散模型是Stable Diffusion的核心，因此深入了解扩散模型的原理显得尤为重要。
扩散模型本质上属于生成模型，生成模型的工作是“将噪声转换为具有代表性的数据样本”
想要了解什么是扩散模型，我们需要先了解一下扩散的概念$\Rightarrow$扩散：由浓度梯度驱动的粒子或分子由高浓度区域向低浓度区域的运动
接下来，我们首先从宏观层面探讨基于扩散的图像生成模型的工作原理。

How Do Diffusion-Based Image Generation Models Work?

扩散模型：存在一系列高斯噪声（$T$轮），在扩散模型的正向扩散过程中将输入图片$x_0$变成纯高斯噪声$x_T$，然后在扩散模型的反向扩散过程中将$x_T$还原回图片$x_0$

上面这张图，展示了扩散模型的正向扩散过程（从右到左）和反向扩散过程（从左到右）
Forward Diffusion：

1. 原始图像$x_0$通过添加高斯噪声而被迭代地缓慢破坏（马尔可夫链）
2. 这个过程是在$T$个时间步长内完成的
3. 时间步$t$的图像：$x_{t-1}+{\epsilon }_{t-1}\to x_t$
4. 在正向扩散过程$x_T$结束时，已经将数据分布转换为高斯分布

Backward Diffusion：

1. 再次以迭代方式去除Forward Diffusion中添加的噪声
2. 反向扩散过程的任务如下：给定时间步$t$和噪声图像$x_t$，预测在时间步$t-1$添加到图像中的噪声，即$x_t\to Model\to ϵ$

在整个扩散的过程中，“图像的维度保持不变”

The Need For Generative Models

扩散模型是一个概率模型，我们可以结合正向扩散过程和反向扩散过程进行理解
在正向扩散过程中，由于是向数据中逐步添加噪声的，所以每个时间步的状态依赖于前一个时间步的状态。因此，前向扩散过程中的状态转移是基于条件概率的，条件概率描述了如何从一个状态转变为另一个状态
在反向扩散过程中，由于是从噪声数据中逐步去除噪声，而且逆向扩散与正向扩散是相反的过程，所以前一个时间步的状态依赖于当前时间步的状态。因此，逆向扩散过程中的状态转移同样基于条件概率
既然涉及到条件概率，那么贝叶斯公式肯定是不可避免的

Bayesian Inference

我们首先考虑这样的一个问题：$X$的概率密度函数为$f(x|\theta )$，观测到一组样本$(x_1,x_2,...,x_n)$，此时需要估计参数$\theta$
当$\hat{\theta }$满足“$\theta =\hat{\theta }$时，该组观测样本$(x_1,x_2,...,x_n)$更容易被观测到”$\Rightarrow$$\hat{\theta }$是$\theta$的极大似然估计值，$\hat{\theta }$使得事件$(x_1,x_2,...,x_n)$发生的可能性最大
极大似然估计的数学描述：$L(\theta |x)=f(x|\theta )=f(x_1,x_2,...,x_n|\theta )={\prod }_{i=1}^{n}f(x_i|\theta )$$\Rightarrow$$\hat{\theta }=argma{x}_{\theta }L(\theta |x)$（如果有时间的话，笔者会专门写一起关于极大似然估计的文章，对极大似然估计进行介绍）
上述介绍的扩散模型和极大似然估计都与条件概率密切相关，而条件概率又与贝叶斯公式息息相关。
贝叶斯公式的数学描述：

• 离散变量：$P(B_i|A)=\frac{P(B_i)P(A|B_i)}{P(A)}=\frac{P(B_i)P(A|B_i)}{{\sum }_{j=1}^{n}P(B_j)P(A|B_j)}$，其中$P(B_i|A)$是后验概率、$P(B_i)$是先验概率、$P(A|B_i)$是似然函数、$P(A)$是边缘概率
• 连续变量：$p(z|x)=\frac{p(z)p(x|z)}{\int p(z)p(x|z)dz}$，其中$p(z|x)$是后验概率、$p(z)$是先验概率、$p(x|z)$是似然函数、$\int p(z)p(x|z)dz$是边缘概率$\iff$$\pi (\theta |x)=\frac{f(x|\theta )\pi (\theta )}{m(x)}=\frac{f(x|\theta )\pi (\theta )}{\int f(x|\theta )\pi (\theta )d(\theta )}$，其中$\pi (\theta )$是参数$\theta$的先验分布，表示对$\theta$的主观认识；$\pi (\theta |x)$是参数$\theta$的后验概率；$m(x)$是观测数据$x$的边际似然函数，边际似然函数是一个用来归一化后验概率分布的因子，确保后验概率的总和（或积分）等于1

贝叶斯估计可以看作是，在假定$\theta$服从$\pi (\theta )$的先验分布的前提下，根据样本信息去校正先验分布，得到后验分布$\pi (\theta |x)$
最大后验估计（Maximum A Posteriori estimation，MAP）：${\hat{\theta }}_{map}=argma{x}_{\theta }\pi (\theta |x)=argma{x}_{\theta }\frac{f(x|\theta )\pi (\theta )}{m(x)}=argma{x}_{\theta }f(x|\theta )\pi (\theta )$，由于$m(x)$和$\theta$无关，因此简化了计算
贝叶斯最大后验就是把极大似然加了一个先验信息

What Are Diffusion Probabilistic Models?

让我们回过头来看什么是扩散概率模型
“我们可以使用马尔可夫链将一种分布逐渐转换为另一种分布”—《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》（ICML2015、CCF-A）
马尔可夫链简单理解就是：“未来独立于过去，而只依赖于当下”（如果有时间的话，笔者会专门写一起关于马尔可夫链的文章，对马尔可夫链进行介绍）
回到扩散模型，扩散模型由两个相反的过程组成，即正向和反向扩散过程

Itsy-Bitsy Mathematical Details Behind Denoising Diffusion Probabilistic Models

箭头上提到了两个术语：

1. $q(x_t|x_{t-1})$：定义了在正向扩散过程中时间步$t$时图像的概率密度函数
2. $p_{\theta }(x_{t-1}|x_t)$：定义了在反向扩散过程中时间步$t-1$时图像的概率密度函数

Mathematical Details Of The Forward Diffusion Process

将前向扩散过程中的分布$q$定义为马尔可夫链（前向过程每个时刻$t$只与时刻$t-1$相关），由下式给出：
$\begin{array}{rlrl}q(x_1,\dots ,x_T|x_0)& :=\prod _{t=1}^{T}q(x_t|x_{t-1})& & (1)\\ q(x_t|x_{t-1})& :=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)& & (2)\end{array}$，其中高斯分布方差的超参数 { β t ∈ ( 0 , 1 ) } t = 1 T \{\beta_{t}\in(0,1)\}_{t=1}^T {βt​∈(0,1)}t=1T​

1. 首先从数据集中获取一张图像$x_0\sim q(x)$，数据集中的图像分布满足$q(x)$
2. 从时间步$1$到$T$的所有中间噪声被称为latents，latents的大小与原始图像相同
3. 正向扩散过程的概率密度函数是正态分布
4. 在每一个时间步$t$，定义图像$x_t$分布的参数：均值$\sqrt{1-{\beta }_t}{x}_{t-1}$、方差${\beta }_{t}I$
5. $\beta$被称为扩散率，使用方差调度器预先计算
6. 在每个时间步对原始图像添加少量的高斯噪声，得到$x_1,x_2,...,x_T$，添加的噪声量由调度器进行调节。随着$t$的增大，$x_t$越来越接近纯噪声，当$T\to \infty$时，$x_T$完全的高斯噪声（与均值系数$\sqrt{1-{\beta }_t}$的选择相关），且实际中${\beta }_t$随着$t$增大是递增的，即${\beta }_1<{\beta }_2<\dots <{\beta }_T$

$\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$：一个多变量的正态分布

• $\sqrt{1-{\beta }_t}{x}_{t-1}$：$x_t$的条件均值，表示$x_t$的期望是$x_{t-1}$的一个加权版本$\Rightarrow$$\sqrt{1-{\beta }_t}$描述了从前一个状态$x_{t-1}$到当前状态$x_t$的直接传递的强度，如果${\beta }_t$接近 0，这意味着$x_{t-1}$与$x_t$有很强的相关性；如果${\beta }_t$接近 1，这意味着$x_{t-1}$与$x_t$的相关性较弱
• ${\beta }_{t}I$：${\beta }_t$控制了新状态$x_t$与旧状态$x_{t-1}$之间的相似度，以及新状态的噪声水平；$I$是单位矩阵

如果从某个分布（高斯分布）中随机采样一个样本，这个过程是无法反向传播梯度的，而这个通过高斯噪声采样得到$x_t$的过程在diffusion中到处都是，因此我们需要通过重参数技巧来使得它可微（如果有时间的话，笔者会专门写一起关于重参数化的文章，对重参数化进行介绍）
最通常的做法是把随机性通过一个独立的随机变量$ϵ$引导过去。举个例子，如果要从高斯分布$z\sim \mathcal{N}(z;{\mu }_{\theta },{\sigma }_{\theta }^2\mathbf{I})$采样一个$z$，可以写成$z={\mu }_{\theta }+{\sigma }_{\theta }\odot ϵ,ϵ\sim \mathcal{N}(0,\mathbf{I})$
上式的$z$依旧是有随机性的， 且满足均值为${\mu }_{\theta }$、方差为${\sigma }_{\theta }^2$的高斯分布，这里的${\mu }_{\theta }$和${\sigma }_{\theta }^2$可以是由参数$\theta$的神经网络推断得到的。整个“采样”过程依旧梯度可导，随机性被转嫁到了$ϵ$上
How do we get image $x_t$ from $x_{t-1}$ and how is noise added at each time step?
从正态分布中对图像$x_t$进行采样：$\begin{array}{r}{x}_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}ϵ\\ \text{;where}ϵ\sim \mathcal{N}(0,I)\end{array}$

1. $ϵ$是从标准高斯分布中随机采样的“噪声”
2. 以上述方式，从$x_0$开始，原始图像从$t=\mathrm{1...}t$迭代破坏

DDPM的作者使用“线性方差调度器”并定义$\beta$在$[0.001，0.02]$范围内，并设置总时间步$T=1000$
每当需要在时间步$t$的潜在样本$x_t$时，我们必须在马尔可夫链中执行$t-1$个步骤

为了解决这个问题，DDPM的作者提出：在这个过程中直接从时间步$0$（即从原始图像）到时间步$t$

at any time $x_t$ can be represented by $x_0$ and $\beta$，我们接下来从数学层面进行推导
首先假设${\alpha }_t=1-{\beta }_t$，并且${\overline{\alpha }}_t={\prod }_{i=1}^T{\alpha }_i$，展开$x_t$可以得到：$\begin{array}{rl}{x}_t& =\sqrt{a_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_1\mathrm{where}{z}_1,z_2,\dots \sim \mathcal{N}(0,\mathbf{I});\\ & =\sqrt{a_t}(\sqrt{a_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{z}_2)+\sqrt{1-{\alpha }_t}{z}_1\\ & =\sqrt{a_t{a}_{t-1}}{x}_{t-2}+(\sqrt{a_t(1-{\alpha }_{t-1})}{z}_2+\sqrt{1-{\alpha }_t}{z}_1)\\ & =\sqrt{a_t{a}_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{z}}_2\mathrm{where}{\overline{z}}_2\sim \mathcal{N}(0,\mathbf{I});\\ & =\dots \\ & =\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}{\overline{z}}_t.\end{array}$
由于独立高斯分布可加性，即$\mathcal{N}(0,{\sigma }_1^2\mathbf{I})+\mathcal{N}(0,{\sigma }_2^2\mathbf{I})\sim \mathcal{N}(0,({\sigma }_1^2+{\sigma }_2^2)\mathbf{I})$，所以$\begin{array}{rl}& \sqrt{a_t(1-{\alpha }_{t-1})}{z}_2\sim \mathcal{N}(0,a_t(1-{\alpha }_{t-1})\mathbf{I})\\ & \sqrt{1-{\alpha }_t}{z}_1\sim \mathcal{N}(0,(1-{\alpha }_t)\mathbf{I})\\ & \sqrt{a_t(1-{\alpha }_{t-1})}{z}_2+\sqrt{1-{\alpha }_t}{z}_1\sim \mathcal{N}(0,[{\alpha }_t(1-{\alpha }_{t-1})+(1-{\alpha }_t)]\mathbf{I})\\ & =\mathcal{N}(0,(1-{\alpha }_t{\alpha }_{t-1})\mathbf{I}).\end{array}$
因此可以混合两个高斯分布得到标准差为$\sqrt{1-{\alpha }_t{\alpha }_{t-1}}$的混合高斯分布，${\overline{z}}_t$仍然是高斯分布，而任意时刻的$x_t$满足$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{a}}_t}{x}_0,(1-{\overline{a}}_t)\mathbf{I})$$\Rightarrow$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$
使用上述公式，可以在马尔可夫链中的任意时间步$t$进行采样

Mathematical Details Of The Reverse Diffusion Process

“在反向扩散过程中，任务是学习有限时间(在$T$时间步内)正向扩散过程的反转”

反向扩散的马尔可夫链从正向过程结束的地方开始，即在时间步$T$：$\begin{array}{rl}q(x_T)& \approx \mathcal{N}(x_t;0,I)\\ p(x_T)& :=\mathcal{N}(x_t;0,I)\end{array}$
反向扩散过程是从纯噪声$x_T$开始的数据样本的所有可能路径的“积分”：$p_{\theta }(x_0):=\int p_{\theta }(x_{0:T})d{x}_{1:T}$$\Rightarrow$$p_{\theta }({\mathbf{x}}_{0;T}):=p({\mathbf{x}}_T){\prod }_{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t),p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t):=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$
如果说正向过程（forward）是加噪的过程，那么反向过程（reverse）就是diffusion的去噪推断过程
无法简单推断$p_{\theta }(x_{t-1}|x_t)$，因此使用深度学习模型（参数$\theta$，目前主流是U-Net+attention的结构）去预测这样的一个反向的分布$p_{\theta }$：$\begin{array}{r}{p}_{\theta }(X_{0:T})=p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t);\\ p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t)).\end{array}$（如果有时间的话，笔者会专门写一起关于UNet的文章，对UNet进行介绍）
虽然我们无法得到逆转后的分布$q(x_{t-1}|x_t)$，但是如果知道$x_0$，是可以通过贝叶斯公式得到$q(x_{t-1}|x_t,x_0)$为：$q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\tilde{\mu }(x_t,x_0),{\tilde{\beta }}_t\mathbf{I})$
在给定条件$C$下的贝叶斯：$P(A|B,C)=\frac{P(B|A,C)\cdot P(A|C)}{P(B|C)}$
推导过程如下：

1. $p(x_{t-1}|x_t,x_0)=\frac{p(x_t|x_{t-1},x_0)\cdot p(x_{t-1}|x_0)}{p(x_t|x_0)}=\frac{p(x_t|x_{t-1})\cdot p(x_{t-1}|x_0)}{p(x_t|x_0)}$由于马尔可夫链的性质，$p(x_t|x_{t-1},x_0)$中的$x_0$可以被直接抹除掉）
2. $p(x_{t-1}|x_0)=\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}}ϵ\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_{t-1}}{x}_0,1-{\bar{\alpha }}_{t-1})$
3. $p(x_t|x_0)=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,1-{\bar{\alpha }}_t)$
4. $p(x_t|x_{t-1})=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}ϵ\sim \mathcal{N}(\sqrt{{\alpha }_t}{x}_{t-1},1-{\alpha }_t)$
5. $p(x_{t-1}|x_t,x_0)\propto exp(-\frac{1}{2}(\frac{(x_t-\sqrt{{\alpha }_t}{x}_{t-1}{)}^2}{{\beta }_t})+\frac{(x_{t-1}-{\sqrt{\alpha }}_{t-1}{x}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}-\frac{(x_t-{\sqrt{\alpha }}_t{x}_0{)}^2}{1-{\bar{\alpha }}_t})$（正态分布：$exp(-\frac{(x-\mu {)}^2}{2{\sigma }^2})=exp(-\frac{1}{2}(\frac{1}{{\sigma }^2}{x}^2-\frac{2\mu }{{\sigma }^2}x+\frac{{\mu }^2}{{\sigma }^2}))$）
6. ${\sigma }^2=\frac{{\beta }_t(1-{\bar{\alpha }}_{t-1})}{{\alpha }_t(1-{\bar{\alpha }}_{t-1})+{\beta }_t}$、$\mu =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0$
7. 但是在反向过程中，$x_0$是未知的，$x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t)$，将$x_0$代入$\mu$中，$\mu =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ)$
8. 因为$ϵ$是个噪音，用神经网络去拟合，表示为${ϵ}_{\theta }(x_t,t)$，从而可以得到要优化的目标，就是拟合出每一步的噪音，$||{ϵ}_t-{ϵ}_{\theta }(x_t,t)|{|}^2$$\Rightarrow$损失函数：噪声之间的均方差$||{ϵ}_t-{ϵ}_{\theta }(x_t,t)|{|}^2$$\Rightarrow$${\mathbb{E}}_{{\mathbf{x}}_0,ϵ\sim \mathcal{N}(0,\mathbf{I})}\left[\Vert ϵ-{ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t\right){\Vert }^2\right]$
9. DDPM最后的采样公式：$x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))+{\sigma }_tϵϵ\sim \mathcal{N}(0,I)$

DDPM论文中${\beta }_t$的生成是将一个给定范围均匀的时间步分成$T$份，然后每个时间步对应其中的某个点：

betas = torch.linspace(start=0.0001, end=0.02, steps=1000)

采用这种方式生成${\beta }_t$致的问题：前向扩散过程中靠后的时间步噪声加的很多，这些噪声在反向生成采样的过程中没有产生太大的贡献，即使跳过也不会对结果产生多大的影响

DDPM的训练过程：随机选择一个训练样本$\to$从时间步$1$到时间步$T$随机抽样一个$t$$\to$随机产生噪声$\to$计算当前带噪声的数据$\to$输入网络预测噪声$\to$计算产生的噪声和预测的噪声的L2损失$\to$计算梯度并更新网络
DDPM的采样过程：从一个随机噪声开始$\to$使用训练好的网络预测噪声$\to$计算条件分布的均值$\to$用均值加标准差乘以一个随机噪声，直至$t=0$完成新样本的生成（最后一步不加噪声）

参考文献

1、An In-Depth Guide to Denoising Diffusion Probabilistic Models – From Theory to Implementation
2、Diffusion Models for Image Generation – A Comprehensive Guide
3、由浅入深了解Diffusion Model
4、The Annotated Diffusion Model
5、DDPM和DDIM公式推导
6、扩散模型之DDPM