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.3 Score Matching
本章迄今讨论的去噪扩散模型与另一类深度生成模型存在紧密关联——这类模型基于得分匹配(Score Matching)理论(Hyvärinen,2005;Song和Ermon,2019)独立发展而来。它们的核心工具是得分函数(Score Function)或Stein得分,其定义为数据向量 xxx 的对数似然函数关于 xxx 的梯度,具体表达式为:
在此需特别强调:该梯度是针对数据向量 xxx 计算的,而非针对任何参数向量。需注意,得分函数 s(x)s(x)s(x) 是一个与 xxx 维度相同的向量值函数,其每个分量 si(x)=∂lnp(x)∂xis_i(x) = \frac{\partial \ln p(x)}{\partial x_i}si(x)=∂xi∂lnp(x) 对应于 xxx 的第 iii 个分量 xix_ixi。例如,若 xxx 为图像,则 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),若它们的得分函数相等(即对所有 xxx 满足 ∇xlnq(x)=∇xlnp(x)\nabla_x \ln q(x) = \nabla_x \ln p(x)∇xlnq(x)=∇xlnp(x)),则对等式两边关于 xxx 积分并取指数后,可得 q(x)=Kp(x)q(x) = K p(x)q(x)=Kp(x),其中 KKK 是与 xxx 无关的常数。因此,若能通过模型 s(x,w)s(x, w)s(x,w) 学习到得分函数,则相当于建模了原始数据密度 p(x)p(x)p(x)(仅相差一个比例常数 KKK)。
20.3.1 Score loss function
要训练这样的模型,我们需要定义一个损失函数,其目标是使模型预测的得分函数 s(x,w)s(x, w)s(x,w) 与生成数据的真实分布 p(x)p(x)p(x) 的得分函数 ∇xlnp(x)\nabla_x \ln p(x)∇xlnp(x) 相匹配。此类损失函数的一个典型例子是模型得分与真实得分之间的期望平方误差,其表达式为:
正如我们在能量模型讨论中提到的,得分函数无需假设相关概率密度已归一化,因为梯度算子会消去归一化常数,这为模型选择提供了极大的灵活性。使用深度神经网络表示得分函数 s(x,w)s(x, w)s(x,w) 时,主要有两种方法:
-
直接输出法:由于得分函数的每个分量 sis_isi 对应数据 xxx 的第 iii 个分量 xix_ixi,因此可设计一个神经网络,使其输出维度与输入维度相同。
-
梯度计算法:得分函数本质上是标量函数(对数概率密度)的梯度,属于更受限的函数类。因此另一种方法是构建一个仅输出标量 ϕ(x)\phi(x)ϕ(x) 的网络,再通过自动微分计算 ∇xϕ(x)\nabla_x \phi(x)∇xϕ(x)。不过,这种方法需要两次反向传播(一次计算 ϕ(x)\phi(x)ϕ(x),一次计算梯度),计算成本更高。
由于上述原因,大多数实际应用采用第一种方法。
20.3.2 Modified score loss
损失函数(20.43)存在的一个关键问题是:我们无法直接对其进行最小化优化,因为真实数据的得分函数 ∇xlnp(x)\nabla_x \ln p(x)∇xlnp(x) 是未知的。我们手中仅有的资源是有限的数据集 D={x1,…,xN}\mathcal{D} = \{x_1, \ldots, x_N\}D={x1,…,xN},通过它可以构造一个经验分布(empirical distribution):
其中 δ(x)\delta(x)δ(x) 是狄拉克δ函数(Dirac delta function),可直观理解为在 x=0x=0x=0 处的一个无限高、无限窄的“尖峰”,其具有以下特性:
由于式(20.44)关于 xxx 不可微,因此无法直接计算其得分函数。为解决这一问题,可通过引入噪声模型对数据点进行**“平滑处理”,从而得到一个光滑且可微的概率密度表示。这种方法被称为帕森窗估计(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 是模型的参数。
- ∇zlnq(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,σ2I)\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) 变为:
这里,求和是对数据集中的 NNN 个数据点 xn\mathbf{x}_nxn 进行的。通过使用经验密度,我们将连续的积分形式转换为了离散的求和形式,这使得在实际应用中,当我们只有有限的数据样本时,能够更方便地计算和优化损失函数。
总体而言,这部分内容展示了如何通过特定的代入操作和经验密度的使用,将损失函数转换为更便于实际应用的形式,为后续的模型训练和优化提供了基础。
图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)上,那么在流形之外的点处,概率密度为零,并且由于 lnp(x)\ln p(x)lnp(x) 无定义,此处分数函数也无定义。
其次,在数据密度较低的区域,由于损失函数(20.43)是按密度加权的,因此对分数函数的估计可能不准确。使用朗之万采样时,不准确的分数函数可能会导致生成较差的轨迹。
第三,即使分数函数的模型准确,如果数据分布由多个不相交的分布混合而成,朗之万过程也可能无法正确采样。
这三个问题都可以通过为核函数(20.48)中使用的噪声方差 σ2\sigma^2σ2 选择一个足够大的值来解决,因为这样会使数据分布变得模糊(平滑)。然而,方差过大会严重扭曲原始分布,这本身就会导致分数函数建模的不准确。可以通过考虑一系列方差值 σ12<σ22<⋯<σT2\sigma^2_1 < \sigma^2_2 < \cdots < \sigma^2_Tσ12<σ22<⋯<σT2(Song和Ermon,2019)来处理这种权衡,其中 σ12\sigma^2_1σ12 足够小,能够准确表示数据分布,而 σT2\sigma^2_TσT2 足够大,可以避免上述问题。
然后,分数网络被修改为将方差作为额外输入 s(x,w,σ2)s(\mathbf{x}, \mathbf{w}, \sigma^2)s(x,w,σ2),并通过使用一个加权和形式的损失函数进行训练,该损失函数是形式为(20.51)的损失函数的加权总和,其中每一项表示相关网络与相应扰动数据集之间的误差。对于数据向量 xn\mathbf{x}_nxn,损失函数的形式则为
其中 λ(i)\lambda^{(i)}λ(i) 是权重系数。我们可以看到,这种训练过程与用于训练分层去噪网络的训练过程完全对应(意思一致、相呼应)。
训练完成后,可按顺序依次针对 i=L,L−1,…,2,1i = L, L - 1, \ldots, 2, 1i=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)是随机的,例如由无穷小的高斯步长给出。这里的参数 ttt 借物理系统的类比,常被称为“时间”。通过对连续时间极限的推导,扩散模型的正向噪声过程(20.3)可表示为(20.55)形式的随机微分方程。
对于随机微分方程(SDE)(20.55),存在相应的逆向随机微分方程(Song等人,2020),由下式给出:
其中,我们认定 ∇zlnp(z)\nabla_{\mathbf{z}} \ln p(\mathbf{z})∇zlnp(z) 为分数函数。需从 t=Tt = Tt=T 到 t=0t = 0t=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),并指出其中 ∇zlnp(z)\nabla_{\mathbf{z}} \ln p(\mathbf{z})∇zlnp(z) 为分数函数,且说明了求解该逆向随机微分方程的时间范围。
- 随机微分方程的数值求解:阐述了数值求解随机微分方程需要对时间变量离散化,介绍了最简单的欧拉 - 玛雅那求解器,以及对于逆向随机微分方程可恢复朗之万方程形式,还提到有更复杂的求解器可采用。
- 确定性过程与常微分方程:说明对于由随机微分方程控制的扩散过程,存在具有相同边际概率密度的确定性过程,由常微分方程描述,并给出了形式为(20.56)的随机微分方程对应的常微分方程(20.57)。
- 常微分方程的优势:指出常微分方程的公式化表述可使用高效自适应步长求解器,减少函数求值次数,还能使概率扩散模型与归一化流模型相关联,用于精确估计对数似然。