Diffusion Models 扩散模型
我们已经了解到,构建强大的生成模型的一种有效方法是:先引入一个关于潜在变量z的分布p(z),然后使用深度神经网络将z变换到数据空间x。由于神经网络具有通用性,能够将简单固定的分布转化为关于x的高度灵活的分布族,因此为p(z)采用如高斯分布N(z|0, I)这类简单固定的分布就足够了。在之前的章节中,我们探讨了多种符合这一框架的模型,这些模型基于生成对抗网络、变分自编码器以及归一化流,在定义和训练深度神经网络方面采用了不同的方法。
在本章中,我们将探讨这一通用框架下的第四类模型,即扩散模型(diffusion models),也被称为去噪扩散概率模型(denoising diffusion probabilistic models,简称DDPMs)(Sohl-Dickstein等人,2015年;Ho、Jain和Abbeel,2020年)。这类模型已成为许多应用领域中最先进的模型。为便于说明,我们将重点讨论图像数据模型,尽管该框架具有更广泛的适用性。其核心思想是,对每张训练图像应用多步加噪过程,将其逐步破坏并最终转化为一个服从高斯分布的样本。这一过程如图20.1所示。随后,训练一个深度神经网络来逆转这一过程;一旦训练完成,该网络就可以从高斯分布中采样作为输入,进而生成新的图像。
扩散模型可被视为一种分层变分自编码器的变体,其中编码器分布是固定的,由加噪过程所定义,而仅需学习生成分布(Luo,2022)。这类模型易于训练,在并行硬件上扩展性良好,且能规避对抗训练中的挑战与不稳定性问题,同时生成结果的质量可与生成对抗网络相媲美甚至更优。然而,由于需要通过解码器网络进行多次前向传播,生成新样本的计算成本可能较高(Dhariwal和Nichol,2021)。
20.1 Forward Encoder 前向编码器
假设我们从训练集中选取一张图像,记为 x,然后对每个像素独立地添加高斯噪声进行混合,得到一个受噪声干扰的图像 z₁,其定义如下:
其中,ε₁ ∼ N(ε₁|0, I),且 β₁ < 1 为噪声分布的方差,β₁ ∈ (0,1):控制噪声强度的超参数(称为噪声的方差系数)。公式(20.1)和(20.3)中系数 √(1 − β₁) 和 √β₁ 的选取,确保了 zₜ 的分布均值相较于 zₜ₋₁ 更接近零,且 zₜ 的协方差矩阵相较于 zₜ₋₁ 更接近单位矩阵。我们可以将变换(20.1)改写为如下形式:
随后,我们通过多次添加独立的加性高斯噪声,重复这一过程,从而生成一系列噪声逐步增强的图像序列 z₂, …, zₜ。需要注意的是,在扩散模型的相关文献中,这些潜在变量有时被记为 x₁, …, xₜ,而观测变量则记为 x₀。为与本书其余部分保持一致,我们采用 z 表示潜在变量、x 表示观测变量的符号体系。每一幅后续图像的生成方式如下:
其中,εₜ ∼ N(εₜ|0, I)。同样地,我们可以将公式(20.3)改写为如下形式:
条件分布序列(20.4)构成一个马尔可夫链,并可用如图20.2所示的概率图模型表示。方差参数 βₜ ∈ (0, 1) 需手动设定,通常按照预设的递增规则选择,使得方差值沿链逐步增大,即满足 β₁ < β₂ < … < βₜ。
这个公式描述的是扩散模型(Diffusion Model)中的前向过程(Forward Process),即如何逐步向原始图像添加噪声,最终将其转化为纯噪声。详细解析公式和方差参数空间门
公式 (20.3)
z t = 1 − β t z t − 1 + β t ϵ t z_t = \sqrt{1 - \beta_t} z_{t-1} + \sqrt{\beta_t} \epsilon_t zt=1−βtzt−1+βtϵt
- z t z_t zt 和 z t − 1 z_{t-1} zt−1:分别表示在时间步 t t t 和 t − 1 t-1 t−1 的图像(或潜在变量)。
- β t \beta_t βt:是一个预先设定的参数,取值范围在 ( 0 , 1 ) (0, 1) (0,1) 之间,控制每一步添加噪声的比例。
- ϵ t \epsilon_t ϵt:是从标准正态分布 N ( ϵ t ∣ 0 , I ) \mathcal{N}(\epsilon_t | 0, \mathbf{I}) N(ϵt∣0,I) 中采样得到的高斯噪声,其中 I \mathbf{I} I 是单位矩阵。
这个公式的含义是,在每一步 t t t,当前的图像 z t z_t zt 是由前一步的图像 z t − 1 z_{t-1} zt−1 加上一定比例的高斯噪声 β t ϵ t \sqrt{\beta_t} \epsilon_t βtϵt 得到的,同时保留了前一步图像的一部分 1 − β t z t − 1 \sqrt{1 - \beta_t} z_{t-1} 1−βtzt−1。
公式 (20.4)
q ( z t ∣ z t − 1 ) = N ( z t ∣ 1 − β t z t − 1 , β t I ) q(z_t | z_{t-1}) = \mathcal{N}(z_t | \sqrt{1 - \beta_t} z_{t-1}, \beta_t \mathbf{I}) q(zt∣zt−1)=N(zt∣1−βtzt−1,βtI)
这个公式是从概率的角度描述了公式 (20.3)。它表示在给定 z t − 1 z_{t-1} zt−1 的条件下, z t z_t zt 的条件概率分布是一个均值为 1 − β t z t − 1 \sqrt{1 - \beta_t} z_{t-1} 1−βtzt−1,协方差矩阵为 β t I \beta_t \mathbf{I} βtI 的高斯分布。
参数 β t \beta_t βt 的设定
参数 β t \beta_t βt 是手动设定的,并且通常按照一个预设的递增规则选择,使得 β 1 < β 2 < ⋯ < β T \beta_1 < \beta_2 < \cdots < \beta_T β1<β2<⋯<βT。这意味着随着时间步 t t t 的增加,添加的噪声比例逐渐增大,最终在 t = T t = T t=T 时,图像几乎完全转化为噪声。
通过这种方式,扩散模型的前向过程逐步将原始图像转化为一个各向同性的高斯分布,为后续的逆向生成过程(即从噪声中恢复图像)奠定了基础。
图20.2 一个表示为概率图模型的扩散过程。原始图像 x \mathbf{x} x由阴影节点表示,因为它是一个观测变量,而受噪声干扰的图像 z 1 , … , z T z_1, \ldots, z_T z1,…,zT被视为潜在变量。噪声过程由前向分布 q ( z t ∣ z t − 1 ) q(z_t|z_{t - 1}) q(zt∣zt−1)定义,可被视为一个编码器。我们的目标是学习一个模型 p ( z t − 1 ∣ z t , w ) p(z_{t - 1}|z_t, \mathbf{w}) p(zt−1∣zt,w),该模型试图逆转这个噪声过程,可被视为一个解码器。正如我们稍后将看到的,条件分布 q ( z t − 1 ∣ z t , x ) q(z_{t - 1}|z_t, \mathbf{x}) q(zt−1∣zt,x)在定义训练过程中起着重要作用。
图的解析
这张图展示了一个用概率图模型表示的扩散过程。
节点和变量:
- 图中左侧的蓝色阴影节点表示原始图像 x \mathbf{x} x,它是一个观测变量。
- 红色节点 z 1 , … , z T z_1, \ldots, z_T z1,…,zT 表示受噪声干扰的图像,被视为潜在变量。
过程描述:
- 前向过程(噪声添加过程):由前向分布 q ( z t ∣ z t − 1 ) q(z_t | z_{t - 1}) q(zt∣zt−1) 定义,可以看作是一个编码器。从原始图像 x \mathbf{x} x 开始,通过一系列步骤逐步添加噪声,最终在 T T T 步后得到一个接近纯噪声的潜在变量 z T z_T zT。图中绿色虚线箭头表示了这一前向过程的方向,即从 z t − 1 z_{t-1} zt−1 到 z t z_t zt 的噪声添加步骤。
- 逆向过程(生成过程):目标是学习一个模型 p ( z t − 1 ∣ z t , w ) p(z_{t - 1} | z_t, \mathbf{w}) p(zt−1∣zt,w),试图逆转这个噪声添加过程,可以看作是一个解码器。图中蓝色实线箭头表示了逆向过程的方向,即从 z t z_t zt 恢复到 z t − 1 z_{t - 1} zt−1 的步骤。
条件分布的作用:
- 条件分布 q ( z t − 1 ∣ z t , x ) q(z_{t - 1} | z_t, \mathbf{x}) q(zt−1∣zt,x) 在定义训练过程中起着重要作用。它结合了观测变量 x \mathbf{x} x 的信息,帮助模型学习如何从噪声潜在变量 z t z_t zt 恢复到前一步的潜在变量 z t − 1 z_{t - 1} zt−1。
总体而言,这张图直观地展示了扩散模型中前向的噪声添加过程和逆向的生成过程,以及各个变量和分布之间的关系。
20.1.1 Diffusion kernel
在给定观测数据向量 x \mathbf{x} x 的条件下,潜在变量的联合分布由下式给出:
如果我们现在对中间变量 z 1 , … , z t − 1 z_1, \ldots, z_{t - 1} z1,…,zt−1 进行边缘化处理(边缘化处理的目的是消除中间变量,从而得到仅关于 z t z_t zt 和 x \mathbf{x} x 的条件概率分布 q ( z t ∣ x ) q(\mathbf{z}_t | \mathbf{x}) q(zt∣x)),就得到了扩散核:
在此我们定义:
我们看到,每个中间分布都有一个简单的闭式高斯表达式,我们可以直接从中进行采样。这在训练去噪扩散概率模型(Denoising Diffusion Probabilistic Models, DDPMs)时将被证明是很有用的,因为它允许使用马尔可夫链中随机选择的中间项来进行高效的随机梯度下降,而无需运行整个链。我们还可以将公式(20.6)写成如下形式
其中再次(有) ϵ t ∼ N ( ϵ t ∣ 0 , I ) \epsilon_t \sim \mathcal{N}(\epsilon_t|0, \mathbf{I}) ϵt∼N(ϵt∣0,I)(这里表示 ϵ t \epsilon_t ϵt服从均值为0、协方差矩阵为单位矩阵 I \mathbf{I} I的多元正态分布)。请注意,此时 ϵ t \epsilon_t ϵt代表添加到原始图像上的总噪声,而不是在马尔可夫链的这一步添加的增量噪声。
经过许多步骤后,图像变得与高斯噪声无法区分,当步数 T T T 趋于无穷大( T → ∞ T \to \infty T→∞)时,我们有(以下情况/结果)
因此,关于原始图像的所有信息都丢失了。公式 (20.3) 中系数 1 − β t \sqrt{1 - \beta_t} 1−βt 和 β t \sqrt{\beta_t} βt 的选择确保了一旦马尔可夫链收敛到一个均值为0、协方差为单位矩阵的分布,进一步的更新将不会改变这一分布。
由于公式 (20.9) 的右侧与 x \mathbf{x} x(或 x,根据上下文确定具体变量)无关,因此可以得出 z T z_T zT(或 z T \mathbf{z}_T zT ,根据上下文确定具体变量)的边缘分布由以下式子给出(或可表示为)
人们通常将马尔可夫链(公式20.4)称为前向过程,它类似于变分自动编码器(VAE)中的编码器,只不过这里(马尔可夫链所代表的过程)是固定的,而非通过学习得到的。然而,需要注意的是,文献中的常用术语与标准化流(normalizing flows)相关文献中通常使用的术语是相反的,在标准化流的文献中,从潜在空间到数据空间的映射被视为前向过程。
20.1.2 Conditional distribution 条件扩散
我们的目标是学会逆转(消除、还原)加噪过程,因此很自然地会考虑条件分布 q ( z t ∣ z t − 1 ) q(z_t | z_{t - 1}) q(zt∣zt−1) 的逆过程,我们可以借助贝叶斯定理将其表示为如下形式
我们可以将边缘分布 q ( z t − 1 ) q(z_{t-1}) q(zt−1) 表示为如下形式
其中, q ( z t − 1 ∣ x ) q(z_{t-1}|\mathbf{x}) q(zt−1∣x) 由条件高斯分布(公式20.6)给出。然而,该分布难以直接处理(无法解析求解),因为我们必须对未知的数据分布 p ( x ) p(\mathbf{x}) p(x) 进行积分。如果我们使用训练数据集中的样本对积分进行近似,则会得到一个复杂的分布,该分布可表示为高斯混合分布。
相反,我们考虑反向分布的条件版本,即在给定数据向量 x \mathbf{x} x 的条件下,定义为 q ( z t − 1 ∣ z t , x ) q(z_{t-1}|z_t, \mathbf{x}) q(zt−1∣zt,x) 的分布。我们很快将会看到,该分布实际上是一个简单的高斯分布。从直觉上看,这是合理的:因为给定一张含噪图像时,很难推测出是哪张低噪声图像生成了它;而如果我们还知道原始(起始)图像,那么问题就会变得简单得多。我们可以利用贝叶斯定理来计算这个条件分布:
现在,我们利用前向过程的马尔可夫性质来进行推导( z t z_t zt只和前一时刻 z t − 1 z{t-1} zt−1有关,与x无关)
其中,等式右侧由(20.4)式给出。作为 z t − 1 z_{t-1} zt−1 的函数,它呈现为二次型的指数形式。(20.13)式分子中的 q ( z t − 1 ∣ x ) q(z_{t-1}|\mathbf{x}) q(zt−1∣x) 项是由(20.6)式给出的扩散核,它同样涉及关于 z t − 1 z_{t-1} zt−1 的二次型指数。由于(20.13)式的分母作为 z t − 1 z_{t-1} zt−1 的函数是常数,我们可以忽略它。因此,我们看到(20.13)式的右侧呈现为高斯分布的形式,并且我们可以使用“配方法”来确定其均值和协方差,具体如下:
上述条件分布的转化,化简解释秘籍——》宝典
20.2 Reverse Decoder
我们已经看到,前向编码器模型是由一系列高斯条件分布 q(zₜ|zₜ₋₁) 定义的,但直接对其进行逆向推导会得到一个难以处理的分布 q(zₜ₋₁|zₜ),因为这需要对所有可能的起始向量 x 的值进行积分,而其分布正是我们希望建模的未知数据分布 p(x)。相反,我们将使用一个由深度神经网络控制的分布 p(zₜ₋₁|zₜ, w) 来学习对反向分布的近似,其中 w 表示网络的权重和偏置。这一反向步骤类似于变分自编码器中的解码器,具体如图20.2所示。网络训练完成后,我们可以从 zₜ 的简单高斯分布中采样,并通过反复应用训练好的网络进行一系列反向采样步骤,从而将其转换为来自数据分布 p(x) 的样本。
图20.4描绘的左边是q(zₜ|zₜ₋₁)分布,拥有更小的方差βt,右侧蓝色曲线为q(zₜ₋₁|zₜ),可以看出近似一个高斯分布。
直观来讲,如果我们将方差控制得较小,使得βt≪1,那么各步骤间潜在向量(latent vector)的变化会相对较小,因此学习逆转该变换过程应该会更容易。更具体地说,如果βt≪1,那么条件分布q(zₜ₋₁|zₜ) 将近似为关于zₜ₋₁的一个高斯分布。
这个结论可以从数学角度进行严谨解释:
1. 数学推导基础:
根据贝叶斯定理,逆向分布可表示为:
q(zₜ₋₁|zₜ) ∝ q(zₜ|zₜ₋₁)q(zₜ₋₁)
2. 关键假设条件:
当βₜ≪1时:
- 前向转移核q(zₜ|zₜ₋₁) = N(zₜ; √(1-βₜ)zₜ₋₁, βₜI) 成为"窄峰"分布
- 在zₜ₋₁的局部邻域内,边缘分布q(zₜ₋₁)近似恒定
3. 分布近似过程:
由于βₜ极小:
√(1-βₜ) ≈ 1 ⇒ q(zₜ|zₜ₋₁) ≈ N(zₜ; zₜ₋₁, βₜI)
这使得乘积q(zₜ|zₜ₋₁)q(zₜ₋₁)中:
- q(zₜ₋₁)在zₜ邻域内基本不变
- 分布形态完全由q(zₜ|zₜ₋₁)的高斯项主导
4. 物理意义:
小βₜ意味着:
- 单步噪声扰动极小
- 逆向过程只需做"微调"
- 局部线性假设成立
5. 对比实验验证:
当βₜ较大时(如图20.3):
- 前向核覆盖范围广
- 可能跨越q(zₜ₋₁)的多个峰
- 导致逆向分布呈多模态
6. 实际应用启示:
这种近似使得:
- 可以用高斯分布参数化逆向过程
- 神经网络只需预测均值/方差
- 大幅降低学习难度
(补充说明:这个性质是DDPM等扩散模型的理论基础,但实际实现时需要权衡βₜ大小与计算成本)
这一点可从公式(20.11)中看出,因为该公式的右侧通过q(zₜ|zₜ₋₁)和q(zₜ₋₁)与zₜ₋₁相关联。如果q(zₜ|zₜ₋₁)是一个足够窄的高斯分布,那么在q(zₜ|zₜ₋₁)具有显著概率质量的区域内,q(zₜ₋₁)的变化将非常小,因此q(zₜ₋₁|zₜ)也将近似为高斯分布。如图20.3和图20.4所示的一个简单例子可以证实这种直觉。然而,由于每一步的方差都很小,我们必须使用大量的步骤来确保通过前向加噪过程得到的最终潜在变量zₜ的分布仍然接近高斯分布,而这会增加生成新样本的成本。在实际应用中,T可能达到数千。
通过围绕点(将 z t z_t zt 视为关于 z t − 1 z_{t-1} zt−1(的函数))对 ln q ( z t − 1 ∣ z t ) \ln q(z_{t-1}|z_t) lnq(zt−1∣zt) 进行泰勒级数展开,我们可以更正式地看出 q ( z t − 1 ∣ z t ) q(z_{t-1}|z_t) q(zt−1∣zt) 将近似为高斯分布。这也表明,当方差较小时,反向分布 q ( z t ∣ z t − 1 ) q(z_t|z_{t-1}) q(zt∣zt−1) 的协方差将接近前向噪声过程 q ( z t − 1 ∣ z t ) q(z_{t-1}|z_t) q(zt−1∣zt)(此处应为 q ( z t ∣ z t − 1 ) q(z_t|z_{t - 1}) q(zt∣zt−1) ,原表述有误)的协方差 β t I \beta_t I βtI。因此,我们使用如下形式的高斯分布对反向过程进行建模:
其中, μ ( z t , w , t ) \mu(z_t, w, t) μ(zt,w,t) 是一个由参数集 w w w 控制的深度神经网络。需注意,该网络将步骤索引 t t t 明确作为输入,以便能够考虑链式过程中不同步骤间方差 β t \beta_t βt 的变化。这使得我们能够使用单个网络来逆转马尔可夫链中的所有步骤,而无需为每个步骤单独学习一个网络。此外,通过在网络中引入更多输出以考虑 z t z_t zt 邻域内分布 q ( z t − 1 ) q(z_{t-1}) q(zt−1) 的曲率,我们还可以学习去噪过程的协方差(Nichol 和 Dhariwal,2021)。在满足输出与输入维度相同这一限制的前提下,用于建模 μ ( z t , w , t ) \mu(z_t, w, t) μ(zt,w,t) 的神经网络在架构选择上具有相当大的灵活性。鉴于这一限制,在图像处理应用中,U-net 架构是一种常见选择。
那么,整体反向去噪过程就呈现出如下(所给出的)马尔可夫链形式:
此处,假设 p ( z T ) p(z_T) p(zT)与 q ( z T ) q(z_T) q(zT)的分布相同,因此其形式为 N ( z T ∣ 0 , I ) \mathcal{N}(z_T|0, I) N(zT∣0,I)(即均值为0、单位协方差矩阵的高斯分布)。一旦模型训练完成,采样过程便十分直接:我们首先从简单的高斯分布 p ( z T ) p(z_T) p(zT)中采样,然后依次从每个条件分布 p ( z t − 1 ∣ z t , w ) p(z_{t-1}|z_t, w) p(zt−1∣zt,w)中顺序采样,最终从 p ( x ) p(x) p(x)。
20.2.1 Training the decoder
接下来,我们需要为神经网络的训练确定一个目标函数。
显而易见的选择是似然函数,对于数据点 x x x,其(似然函数)表达式为
其中, p ( x , z 1 , … , z T ∣ w ) p(x, z_1, \ldots, z_T | w) p(x,z1,…,zT∣w) 由式(20.19)定义。这是通用隐变量模型(16.81)的一个实例,其中隐变量为 z = ( z 1 , … , z T ) z = (z_1, \ldots, z_T) z=(z1,…,zT),观测变量为 x x x。需注意,与归一化流(normalizing flows)的情况相同,所有隐变量的维度均与数据空间一致,而变分自编码器(variational autoencoders)或生成对抗网络(generative adversarial networks)则并非如此。从式(20.20)可以看出,似然函数涉及对所有可能的噪声样本轨迹进行积分,这些轨迹可能生成观测到的数据点。由于式(20.20)中的积分涉及对高度复杂的神经网络函数进行积分,因此这些积分是难以处理的。
20.2.2 Evidence lower bound 证据下界 ( ELBO)
由于精确的似然函数难以计算(不可处理),我们可以采用与变分自编码器类似的方法,最大化对数似然的一个下界——即证据下界(ELBO)。此处,我们针对第16.3节所述的扩散模型重新推导该下界。对于任意选择的分布 q ( z ) q(z) q(z),以下关系始终成立:
b站学习笔记
其中, L \mathcal{L} L 为证据下界(Evidence Lower Bound),亦称变分下界(Variational Lower Bound),其表达式为:
而两个概率密度函数 f ( z ) f(z) f(z) 和 g ( z ) g(z) g(z) 之间的 Kullback-Leibler 散度(KL 散度),记作 KL ( f ∥ g ) \text{KL}(f \| g) KL(f∥g),其定义为:
为验证关系式(20.21),首先需注意,根据概率的乘法规则(乘积法则),我们有:
这里是x,z同时发生的概率等于在x发生的条件下z发生的概率乘以x发生的概率。
将(20.24)代入(20.22),并利用(20.23),可以得到(20.21)。Kullback–Leibler 散度(KL 散度)具有 K L ( ⋅ ∥ ⋅ ) ≥ 0 KL(\cdot\|\cdot) \geq 0 KL(⋅∥⋅)≥0 的性质,由此可得:
由于对数似然函数难以直接计算(不可处理),我们通过最大化下界 L ( w ) \mathcal{L}(w) L(w) 来训练神经网络。
为实现这一目标,我们首先推导出扩散模型下界的显式表达式。在定义该下界时,我们可以自由选择任意满足概率分布条件的 q ( z ) q(z) q(z) 形式,即需满足非负性且归一化为1。在变分自编码器等ELBO应用场景中,我们通常选择具有可调参数的 q ( z ) q(z) q(z) 形式 (如深度神经网络参数化),并同时针对这些参数以及联合分布 p ( x , z ∣ w ) p(x,z|w) p(x,z∣w) 的参数最大化ELBO。通过优化 q ( z ) q(z) q(z) 的分布,可使下界更紧,从而使 p ( x , z ∣ w ) p(x,z|w) p(x,z∣w) 的参数优化更接近最大似然估计。然而在扩散模型中,我们直接采用由马尔可夫链(20.5)定义的固定分布 q ( z 1 , . . . , z T ∣ x ) q(z_1,...,z_T|x) q(z1,...,zT∣x) 作为 q ( z ) q(z) q(z),因此唯一可调参数仅来自反向马尔可夫链的模型 p ( x , z 1 , . . . , z T ∣ w ) p(x,z_1,...,z_T|w) p(x,z1,...,zT∣w)。值得注意的是,我们正是利用了 q ( z ) q(z) q(z) 形式选择的灵活性,使其能够显式依赖于观测数据 x x x。
因此,我们使用公式(20.5)来替换(即代入到)公式(20.21)中的 q ( z 1 , … , z T ∣ x ) q(z_1, \ldots, z_T | x) q(z1,…,zT∣x),同样地,我们使用公式(20.19)来替换 p ( x , z 1 , … , z T ∣ w ) p(x, z_1, \ldots, z_T | w) p(x,z1,…,zT∣w)。这使得我们能够将证据下界(ELBO)写成特定的形式。
公式(20.26)右侧的第一项 ln p ( z T ) \ln p(z_T) lnp(zT) 对应固定分布 N ( z T ∣ 0 , I ) \mathcal{N}(z_T | 0, I) N(zT∣0,I),由于该分布不含可训练参数,因此作为固定加性常数可从证据下界(ELBO)中直接移除。同理,第三项 − ln q ( z 1 ∣ x ) -\ln q(z_1|x) −lnq(z1∣x) 与模型参数 w w w 无关,故同样可省略。
公式(20.26)右侧的第四项对应于变分自编码器中的重建项。可以通过从由(20.2)所定义的关于 z 1 z_1 z1 的分布中抽取样本,并通过蒙特卡洛估计来近似期望 E q [ ⋅ ] \mathbb{E}_q[\cdot] Eq[⋅],从而对这一项进行计算,以便(后续处理或优化)。
其中 z 1 ( l ) ∼ N ( z 1 ∣ 1 − β 1 x , β 1 I ) z_1^{(l)} \sim \mathcal{N}(z_1 | \sqrt{1 - \beta_1}x, \beta_1 I) z1(l)∼N(z1∣1−β1x,β1I)。与变分自编码器(VAEs)不同,我们无需通过采样值反向传播误差信号,因为 q q q 分布是固定的,因此此处无需采用重参数化技巧。
这样便只剩下公式(20.26)右侧的第二项,它由一系列项相加组成,其中每一项都依赖于相邻的两个潜在变量值 z t − 1 z_{t-1} zt−1 和 z t z_t zt。在之前推导扩散核(20.6)时,我们曾指出可以直接从高斯分布 q ( z t − 1 ∣ x ) q(z_{t-1}|x) q(zt−1∣x) 中采样,然后利用(20.4)式(该式同样对应高斯分布)进一步生成对应的 z t z_t zt 样本。尽管从理论上讲,当样本数量趋于无穷大时,这种基于成对样本的估计方法是正确的,但在实际应用中,由于成对采样会导致估计值噪声极大、方差过高,因此需要采集远超必要数量的样本才能获得可靠结果。为此,我们通过数学变换将证据下界(ELBO)改写为每项仅需采样一个值的估计形式。
20.2.3 Rewriting the ELBO 重写证据下界(ELBO)
基于我们之前对变分自编码器中证据下界(ELBO)的讨论,此处我们的目标是将ELBO用KL散度(Kullback–Leibler散度)的形式表示,进而推导出其闭式解。在扩散模型中,神经网络用于建模反向过程的条件分布 p ( z t − 1 ∣ z t , w ) p(z_{t-1}|z_t, w) p(zt−1∣zt,w),而前向过程的 q q q 分布则表示为 q ( z t ∣ z t − 1 , x ) q(z_t|z_{t-1}, x) q(zt∣zt−1,x)。因此,我们利用贝叶斯定理对条件分布的方向进行转换,具体表达式为:
这使我们能够将公式(20.26)中的第二项改写为如下形式:
公式(20.30)右侧的第二项与模型参数 w w w 无关,因此可以忽略。将(20.30)代入(20.26)后,我们得到
最后,我们可以将(20.31)改写为如下形式:
其中,我们对第一项中关于 q ( z 1 , … , z T ∣ x ) q(z_1, \ldots, z_T | x) q(z1,…,zT∣x) 的期望进行了简化,因为在被积函数中仅出现潜在变量 z 1 z_1 z1。因此,在由(20.27)定义的期望中,所有条件分布的积分结果均为1,最终仅保留对 z 1 z_1 z1 的积分。类似地,在第二项中,每个积分仅涉及两个相邻的潜在变量 z t − 1 z_{t-1} zt−1 和 z t z_t zt,其余变量均可被积分消去。
边界(20.32)与变分自编码器(VAE)的证据下界(ELBO,如公式19.14所示)高度相似,区别在于当前模型包含多级编码器与解码器结构。其中,重构项通过提升观测数据样本的生成概率实现优化,其训练方式与VAE中对应项一致,可采用采样近似(20.28)进行估计。而(20.32)中的一致性项定义于成对高斯分布之间,因此可推导出闭式解,具体如下:分布 q ( z t − 1 ∣ z t , x ) q(z_{t-1}|z_t, x) q(zt−1∣zt,x) 由(20.15)给出,分布 p ( z t − 1 ∣ z t , w ) p(z_{t-1}|z_t, w) p(zt−1∣zt,w) 由(20.18)给出,因此两者的KL散度可表示为
其中, m t ( x , z t ) m_t(x, z_t) mt(x,zt) 由公式(20.16)定义,且所有与网络参数 w w w 无关的加性项均被合并至常数项中(该常数项在训练过程中不起作用)。公式(20.32)中的每个一致性项均包含一个关于 z t z_t zt 的剩余积分,其权重为 q ( z t ∣ x ) q(z_t|x) q(zt∣x)。该积分可通过从 q ( z t ∣ x ) q(z_t|x) q(zt∣x) 中采样进行近似,而这一采样过程可利用扩散核(20.6)高效实现。
我们发现,KL散度(20.33)呈现为简单的平方损失函数形式。由于我们通过调整网络参数来最大化(20.32)中的证据下界(ELBO),而ELBO中KL散度项前带有负号,因此这一优化过程等价于最小化该平方误差。
20.2.4 Predicting the noise
为实现更高质量的生成效果,可对神经网络的作用进行改进:不再让其在马尔可夫链的每一步预测去噪后的图像,而是改为预测原始图像到当前带噪图像过程中所添加的总噪声分量(Ho、Jain与Abbeel,2020)。为此,我们首先对公式(20.8)进行变形推导,得到:
若将上述结果代入公式(20.16),我们可将反向条件分布 q ( z t − 1 ∣ z t , x ) q(z_{t-1}|z_t, x) q(zt−1∣zt,x) 的均值 m t ( x , z t ) m_t(x, z_t) mt(x,zt) 用原始数据向量 x x x 和噪声 ϵ \epsilon ϵ 重新表示为:
类似地,我们不再使用预测去噪图像的神经网络 μ ( z t , w , t ) \mu(z_t, w, t) μ(zt,w,t),而是引入一个神经网络 g ( z t , w , t ) g(z_t, w, t) g(zt,w,t),其目标是预测从原始数据 x x x 生成 z t z_t zt 时所添加的总噪声。通过与推导(20.35)相同的步骤,可证明这两个网络函数满足如下关系:
现在,我们可以将公式(20.35)和(20.36)代入公式(20.33),得到:
其中在最后一行中,我们利用公式(20.8)对 z t z_t zt 进行了替换。
ELBO(20.32)中的重构项可通过(20.28)结合采样得到的 z 1 z_1 z1 值进行近似。当生成模型 p ( x ∣ z , w ) p(x|z, w) p(x∣z,w) 采用(20.18)的形式时,该重构项可具体表示为:
若将(20.36)中的 μ ( z 1 , w , 1 ) \mu(z_1, w, 1) μ(z1,w,1) 表达式代入,并结合(20.1)对 x x x 进行替换,同时利用由(20.7)推导出的 α 1 = ( 1 − β 1 ) \alpha_1 = (1 - \beta_1) α1=(1−β1) 关系,最终可得:
这与 t = 1 t=1 t=1 时的特殊情形下公式(20.37)的表达式完全一致,因此重构项与一致性项可以合并处理。
Ho、Jain与Abbeel(2020)通过实验发现,若直接省略公式(20.37)前的系数 β t 2 ( 1 − α t ) ( 1 − β t ) \frac{\beta_t}{2(1-\alpha_t)(1-\beta_t)} 2(1−αt)(1−βt)βt(即让马尔可夫链的所有步骤具有相同权重),模型性能可进一步提升。将这一简化后的(20.37)式代入(20.33)式,可得到如下形式的训练目标函数:
公式(20.40)右侧的平方误差项具有非常直观的解释:对于马尔可夫链中的给定步骤 t t t 和给定训练数据点 x x x,我们首先采样一个噪声向量 ϵ t \epsilon_t ϵt,并利用该噪声生成对应的带噪隐变量 z t z_t zt。此时损失函数即为预测噪声与实际噪声之间的平方差。需要注意的是,神经网络 g ( ⋅ , ⋅ , ⋅ ) g(\cdot,\cdot,\cdot) g(⋅,⋅,⋅) 预测的是原始数据向量 x x x 上累积添加的总噪声,而不仅仅是步骤 t t t 中新增的噪声分量。
在使用随机梯度下降法时,我们会从训练集中随机选取一个数据点 x x x,并计算损失函数相对于网络参数的梯度向量。此外,对于每个这样的数据点,我们会随机选择马尔可夫链中的一个步骤 t t t,而非对公式(20.40)中所有 t t t 的求和项逐一计算误差。这些梯度会通过小批量数据样本进行累积,随后用于更新网络权重。
还需注意,该损失函数天然具备数据增强的效果——因为每次使用某个特定训练样本 x x x 时,都会将其与新采样的噪声 ϵ t \epsilon_t ϵt 相结合。以上讨论均针对训练集中的单个数据点 x x x,其对应的梯度计算过程如算法20.1所示。
20.2.5 Generating new samples
当网络训练完成后,我们可以通过以下步骤在数据空间中生成新样本:首先从高斯分布 p ( z T ) p(z_T) p(zT) 中采样初始隐变量 z T z_T zT,然后沿马尔可夫链逐步进行去噪。具体到第 t t t 步的去噪样本 z t z_t zt,生成前一步样本 z t − 1 z_{t-1} zt−1 的过程分为三步:
- 将 z t z_t zt、网络权重 w w w 和时间步 t t t 输入神经网络,计算输出 g ( z t , w , t ) g(z_t, w, t) g(zt,w,t);
- 利用公式(20.36)根据网络输出计算均值 μ ( z t , w , t ) \mu(z_t, w, t) μ(zt,w,t);
- 从条件分布 p ( z t − 1 ∣ z t , w ) = N ( z t − 1 ∣ μ ( z t , w , t ) , β t I ) p(z_{t-1}|z_t, w) = \mathcal{N}(z_{t-1}|\mu(z_t, w, t), \beta_t I) p(zt−1∣zt,w)=N(zt−1∣μ(zt,w,t),βtI) 中采样 z t − 1 z_{t-1} zt−1,即通过向均值 μ ( z t , w , t ) \mu(z_t, w, t) μ(zt,w,t) 添加方差为 β t \beta_t βt 的缩放噪声实现。
其中 ϵ ∼ N ( ϵ ∣ 0 , I ) \epsilon \sim \mathcal{N}(\epsilon|0, I) ϵ∼N(ϵ∣0,I)。需注意,神经网络 g ( ⋅ , ⋅ , ⋅ ) g(\cdot,\cdot,\cdot) g(⋅,⋅,⋅) 预测的是从原始数据向量 x x x 生成 z t z_t zt 时所添加的总噪声,但在采样过程中,我们仅从 z t − 1 z_{t-1} zt−1 中减去该噪声的 β t 1 − α t \frac{\beta_t}{\sqrt{1-\alpha_t}} 1−αtβt 分量,随后再添加方差为 β t \beta_t βt 的新噪声以生成 z t − 1 z_{t-1} zt−1。在最终生成合成数据样本 x x x 时,由于目标是输出无噪声结果,因此不再添加额外噪声。整个采样流程总结于算法20.2。
扩散模型在数据生成方面的主要缺点在于:生成样本时需要多次顺序遍历训练好的神经网络,这一过程计算成本较高。为加速采样,一种可行方法是:先将去噪过程转化为关于连续时间的微分方程,再利用高效的离散化数值方法快速求解该方程。
本章中我们假设数据和隐变量均为连续型变量,因此可采用高斯噪声模型进行建模。不过,扩散模型同样可推广至离散空间(如Austin等人,2021),例如在生成候选药物分子的场景中,其部分生成过程需从化学元素子集中选择原子类型。
我们已发现,扩散模型因需按顺序逆向执行可能包含数百甚至数千步的噪声过程,导致计算成本高昂。为此,Song、Meng与Ermon(2020)提出了一种相关技术——去噪扩散隐式模型(Denoising Diffusion Implicit Models),该技术通过放宽噪声过程的马尔可夫假设,同时保持原有训练目标函数不变,从而在采样阶段实现1-2个数量级的加速,且生成的样本质量未受显著影响。
20.3 Score Matching
本章迄今讨论的去噪扩散模型与另一类深度生成模型存在紧密关联——这类模型基于得分匹配(Score Matching)理论(Hyvärinen,2005;Song和Ermon,2019)独立发展而来。它们的核心工具是得分函数(Score Function)或Stein得分,其定义为数据向量 x x x 的对数似然函数关于 x x x 的梯度,具体表达式为:
在此需特别强调:该梯度是针对数据向量 x x x 计算的,而非针对任何参数向量。需注意,得分函数 s ( x ) s(x) s(x) 是一个与 x x x 维度相同的向量值函数,其每个分量 s i ( x ) = ∂ ln p ( x ) ∂ x i s_i(x) = \frac{\partial \ln p(x)}{\partial x_i} si(x)=∂xi∂lnp(x) 对应于 x x x 的第 i i i 个分量 x i x_i xi。例如,若 x x x 为图像,则 s ( x ) s(x) s(x) 也可表示为同尺寸的图像,其中每个像素对应原图像的梯度值。图20.5展示了一个二维概率密度函数及其对应的得分函数示例。
图20.5 分数函数的图示,展示了二维空间中的一种分布,该分布由作为热图表示的高斯混合组成,以及由(20.42)定义的相应分数函数,以向量的形式绘制在规则的x值网格上。
# 图表解释
这张图展示了一个二维空间中的分布情况。图中主要包含两个部分的信息:
- **热图部分**:背景颜色表示了一个由多个高斯分布混合而成的概率分布。颜色越亮(如黄色、橙色区域),表示该区域的概率密度越高;颜色越暗(如黑色区域),表示概率密度越低。这种热图的可视化方式有助于直观地理解分布的形状和集中区域。
- **向量部分**:在规则的x值网格上绘制了向量,这些向量代表了由公式(20.42)定义的分数函数。向量的方向和长度传达了分数函数在不同位置的信息,可能用于指示某种梯度、方向场或其他与分布相关的特性。
总体而言,这张图通过热图和向量的结合,提供了对二维空间中高斯混合分布及其相关分数函数的直观可视化,有助于理解分布的结构和分数函数的行为。
要理解得分函数为何有用,可考虑两个概率密度函数 q ( x ) q(x) q(x) 和 p ( x ) p(x) p(x),若它们的得分函数相等(即对所有 x x x 满足 ∇ x ln q ( x ) = ∇ x ln p ( x ) \nabla_x \ln q(x) = \nabla_x \ln p(x) ∇xlnq(x)=∇xlnp(x)),则对等式两边关于 x x x 积分并取指数后,可得 q ( x ) = K p ( x ) q(x) = K p(x) q(x)=Kp(x),其中 K K K 是与 x x x 无关的常数。因此,若能通过模型 s ( x , w ) s(x, w) s(x,w) 学习到得分函数,则相当于建模了原始数据密度 p ( x ) p(x) p(x)(仅相差一个比例常数 K K K)。
20.3.1 Score loss function
要训练这样的模型,我们需要定义一个损失函数,其目标是使模型预测的得分函数 s ( x , w ) s(x, w) s(x,w) 与生成数据的真实分布 p ( x ) p(x) p(x) 的得分函数 ∇ x ln p ( x ) \nabla_x \ln p(x) ∇xlnp(x) 相匹配。此类损失函数的一个典型例子是模型得分与真实得分之间的期望平方误差,其表达式为:
正如我们在能量模型讨论中提到的,得分函数无需假设相关概率密度已归一化,因为梯度算子会消去归一化常数,这为模型选择提供了极大的灵活性。使用深度神经网络表示得分函数 s ( x , w ) s(x, w) s(x,w) 时,主要有两种方法:
直接输出法:由于得分函数的每个分量 s i s_i si 对应数据 x x x 的第 i i i 个分量 x i x_i xi,因此可设计一个神经网络,使其输出维度与输入维度相同。
梯度计算法:得分函数本质上是标量函数(对数概率密度)的梯度,属于更受限的函数类。因此另一种方法是构建一个仅输出标量 ϕ ( x ) \phi(x) ϕ(x) 的网络,再通过自动微分计算 ∇ x ϕ ( x ) \nabla_x \phi(x) ∇xϕ(x)。不过,这种方法需要两次反向传播(一次计算 ϕ ( x ) \phi(x) ϕ(x),一次计算梯度),计算成本更高。
由于上述原因,大多数实际应用采用第一种方法。
20.3.2 Modified score loss
损失函数(20.43)存在的一个关键问题是:我们无法直接对其进行最小化优化,因为真实数据的得分函数 ∇ x ln p ( x ) \nabla_x \ln p(x) ∇xlnp(x) 是未知的。我们手中仅有的资源是有限的数据集 D = { x 1 , … , x N } \mathcal{D} = \{x_1, \ldots, x_N\} D={x1,…,xN},通过它可以构造一个经验分布(empirical distribution):
其中 δ ( x ) \delta(x) δ(x) 是狄拉克δ函数(Dirac delta function),可直观理解为在 x = 0 x=0 x=0 处的一个无限高、无限窄的“尖峰”,其具有以下特性:
由于式(20.44)关于 x x x 不可微,因此无法直接计算其得分函数。为解决这一问题,可通过引入噪声模型对数据点进行**“平滑处理”,从而得到一个光滑且可微的概率密度表示。这种方法被称为帕森窗估计(Parzen Estimator)或核密度估计(Kernel Density Estimator),其定义为:
其中 q ( z ∣ x , σ ) q(z|x, \sigma) q(z∣x,σ) 是噪声核函数。一种常用的核函数选择是高斯核(Gaussian Kernel)**,即:
通过引入噪声模型和高斯核函数,我们可以将离散的经验分布转换为连续且可微的密度估计,从而能够计算分数函数并进行后续的优化和建模工作。这种方法在处理有限数据集和构建可微的概率密度模型时非常有用。
此时我们不再直接最小化损失函数(20.43),而是改用针对平滑后的帕森密度(Parzen density)的对应损失函数,其形式为:
一个关键的结果是,把公式(20.47)代入到公式(20.49)中,能够将损失函数重新写成由(Vincent, 2011)给出的等效形式。这里的核心在于通过特定的代入操作,对损失函数进行形式上的转换,以便更好地进行分析和处理。
等效形式的损失函数
改写后的损失函数 J ( w ) J(\mathbf{w}) J(w) 的表达式为:
这个表达式是损失函数在一般情况下的等效形式。其中:
- s ( z , w ) \mathbf{s}(\mathbf{z}, \mathbf{w}) s(z,w) 是我们学习的分数函数模型, w \mathbf{w} w 是模型的参数。
- ∇ z ln q ( z ∣ x , σ ) \nabla_{\mathbf{z}} \ln q(\mathbf{z}|\mathbf{x}, \sigma) ∇zlnq(z∣x,σ) 是条件概率分布 q ( z ∣ x , σ ) q(\mathbf{z}|\mathbf{x}, \sigma) q(z∣x,σ) 的对数关于 z \mathbf{z} z 的梯度,也就是分数函数。
- q ( z ∣ x , σ ) q(\mathbf{z}|\mathbf{x}, \sigma) q(z∣x,σ) 是噪声核函数,通常为高斯核函数 N ( z ∣ x , σ 2 I ) \mathcal{N}(\mathbf{z}|\mathbf{x}, \sigma^2 \mathbf{I}) N(z∣x,σ2I)。
- p ( x ) p(\mathbf{x}) p(x) 是原始数据的概率密度函数。
- 整个积分表达式衡量了学习到的分数函数模型与真实的分数函数之间的差异,通过积分在所有可能的 x \mathbf{x} x 和 z \mathbf{z} z 上进行加权平均。
使用经验密度代入后的损失函数
如果我们使用经验密度(20.44)来代替 p ( x ) p(\mathbf{x}) p(x),即用有限的数据集来近似原始数据的分布,那么损失函数 J ( w ) J(\mathbf{w}) J(w) 变为:
这里,求和是对数据集中的 N N N 个数据点 x n \mathbf{x}_n xn 进行的。通过使用经验密度,我们将连续的积分形式转换为了离散的求和形式,这使得在实际应用中,当我们只有有限的数据样本时,能够更方便地计算和优化损失函数。
总体而言,这部分内容展示了如何通过特定的代入操作和经验密度的使用,将损失函数转换为更便于实际应用的形式,为后续的模型训练和优化提供了基础。
图20.6 使用由(14.61)定义的朗之万动力学(Langevin dynamics)对图20.5中所示的分布进行采样得到的轨迹示例,展示了三条均从绘图中心开始的轨迹。
### 详细解析
1. **朗之万动力学(Langevin dynamics)**
- 朗之万动力学是一种用于描述随机过程的物理模型,常用于从复杂的概率分布中进行采样。在机器学习和统计物理中,它被广泛应用于模拟分子运动、优化算法以及概率密度估计等任务。公式(14.61)定义了具体的朗之万动力学方程,虽然图中未给出该公式的具体内容,但通常朗之万方程会包含确定性力和随机力两部分,用于描述系统在随机噪声影响下的演化过程。
2. **对图20.5中分布的采样**
- 图20.5展示了一个二维空间中的分布,该分布由高斯混合表示,并通过热图和向量进行了可视化。这里的采样是指根据图20.5所表示的概率分布生成样本点,而朗之万动力学提供了一种从该分布中生成样本的动态方法。通过模拟朗之万动力学过程,可以生成一系列的样本点,这些样本点的轨迹反映了系统在概率空间中的演化路径。
3. **轨迹示例**
- 图20.6展示了三条采样轨迹,这些轨迹都从绘图的中心开始。每条轨迹由一系列的点组成,点之间的连线表示系统在不同时间步的演化路径。轨迹上的箭头可能表示系统在每个时间步的移动方向,反映了朗之万动力学中确定性力和随机力的综合作用。
- 从图中可以看出,不同的轨迹在演化过程中会受到随机噪声的影响,从而呈现出不同的路径。尽管起始点相同,但由于随机力的存在,轨迹会逐渐分散,覆盖概率分布的不同区域。这展示了朗之万动力学在探索复杂概率分布空间时的随机性和多样性。
总体而言,图20.6通过具体的轨迹示例,直观地展示了如何使用朗之万动力学从给定的概率分布中进行采样,以及采样过程中系统的动态演化特征。
对于高斯帕曾核(Gaussian Parzen kernel)(公式20.48),分数函数变为:
其中 ϵ = z − x \boldsymbol{\epsilon} = \mathbf{z} - \mathbf{x} ϵ=z−x 是从 N ( z ∣ 0 , I ) \mathcal{N}(\mathbf{z}|\mathbf{0}, \mathbf{I}) N(z∣0,I) 中抽取的。如果我们考虑特定的噪声模型(公式20.6),则得到:
因此,我们可以看到分数损失(公式20.50)衡量的是神经网络预测与噪声 ϵ \boldsymbol{\epsilon} ϵ 之间的差异。所以,这个损失函数与去噪扩散模型中使用的(公式20.37)形式具有相同的最小值,分数函数 s ( z , w ) \mathbf{s}(\mathbf{z}, \mathbf{w}) s(z,w) 起着与噪声预测网络 g ( z , w ) \mathbf{g}(\mathbf{z}, \mathbf{w}) g(z,w) 相同的作用,只是有一个常数缩放因子 − 1 / 1 − α t -1/\sqrt{1 - \alpha_t} −1/1−αt(Song和Ermon,2019)。最小化(公式20.50)被称为去噪分数匹配,我们可以看到它与去噪扩散模型的紧密联系。关于如何选择噪声方差 σ 2 \sigma^2 σ2 的问题仍然存在,我们很快会回到这个问题。
在训练了一个基于分数的模型后,我们需要抽取新的样本。朗之万动力学非常适合基于分数的模型,因为它基于分数函数,因此不需要归一化的概率分布,如图20.6所示。
解释
- 分数函数的推导
- 首先,针对高斯帕曾核,通过对数条件概率关于 z \mathbf{z} z 求梯度,得到了分数函数的表达式(公式20.52),其中 ϵ \boldsymbol{\epsilon} ϵ 是从标准正态分布中抽取的。
- 当考虑特定的噪声模型时,分数函数的表达式变为(公式20.53),这里引入了与噪声模型相关的参数 α t \alpha_t αt。
- 分数损失与去噪扩散模型的联系
- 分数损失(公式20.50)的作用是衡量神经网络预测和噪声之间的差异。
- 该损失函数与去噪扩散模型中的某个形式(公式20.37)具有相同的最小值,基于分数的模型中的分数函数 s ( z , w ) \mathbf{s}(\mathbf{z}, \mathbf{w}) s(z,w) 和去噪扩散模型中的噪声预测网络 g ( z , w ) \mathbf{g}(\mathbf{z}, \mathbf{w}) g(z,w) 功能相似,只是存在一个常数缩放因子。
- 最小化分数损失的过程被称为去噪分数匹配,这显示了基于分数的模型与去噪扩散模型之间的紧密联系。
- 样本抽取与朗之万动力学
- 训练好基于分数的模型后,需要抽取新样本。
- 朗之万动力学适用于基于分数的模型,因为它直接利用分数函数,不需要归一化的概率分布,图20.6展示了使用朗之万动力学进行采样的轨迹示例。
- 未解决的问题
- 文中指出,关于如何选择噪声方差 σ 2 \sigma^2 σ2 的问题仍然存在,后续会进一步讨论这个问题。
20.3.3 Noise variance 噪声方差
我们已经了解了如何从一组训练数据中学习分数函数,以及如何使用朗之万采样(Langevin sampling)从学习到的分布中生成新样本。然而,我们可以发现这种方法存在三个潜在问题(Song和Ermon,2019;Luo,2022)。
首先,如果数据分布位于一个维度低于数据空间的流形(manifold)上,那么在流形之外的点处,概率密度为零,并且由于 ln p ( x ) \ln p(x) lnp(x) 无定义,此处分数函数也无定义。
其次,在数据密度较低的区域,由于损失函数(20.43)是按密度加权的,因此对分数函数的估计可能不准确。使用朗之万采样时,不准确的分数函数可能会导致生成较差的轨迹。
第三,即使分数函数的模型准确,如果数据分布由多个不相交的分布混合而成,朗之万过程也可能无法正确采样。
这三个问题都可以通过为核函数(20.48)中使用的噪声方差 σ 2 \sigma^2 σ2 选择一个足够大的值来解决,因为这样会使数据分布变得模糊(平滑)。然而,方差过大会严重扭曲原始分布,这本身就会导致分数函数建模的不准确。可以通过考虑一系列方差值 σ 1 2 < σ 2 2 < ⋯ < σ T 2 \sigma^2_1 < \sigma^2_2 < \cdots < \sigma^2_T σ12<σ22<⋯<σT2(Song和Ermon,2019)来处理这种权衡,其中 σ 1 2 \sigma^2_1 σ12 足够小,能够准确表示数据分布,而 σ T 2 \sigma^2_T σT2 足够大,可以避免上述问题。
然后,分数网络被修改为将方差作为额外输入 s ( x , w , σ 2 ) s(\mathbf{x}, \mathbf{w}, \sigma^2) s(x,w,σ2),并通过使用一个加权和形式的损失函数进行训练,该损失函数是形式为(20.51)的损失函数的加权总和,其中每一项表示相关网络与相应扰动数据集之间的误差。对于数据向量 x n \mathbf{x}_n xn,损失函数的形式则为
其中 λ ( i ) \lambda^{(i)} λ(i) 是权重系数。我们可以看到,这种训练过程与用于训练分层去噪网络的训练过程完全对应(意思一致、相呼应)。
训练完成后,可按顺序依次针对 i = L , L − 1 , … , 2 , 1 i = L, L - 1, \ldots, 2, 1 i=L,L−1,…,2,1 的各个模型,运行若干步朗之万采样来生成样本。这种技术被称为退火朗之万动力学(annealed Langevin dynamics),其原理与用于从去噪扩散模型中采样的算法20.2类似。
20.3.4 Stochastic differential equations 随机微分方程
我们已看到,在为扩散模型构建噪声过程时,采用大量步骤(通常多达数千步)是很有帮助的。因此,很自然地我们会思考:若像引入神经微分方程时对无限深层神经网络所做的那样,考虑无限步数的极限情况,会发生什么?在求取这种极限时,我们需要确保每一步的噪声方差(见18.3.1节) β t \beta_t βt 随步长的减小而变小。这促使我们将扩散模型在连续时间下的形式表述为随机微分方程(Stochastic Differential Equations,SDEs)(Song等人,2020)。如此一来,去噪扩散概率模型和分数匹配模型都可被视为连续时间随机微分方程的一种离散化形式。
我们可以将一般的随机微分方程(SDE)写成对向量 z \mathbf{z} z 的无穷小更新的形式,即
其中,漂移项(drift term)如同常微分方程(ODE)中那样是确定性的,但扩散项(diffusion term)是随机的,例如由无穷小的高斯步长给出。这里的参数 t t t 借物理系统的类比,常被称为“时间”。通过对连续时间极限的推导,扩散模型的正向噪声过程(20.3)可表示为(20.55)形式的随机微分方程。
对于随机微分方程(SDE)(20.55),存在相应的逆向随机微分方程(Song等人,2020),由下式给出:
其中,我们认定 ∇ z ln p ( z ) \nabla_{\mathbf{z}} \ln p(\mathbf{z}) ∇zlnp(z) 为分数函数。需从 t = T t = T t=T 到 t = 0 t = 0 t=0 逆向求解由(20.55)给出的随机微分方程。
为了数值求解随机微分方程,我们需要对时间变量进行离散化。最简单的方法是使用固定且等间距的时间步长,这被称为欧拉 - 玛雅那(Euler - Maruyama)求解器。对于逆向随机微分方程,我们则恢复一种朗之万方程的形式。然而,可以采用更复杂的求解器,它们使用更灵活的离散化形式(Kloeden和Platen,2013)。
对于所有由随机微分方程控制的扩散过程,存在一个由常微分方程(ODE)描述的相应确定性过程,其轨迹具有与随机微分方程相同的边际概率密度 p ( z ∣ t ) p(\mathbf{z}|t) p(z∣t)(Song等人,2020)。对于形式为(20.56)的随机微分方程,相应的常微分方程由下式给出:
常微分方程的公式化表述允许使用高效的自适应步长求解器,从而大幅减少函数求值的次数。此外,它使得概率扩散模型与归一化流模型相关联,从中可以使用变量变换公式(18.1)来提供对数似然的精确估计。
解释
这段内容主要围绕随机微分方程(SDE)及其相关概念展开:
- 逆向随机微分方程:介绍了与给定随机微分方程(20.55)对应的逆向随机微分方程(20.56),并指出其中 ∇ z ln p ( z ) \nabla_{\mathbf{z}} \ln p(\mathbf{z}) ∇zlnp(z) 为分数函数,且说明了求解该逆向随机微分方程的时间范围。
- 随机微分方程的数值求解:阐述了数值求解随机微分方程需要对时间变量离散化,介绍了最简单的欧拉 - 玛雅那求解器,以及对于逆向随机微分方程可恢复朗之万方程形式,还提到有更复杂的求解器可采用。
- 确定性过程与常微分方程:说明对于由随机微分方程控制的扩散过程,存在具有相同边际概率密度的确定性过程,由常微分方程描述,并给出了形式为(20.56)的随机微分方程对应的常微分方程(20.57)。
- 常微分方程的优势:指出常微分方程的公式化表述可使用高效自适应步长求解器,减少函数求值次数,还能使概率扩散模型与归一化流模型相关联,用于精确估计对数似然。
20.4 Guided Diffusion
到目前为止,我们将扩散模型视为一种表示无条件密度 p ( x ) p(x) p(x) 的方法,该密度是从独立采自 p ( x ) p(x) p(x) 的一组训练样本 x 1 , … , x N x_1, \ldots, x_N x1,…,xN 中学习得到的。模型训练完成后,我们可以从该分布中生成新的样本。我们已经在图1.3中看到了从面部图像深度生成模型进行无条件采样的例子,在那个例子中,使用的是生成对抗网络(GAN)模型。
然而,在许多应用场景中,我们希望从条件分布 p ( x ∣ c ) p(x|c) p(x∣c) 中进行采样,其中条件变量 c c c 可以是类别标签,或是图像所需内容的文本描述等。这也构成了图像超分辨率、图像修复、视频生成等众多应用的基础。实现这一目标最简单的方法是将 c c c 作为去噪神经网络 g ( z , w , t , c ) g(\mathbf{z}, \mathbf{w}, t, c) g(z,w,t,c) 的额外输入,然后使用匹配的样本对 { x n , c n } \{\mathbf{x}_n, c_n\} {xn,cn} 对网络进行训练。但这种方法的主要局限在于,网络可能对条件变量赋予的权重不足,甚至完全忽略它们。因此,我们需要一种方法来控制赋予条件信息的权重,并在样本多样性与条件信息匹配度之间进行权衡。这种增强条件信息匹配度的额外约束被称为引导(guidance)。根据是否使用单独的分类器模型,引导方法主要分为两类。
20.4.1 Classifier guidance
假设存在一个已训练好的分类器 p ( c ∣ x ) p(\mathbf{c}|\mathbf{x}) p(c∣x),并且从分数函数的角度考虑一个扩散模型。利用贝叶斯定理,我们可以将条件扩散模型的分数函数写成如下形式:
其中,由于 p ( c ) p(\mathbf{c}) p(c) 与 x \mathbf{x} x 无关,我们使用了 ∇ x ln p ( c ) = 0 \nabla_{\mathbf{x}} \ln p(\mathbf{c}) = 0 ∇xlnp(c)=0 这一条件。(20.58) 右侧的第一项是常见的无条件分数函数,而第二项将去噪过程推向在分类器模型下使给定标签 c \mathbf{c} c 的概率最大化的方向(Dhariwal和Nichol,2021)。通过引入一个超参数 λ \lambda λ(称为引导尺度),可以控制分类器梯度所赋予的权重,从而控制分类器的影响。此时用于采样的分数函数变为:
如果 λ = 0 \lambda = 0 λ=0,我们得到原始的无条件扩散模型;而如果 λ = 1 \lambda = 1 λ=1,我们得到与条件分布 p ( x ∣ c ) p(\mathbf{x}|\mathbf{c}) p(x∣c) 对应的分数。对于 λ > 1 \lambda > 1 λ>1,强烈建议模型尊重条件标签,可以使用较大的 λ \lambda λ 值,例如 λ = 10 \lambda = 10 λ=10。然而,这样做是以牺牲样本的多样性为代价的,因为模型更倾向于分类器能够正确分类的“简单”样本。
基于分类器的引导方法存在一个问题,即必须训练一个单独的分类器。此外,这个分类器需要能够对具有不同程度噪声的样本进行分类,而标准分类器是在干净样本上训练的。因此,我们转向一种避免使用单独分类器的替代方法。
解释
这段内容主要围绕基于分类器的扩散模型引导方法展开:
- 分数函数推导:利用贝叶斯定理,从已训练好的分类器 p ( c ∣ x ) p(\mathbf{c}|\mathbf{x}) p(c∣x) 出发,推导出条件扩散模型的分数函数 ∇ x ln p ( x ∣ c ) \nabla_{\mathbf{x}} \ln p(\mathbf{x}|\mathbf{c}) ∇xlnp(x∣c),并将其分解为无条件分数函数 ∇ x ln p ( x ) \nabla_{\mathbf{x}} \ln p(\mathbf{x}) ∇xlnp(x) 和与分类器相关的项 ∇ x ln p ( c ∣ x ) \nabla_{\mathbf{x}} \ln p(\mathbf{c}|\mathbf{x}) ∇xlnp(c∣x) 两部分。
- 引导尺度 λ \lambda λ 的作用:引入超参数 λ \lambda λ(引导尺度)来控制分类器梯度对去噪过程的影响权重。通过调整 λ \lambda λ 的值,可以平衡样本多样性与对条件标签的遵循程度。
- 不同 λ \lambda λ 值的效果:当 λ = 0 \lambda = 0 λ=0 时,模型为无条件扩散模型; λ = 1 \lambda = 1 λ=1 时得到与条件分布对应的分数; λ > 1 \lambda > 1 λ>1 时模型更倾向于遵循条件标签,但会牺牲样本多样性。
- 基于分类器引导方法的局限性:指出该方法需要训练一个单独的分类器,且该分类器需能处理带噪声的样本,而标准分类器通常在干净样本上训练,因此考虑采用避免使用单独分类器的替代方法。
20.4.2 Classifier-free guidance
如果我们用(20.58)来替换(20.59)中的 ∇ x ln p ( c ∣ x ) \nabla_{\mathbf{x}} \ln p(\mathbf{c}|\mathbf{x}) ∇xlnp(c∣x),我们可以将分数函数写成如下形式
当 0 < λ < 1 0 < \lambda < 1 0<λ<1 时,该形式表示条件对数密度 ln p ( x ∣ c ) \ln p(\mathbf{x}|\mathbf{c}) lnp(x∣c) 与无条件对数密度 ln p ( x ) \ln p(\mathbf{x}) lnp(x) 的一种凸组合。而当 λ > 1 \lambda > 1 λ>1 时,无条件分数项的贡献变为负值,这意味着模型会主动降低生成忽略条件信息样本的概率,转而倾向于生成符合条件信息的样本。
此外,我们无需分别训练用于建模 p ( x ∣ c ) p(\mathbf{x}|\mathbf{c}) p(x∣c) 和 p ( x ) p(\mathbf{x}) p(x) 的独立网络,而是可以训练一个单一的条件模型。具体做法是,在训练过程中,以一定概率(通常约为 10% - 20%)将条件变量 c \mathbf{c} c 设置为空值,例如 c = 0 \mathbf{c} = 0 c=0。这样一来, p ( x ) p(\mathbf{x}) p(x) 就可以用 p ( x ∣ c = 0 ) p(\mathbf{x}|\mathbf{c} = 0) p(x∣c=0) 来表示。这在一定程度上类似于神经网络中的“Dropout”技术,只不过在这里,是随机将一部分训练向量的条件输入整体置为零。
训练完成后,分数函数(20.60)会被用于增强条件信息的权重。在实际应用中,无分类器引导(classifier - free guidance)相较于基于分类器的引导(classifier guidance),能产生质量更高的结果(Nichol等人,2021;Saharia等人,2022)。原因在于,分类器 p ( c ∣ x ) p(\mathbf{c}|\mathbf{x}) p(c∣x) 只要能对 c \mathbf{c} c 做出准确预测,就可能会忽略输入向量 x \mathbf{x} x 的大部分信息;而无分类器引导是基于条件密度 p ( x ∣ c ) p(\mathbf{x}|\mathbf{c}) p(x∣c),该密度必须对 x \mathbf{x} x 的所有方面都赋予高概率。
文本引导的扩散模型可以借助大语言模型的技术,使条件输入成为一个通用的文本序列(即提示词 prompt),而不仅仅是从预定义的类别标签集中进行选择。这使得文本输入可以通过两种方式影响去噪过程:其一,将基于Transformer的语言模型的内部表示与去噪网络的输入进行拼接;其二,允许去噪网络中的交叉注意力层关注文本标记序列。以文本提示为条件的无分类器引导如图20.7所示。
图20.7 展示了使用条件文本“一只吃竹子的熊猫的彩色玻璃窗”通过名为GLIDE的模型生成的扩散模型的无分类器引导示例。左侧的示例是在λ = 0(无引导,仅为普通条件模型)的情况下生成的,而右侧的示例是在λ = 3的情况下生成的。[经Nichol等人(2021)许可使用]
解释:
这张图展示了无分类器引导(classifier - free guidance)在扩散模型中的应用效果。通过给定的条件文本“一只吃竹子的熊猫的彩色玻璃窗”,模型生成了不同的图像。左侧图像是在没有引导(λ = 0)的情况下,由普通条件模型生成的,图像的多样性和与条件的契合度可能相对有限。而右侧图像是在λ = 3的引导下生成的,可以看到图像更加符合条件文本的描述,在风格和内容上都更贴近“彩色玻璃窗”以及“吃竹子的熊猫”这一主题,说明无分类器引导能够增强模型对条件信息的利用,生成更符合预期的图像。
条件扩散模型的另一个应用是图像超分辨率,该技术可将低分辨率图像转换为相应的高分辨率图像。这本质上是一个逆问题,因为给定一张低分辨率图像,可能对应多张高分辨率图像。通过将低分辨率图像作为条件变量,对从高斯分布中采样得到的高分辨率样本进行去噪处理,即可实现图像超分辨率(Saharia, Ho等人,2021)。该方法的应用示例见图20.8。此类模型可以级联使用以实现超高分辨率(Ho等人,2021),例如从64×64分辨率提升到256×256,再从256×256提升到1024×1024。每个阶段通常采用U - net架构,且每个U - net都以前一阶段的最终去噪输出作为条件。
图20.8 扩散模型生成的低分辨率图像以及相应高分辨率图像的示例。上行展示了一张16×16的输入图像、对应的128×128输出图像,以及生成输入图像所用的原始图像。下行展示了一张64×64的输入图像、对应的256×256输出图像,以及用于对比的原始图像。[经Saharia、Ho等人(2021)许可使用]
这种级联方式也可应用于图像生成扩散模型,即先在较低分辨率下对图像进行去噪处理,随后利用单独的网络(该网络也可将文本提示作为输入)对去噪结果进行上采样,从而得到最终的高分辨率输出(Nichol等人,2021;Saharia等人,2022)。由于去噪过程可能需要数百次通过网络,与直接在高维空间中操作相比,这种方法能显著降低计算成本。需注意的是,这些方法仍直接在图像空间内进行操作,只不过是在较低分辨率下开展。
另一种解决在高分辨率图像空间中直接应用扩散模型计算成本高昂问题的方法,是潜在扩散模型(Latent Diffusion Models)(Rombach等人,2021)。该方法首先在无噪声图像(见第19.1节)上训练一个自动编码器,以获取图像的低维表示,之后该自动编码器保持固定。接着,训练一个U-net架构在低维空间中进行去噪处理,不过这个低维空间本身并不能直接解读为图像。最后,利用固定自动编码器网络的后半部分(即解码器部分),将去噪后的低维表示映射到高分辨率图像空间。这种方法能更高效地利用低维空间,使其专注于图像语义信息,而将生成清晰高分辨率图像的任务交给解码器,由解码器根据去噪后的低维表示来创建对应的高分辨率图像。
条件图像生成还有许多其他应用,包括图像修复(inpainting)、图像外扩(un - cropping)、图像复原(restoration)、图像变形(image morphing)、风格迁移(style transfer)、图像上色(colourization)、去模糊(de - blurring)以及视频生成(video generation)(Yang、Srivastava 和 Mandt,2022)。图像修复的一个示例见图20.9。
图20.9 图像修复示例,左侧为原始图像,中间为移除了部分区域的图像,右侧为经过修复的图像。[经Saharia、Chan、Chang等人(2021)许可使用]