Skip to content

深入理解扩散模型

1. 引言:生成模型

给定从感兴趣分布中观察到的样本 x,生成模型的目标是学习其真实数据分布 p(x)。一旦学习完成,就可以从近似模型中生成新样本,并在某些设定下评估观测数据的可能性。

当前文献中主要的生成模型方向包括:

  • 生成对抗网络(GANs):通过生成器与判别器的对抗训练学习复杂分布的采样过程。
  • 基于似然的模型:直接最大化样本数据的似然,包括自回归模型、归一化流模型和变分自编码器(VAEs)。
  • 基于能量的模型(EBMs):将分布表示为灵活的能量函数并经归一化得到。
  • 基于分数的生成模型:不直接建模概率密度,而是学习其对数密度的梯度(即分数函数)。

扩散模型与上述方向都有密切联系:它既可以从变分推断和证据下界(ELBO)的角度理解,也可以从分数函数和朗之万动力学的角度理解。本文将从 ELBO 与 VAE 出发,逐步推导变分扩散模型(VDM)的数学原理,并介绍条件生成与无分类器引导等关键技术。

2. 背景:ELBO、VAE 和层次 VAE

对于许多模态,可以认为观察到的数据由一个相关的不可见潜在变量表示或生成,我们可以通过随机变量 z来表示。表达这一想法的最好直觉是通过柏拉图的洞穴寓言来理解。在寓言中,一群人一生中都被锁在一个洞穴里,只能看到投射在他们面前墙上的二维阴影,这些阴影是由在火前经过的不可见的三维物体产生的。对于这些人来说,他们观察到的一切都是由他们永远无法看到的更高维度的抽象概念决定的。

Untitled

洞穴寓言故事大致描述,把一群人关起来,如图所示,他们永远只能看岩壁上的二维画面,这个二维画面的影像是外部三维世界的物体经过一个固定的火把后在岩壁上的投影。岩壁上二维的画面显然是由外部三维实体决定的。进一步拓展思路,很有可能我们在现实世界中观察到的很多现象(数据)是从更高维度投影过来。

类似地,我们在现实世界中遇到的物体也可能是由某些更高层次的表示生成的;例如,这些表示可能封装了颜色、大小、形状等抽象属性。然后,我们观察到的可以被解释为这些抽象概念的三维投影或实例化,就像洞穴居民观察到的实际上是三维物体的二维投影一样。尽管洞穴居民永远看不到(甚至完全理解不了)隐藏的物体,他们仍然可以推断出关于它们的信息;类似地,我们可以近似潜在表示来描述我们观察到的数据。

尽管柏拉图的洞穴寓言阐释了潜在变量作为可能不可观察的表示来决定观察结果的背后思想,但这个类比的一个缺点是,在生成建模中,我们通常寻求学习低于观察维度的潜在表示,而不是高于观察维度的。这是因为试图学习一个比观察结果更高维度的表示,如果没有强有力的先验知识是没有用的。另一方面,学习低维潜在变量也可以看作是一种压缩形式,并且有可能揭示出描述观察结果的语义上有意义的结构。

2.1 证据下界(Evidence Lower Bound, ELBO)

数学上,我们可以想象潜在变量和我们观察到的数据由一个联合分布 p(x,z) 来建模。生成建模中称为“基于似然的”方法之一是学习一个模型来最大化所有观测到的x的似然p(x)。有两种方法可以操作这个联合分布来恢复我们观测数据的似然p(x), 我们可以明确地对潜在变量z积分,计算边缘分布p(x)

p(x)=p(x,z)dz

或者,我们也可以使用概率链式法则:

p(x)=p(x,z)p(zx)

直接计算和最大化似然 p(x)是困难的,因为它要么涉及将所有潜在变量z积分出来(方程 (1) ), 这对于复杂模型来说是不可行的;要么涉及获得真实的潜在编码器p(z|x)(方程 (2) )。

然而,使用这两个方程,我们可以推导出一个称为证据下界(ELBO)的项,顾名思义,它是证据的一个下界。在这种情况下,证据被量化为观测数据的对数似然。然后,最大化 ELBO 成为一个代理目标,用于优化潜在变量模型;在最好的情况下,当 ELBO 被正确地参数化并且被完美优化时,它与证据完全等价。

正式地,证据下界 ELBO 的方程是:

Eqϕ(zx)[logp(x,z)qϕ(zx)]

为了明确与证据的关系,我们可以数学上写为:

logp(x)Eqϕ(zx)[logp(x,z)qϕ(zx)]

这里,qϕ(z|x)是一个灵活的近似变分分布,其参数ϕ是我们要寻求优化的。直观上,它可以被看作是一个可参数化的模型,用于学习估计给定观测x的潜在变量的真实分布;换句话说,它寻求近似真实的后验p(z|x)。当我们通过调整参数ϕ来增加下界以最大化 ELBO 时,我们获得了可以用来建模真实数据分布并从中采样的组件,从而学习一个生成模型。

变分分布:

变分分布是变分推断中用于近似复杂后验分布的简化模型。它通过优化一个参数化的分布族来捕捉原始分布的关键特征,通常选择易于计算且足够灵活的分布作为变分分布族(如最为常见的正态分布),并通过最小化与真实后验分布的 KL 散度来调整参数。当这两个分布差距很小时,就可以用新分布去代替实际要求的分布了。这种方法使得原本在计算上不可行的复杂概率分布推断变得可行。

2.2 证据下界(ELBO)的推导

让我们从方程(1)开始推导证据下界 ELBO:

logp(x)=logp(x,z)dz=logp(x,z)qϕ(zx)qϕ(zx)dz=logEqϕ(zx)[p(x,z)qϕ(zx)]Eqϕ(zx)[logp(x,z)qϕ(zx)]

(6)->(7): 将对随机变量的积分转换为对其概率分布的期望值,意味着我们不是直接计算随机变量在所有可能取值上的积分,而是计算这个随机变量按照其概率分布取值时的平均值。这种转换是基于期望值的定义,即一个随机变量的期望值是它所有可能取值的概率加权和。

在变分推断中,由于直接计算(5)这个积分可能非常复杂,因此利用期望的性质来简化问题。期望的性质允许我们将对随机变量的积分转换为对其概率分布的期望值,进一步可以通过蒙特卡洛采样方法来进行估计。

(7)->(8):Jensen 不等式是一个凸函数性质的结果,对于任意凹函数 f(如log)和概率分布P,以下不等式恒成立:logEP[X]EP[logX]

在这个推导中,我们通过应用 Jensen 不等式 直接得到了我们的下界。然而,这并没有提供关于实际发生的情况的有用信息;关键的是,这个证明没有给出为什么 ELBO 实际上是证据下界的直觉,因为 Jensen 不等式 把它模糊处理了。此外,仅仅知道 ELBO 是数据的下界,并没有真正告诉我们为什么我们想要将其最大化作为目标。为了更好地理解证据和 ELBO 之间的关系,让我们进行另一个推导,这次使用方程(2):

logp(x)=logp(x)qϕ(zx)dz=qϕ(zx)logp(x)dz=Eqϕ(zx)[logp(x)]=Eqϕ(zx)[logp(x,z)p(zx)]=Eqϕ(zx)[logp(x,z)qϕ(zx)p(zx)qϕ(zx)]=Eqϕ(zx)[logp(x,z)qϕ(zx)]+Eqϕ(zx)[logqϕ(zx)p(zx)]=Eqϕ(zx)[logp(x,z)qϕ(zx)]+DKL(qϕ(zx)∣∣p(zx))Eqϕ(zx)[logp(x,z)qϕ(zx)]

KL 散度:

KL 散度也称为相对熵(Relative Entropy),是衡量两个概率分布差异的一种方法。它是两个概率分布 P 和 Q 之间的非对称距离度量,定义为:DKL(P||Q)=xP(x)log(P(x)Q(x)),其中,P 是数据的真实分布,而 Q 是模型或估计分布。KL 散度的值总是非负的,当且仅当 P 和 Q 完全相同时,KL 散度为零。

从这个推导中,我们清楚地从方程 (15) 观察到证据等于 证据下界 ELBO 加上近似后验 qϕ(z|x)和真实后验p(z|x)之间的 KL 散度。实际上,正是这个 KL 散度项在第一个推导的方程(8)中通过 Jensen 不等式 被神奇地移除了。理解这个项是理解证据下界 ELBO 和证据之间关系的关键,也是理解为什么优化 ELBO 是一个合适的目标的原因。

首先,我们现在知道为什么 ELBO 确实是一个下界:

证据和 ELBO 之间的差是一个严格非负的 KL 项,因此 ELBO 的值永远不会超过证据。

其次,我们探讨为什么寻求最大化 ELBO。

在我们引入了想要建模的潜在变量 z之后,我们的目标是学习描述我们观测数据的潜在结构。换句话说,我们想要优化我们的变分后验qϕ(z|x)的参数,以便它完全匹配真实的后验分布p(z|x),这通过最小化它们的 KL 散度(理想情况下为零)来实现的。不幸的是,直接最小化这个 KL 散度项是不可行的,因为我们无法访问真实的p(z|x)分布。

但是,请观察,在方程(15)的左侧,我们数据的似然性(因此我们的证据项 logp(x))相对于ϕ总是一个常数,因为它是通过从联合分布p(x,z)中边缘化所有潜在z来计算的,根本不依赖于ϕ。由于证据下界 ELBO 和 KL 散度项加起来是一个常数,对证据下界 ELBO 项的任何最大化都必然引起 KL 散度项的等量最小化。

因此,证据下界 ELBO 可以被最大化,作为学习完美地模拟真实潜在后验分布的代理目标;我们越优化 ELBO,我们的近似后验就越接近真实后验。此外,一旦训练完成,ELBO 还可以用来估计观测或生成数据的可能性,因为它被学习来近似模型证据 logp(x)

2.3 变分自编码器(Variational Autoencoders)

VAE:

在 VAE 的架构中,设计两个网络,一个网络做 Encoder,将观察的数据 x映射到隐变量z上,另一个网络做 Decoder,从采样的z中解码回x。从实践的角度,VAE 需要训练两套参数θϕ,即qϕ(z|x)pθ(x|z) 分别对应 encoder 和 decoder。

VAE 的目标函数采用的是最大似然(likelihood-based)的路线,即对观察到的数据 x 计算最大p(x),借助数学工具 ELBO(Evidence Lower Bound),可以得到最大似然分布的一个下界。

在变分自编码器(VAE)的默认方程中,我们直接最大化证据下界(ELBO)。这种方法是变分的,因为我们在由 ϕ参数化的潜在后验分布族中优化最佳的qϕ(z|x)。它被称为自编码器,因为它让人想起传统的自编码器模型,其中输入数据被训练在经历一个中间瓶颈表示步骤后进行预测。为了明确这种联系,让我们进一步剖析 ELBO 项:

Eqϕ(zx)[logp(x,z)qϕ(zx)]=Eqϕ(zx)[logpθ(xz)p(z)qϕ(zx)]=Eqϕ(zx)[logpθ(xz)]+Eqϕ(zx)[logp(z)qϕ(zx)]=Eqϕ(zx)[logpθ(xz)]reconstruction termDKL(qϕ(zx)∣∣p(z))prior matching term

在这个情况下,我们学习了一个中间瓶颈分布 qϕ(z|x),它可以被看作是一个编码器;将输入转换为可能的潜在分布。同时,我们学习了一个确定性函数pθ(x|z)将给定的潜在向量z转换为观测x,这可以被解释为解码器。

Visualizing a Variational Autoencoder

可视化普通变分自编码器。通过重参数化技巧联合学习潜在编码器和解码器 [6, 7]

方程(19)中的两项各自有直观的描述:

  • 第一项测量了解码器从我们的变分分布中重构的似然性;这确保了学习到的分布是模拟有效的潜在变量,原始数据可以从中再生。
  • 第二项测量了学习到的变分分布与我们对潜在变量持有的先验信念的相似性。最小化这个项鼓励编码器实际学习一个分布,而不是塌缩成一个狄拉克δ函数。

通俗来说,第一项要求隐变量尽可能无损的保留原始数据的信息,即 Decoder 的输出尽量接近原图;第二项让 隐变量的先验分布和后验分布尽可能接近,即 Encoder 输出的分布要接近真实分布 p(𝑧).

因此,最大化 ELBO 等同于最大化其第一项并最小化其第二项。

VAE 的一个决定性特征是 ELBO 是如何联合优化参数 ϕθ。VAE 的编码器qϕ(z|x)通常选择为具有对角线协方差的多元高斯分布进行建模,而先验pθ(z)通常选择为标准多元高斯分布:

qϕ(zx)=N(z;μϕ(x),σϕ2(x)I)p(z)=N(z;0,I)

这时,证据下界 ELBO 中的 KL 散度项能够直接计算解析解,而重建项可以使用蒙特卡洛估计来近似。我们的目标可以被重写为:

argmaxϕ,θEqϕ(zx)[logpθ(xz)]DKL(qϕ(zx)∣∣p(z))argmaxϕ,θl=1Llogpθ(xz(l))DKL(qϕ(zx)∣∣p(z))

其中潜在变量 {z(l)}l=1L是从编码分布qϕ(z|x)中采样的,对于数据集中的每个观测x都是如此。然而,在这个默认设置中出现了一个问题:我们计算损失的每个z(l)都是通过一个随机采样过程生成的,这通常是不可微分的。幸运的是,当qϕ(z|x)被设计为模拟某些分布(包括多元高斯分布)时,这个问题可以通过重参数化技巧来解决。

重参数化技巧

重参数化技巧是机器学习和深度学习中用于处理随机梯度问题的一种方法。在许多模型中,特别是在生成模型如变分自编码器(VAE)中,我们经常需要从概率分布中采样,而这些采样操作本身是随机的,不可导的。这就意味着我们不能直接通过这些随机采样来计算梯度,并用它们来更新模型参数。

为了解决这个问题,重参数化技巧通过引入一个新的随机变量,并将其作为一个噪声变量来重新表达原始的随机采样。这个新的表达式是原始随机变量的一个确定性函数,这意味着它完全依赖于模型参数和一个新的独立的随机变量。这样,我们就可以通过这个新的随机变量来计算原始随机变量的期望值,而这个期望值是可导的,因为它不再依赖于模型参数的随机选择。

具体来说,假设我们有一个随机变量 Z,它是通过从某个分布 P(Z)中采样得到的。在没有重参数化技巧的情况下,我们不能直接对Z进行梯度下降优化。但是,如果我们能够找到一种方法,使得 Z 可以被表达为一个确定性函数 g 和一个独立的随机变量ϵ的组合,即Z=g(ϵ;θ),其中θ是模型参数,那么我们就可以使用ϵ来计算 Z 的期望值,并通过ϵ的分布来计算梯度。这种方法的关键优势在于,尽管 Z 本身是随机的,但g(ϵ;θ)是确定性的,并且可以通过ϵ来计算梯度。因此,我们可以使用标准的梯度下降算法来优化模型参数θ,即使整个表达式涉及到随机变量。

在 VAE 的情况下,重参数化技巧允许我们将编码器输出的随机潜在变量 Z重写为一个关于模型参数ϕ的确定性函数,同时依赖于一个独立的随机变量。这样,我们就可以通过梯度下降来优化编码器和解码器的参数,同时保持生成的数据与真实数据的相似性。

例如,从正态分布 xN(x;μ,σ2)中采样的随机变量,其中μ是任意的均值,σ2是方差,可以被重写为:

x=μ+σϵwithϵN(ϵ;0,I)

换句话说,任意的高斯分布可以被解释为标准高斯(其中 ϵ是一个样本),通过添加将均值从零移动到目标均值μ,以及通过目标方差σ2拉伸方差。因此,通过重参数化技巧,从任意高斯分布中采样可以通过从标准高斯中采样,将结果乘以目标标准差,然后通过目标均值进行平移来执行。

在 VAE 中,每个潜在变量 z被因此可以被计算为输入x和辅助噪声变量ϵ的确定性函数:

z=μϕ(x)+σϕ(x)ϵwithϵN(ϵ;0,I)

在重参数化的 z版本下,可以计算相对于ϕ的梯度,以期望优化μϕσϕ。因此,VAE 利用重参数化技巧和蒙特卡洛估计联合优化证据下界 ELBO。其中表示逐元素乘积。

训练完 VAE 后,可以通过直接从潜在空间 p(z)中采样,然后将其通过解码器来生成新数据。当潜在变量z的维度小于输入x的维度时,变分自编码器特别有趣,因为那时我们可能学习紧凑、有用的表示。此外,当学习到一个语义上有意义的潜在空间时,潜在向量可以在传递给解码器之前进行编辑,以更精确地控制生成的数据。

2.4 分层变分自编码器(Hierarchical Variational Autoencoders)

分层变分自编码器(HVAE)是变分自编码器(VAE)的推广,可扩展到潜在变量的多个层次结构。在这个框架下,潜在变量本身被视为由其他更高级别的、更抽象的潜在变量生成。就像我们将三维观察对象视为从更高层次的抽象潜在对象生成一样,柏拉图洞穴中的人们将三维对象视为生成二维观察的潜在对象。因此,从柏拉图洞穴居民的视角来看,他们的观察可以被视为由两个或更多层次的潜在层次结构建模。

Visualizing a Hierarchical VAE

具有 T 个分层潜在变量的马尔可夫分层变分自编码器。生成过程被建模为马尔可夫链,其中每个潜在变量zt仅从前一个潜在变量zt+1生成。

与具有 T 层层次结构的一般分层变分自编码器(HVAE)不同, 在一般的 HVAE 中,每个潜在变量被允许依赖于之前所有的潜在变量。然而,在本研究中,我们关注的是一个特殊案例,我们称之为马尔可夫分层变分自编码器(MHVAE)。在 MHVAE 中,生成过程是一个马尔可夫链,也就是说,层次结构中的每个转换都是马尔可夫的,其中解码每个潜在变量 zt仅依赖于前一个潜在变量zt+1。这可以被看作是简单地将变分自编码器(VAEs)堆叠在一起,如图 2 所示;描述这个模型的另一个恰当术语是递归变分自编码器(Recursive VAE)。数学上,我们表示马尔可夫分层变分自编码器的联合分布和后验分布为:

p(x,z1:T)=p(zT)pθ(xz1)t=2Tpθ(zt1zt)

其后验表示为:

qϕ(z1:Tx)=qϕ(z1x)t=2Tqϕ(ztzt1)

然后,我们可以很容易地将证据下界 ELBO 扩展为:

logp(x)=logp(x,z1:T)dz1:T=logp(x,z1:T)qϕ(z1:Tx)qϕ(z1:Tx)dz1:T=logEqϕ(z1:Tx)[p(x,z1:T)qϕ(z1:Tx)]Eqϕ(z1:Tx)[logp(x,z1:T)qϕ(z1:Tx)]

然后,我们可以将我们的联合分布(方程(23))和后验(方程(24))代入方程(28),以产生另一种形式:

Eqϕ(z1:Tx)[logp(x,z1:T)qϕ(z1:Tx)]=Eqϕ(z1:Tx)[logp(zT)pθ(xz1)t=2Tpθ(zt1zt)qϕ(z1x)t=2Tqϕ(ztzt1)]

正如我们将在下面探讨变分扩散模型时所指出的,这个目标可以进一步分解为可解释的组成部分。

3. 变分扩散模型(Variational Diffusion Models)

基本概念:

业界针对 VAE 的问题做了很多尝试,Diffusion Model 可以算作是其中一种比较成功的尝试。以图片生成举例,VAE 想一步到位,同时训练 Encoder 和 Decoder,将原始图片映射到隐变量中,又能解码回来,端到端地学习整个过程太难了。有没有什么办法可以降低问题的难度?一个常见的思路是将困难的问题分而治之,即能不能只训练 Encoder 或 Decoder?能不能把问题分解成一个个子问题?

变分扩散模型(VDM)可以简单地看作是一个具有三个关键限制的马尔可夫分层变分自编码器(Markovian Hierarchical Variational Autoencoder, MHVAE):

  1. 潜在维度完全等于数据维度。
  2. 每个时间步的潜在编码器的结构不是学习得到的,它被预定义为线性高斯模型。换句话说,它是以前一时间步的输出为中心的高斯分布。
  3. 潜在编码器的高斯分布参数随时间变化,使得最终时间步 T 的潜在分布是标准高斯分布。

此外,我们明确保持了从标准马尔可夫分层变分自编码器继承的分层转换之间的马尔可夫性质。让我们进一步阐述这些假设的含义。

根据第一个限制,现在我们可以将真实的数据样本和潜在变量表示为 xt,其中t=0表示真实的数据样本,t[1,T]表示具有层次索引t的相应潜在变量。变分扩散模型的后验与 MHVAE 的后验相同(方程 24),现在可以重写为:

q(x1:Tx0)=t=1Tq(xtxt1)

根据第二个假设,我们知道每个潜在变量在编码器中的分布是一个以它前一个层次潜在变量为中心的高斯分布。与 MHVAE 不同,在变分扩散模型中,每个时间步 t的编码器的结构不是学习得到的,它被固定为线性高斯模型,其中均值和标准差可以预先设置为超参数,或者作为参数学习。

我们用均值 μt(xt)=αtxt1和方差Σt(xt)=(1αt)I来参数化高斯编码器,其中系数的形式被选择为保持潜在变量的方差在相似的尺度上;换句话说,编码过程是方差保持的, 请注意,允许使用其他高斯参数化方法,并且可以得到类似的推导结果。主要的收获是αt是一个(可能是可学习的)系数,它可以根据层次深度t变化,以增加灵活性。数学上,编码器的转换表示为:

q(xtxt1)=N(xt;αtxt1,(1αt)I)

根据第三个假设,我们知道 αt按照固定或可学习的时间表随时间演变,使得最终潜在的分布p(xT)是标准高斯分布。然后,我们可以更新马尔可夫分层变分自编码器(MHVAE)的联合分布(方程 23),将变分扩散模型(VDM)的联合分布写为:

p(x0:T)=p(xT)t=1Tpθ(xt1xt)where,p(xT)=N(xT;0,I)

总的来说,这组假设描述了图像输入随时间逐步噪声化的过程:我们通过逐步添加高斯噪声逐渐破坏图像,直到最终它完全变成纯高斯噪声。这个过程在图 3 中以视觉方式呈现。

Visualizing a Variational Diffusion Model

变分扩散模型的视觉表示。输入随着时间的推移稳定地产生噪声,直到它变得与高斯噪声相同;扩散模型学习逆向过程。

请观察,我们的编码器分布 q(xt|xt1)已经不再由ϕ参数化,因为它们完全被建模为在每个时间步骤具有确定的均值和方差参数的高斯分布。因此,在变分扩散模型中,我们只对条件分布pθ(xt1|xt)感兴趣,这样我们就可以模拟新数据。在优化变分扩散模型之后,采样过程就像从p(xT)中采样高斯噪声,然后迭代地运行去噪转换pθ(xt1|xt) 进行T步来生成一个新的x0

3.1 第一种常见解释

像任何分层变分自编码器 HVAE 一样,变分扩散模型可以通过最大化证据下界 ELBO 来优化,它可以被推导为:

logp(x)Eq(x1:Tx0)[logp(x0:T)q(x1:Tx0)]=Eq(x1x0)[logpθ(x0x1)]reconstruction termEq(xT1x0)[DKL(q(xTxT1)∣∣p(xT))]prior matching termt=1T1Eq(xt1,xt+1x0)[DKL(q(xtxt1)∣∣pθ(xtxt+1))]consistency term

证明过程如下:

Deriving the ELBO of a VDM in terms of Consistency terms

注:

(34):边缘分布的概念

(36):期望的定义

(37):Jensen 不等式

(42)-(43):求和的期望(期望对求和的操作)与期望的求和(求和对期望的操作)在数学上是相等的。这是由于期望的线性属性所保证的。

(43)-(44):根据表达式中的下标,对下标具体化

(44)-(45):这里解释其中第二项 Eq(xT1,xT|x0)[logp(xT)q(xT|xT1)]=Eq(xT1|x0)[DKL(q(xT|xT1)p(xT))],第三项同理,期望是在q(xT1,xT|x0)这个联合分布下计算的,意味着考虑了xT1xT的所有可能的组合。而Eq(xT1|x0)[DKL(q(xT|xT1)p(xT))]是首先对xT1求期望,然后计算给定xT1时,xT的分布q(xT|xT1)与其先验p(xT)之间的 KL 散度的负值。这两个表达式等价的原因如下:

  • 期望的性质:由于 KL 散度内部本身就是一个对数期望形式,这个期望的外部和内部可以合并,这是因为期望的计算是线性的,可以跨越概率分布的边界。

  • KL 散度的定义:上述的 KL 散度定义直接对应于你的第一个表达式的内部。KL 散度衡量的是在特定条件下,两个分布的相对熵。 所以,当从 q(xT1,xT|x0)这个条件分布计算期望时,实际上是在遍历所有xT1xT的组合,并评估q(xT|xT1)p(xT)之间的相对熵。当你把这个过程表示为 KL 散度,并注意到 KL 散度本身就是期望的形式时,两边就能够匹配上,并通过将 KL 散度的定义带入,完成从左侧到右侧的转换。

最后三行公式是对变分扩散模型(VDM)中的证据下界(ELBO)进行分解和解释的过程。ELBO 是变分推断中的一个重要概念,用于衡量模型对数据的拟合程度。ELBO 被分解为三个主要部分:重构项(reconstruction term)、先验匹配项(prior matching term)和一致性项(consistency term)。下面是各个组成部分的解释:

  1. 重建项: Eq(x1|x0)[logpθ(x0|x1)],这一项计算的是给定第一步潜在变量x1时,预测原始数据x0的对数概率。这一项反映了模型重构数据的能力,与普通 VAE 中的重构项类似。
  2. 先验匹配项: Eq(xT1|x0)[KL(q(xT|xT1)p(xT))];这一项计算的是最终潜在变量xT的分布与高斯先验分布之间的对数比。当最终潜在变量的分布在所有步骤后与高斯先验分布匹配时,这一项的值最小。 这一项不需要优化,因为它没有可训练的参数;此外,由于我们假设了一个足够大的 T,使得最终分布是高斯分布,这个项实际上变成了零。
  3. 一致性项: Eq(xt1,xt,xt+1|x0)[KL(q(xt|xt1)pθ(xt|xt+1))], 它努力使xt的分布在两个方向上保持一致,即从更噪声的图像到更清晰的图像的去噪步骤和从更清晰的图像到更噪声的图像的加噪步骤相匹配。这种一致性在数学上通过 KL 散度反映出来。当两个去噪步骤尽可能匹配时,这一项被最小化,如方程 31 中定义的高斯分布q(xt|xt1)

通俗地讲:重构项度量了从第一步重构到原始输入的部分,和 VAE 比较像,保证了模型的生成能力;先验匹配项让最后一步的输出尽可能接近标准高斯分布;前后一致项保证了中间每一个时间步,在经过模型还原后,要和前向扩散过程保持一致。

这些项的和构成了 VDM 的 ELBO,它是优化过程中的目标函数。通过最大化 ELBO,我们可以调整模型参数 θ来提高模型对数据的拟合程度,同时确保生成过程的一致性。这个过程涉及到对联合分布和条件分布的积分进行计算,通常需要使用蒙特卡洛估计和重参数化技巧来近似这些期望值。通过这种方式,我们可以有效地训练 VDM 模型,使其能够生成与训练数据相似的新样本。

Deriving the ELBO of a VDM in terms of Consistency terms

图 4:根据第一个推导,变分扩散模型(VDM)可以通过确保对于每一个中间潜在变量 xt,来自上层潜在变量的后验分布pθ(xt|xt+1)与来自下层潜在变量的高斯噪声破坏的分布q(xt|xt1)相匹配来优化。在这个图中,对于每一个中间潜在变量xt,最小化由粉色和绿色箭头所表示的分布之间的差异。

在图 4 中,我们以视觉方式展示了 ELBO 的这种解释;对于每个中间的 xt,我们最小化由粉色和绿色箭头表示的分布之间的差异。在完整的图中,每条粉色箭头也必须从x0开始,因为它也是一个条件项。

优化一个变分扩散模型(VDM)的成本主要由第三项主导,因为我们必须要在所有时间步骤 t 上进行优化。在这种推导下,ELBO 的所有项都被计算为期望,因此可以使用蒙特卡洛估计来近似。然而,实际上使用我们刚刚推导出的项来优化 ELBO 可能是次优的;因为一致性项是作为两个随机变量 {xt1,xt+1}的期望来计算的,对于每个时间步,它的蒙特卡洛估计的方差可能会比使用每个时间步一个随机变量估计的项更高。由于它是通过对T1个一致性项求和来计算的,对于大的T值,ELBO 的最终估计值可能具有高方差。

注:原先推导出来一致性项是作为两个随机变量 {xt1,xt+1}的期望值来计算的。当使用蒙特卡洛方法来近似这个期望值时,通常是通过从这些随机变量的分布中采样来进行的。由于每次采样都是独立的,当同时对两个随机变量进行采样时,这些采样点的组合将形成一个联合分布,其方差可能会比单个随机变量的方差更大。

在蒙特卡洛估计中,方差是衡量估计值稳定性的一个指标。高方差意味着估计值在不同的采样下会有较大的波动,这可能导致优化过程中的不稳定性。在 VDM 的情况下,由于一致性项是通过累加 T-1 个这样的 KL 散度项来计算的,如果 T 很大,那么最终的 ELBO 估计值的方差可能会非常高,这会影响优化算法的性能和收敛速度。

为了解决这个问题,我们可以尝试推导出一个新的 ELBO 形式,其中每个项都是作为仅对一个随机变量的期望值来计算的。这样做可以降低蒙特卡洛估计的方差,从而提高优化过程的稳定性和效率。通过这种方式,我们可以更准确地估计 ELBO,并更有效地优化模型参数。

让我们尝试推导出一个 ELBO 的形式,其中每个项都是对最多一个随机变量的期望进行计算。关键的洞察是我们可以重写编码器转换为 q(xt|xt1)=q(xt|xt1,x0),其中额外的条件项由于马尔可夫性质而变得多余。然后,根据贝叶斯规则,我们可以将每个转换重写为:

q(xtxt1,x0)=q(xt1xt,x0)q(xtx0)q(xt1x0)

利用这个新方程,ELBO 可以推导出为以下形式:

logp(x)Eq(x1:Tx0)[logp(x0:T)q(x1:Tx0)]=Eq(x1x0)[logpθ(x0x1)]reconstruction termDKL(q(xTx0)∣∣p(xT))prior matching termt=2TEq(xtx0)[DKL(q(xt1xt,x0)∣∣pθ(xt1xt))]denoising matching term

证明过程如下:

根据去噪匹配项推导 VDM 的 ELBO

因此,我们已经成功地推导出了 ELBO 新的解释,ELBO 可以分解为单个项的总和,这些项是一次最多一个随机变量的期望,因此具有较低的方差。这个方程还有一个优雅的解释,当检查每个单独的项时,可以揭示出来:

  1. Eq(x1|x0)[logpθ(x0|x1)]可以被解释为一个重建项;像 VAE 的 ELBO 中的类似项一样,这个项可以使用蒙特卡洛估计来近似和优化。

  2. DKL(q(xTx0)∣∣p(xT))表示最终噪声输入的分布与标准高斯先验的接近程度。它没有可训练的参数,并且在我们的假设下也等于零。

  3. Eq(xtx0)[DKL(q(xt1xt,x0)∣∣pθ(xt1xt))]是一个去噪匹配项。我们学习期望的去噪转换步骤pθ(xt1|xt)作为可处理的、真实的去噪转换步骤q(xt1|xt,x0)的近似。由于q(xt1|xt,x0)转换步骤可以作为真实的信号,因为它定义了如何使用最终完全去噪的图像x0来去噪噪声图像xt。因此,当两个去噪步骤尽可能匹配时,这个项被最小化,通过它们之间的 KL 散度来衡量。

下图描述了此 ELBO 分解的直观解释:

Deriving the ELBO of a VDM in terms of Denoising Matching terms

变分扩散模型(VDM)也可以通过学习每个独立潜在变量的去噪步骤来优化,方法是将其与可计算的、可靠的真实去噪步骤相匹配。这再次通过视觉上将绿色箭头所代表的分布与粉色箭头所代表的分布相匹配来表示。这里运用了艺术上的自由;在完整的画面中,每个粉色箭头也必须源自真实图像,因为真实图像也是一个条件项。

作为旁注,我们观察到,在 ELBO 的推导过程中,只使用了马尔可夫假设;因此,这些方程对于任何任意的马尔可夫 HVAE 都是成立的。此外,当我们设置 T=1时,变分扩散模型的两种 ELBO 解释对于标准 VAE 的 ELBO 方程(方程 19)是完全相同的。

在这种 ELBO 的推导中,大部分优化成本再次由求和项占据,它支配了重建项。虽然每个 KL 散度项 DKL(q(xt1xt,x0)∣∣pθ(xt1xt))对于任意复杂的马尔可夫 HVAEs 中的任意后验来说都是难以最小化的,因为同时学习编码器的复杂性增加了,但在变分扩散模型(VDM)中,我们可以利用高斯转换假设来使优化变得可行。通过贝叶斯规则,我们有:

q(xt1xt,x0)=q(xtxt1,x0)q(xt1x0)q(xtx0)

由于我们已经知道根据我们对编码器转换的假设(方程 31)有: q(xt|xt1,x0)=q(xt|xt1)=N(xt;αtxt1,(1αt)I),剩下的的工作是推导出q(xt|x0)q(xt1|x0)的形式。幸运的是,通过利用变分扩散模型的编码器转换是线性高斯模型这一事实, 这些也可以变得可行。回想一下,在重参数化技巧下,样本xtq(xt|xt1)可以被重写为:

xt=αtxt1+1αtϵwithϵN(ϵ;0,I)

类似地,样本 xt1q(xt1|xt2)可以被重写为:

xt1=αt1xt2+1αt1ϵwithϵN(ϵ;0,I)

然后,q(xt|x0)的形式可以通过重复应用重新参数化技巧来递归推导出来。假设我们有2T个随机噪声变量{ϵt,ϵt}t=0TiidN(ϵ;0,I)。那么,对于任意样本xtq(xt|x0),我们可以重写为:

xt=αtxt1+1αtϵt1=αt(αt1xt2+1αt1ϵt2)+1αtϵt1=αtαt1xt2+αtαtαt1ϵt2+1αtϵt1=αtαt1xt2+αtαtαt12+1αt2ϵt2=αtαt1xt2+αtαtαt1+1αtϵt2=αtαt1xt2+1αtαt1ϵt2==i=1tαix0+1i=1tαiϵ0=α¯tx0+1α¯tϵ0N(xt;α¯tx0,(1α¯t)I)

其中,在方程 44 中,我们利用了两个独立的高斯随机变量之和仍然是高斯分布的事实,其均值为两个均值之和,方差为两个方差之和。将 1αtϵt1解释为从高斯N(0,(1αt)I)中抽取的样本,以及αtαtαt1ϵt2解释为从高斯N(0,(αtαtαt1)I)中抽取的样本,我们可以将它们的和视为从高斯N(0,(1αtαt1)I)中抽取的随机变量N(0,(1αt+αtαtαt1)I)=N(0,(1αtαt1)I)。这个分布的样本可以使用重参数化技巧表示为1αtαt1ϵt2,如方程 46 所示。

因此,我们已经推导出了 q(xt|x0)的高斯形式。这个推导可以修改得到描述q(xt1|x0)的高斯参数化。现在,知道了q(xt|x0)q(xt1|x0)的形式,我们可以通过将它们代入贝叶斯规则扩展来计算q(xt1|xt,x0)的形式:

q(xt1xt,x0)=q(xtxt1,x0)q(xt1x0)q(xtx0)=N(xt;αtxt1,(1αt)I)N(xt1;α¯t1x0,(1α¯t1)I)N(xt;α¯tx0,(1α¯t)I)N(xt1;αt(1α¯t1)xt+α¯t1(1αt)x01α¯tμq(xt,x0),(1αt)(1α¯t1)1α¯tIΣq(t))

证明过程如下:

推导每个中间噪声潜伏的分布

其中在第 5 行中,C(xt,x0)是一个关于xt1的常数项,它是xt,x0α值的组合;该项在最后一行中被隐式地返回。

因此,我们已经展示了在每一步,xt1q(xt1|xt,x0)都是正态分布的,其均值μq(xt,x0)xtx0的函数,方差Σq(t)α系数的函数。这些α系数在每个时间步是已知并固定的;它们要么被永久地设置为超参数,要么被视为一个网络的当前推理输出,该网络旨在对它们进行建模。根据方程 53,我们可以将我们的方差方程重写为Σq(t)=σq2(t)I,其中:

σq2(t)=(1αt)(1α¯t1)1α¯t

为了尽可能地将近似去噪转换步骤 pθ(xt1|xt)与真实去噪转换步骤q(xt1|xt,x0)相匹配,我们也可以将其建模为高斯分布。此外,由于所有的α项在每个时间步都是已知并固定的,我们可以立即构造近似去噪转换步骤的方差为Σq(t)=σq2(t)I。我们必须将均值μθ(xt,t)参数化为xt的函数,因为pθ(xt1|xt)不依赖于x0

回想一下,两个高斯分布之间的 KL 散度是:

DKL(N(x;μx,Σx)∣∣N(y;μy,Σy))=12[logΣyΣxd+tr(Σy1Σx)+(μyμx)TΣy1(μyμx)]

在我们的例子中,我们可以将两个高斯的方差设置为完全匹配,优化 KL 散度项简化为最小化两个分布的均值之间的差异:

argminθDKL(q(xt1xt,x0)∣∣pθ(xt1xt))=argminθDKL(N(xt1;μq,Σq(t))∣∣N(xt1;μθ,Σq(t)))=argminθ12σq2(t)[μθμq22]

证明过程如下:

推导两个高斯函数之间的 KL

其中,我们已经将 μq(xt,x0)简写为μq,将μθ(xt,t)简写为μθ,以便于书写。换句话说,我们想要优化一个μθ(xt,t)来匹配μq(xt,x0),根据我们推导出的公式 53,它的形式是:

μq(xt,x0)=αt(1α¯t1)xt+α¯t1(1αt)x01α¯t

由于 μθ(xt,t)也依赖于xt,我们可以通过将其设置为以下形式来密切匹配μq(xt,x0)

μθ(xt,t)=αt(1α¯t1)xt+α¯t1(1αt)x^θ(xt,t)1α¯t

其中 x^θ(xt,t)是一个神经网络,它的目标是从噪声图像xt和时间索引t预测原始图像x0。然后,优化问题简化为:

argminθDKL(q(xt1xt,x0)∣∣pθ(xt1xt))=argminθDKL(N(xt1;μq,Σq(t))∣∣N(xt1;μθ,Σq(t)))=argminθ12σq2(t)α¯t1(1αt)2(1α¯t)2[x^θ(xt,t)x022]

证明过程如下:

推导图像预测目标

因此,优化一个变分扩散模型归结为学习一个神经网络,该神经网络从任意噪声版本来预测原始图像 x0。此外,通过最小化我们推导出的 ELBO 目标中的所有噪声水平的求和项(方程 38),可以通过最小化所有时间步骤的期望来近似:

argminθt=2TEq(xtx0)[DKL(q(xt1xt,x0)∣∣pθ(xt1xt))]=argminθEtU{2,T}[Eq(xtx0)[12σq2(t)α¯t1(1αt)2(1α¯t)2[x^θ(xt,t)x022]]]

然后可以使用随时间步长的随机样本进行优化。

3.2 学习扩散噪声参数

让我们探讨如何联合学习变分扩散模型(VDM)的噪声参数。一种潜在的方法是使用带有参数 η的神经网络α^η(t)来对αt进行建模。然而,这种方法效率低下,因为在每个时间步t都必须多次执行推断以计算α¯t。尽管通过缓存可以缓解这种计算成本,我们还可以推导出另一种学习扩散噪声参数的方法。通过将我们的方差方程(方程 54)代入我们推导出的每个时间步的目标方程 (方程 61),我们可以简化:

12σq2(t)α¯t1(1αt)2(1α¯t)2[x^θ(xt,t)x022]=12(α¯t11α¯t1α¯t1α¯t)[x^θ(xt,t)x022]

证明过程如下:

使用 SNR 项导出目标

回忆一下方程 50 , q(xtx0)是一个高斯分布, 形式为N(xt;α¯tx0,(1α¯t)I)。然后,遵循信噪比 (SNR)的定义:μ2σ2,我们可以将每个时间步 t 的信噪比写为:

SNR(t)=α¯t1α¯t

然后,我们推导出的方程 63(和方程 61) 可以简化为:

12σq2(t)α¯t1(1αt)2(1α¯t)2[x^θ(xt,t)x022]=12(SNR(t1)SNR(t))[x^θ(xt,t)x022]

顾名思义,信噪比(SNR)代表了原始信号与存在的噪声量之间的比率;较高的 SNR 表示更多的信号,较低的 SNR 表示更多的噪声。在扩散模型中,我们要求 SNR 随着时间步 t的增加而单调递减;这形式化了扰动输入xt随着时间的推移变得越来越嘈杂的概念,直到在t=T时完全变成标准高斯噪声。

根据方程(65)中目标的简化,我们可以直接使用神经网络对每个时间步的 SNR 进行参数化,并与扩散模型一起联合学习[10]。由于 SNR 必须随时间单调递减,我们可以将其表示为:

SNR(t)=exp(ωη(t))

其中 wn(t)被建模为具有参数n的单调递增神经网络。对wn(t)取反得到一个单调递减的函数,而指数函数则迫使结果项为正。注意,方程(62)中的目标现在也必须对n进行优化:

argminθ,ηt=2TEq(xtx0)[DKL(q(xt1xt,x0)∣∣pθ(xt1xt))]=argminθ,ηEtU{2,T}[Eq(xtx0)[12(SNR(t1)SNR(t))[x^θ(xt,t)x022]]]

通过将方程(66)中的 SNR 参数化与方程(64)中的 SNR 定义结合起来,我们还可以显式地推导出 αt以及1αt的优雅形式:

α¯t1α¯t=exp(ωη(t))α¯t=sigmoid(ωη(t))1α¯t=sigmoid(ωη(t))

这些项对于各种计算是必需的;例如,在优化过程中,它们被用来使用重新参数化技巧从输入 x0创建任意噪声的xt,如在方程 49 中推导的。

3.3 三种等效的解释 (Three Equivalent Interpretations)

3.3.1 第二种常见解释

正如我们之前证明的,可以通过简单地学习一个神经网络来预测原始自然图像 x0, 以任意噪声版本xt及其时间索引t作为输入来训练变分扩散模型。然而,x0有两个其他的等效参数化形式,这导致了变分扩散模型的另外两种解释。

首先,可以利用重参数化技巧。在推导 q(xt|x0)的形式时,我们可以重新排列方程(49):

x0=xt1α¯tϵ0α¯t

将这个代入我们之前推导出的真实去噪转换均值 μq(xt,x0),我们可以重新推导如下:

μq(xt,x0)=αt(1α¯t1)xt+α¯t1(1αt)x01α¯t=1αtxt1αt1α¯tαtϵ0

因此,我们可以将近似去噪转换均值 μθ(xt,t)设置为:

μθ(xt,t)=1αtxt1αt1α¯tαtϵ^θ(xt,t)

相应的优化问题变为:

argminθDKL(q(xt1xt,x0)∣∣pθ(xt1xt))=argminθDKL(N(xt1;μq,Σq(t))∣∣N(xt1;μθ,Σq(t)))=argminθ12σq2(t)(1αt)2(1α¯t)αt[ϵ0ϵ^θ(xt,t)22]

其中,ϵ^θ(xt,t)是一个神经网络,它学习预测决定源噪声ϵ0N(ϵ;0,I),这个噪声决定了从x0xt的转换。因此,我们已经展示了通过学习变分扩散模型预测原始图像x0是和学习预测噪声是等价的;然而,实证上,一些研究工作发现预测噪声可以带来更好的性能

3.3.2 第三种常见解释

为了推导出变分扩散模型的第三种常见解释,我们引用 Tweedie 方程。Tweedie 方程指出,给定从指数族分布中抽取的样本,其真实均值可以通过样本的最大似然估计(即经验均值)加上涉及估计分数的一些校正项来估计。在只有一个观测样本的情况下,经验均值就是样本本身。它通常用于减少样本偏差;如果观测样本都位于潜在分布的一端,那么负得分会变得很大,并将原始的最大似然估计的样本向真实均值进行修正。

数学上,对于一个高斯变量 zN(z;μz,Σz),Tweedie 方程表明:E[μzz]=z+Σzzlogp(z)

在这种情况下,我们将其应用于预测给定其样本的 xt的真实后验均值。从方程 50,我们知道q(xt|x0)=N(xt;αtx0,(1αt)I)。然后,根据 Tweedie 方程,我们有:

E[μxtxt]=xt+(1α¯t)xtlogp(xt)

其中我们为了简化符号,将 xtlogp(xt)写为logp(xt)。根据 Tweedie 方程,xt生成的真实均值μxt=αtx0被定义为:

α¯tx0=xt+(1α¯t)logp(xt)x0=xt+(1α¯t)logp(xt)α¯t

然后,我们可以将方程 133 再次代入我们的真实去噪转换均值 μq(xt,x0)并推导出一个新的形式:

μq(xt,x0)=αt(1α¯t1)xt+α¯t1(1αt)x01α¯t=1αtxt+1αtαtlogp(xt)

证明过程如下:

使用分数项导出真实的去噪平均值

因此,我们也可以将近似去噪转换均值 μθ(xt,t)设置为:

μθ(xt,t)=1αtxt+1αtαtsθ(xt,t)

相应的优化问题变为:

argminθDKL(q(xt1xt,x0)∣∣pθ(xt1xt))=argminθDKL(N(xt1;μq,Σq(t))∣∣N(xt1;μθ,Σq(t)))=argminθ12σq2(t)(1αt)2αt[sθ(xt,t)logp(xt)22]

其中,sθ(xt,t)是一个神经网络,它学习预测任意噪声水平txt的得分函数xtlogp(xt),即对数概率密度在数据空间中的梯度。

敏锐的读者会观察到得分函数 logp(xt)在形式上与源噪声ϵ0非常相似。这可以通过将 Tweedie 方程(方程 79)与重新参数化技巧(方程 71)结合起来明确展示:

x0=xt+(1α¯t)logp(xt)α¯t=xt1α¯tϵ0α¯t(1α¯t)logp(xt)=1α¯tϵ0logp(xt)=11α¯tϵ0

事实证明,这两项之间存在一个随时间变化的常数因子!得分函数测量了在数据空间中如何移动以最大化对数概率;直观上,由于源噪声被添加到自然图像中以破坏它,向相反方向移动“去噪”图像,并且将是增加随后对数概率的最佳更新。我们的数学证明只是证实了这种直觉;我们已经明确展示了学习建模得分函数等同于建模源噪声的负值(乘以一个随时间变化的常数因子)。

因此,我们已经推导出了优化变分扩散模型的三种等效目标, 学习一个神经网络来:

  • 预测任意噪声版本的原始图像 x0
  • 预测任意噪声图像的源噪声 ϵ0
  • 预测任意噪声水平下的图像得分 logp(xt)

变分扩散模型可以通过随机采样时间步 t并最小化预测与真实目标的范数来可扩展地训练。

4. 分数生成模型(Score-based Generative Models)

我们已经展示了变分扩散模型可以通过简单地优化一个神经网络 sθ(xt,t)来预测xt在任意噪声水平下的得分函数logp(xt)来学习。然而,在我们的推导中,分数项是通过应用 Tweedie 方程得出的;这并不一定为我们提供了关于得分函数究竟是什么或者为什么值得建模的深刻直觉或洞察。幸运的是,我们可以向另一类生成模型——基于分数的生成模型(Score-based Generative Models)来获得这种直觉。结果表明,我们可以展示我们之前推导出的 VDM 公式与等价的基于得分的生成模型公式是一致的,这使我们能够灵活地在这两种解释之间随意切换。

为了开始理解优化得分函数的意义,我们绕道并重新审视基于能量的模型(energy-based models)。任意灵活的概率分布可以写成以下形式:

pθ(x)=1Zθefθ(x)

其中 fθ(x)是一个任意灵活的、可参数化的函数,称为能量函数,通常由神经网络建模,Zθ是一个规范化常数,以确保积分pθ(x)dx=1。学习这样一个分布的一种方法是最大似然;然而,这需要能够灵活地计算规范化常数Zθ=efθ(x)dx,对于复杂fθ(x)函数来说可能是不可能的。

避免计算或建模规范化常数的一种方法是使用一个神经网络 sθ(x)来学习分布p(x)的得分函数logp(x)而不是能量函数。这是由以下观察所激发的:对上面方程的两边取对数后进行微分得到:

xlogpθ(x)=xlog(1Zθefθ(x))=xlog1Zθ+xlogefθ(x)=xfθ(x)sθ(x)

这可以自由地表示为一个神经网络,而不涉及任何规范化常数。得分模型可以通过最小化与真实得分函数的 Fisher 散度来优化:

Ep(x)[sθ(x)logp(x)22]

得分函数代表什么?对于每一个 x,对其对数似然关于x取梯度,本质上描述了数据空间中移动的方向,以进一步增加其似然。

直观地说,得分函数定义了整个数据 x​ 所占据的空间上的一个向量场,指向模式。在图 6 的右图中被视觉化地表示出来。

这是使用 Langevin 动力学生成的三个随机采样轨迹的可视化,所有轨迹都从相同的初始点开始,针对的是高斯混合模型。左侧的图形在三维等高线上绘制了这些采样轨迹,而右侧的图形则将采样轨迹与真实得分函数相对比。由于 Langevin 动力学采样过程中的随机噪声项,我们可以从相同的初始点生成来自不同模式的样本;如果没有这个噪声项,从固定点开始的采样将总是在每次试验中确定性地沿着得分走向同一个模式。

然后,通过学习真实数据分布的得分函数,我们可以通过从同一空间中的任意点开始,迭代地跟随得分直到到达一个模式来生成样本。这种采样过程被称为 Langevin 动力学,数学上描述为:

xi+1xi+clogp(xi)+2cϵ,i=0,1,...,K
Visualizing sampled paths from Langevin Dynamics against ground-truth scores

其中 x0是从先验分布(如均匀分布)中随机采样的,ϵN(ϵ;0,I)是一个额外的噪声项,以确保生成的样本不会总是坍缩到一个模式上,而是在其周围徘徊以实现多样性。此外,由于学习到的得分函数是确定性的,涉及噪声项的采样为生成过程增加了随机性,使我们能够避免确定性的轨迹。当采样从位于多个模式之间的位置初始化时,这特别有用。Langevin 动力学采样和噪声项的好处在图 6 中以视觉方式展示。

观察到,方程(93)中的目标依赖于我们能够访问真实得分函数,对于像自然图像这样的复杂分布,我们是无法获得的。幸运的是,已经推导出了称为得分匹配(score matching)的替代技术,可以在不知道真实得分的情况下最小化这个 Fisher 散度,并且可以使用随机梯度下降来优化。

总的来说,学习将分布表示为得分函数,并使用它通过马尔可夫链蒙特卡洛技术(如 Langevin 动力学)生成样本,被称为基于分数的生成建模(Score-based Generative Modeling)。

标准得分匹配存在三个主要问题,这些问题由 Song 和 Ermon 详细说明[3]。首先,当 x位于高维空间中的低维流形上时,得分函数是未定义的。这可以从数学上看出;所有不在低维流形上的点都有零概率,其对数是未定义的。这在尝试学习自然图像上的生成模型时特别不方便,因为已知自然图像位于整个空间的低维流形上。其次,通过标准得分匹配训练出的估计得分函数在低密度区域将不准确。这从我们在方程 93 中最小化的目标中可以看出。因为它是对p(x)的期望,并且明确地在来自它的样本上进行训练,模型将不会为很少看到或未见过的示例提供准确的学习信号。这是有问题的,因为我们的采样策略涉及从高维空间中的随机位置开始,这很可能是随机噪声,然后根据学习到的得分函数移动。由于我们正在跟随一个嘈杂或不准确的得分估计,最终生成的样本也可能是次优的,或者需要更多的迭代才能收敛到准确的输出。

最后,即使使用真实得分执行,Langevin 动力学采样也可能不会混合。假设真实数据分布是两个不相交分布的混合

p(x)=c1p1(x)+c2p2(x)

然后,当计算得分时,这些混合系数丢失了,因为对数运算将系数从分布中分离出来,梯度运算将其置零。为了可视化这一点,请注意上图中显示的真实得分函数对三个分布之间的不同权重是不可知的;从所描述的初始点开始的 Langevin 动力学采样到达每个模式的概率大致相等,尽管右下角的模式在实际的高斯混合中具有更高的权重。

结果表明,通过向数据添加多个水平的高斯噪声,可以同时解决这三个缺点。首先,由于高斯噪声分布的支持是整个空间,扰动的数据样本将不再局限于低维流形。其次,添加大量的高斯噪声将增加每个模式在数据分布中覆盖的区域,从而在低密度区域增加更多的训练信号。最后,添加具有逐渐增加的方差的多个水平的高斯噪声将导致中间分布,这些分布尊重真实的混合系数。

正式地,我们可以选择一个正噪声水平序列 {σt}t=1T并定义一系列逐渐受到干扰的数据分布:

pσt(xt)=p(x)N(xt;x,σt2I)dx

然后,使用得分匹配学习神经网络 sθ(x,t)来同时学习所有噪声水平的得分函数:

argminθt=1Tλ(t)Epσt(xt)[sθ(x,t)logpσt(xt)22]

其中 λ(t)是一个正的权重函数,它依赖于噪声水平t。注意,这个目标几乎完全匹配我们在方程 84 中推导出的用于训练变分扩散模型的目标。此外,作者提出了退火 Langevin 动力学采样作为一种生成过程,其中样本是通过按顺序运行 Langevin 动力学来产生的,每个t=T,T1,,2,1。初始化是从某个固定的先验(如均匀分布)中选择的,每个后续的采样步骤从上一个模拟的最终样本开始。因为噪声水平在时间步t上稳步下降,并且我们随着时间的推移减小步长,样本最终会收敛到一个真实模式。这与在变分扩散模型的马尔可夫 HVAE 解释中执行的采样程序直接类似,其中随机初始化的数据向量在降低噪声水平的过程中被迭代地细化。

因此,我们在训练目标和采样程序上建立了变分扩散模型和基于得分的生成模型之间的明确联系。

一个问题是如何将扩散模型自然地推广到无限数量的时间步。在马尔可夫 HVAE 的视角下,这可以被解释为将层级的数量扩展到无穷 T。从等效的基于得分的生成模型的角度来表示这一点更清晰;在无限数量的噪声尺度下,图像在连续时间上的扰动可以被表示为一个随机过程,因此可以用随机微分方程(SDE)来描述。然后通过逆转 SDE 来执行采样,这自然需要在每个连续值的噪声水平上估计得分函数[12]。SDE 的不同参数化本质上描述了不同的扰动方案,使得灵活地建模噪声过程成为可能。

5. 指导(Guidance)

到目前为止,我们只关注了对数据分布 p(x)的建模。然而,我们通常也对学习条件分布p(x|y)感兴趣,这将使我们能够通过条件信息y明确控制我们生成的数据。这构成了图像超分辨率模型的支柱,例如级联扩散模型,以及最先进的图像文本模型的关键,例如 DALL-E 2。

自然地添加条件信息的方法是在每次迭代中,简单地将其与时间步信息一起加入。回想我们方程 32 中的联合分布:

p(x0:T)=p(xT)t=1Tpθ(xt1xt)

然后,要将其变成一个条件扩散模型,我们可以简单地在每个转换步骤中添加任意的条件信息 y

p(x0:Ty)=p(xT)t=1Tpθ(xt1xt,y)

例如,y可以是图像-文本生成中的文本编码,或者是执行超分辨率的低分辨率图像。因此,我们可以像以前一样学习变分扩散模型的核心神经网络,通过预测x^θ(xt,t,y)x0ϵ^θ(xt,t,y)ϵ0,或sθ(xt,t,y)logp(xt|y)来实现每种所需的解释和实现。

这种基本公式的一个缺点是,以这种方式训练的条件扩散模型可能会学会忽略或淡化任何给定的条件信息。因此,提出了指导作为一种更明确地控制模型给予条件信息权重的方法,代价是牺牲样本多样性。两种最流行的形式的指导被称为分类器指导(Classifier Guidance)和无分类器指导(Classifier-Free Guidance)。

5.1 分类器指导(Classifier Guidance)

让我们从扩散模型的基于得分的方程开始,我们的目标是在任意噪声水平 t下学习条件模型的得分logp(xt|y)。为了简洁起见,xt的简写。根据贝叶斯定理,我们可以推导出以下等价形式:

logp(xty)=log(p(xt)p(yxt)p(y))=logp(xt)+logp(yxt)logp(y)=logp(xt)unconditional score+logp(yxt)adversarial gradient

这里, 我们利用了相对于 xtlogp(y)的梯度为零的事实。

我们最终推导出的结果可以解释为学习一个无分类器得分函数与一个分类器 p(y|xt)的对抗梯度的组合。因此,在分类器指导中,无分类器扩散模型的得分像之前推导的那样学习,同时学习一个分类器,该分类器接收任意噪声的xt并尝试预测条件信息y。然后,在采样过程中,用于退火 Langevin 动力学的整体条件得分函数被计算为无分类器得分函数和噪声分类器的对抗梯度的总和。

为了引入细粒度控制,以鼓励或阻止模型考虑条件信息,分类器指导通过一个 γ超参数项来缩放噪声分类器的对抗梯度。在分类器指导下学习得分函数可以总结为:

logp(xty)=logp(xt)+γlogp(yxt)

直观地说:

  • γ=0时,条件扩散模型完全忽略条件信息,
  • γ很大时,条件扩散模型学会产生严重依赖于条件信息的样本。这将以牺牲样本多样性为代价,因为它只会产生容易从提供的条件下重新生成的样本,即使在噪声水平下也是如此。

分类器指导的一个缺点是它依赖于单独学习分类器。由于分类器必须处理任意噪声的输入,这超出了大多数现有的预训练分类模型的优化范围,因此它必须与扩散模型一起学习。

5.2 无分类器指导(Classifier-Free Guidance)

在无分类器指导中,作者放弃了训练单独的分类器模型,转而使用一个无分类器扩散模型和一个条件扩散模型。为了推导无分类器指导下的得分函数,我们可以先重新排列方程(101):

logp(yxt)=logp(xty)logp(xt)

然后,将这个代入方程(102),我们得到:

logp(xty)=logp(xt)+γ(logp(xty)logp(xt))=logp(xt)+γlogp(xty)γlogp(xt)=γlogp(xty)conditional score+(1γ)logp(xt)unconditional score

再一次,γ是一个控制我们学习的条件模型对条件信息关注程度的项。

  • γ=0时,学习的条件模型完全忽略条件器并学习一个无分类器扩散模型。
  • γ=1时,模型明确学习没有指导的普通条件分布。
  • γ>1时,扩散模型不仅优先考虑条件得分函数,而且还向远离无分类器得分函数的方向移动。换句话说,它减少了生成不使用条件信息的样本的概率,而有利于明确使用条件信息的样本。这也通过牺牲样本多样性为代价,增加了生成与条件信息准确匹配的样本的可能性。

由于学习两个单独的扩散模型是昂贵的,我们可以将条件和无分类器扩散模型一起作为一个单一的条件模型学习;无分类器扩散模型可以通过将条件信息替换为固定常数值(如零)来查询,这本质上是在条件信息上执行随机 dropout。无分类器指导之所以优雅,是因为它使我们能够更精细地控制我们的条件生成过程, 同时除了训练扩散模型之外不需要任何其他操作。

总结

让我们回顾一下这篇博客文章中的发现。首先,我们将变分扩散模型(Variational Diffusion Models)推导为马尔可夫分层变分自编码器(Markovian Hierarchical Variational Autoencoder)的一个特例,其中三个关键假设使得对 ELBO 的计算可行并且可扩展优化成为可能。然后,我们证明了优化变分扩散模型归结为学习一个神经网络来预测以下三个潜在目标之一:

  • 从任何随机噪声版本中预测原始源图像,
  • 从任何随机噪声图像中预测原始源噪声,
  • 在任意噪声水平下预测噪声图像的得分函数。

然后,我们深入探讨了学习得分函数的含义,并通过 Tweedie 方程明确地将扩散模型的变分视角与基于分数的生成建模视角联系起来。最后,我们介绍了如何通过指导使用扩散模型学习条件分布。

尽管扩散模型在生成建模方面取得了令人难以置信的成功,但仍有一些缺点需要考虑,这些缺点是未来工作令人兴奋的方向:

  • 我们人类似乎不太可能以这种方式自然地建模和生成数据;我们不会将样本生成为随机噪声,然后逐步去噪。
  • 变分扩散模型(VDM)不产生可解释的潜在表示。而变分自编码器(VAE)有望通过优化其编码器学习到一个结构化的潜在空间,但在变分扩散模型中,每个时间步的编码器已经被给定为线性高斯模型,并且不能灵活地优化。因此,中间潜在表示被限制为原始输入的噪声版本。
  • 潜在变量被限制为与原始输入相同的维度,进一步阻碍了学习有意义、压缩的潜在结构的努力。
  • 采样是一个昂贵的过程,因为在两种方程下都必须运行多个去噪步骤。
  • 为了确保最终潜在表示完全是高斯噪声,时间步的数量通常非常大;在采样期间,我们必须迭代所有这些时间步来生成一个样本。

最后需要注意的是,扩散模型的成功突显了分层变分自编码器(Hierarchical VAEs)作为生成模型的强大能力。我们已经展示了,当我们推广到无限的潜在层次时,即使编码器是平凡的,潜在维度是固定的,且假设了马尔可夫转换,我们仍然能够学习强大的数据模型。这表明,在一般情况下,深度 HVAEs 可以取得进一步的性能提升,其中可以潜在地学习复杂的编码器和语义上有意义的潜在空间。

6. Reference

Maintained by Robin