Improving Adversarial Energy-Based Model via Diffusion Process
公和众和号:EDPJ(进 Q 交流群:922230617 或加 VX:CV_EDPJ 进 V 交流群)
目录
0. 摘要
2. 相关工作
3. 去噪扩散对抗 EBM
3.1. 具有去噪扩散过程的对抗 EBM
3.2. 对 pϕ 的最小化
3.3. 相对于 Eθ 的最大化步骤
4. 实验
附录
B.2 网络结构
0. 摘要
生成模型展示了强大的生成能力,而高效的似然估计却较少探索。基于能量的模型(Energy-based model,EBM)定义了一个灵活的能量函数,以有效地参数化未归一化的密度,但因难以训练而臭名昭著。对抗式 EBM 引入了一个生成器来形成一个最大值最小化的训练游戏,以避免传统 EBM 中使用的昂贵的马尔可夫链蒙特卡罗(Markov chain Monte Carlo,MCMC)采样,但对抗式 EBM 与其他强大的生成模型之间仍存在明显的差距。受扩散模型的启发,我们将 EBM 嵌入到每个去噪步骤中,将一个长时间生成的过程分成几个较小的步骤。此外,我们采用对称的 Jeffrey 散度,并引入一个变分后验分布来训练生成器,以解决对抗式 EBM 中存在的主要挑战。我们的实验显示,与现有的对抗式 EBM 相比,在生成方面取得了显著的改进,同时为高效的密度估计提供了一个有用的能量函数。
2. 相关工作
(2017|ICLR,EBGAN,AE 鉴别器,正则化)基于能量的 GAN
EBM 在机器学习中有着悠久的历史,可以追溯到 Hopfield 网络(Hopfield, 1982)和 Boltzmann 机器(Hinton & Sejnowski, 1983)。近年来,深度 EBM 在特别是图像生成方面变得越来越受欢迎,可以广泛地分为两类。
- 基于 MCMC 的方法在训练过程中(Du & Mordatch, 2019; Arbel等,2021)或采样过程中(Song等,2020; Song & Ermon, 2019)模拟马尔可夫链。这可能会很昂贵、慢和难以控制。
- 协作学习(Cooperative learning)方法(Xie等,2020)联合训练一个自顶向下的生成器和一个用于 MCMC 教学的能量函数,使用生成器作为 Langevin 采样的快速初始化器以减轻 MCMC 的负担。
- 然而,它们仍然效率低下。
另一类是对抗性 EBM(Zhai等,2016; Geng等,2021; Kumar等,2019; Zhao等,2017),它们引入了一个 “生成器” 来形成能量函数和引入的生成器之间的最大值最小化游戏。Han 等(2019; 2020)和 Kan 等(2022)分别从三角形散度和双层优化问题的角度解释了对抗性 EBM。对抗性EBM 继承了 GAN 的优势并避免了 MCMC 采样,但它们也带来了不稳定训练的风险。
原始的基于扩散的模型(Ho等,2020; Sohl-Dickstein等,2015)从正向过程中学习有限时间的反向扩散,并在图像生成方面表现出很强的性能,但采样过程很慢。一些作品(Song等,2021a; Lu等,2022b; Bao等,2022)改进了无需训练的采样器,并显著提高了从 1000 步采样到仅几步的采样速度。DDGAN(Xiao等,2022)提出使用多模态条件 GAN 模型化每个去噪步骤,解决了快速采样和保持强大的模式覆盖和样本质量之间的生成学习三难问题。然而,它未能提供密度估计,这对许多任务来说至关重要。
尽管一些基于扩散的方法(Kingma等,2021; Song等,2021b; Lu等,2022a)使用分数函数或变分下界来估计数据分布的密度,但这些方法往往复杂且低效,这正是基于能量的模型擅长的地方。DRL(Gao等,2021)试图将 EBM 与基于扩散的模型结合起来,并在几个图像数据集上获得了最先进的生成性能,但它仍依赖于效率低下的 MCMC 采样。
3. 去噪扩散对抗 EBM
基础。令 x ∼ q(x) 为来自潜在数据分布的训练样本。一个 EBM 被指定为一个能量函数 Eθ:X → R,它由参数 θ 参数化,从 Gibbs 分布定义了 X 上的概率分布:
其中 Zθ 是归一化常数(或分区函数)。原则上,对于合适选择的 Eθ,可以用这种方式描述任何密度。我们通过最大似然估计(MLE)来学习 EBM,即我们寻求最大化数据对数似然的 θ*:
训练 EBM 的基本挑战是归一化常数缺乏闭合形式的表达式。解决这个问题的一种常见方法是直接逼近 L(θ) 的梯度,并应用基于梯度的优化:
估计上述方程需要从能量分布 pθ(x) 进行 MCMC 采样,这既不稳定又昂贵。对抗式 EBM(Grathwohl等,2021;Geng等,2021;Han等,2019;Kan等,2022)通过引入一个变分分布pϕ 近似 pθ 来避免 MCMC 。因此,可以形成一个最大值最小化来交替优化两个对抗步骤:
第一步:最小化具有特定合适的散度D 的 D(pϕ, pθ)
第二步:
这种两步训练策略可以解释为一个双层优化问题(Liu等,2022),其中最小化步骤是一个较低级别(LL)子问题,而最大化步骤则是一个较高级别(UL)子问题。 Geng 等人(2021年)证明了,如果 pϕ 无法被优化,它会有使得上界最大化的风险,与复杂和稀疏数据集中的主流生成模型相比,会出现明显的性能差距。
3.1. 具有去噪扩散过程的对抗 EBM
为了使对抗 EBM 可扩展到复杂的分布,我们借鉴了基于扩散的框架,将生成过程分成多个步骤。对于每一步,我们只需要学习一个条件分布,而不是一个复杂的边缘分布。类似于 DDPM(Ho等,2020),我们还使用了一个马尔可夫链来进行扩散过程。具体而言,对于正向过程,我们使用预定义的方差调度 βt,在 T 步中逐渐向真实数据 x0 ∼ q(x0) 添加噪声:
然后我们可以得到去噪分布 q(x_(t−1)| xt) = q(x_(t−1))q(xt | x_(t−1))/ q(xt) 作为我们要学习的目标分布。在基于扩散的模型中,常规的训练目标是最小化每个 t 时真实去噪分布 q(x_(t−1)| xt) 和模拟去噪分布 pθ(x_(t−1)| xt) 之间的 KL 散度,即
受到 DRL(Gao等,2021)的启发,我们设计了一系列条件 EBM,使用在 t 时刻条件化的能量函数来定义我们模拟的去噪分布,
其中 Eθ(xt, t) : R^D × R → R 是由 θ 参数化的能量函数。因此,L 可以简化并通过能量函数来表达:
很容易看出,当公式 (7) 被优化时,pθ(xt) = exp (Eθ(xt,t)) / Z_θ,t = q(xt)(参见附录)。
由于 ˜Zθ,t 是棘手的,如果我们遵循常见的选择来计算目标函数的梯度:
第二项必须通过从分布 q(xt)pθ(x_(t−1) | xt) 进行 MCMC 采样来估计(Gao等,2021),这在训练过程中可能具有计算上的要求和不稳定性。因此,我们通过引入一个由 ϕ 参数化的变分条件分布pϕ(x_(t−1) | xt) 来采用对抗 EBM,形成一个两步最大最小化博弈:
这个对抗训练策略交替优化 pϕ 和 Eθ,一边固定另一边进行优化。这为能量基模型的训练提供了一种可行且无需 MCMC 的方法。我们简单地将我们提出的方法称为基于能量去噪扩散对抗模型(Denoising Diffusion Adversarial Energy-Based Model,DDAEBM)。
3.2. 对 pϕ 的最小化
首先,我们需要对引入的变分条件分布 pϕ(x_(t−1) | xt) 进行具体形式的定义。DDIM(Song 等,2021a)和 DDPM(Ho 等,2020)将 pϕ(x_(t−1) | xt) 定义为一个高斯分布,以匹配高斯后验分布 q(x_(t−1) | xt, x0)。然而,正如在 DRL 中所示,当 βt 很小时,这种正态分布的近似只有在噪声步数很少时才准确,当去噪分布可能复杂且多模态时可能不合理。因此,我们采用重新参数化技巧定义 pϕ(x_(t−1) | xt),即
其中 pϕ (x0 | xt) 是由生成器 Gϕ (xt, z, t):R^D × R^d × R → RD 强加的隐式分布。这里,p(z) 是一个 d 维的潜变量,遵循标准高斯分布 N(z; 0, I),而 q (x_(t−1) | xt, x0) 是在正向扩散过程中定义的后验分布。值得注意的是,这个定义也被 Xiao 等人(2022)探索过,已经验证在拟合复杂和多模态条件分布方面是有效的,尽管他们没有利用基于能量的模型结构来描述数据分布。
到目前为止,KL 散度是基于能量的对抗模型训练的最常见的散度 D 的选择(Geng 等,2021;Grathwohl 等,2021;Zhai 等,2016)。然而,鉴于 KL 散度是一个非对称的度量,它迫使 pϕ(x_(t−1) | xt) 追逐 pθ(x_(t−1) | xt) 的主要模式(mode)。因此,仅仅依赖 KL 散度作为我们的目标可能不足以使生成器有效地捕捉能量分布(Geng 等,2021)。遵循 Kan 等人(2022),我们选择一个对称的 Jeffrey 散度,将 KL 散度和反 KL 散度整合到我们的目标中:
我们将我们的理想目标分为两个项进行单独处理:
找出理想目标第一个项的边界。为了最小化我们的理想目标,我们首先处理等式 (11)。可以通过省略一些无关的项并得到简化:
计算生成分布的熵始终是 EBM 中的一个挑战性任务。已经提出了几种方法(Kumar 等,2019;Grathwohl 等,2021;Geng 等,2021)用于有效的熵近似,但这些方法要么是时间要么是内存密集型。我们需要计算条件熵:
H[pϕ(z | x_(t−1), xt)] 是难以计算的,而 H[pϕ(x_(t−1), z|xt)] 是一个可以忽略的常数(见附录)。最小化等式 (13) 可能会有问题,因为我们无法访问 H[pϕ(z|x_(t−1), xt)],而是通过引入一个近似的高斯后验 qψ(z|x_(t−1), xt) 来最小化一个变分上界。这里,均值和方差通过一个由 ψ 参数化的编码器表示,输入为 x_(t−1)、xt 和 t。我们可以轻松地得到:
用等式 (6) 表示 log pθ(x_(t−1) | xt),并将等式 (16) 应用于等式 (13),可以通过其简化的变分上界最小化第一个项 Eq. (11):
找出理想目标第二项的边界。我们现在继续最小化我们的理想目标,处理第二项方程(12)。通过利用与我们引入的高斯后验分布 qψ(z| x_(t−1), xt) 相关的证据下界(ELBO),我们也可以对这一项进行上界界定,因为
我们可以通过应用方程(18)来得到方程(12)的上界。
对方程(19)的蒙特卡罗近似需要从 q(xt)pθ(x_(t−1) | xt) 中进行采样。由于 pθ(x_(t−1) | xt) 被设计为拟合 q(x_(t−1) | xt),我们使用 q(xt−1, xt) 的样本,并使用重要性比(importance
ratio)来替换 q(xt)pθ(x_(t−1) | xt) 中的样本,就像 BiDVL(Kan等人,2022年)中所做的那样:
这里的重要性比受到 ACT(Kong等人,2023年)的启发而设计:
其中,w 和 w_mid 是基于数据集选择的超参数。t'(·) 是指方差保持(Variance Preserving,VP)SDE 中 xt 的时间的函数(Song等,2021c),因为我们的扩散过程可以看作是连续时间 VP SDE的离散化(见附录)。实施最小化步骤的最终整体目标是相对于 pϕ 和 qψ 的方程 (17) 和方程 (19)的总和。
3.3. 相对于 Eθ 的最大化步骤
在相对于 pϕ 和 qψ 的最小化步骤之后,我们大致假设 pϕ(x_(t−1) | xt) = pθ(x_(t−1) | xt)。接下来,我们应该通过最大化能量函数来优化:
与大多数对抗性 EBM 类似,我们还发现通过为能量的梯度添加 ℓ2 正则化器有助于稳定训练:
其中,γ 是正则化系数。这种正则化是一种梯度惩罚,有助于减轻由于最小化步骤中训练不足而导致的训练不稳定(Grathwohl等人,2021年;Kumar等人,2019年)。