Diffusion Model.

扩散模型 (Diffusion Model)是一类受到非平衡热力学 (non-equilibrium thermodynamics)启发的深度生成模型。这类模型首先定义前向扩散过程的马尔科夫链 (Markov Chain)，向数据中逐渐地添加随机噪声；然后学习反向扩散过程，从噪声中构造所需的数据样本。扩散模型也是一类隐变量模型，其隐变量通常具有较高的维度（与原始数据相同的维度）。

扩散模型的主要优点：

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

扩散模型的主要缺点：

1. 采样速度慢，单次生成需要$T$步采样；
2. 没有编码能力，无法编辑隐空间。

本文目录：

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

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$~$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$项：

其中$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{\epsilon}_\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{\epsilon}_\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)扩散模型。

simple_losses = F.l1_loss(pred_noise, noise, reduction = 'none')

⚪ 建模$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）时间连续型扩散模型的各种变体

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)\)即可实现条件控制生成。

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

（2）事前训练 Classifier-Free

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

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

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

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

对应训练的损失函数：

⭐ 参考文献

• 生成扩散模型漫谈(苏剑林)：介绍扩散模型的中文系列博客。
• 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)改进的去噪扩散概率模型。
• Diffusion Models Beat GANs on Image Synthesis：(arXiv2105)在图像合成任务上扩散模型超越了生成对抗网络。
• More Control for Free! Image Synthesis with Semantic Diffusion Guidance：(arXiv2112)基于语义扩散引导的图像合成。
• Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models：(arXiv2201)Analytic-DPM：扩散概率模型中最优反向方差的分析估计。
• Estimating the Optimal Covariance with Imperfect Mean in Diffusion Probabilistic Models：(arXiv2206)扩散概率模型中具有不准确均值的最优协方差估计。
• Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise：(arXiv2208)Cold Diffusion：反转任意无噪声的图像变换。
• Poisson Flow Generative Models：(arXiv2209)泊松流生成模型。