Diffusion Model.

扩散模型 (Diffusion Model)是一类受到非平衡热力学 (non-equilibrium thermodynamics)启发的深度生成模型。这类模型首先定义前向扩散过程的马尔科夫链 (Markov Chain)，向原始图像中逐渐地添加随机噪声，直至把原始图像破坏成一张完全随机的噪声图像；然后学习反向扩散过程，从噪声图像中不断减去添加噪声（由神经网络预测），最终会还原为清晰的原始图像。扩散模型也是一类隐变量模型，其隐变量通常具有较高的维度（与原始数据相同的维度）。

扩散模型的主要优点：

• 目标函数为回归损失，训练过程平稳，容易训练；
• 与像素顺序无关的逐级自回归过程，图像生成质量高。

扩散模型的主要缺点：

• 采样速度慢，单次生成需要$T$步采样；
• 没有低维空间的编码能力，无法表征和编辑隐空间。

本文目录：

1. 时间离散型扩散模型的基本原理
2. 时间连续型扩散模型的基本原理
3. 条件扩散模型
4. 生成扩散模型的优化

1. 时间离散型扩散模型的基本原理

本节介绍时间离散型扩散模型的基本原理，主要思路如下：

1. 定义前向扩散过程：\(q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)\)
2. 解析地推导：\(q\left(\mathbf{x}_t \mid \mathbf{x}_{0}\right)\)
3. 解析地推导：\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t,\mathbf{x}_{0}\right)\)
4. 近似反向扩散过程：\(p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)\)

（1）前向扩散过程 forward diffusion process

给定从真实数据分布$q(\mathbf{x})$中采样的数据点$\mathbf{x}_0 \sim q(\mathbf{x})$，前向扩散过程定义为逐渐向样本中添加高斯噪声$\boldsymbol{\epsilon}$（共计$T$步），从而产生一系列噪声样本$\mathbf{x}_1,…,\mathbf{x}_T$。噪声的添加程度是由一系列前向方差(forward variances)系数\(\{\beta_t\in (0,1)\}_{t=1}^T\)控制的。

在前向扩散过程中，数据样本$\mathbf{x}_0$逐渐丢失其具有判别性的特征，破坏程度由$\beta_t$控制。使用重参数化技巧，可以采样任意时刻$t$对应的噪声样本$\mathbf{x}_t$。若记$\alpha_t = 1- \beta_t$，则有：

通常在前向扩散过程中会逐渐增大添加噪声的程度，即$0 \approx \beta_1<\beta_2<\cdots < \beta_T$；因此有$\bar{\alpha}_1>\cdots>\bar{\alpha}_T \approx 0$。当$T \to \infty$时，$\mathbf{x}_T$等价于一个各向同性的高斯分布。

⚪ 讨论：与随机梯度朗之万动力学的联系

朗之万动力学 (Langevin dynamics)用于对分子系统进行统计建模。结合随机梯度下降，随机梯度朗之万动力学能够从概率密度$p(\mathbf{x})$中采样，只需要使用更新过程的马尔科夫链中的梯度$\nabla_\mathbf{x} \log p(\mathbf{x})$：

其中$\delta$是更新步长。当$T \to \infty, \epsilon \to 0$，$\mathbf{x}_T$等价于真实概率密度$p(\mathbf{x})$。与标准SGD相比，随机梯度朗之万动力学在参数更新中加入高斯噪声，以避免其崩溃到局部极小值。

（2）反向扩散过程 reverse diffusion process

如果能够求得前向扩散过程\(q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)\)的逆过程\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)\)，则能够从高斯噪声输入\(\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\)出发构造真实样本。注意到当$\beta_t$足够小时，\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)\)也近似服从高斯分布。然而直接估计\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)\)是相当困难的，我们在给定数据集的基础上通过神经网络学习条件概率\(p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)\)：

注意到如果额外引入条件\(\mathbf{x}_0\)，则\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t},\mathbf{x}_0\right)\)是可解的：

其中\(C\left(\mathbf{x}_t, \mathbf{x}_0\right)\)是与\(\mathbf{x}_{t-1}\)无关的项。因此\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t},\mathbf{x}_0\right)\)也服从高斯分布：

注意到$\mathbf{x}_t=\sqrt{\bar{\alpha}_t} \mathbf{x}_0+\sqrt{1-\bar{\alpha}_t} \boldsymbol{\epsilon}$，因此把$\mathbf{x}_0=(\mathbf{x}_t-\sqrt{1-\bar{\alpha}_t} \boldsymbol{\epsilon})/\sqrt{\bar{\alpha}_t}$代入$\tilde{\boldsymbol{\mu}}_t$可得：

⚪ 理解扩散模型的反向过程

扩散模型的反向过程是通过多步迭代来逐渐生成逼真的数据，其关键是设计概率分布\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)\)。一般地，有：

对\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)\)的基本要求也是便于采样，落脚到\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_0,\mathbf{x}_t\right)\)和\(q\left(\mathbf{x}_0 \mid \mathbf{x}_t\right)\)便于采样。这样一来，就可以通过下述流程完成\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)\)的采样：

根据该分解过程，扩散模型每一步的采样\(\mathbf{x}_t \to \mathbf{x}_{t-1}\)实际上包含了两个子步骤：

1. 预估：由\(q\left(\mathbf{x}_0 \mid \mathbf{x}_t\right)\)对\(\mathbf{x}_0\)做一个简单的估计；
2. 修正：由\(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_0,\mathbf{x}_t\right)\)根据估计结果，将估值进行一定程度的修正。

扩散模型的反向过程就是一个反复的“预估-修正”过程，将原本难以一步到位的生成\(\mathbf{x}_t \to \mathbf{x}_0\)分解成逐步推进的过程，并且每一步都进行了数值修正。

（3）目标函数

扩散模型的目标函数为最小化\(p_\theta\left(\mathbf{x}_{0}\right)\)的负对数似然\(\log p_\theta\left(\mathbf{x}_0\right)\)：

可以构造负对数似然的负变分下界 (variational lower bound)：

为了把变分下界公式中的每个项转换为可计算的，可以将上述目标进一步重写为几个KL散度项和熵项的组合：

至此，扩散模型的目标函数（负变分下界）可以被分解为$T+1$项：

其中$L_T$是一个常数（$q$不包含可学习参数$\theta$, $\mathbf{x}_T$是高斯噪声），在训练时可以被省略；$L_0$可以通过一个离散解码器建模；而$L_t$计算了两个高斯分布的KL散度，可以得到闭式解。

⚪ 建模$L_t$

两个正态分布\(\mathcal{N}(\mu_1,\sigma_1^2),\mathcal{N}(\mu_2,\sigma_2^2)\)的KL散度计算为$\log \frac{\sigma_2}{\sigma_1}+\frac{\sigma_1^2+(\mu_1-\mu_2)^2}{2\sigma_2^2}-\frac{1}{2}$。根据之前的讨论，我们有：

通过神经网络建模\(\boldsymbol{\mu}_\theta\left(\mathbf{x}_t, t\right)\)和\(\boldsymbol{\Sigma}_\theta\left(\mathbf{x}_t, t\right)\)，并最小化两个分布的KL散度，即可完成扩散模型的训练过程。

def normal_kl(mean1, logvar1, mean2, logvar2):
    """
    KL divergence between normal distributions parameterized by mean and log-variance.
    """
    return 0.5 * (-1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * torch.exp(-logvar2))

在实践中，\(\boldsymbol{\Sigma}_\theta\left(\mathbf{x}_t, t\right)=\sigma_t^2I\)通常直接人为指定，而不视为可训练参数，以减小训练难度。而\(\boldsymbol{\mu}_\theta\left(\mathbf{x}_t, t\right)\)包含两个输入\(\mathbf{x}_t, t\)，即原则上对于每个不同的$t$都应构造一个不同的模型；实践中共享所有模型的参数，把$t$作为条件传入。

不妨把\(\boldsymbol{\mu}_\theta\left(\mathbf{x}_t, t\right)\)进一步表示为\(\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t, t\right)\)的函数：

则损失$L_t$可以被表示为\(\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t, t\right)\)的函数：

因此在扩散模型中，神经网络学习的目标是在每一步前向扩散过程中加入输入样本的噪声，并尝试在反向扩散过程中去除该噪声；因此扩散模型也被称作去噪(denoising)扩散模型。

⚪ 建模$L_0$

损失项$L_0$表示为：

其中\(\mathbf{x}_0\)表示原始图像（每个像素的取值从\(\{0,1,...,255\}\)归一化到$[-1,1]$）。因此\(p_\theta\left(\mathbf{x}_0 \mid \mathbf{x}_1\right) = \mathcal{N}\left(\mathbf{x}_{0} ; \boldsymbol{\mu}_\theta\left(\mathbf{x}_1, 1\right), \sigma_1^2 \mathbf{I}\right)\)实现了从\(\mathbf{x}_1\)的值域到$[-1,1]$的线性映射，可以通过一个离散解码器建模（由于初始噪声很小，因此\(\mathbf{x}_1\)的值域也落在$[-1,1]$附近，将其均分为$256$份）：

则$L_0$计算为上述分布的负对数似然（NLL）：

def approx_standard_normal_cdf(x):
    return 0.5 * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * (x ** 3))))

def discretized_gaussian_log_likelihood(x, means, log_scales, thres = 0.999):
    """
    计算 -log p_{\theta}(x_0 | x_1)
    """
    assert x.shape == means.shape == log_scales.shape

    centered_x = x - means
    inv_stdv = torch.exp(-log_scales)
    plus_in = inv_stdv * (centered_x + 1. / 255.)
    cdf_plus = approx_standard_normal_cdf(plus_in)
    min_in = inv_stdv * (centered_x - 1. / 255.)
    cdf_min = approx_standard_normal_cdf(min_in)
    log_cdf_plus = math.log(cdf_plus)
    log_one_minus_cdf_min = math.log(1. - cdf_min)
    cdf_delta = cdf_plus - cdf_min

    log_probs = torch.where(x < -thres,
        log_cdf_plus,
        torch.where(x > thres,
            log_one_minus_cdf_min,
            math.log(cdf_delta)))

    return -log_probs

注意到$\beta_1 \approx 0$，即第一次前向扩散通常设置小到可以忽略的噪声，因此在实践中通常认为\(\boldsymbol{\mu}_\theta\left(\mathbf{x}_1, 1\right)\)是没有噪声的（$\sigma_1 =0$）；此时$L_0$与训练过程不相关，\(p_\theta\left(\mathbf{x}_0 \mid \mathbf{x}_1\right)\)是一个确定性的变换：

（4）采样过程

通过训练，我们得到了反向扩散过程的近似条件分布：

从高斯噪声中随机采样\(\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}),\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\)，则扩散模型的采样过程表示为：

（5）时间离散型扩散模型的各种变体

2. 时间连续型扩散模型的基本原理

可以把扩散模型理解为一个在时间上连续的变换过程，并用随机微分方程（Stochastic Differential Equation，SDE）来描述。

（1）前向扩散SDE

前向扩散过程可以用SDE描述为：

其中\(\mathbf{w}\)是标准维纳过程；\(\mathbf{f}_t(\cdot )\)是一个向量函数，被称为\(\mathbf{x}(t)\)的漂移系数(drift coefficient)。\(g(\cdot )\)是一个标量函数，被称为\(\mathbf{x}(t)\)的扩散系数(diffusion coefficient)。

前向扩散SDE也可以等价地写成以下差分方程的形式：

或写作概率模型：

（2）反向扩散SDE

反向扩散SDE旨在求解\(p(\mathbf{x}_{t} \mid \mathbf{x}_{t+\Delta t})\)。根据贝叶斯定理：

通常$\Delta t$比较小，因此有泰勒展开：

代入上式并配方得：

上式也可写作差分方程：

取$\Delta t \to 0$，得到反向过程对应的SDE：

（3）得分匹配

前向和反向扩散过程的SDE：

也可以等价地写成差分形式：

如果进一步知道\(\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)\)，就可以通过反向SDE完成生成过程。

考虑到在离散型的扩散模型中，通常会为\(p(\mathbf{x}_t \mid \mathbf{x}_0)\)设计具有解析解的形式。此时\(\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)\)表示为：

上式表示\(\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)\)计算为\(\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t \mid \mathbf{x}_0)\)在分布\(p(\mathbf{x}_t \mid \mathbf{x}_0)\)上的加权平均。

我们希望用神经网络学一个函数\(s_θ(\mathbf{x}_t,t)\)，使得它能够直接计算\(\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)\)。则\(s_θ(\mathbf{x}_t,t)\)应当也能表示为\(\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t \mid \mathbf{x}_0)\)在分布\(p(\mathbf{x}_t \mid \mathbf{x}_0)\)上的加权平均，或者等价地写成如下损失：

上式被称为(条件)得分匹配 (score matching)损失。

（4）连续型扩散模型的一般流程

构造连续型扩散模型的一般流程：

① 通过随机微分方程定义前向扩散过程：

② 求\(p(\mathbf{x}_t \mid \mathbf{x}_0)\)的表达式；

③ 通过得分匹配损失训练\(s_θ(\mathbf{x}_t,t)\)：

④ 通过随机微分方程实现反向扩散过程：

（5）由SDE到ODE

扩散模型对应了一个随机微分方程，实际上也隐含了一个常微分方程 (ordinary differential equation, ODE)。相比于SDE，ODE用于描述确定性系统的动态行为，不涉及随机性，具有更好的数值稳定性，并且可以实现加速采样。

（6）时间连续型扩散模型的各种变体

3. 条件扩散模型

条件扩散模型是指根据给定的条件来控制生成结果，即在采样过程\(p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)\)中引入输入条件\(\mathbf{y}\)，使得采样过程变为\(p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t,\mathbf{y}\right)\)。条件控制生成的方式可分为两种：事后修改（Classifier-Guidance）和事前训练（Classifier-Free）。

（1）事后修改 Classifier-Guidance

事后修改是指在已经训练好的无条件扩散模型的基础上引入一个可训练的分类器（Classifier），用分类器来调整生成过程以实现控制生成。这类模型的训练成本比较低，但是采样成本会高一些，而且难以控制生成图像的细节。

⚪ 离散形式的事后修改

为了重用训练好的模型\(p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)\)，根据贝叶斯定理：

为了进一步得到可采样的近似结果，在\(\mathbf{x}_{t-1}=\boldsymbol{\mu}_\theta\left(\mathbf{x}_t, t\right)\)处考虑泰勒展开：

并注意到反向扩散过程的建模：

因此有：

则\(p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t,\mathbf{y}\right)\)近似服从正态分布：

因此条件扩散模型\(p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t,\mathbf{y}\right)\)的采样过程为：

⚪ 连续形式的事后修改

对于前向扩散SDE：

对应的反向扩散SDE为：

条件扩散模型需要把\(p_t(\mathbf{x})\)替换为\(p_t(\mathbf{x} \mid \mathbf{y})\)。根据贝叶斯定理：

在一般的参数化下有：

因此：

因此只需要用\(\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t, t\right)-\sqrt{1-\bar{\alpha}_t}\nabla_{\mathbf{x}} \log p_t(\mathbf{y}\mid \mathbf{x})\)替换\(\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t, t\right)\)即可实现条件控制生成：

为了控制事后修改的程度，引入分类器引导尺度 （classifier guidance scale） $w$:

分类器引导尺度$w$越大，模型倾向于生成更具有确定性的结果。

⚪ 基于事后修改的条件扩散模型

（2）事前训练 Classifier-Free

事前训练是指在扩散模型的训练过程中引入条件信号。这类模型的主要缺点是重新训练扩散模型的成本较大，在预算充裕的前提下，这类模型生成的图像通常具有较好的细节。

注意到反向扩散过程的建模：

引入输入条件\(\mathbf{y}\)后，定义反向扩散过程：

一般把\(\boldsymbol{\mu}_\theta\)表示为\(\boldsymbol{\epsilon}_\theta\)的函数：

对应训练的损失函数：

此时应把噪声估计器\(\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t,\mathbf{y}, t\right)\)建模为条件输入模型，如通过交叉注意力机制实现。

⚪ 讨论：事后修改与事前训练的关系

注意到：

对于事后修改的条件扩散模型，其噪声估计器$\boldsymbol{\epsilon}_\theta$表示为：

向噪声估计器$\boldsymbol{\epsilon}_\theta$中引入条件$\mathbf{y}$，有：

此时$s = 1+w\geq 1$被称作无分类器引导尺度（classifier-free guidance scale, CFG scale）。因此通过在有标签和无标签的扩散模型的预测之间进行插值，可以实现类似于事后修改使用分类器梯度将样本引导至标签的效果。

注意到\(\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t, t,\mathbf{y}\right)\)和\(\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t, t\right)\)可以用同一个模型表示：\(\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t, t\right)=\boldsymbol{\epsilon}_\theta\left(\mathbf{x}_t, t,\mathbf{y}=\phi\right)\)；因此在训练时以一定的概率丢弃条件，使得模型同时学习两种情况。

⚪ 基于事前训练的条件扩散模型

5. 生成扩散模型的优化

（1）高分辨率图像生成

⚪ DALL-E 2

DALL-E 2模型通过CLIP模型找到文本特征与图像特征的对应关系，然后训练一个prior模型将文本特征转换为图像特征，再训练一个decoder模型（条件扩散模型GLIDE）根据图像特征生成图像。

⚪ Imagen

Imagen模型采用仅在文本数据上训练的大型冻结语言模型T5-XXL编码文本嵌入序列，并将其为条件输入64×64文本图像扩散模型与两个256×256和1024×1024的超分辨率扩散模型。

⚪ Latent Diffusion Model (LDM)

隐扩散模型（latent diffusion model）没有直接在高维图像空间中执行扩散，而是首先使用变分自编码器把图像压缩到隐空间，再在隐空间中构造扩散过程。

⚪ SDXL

SDXL模型是一个二阶段的级联扩散模型，包括Base模型和Refiner模型。其中Base模型是基本的文生图模型；在Base模型之后级联的Refiner模型是图生图模型，用于对Base模型生成的图像隐特征进行精细化。

（2）采样加速

⚪ Score identity Distillation (SiD)

SiD提出了一种将扩散模型蒸馏为单步生成模型的方案：通过学生模型$x=g_\theta(z)$采样数据并训练扩散模型$\epsilon_\psi(x_t,t)$，并最小化其与教师扩散模型$\epsilon_\phi(x_t,t)$的差异。上述两个目标交替优化：

⭐ 参考文献

• 生成扩散模型漫谈(苏剑林)：介绍扩散模型的中文系列博客。
• What are Diffusion Models?(Lil’Log)：一篇介绍扩散模型的英文博客。
• denoising-diffusion-pytorch(GitHub)：Implementation of Denoising Diffusion Probabilistic Model in Pytorch。
• Denoising Diffusion Probabilistic Models：(arXiv2006)DDPM：去噪扩散概率模型。
• Denoising Diffusion Implicit Models：(arXiv2010)DDIM：去噪扩散隐式模型。
• Score-Based Generative Modeling through Stochastic Differential Equations：(arXiv2011)基于得分匹配的随机微分方程生成式建模。
• Improved Denoising Diffusion Probabilistic Models：(arXiv2102)改进的去噪扩散概率模型。
• Variational Diffusion Models：(arXiv2105)在图像合成任务上扩散模型超越了生成对抗网络。
• Diffusion Models Beat GANs on Image Synthesis：(arXiv2107)变分扩散模型。
• More Control for Free! Image Synthesis with Semantic Diffusion Guidance：(arXiv2112)基于语义扩散引导的图像合成。
• High-Resolution Image Synthesis with Latent Diffusion Models：(arXiv2112)通过隐扩散模型实现高分辨率图像合成。
• GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models：(arXiv2112)GLIDE：通过文本引导的扩散模型实现真实图像生成与编辑。
• Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models：(arXiv2201)Analytic-DPM：扩散概率模型中最优反向方差的分析估计。
• Hierarchical Text-Conditional Image Generation with CLIP Latents：(arXiv2204)通过CLIP隐特征实现层次化文本条件图像生成。
• Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding：(arXiv2205)通过深度语言理解实现真实文本图像扩散模型。
• Estimating the Optimal Covariance with Imperfect Mean in Diffusion Probabilistic Models：(arXiv2206)扩散概率模型中具有不准确均值的最优协方差估计。
• Classifier-Free Diffusion Guidance：(arXiv2207)无分类器引导的条件扩散模型。
• Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise：(arXiv2208)Cold Diffusion：反转任意无噪声的图像变换。
• Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow：(arXiv2209)通过整流流实现数据的生成与转换。
• Poisson Flow Generative Models：(arXiv2209)泊松流生成模型。
• SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis：(arXiv2307)SDXL：改进高分辨率图像合成的隐扩散模型。
• Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation：(arXiv2404)得分恒等蒸馏：单步生成的预训练扩散模型的指数级快速蒸馏。