Improving Diffusion-Based Generative Models via Approximated Optimal Transport
公和众和号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）
目录
0. 摘要
我们介绍了近似最优输运（Approximated Optimal Transport，AOT）技术，这是一种新颖的扩散型生成模型训练方案。我们的方法旨在近似和整合最优输运到训练过程中，显著提高了扩散模型准确估计去噪器输出的能力。这种改进导致扩散模型的 ODE 轨迹具有较低的曲率，在采样过程中减少了截断误差。通过在训练中采用 AOT，我们实现了更优秀的图像质量和减少的采样步骤。具体来说，在无条件生成和有条件生成中，我们分别仅使用 27 NFEs 和 29 NFE 就实现了 1.88 和 1.73 的 FID 分数。此外，当使用 AOT 训练鉴别器用于指导时，我们分别为无条件生成和有条件生成建立了新的最先进的 FID 分数，每个分数为 1.68 和 1.58，每个都使用 29 NFE。这个结果证明了 AOT 在提升扩散模型性能方面的有效性。
1. 简介
EDM [Karras等，2022年]，以其有效的图像合成能力而闻名，观察到具有直线 ODE 轨迹的模型可以在较少的迭代次数内合成高质量的图像。这一观察通过减少 ODE 的曲率和截断误差的技术显著改进了扩散模型。因此，这些模型在比以往更少的步骤中实现了最先进的性能，合成了高质量的图像。此外，几项研究 [Kim等，2023a；Song等，2023；Zhang等，2023；Kim等，2023b] 通过调整其方法以适应表现良好的预训练 EDM 模型，展示了高性能。
尽管有所进展，扩散模型仍然在其 ODE 中表现出高曲率的间隔，这可能是由于训练过程中信息熵的增加所致。虽然 Flow Matching（一类与扩散模型不同的生成模型系列）和相关研究 [Lipman等，2023年；Tong等，2023年；Pooladian等，2023年] 已经提出了使用最优输运 [Villani，2003年] 解决曲线 ODE 轨迹的方法，但因为扩散模型的结构，将这些方法应用于扩散模型会出现计算效率问题。
2. 基础
2.1 扩散与评分模型
扩散模型 [Ho等，2020年] 通过反向扩散过程生成图像，从一个嘈杂的图像中估计然后减少噪声。这种迭代去噪过程使得这些模型能够从纯噪声中合成图像。Song 和 Ermon [2019年] 引入了评分模型，该模型学习对数似然梯度来合成图像。随后，基于评分的生成模型（Score-based Generative Models，SGM）[Song等，2021年] 将评分模型集成到扩散过程中进行图像生成。他们确定了扩散模型处理的噪声与评分之间的对应关系。这些模型使用这个评分和数值积分来求解随机微分方程（SDE）或常微分方程（ODE）进行图像合成。SDE的前向和后向过程可以描述如下：
其中 f(x, t) 和 g(t) 作为系数，w 和 ¯w 分别对应于前向和后向过程相关的维纳过程。一个被表示为 sθ 的基于评分的模型学习数据分布的梯度 ∇_x log p_t (x)。Song 等人 [2021年] 提出了概率流ODE（Probability Flow ODE，PFODE），它利用了与 SDE 过程中使用的相同的评分模型进行隐式图像生成。ODE的反向过程的方程如下：
2.2 EDM
EDM [Karras等，2022年] 为扩散模型带来了几项增强功能，使其能够以比先前模型更少的采样步骤合成高质量图像。这一进步构成了我们实验方法的基础，利用了 EDM 的图像生成效率。 模型训练 EDM 与传统的扩散模型相一致，优化以下目标函数：
其中，Dθ 表示扩散模型，p(y) 表示数据分布，p(σ) 表示 σ 的采样分布，λ(σ) 表示损失权重，ϵ 是随机噪声。对于这些分布和函数的选择显著增强了扩散模型的整体性能。
ODE曲率调度。EDM 确定了 ODE 轨迹曲率降低与采样步骤要求之间的相关性。随着轨迹曲率的降低，切线方向的一致性增加，从而在应用高步幅时减少了截断误差。它提出了利用这种简单的 ODE 的噪声调度，如图 2 所示。EDM 选择其噪声调度和采样方法以最小化 ODE 轨迹中的曲率。
高效采样。EDM 引入了一种优越的采样方法，利用了 Heun's 方法 [Ascher和Petzold，1998年] 和一种新颖的时间步长选择策略。这种方法减少了截断误差，并允许在高质量图像生成中进行更少的采样步骤。时间步长根据以下方程选择：
其中，n 代表步数，而 ti 表示第 i 步中选定的时间步长值。在这个方程中，ρ 调整步幅的大小。较高的 ρ 意味着在 σ_max 附近使用宽步幅，在 σ_min 附近使用窄步幅来优化生成性能；EDM 在其设置中使用了 ρ 为 7。考虑到 Heun 方法是二阶方法，使用 EDM 进行采样时的 NFE 为 2n-1。
2.3 基于 ODE 的模型的最优输运
流匹配（Flow Matching） [Lipman等，2023年] 是一类最近的基于 ODE 的生成模型。这些模型通过结合训练的向量场的向量来合成图像，类似于扩散模型。流匹配的优点是这些模型是无需模拟的。考虑到向量场是训练目标，这些模型可以选择他们打算模拟的目标。
Tong 等人 [2023年] 引入了最优输运条件流匹配（Optimal Transport Conditional Flow Matching，OT-CFM），它利用最优输运 [Villani，2003年] 通过直线化向量场轨迹来增强性能。这种方法与 EDM 方法类似，旨在减少 ODE 轨迹的曲率。几项研究 [Pooladian等，2023年；Liu等，2023年] 已经对直线化轨迹进行了调查。
两个分布的最优输运可以使用 2-Wasserstein 距离 [Villani等人，2009年；Arjovsky等人，2017年] 来确定，表示为以下方程：
其中 Π 表示两个概率分布 q0 和 q1 的联合概率矩阵。c 是表示两个输入之间距离的成本函数。
考虑到 ut 表示流匹配的向量场，成本函数 c 表示欧氏距离，2-Wasserstein 距离的动态形式表示如下：
其中，pt 表示 ODE 过程中特定时间 t 的 x 的分布。
这个动态形式（方程7）说明了为什么将最优输运技术应用于扩散模型是具有挑战性的。正如 Tong 等人 [2023年] 所述，计算最优输运概率需要在所有时间步长上进行积分。在无需模拟的模型，如流匹配中，这个方面并不重要。然而，在基于扩散的生成模型的情况下，对所有时间步长进行模型评估是计算 Wasserstein 距离所必需的。这一方面导致了在将最优输运与扩散模型结合时的计算效率低下。
2.4 Hungarian 算法
分配问题。分配问题 [Kuhn，1955年] 涉及找到代理和任务之间的最优匹配，旨在最小化相关成本。在这个背景下，每个代理都有与执行每个任务相关的唯一成本，并且每个代理被分配一个不重复的任务。分配问题可以在有限分布设置的最优输运背景下解释。
Hungarian 算法。Hungarian算法 [Kuhn，1955年] 是一种搜索分配问题解决方案的算法，其时间复杂度为 O(n3)，其中 n 表示代理或任务的数量。该算法利用了以下特性：当从与同一代理或任务相关的每个成本中加或减精确值时，最优解保持不变。通过上述技术的迭代，算法找到最优解，确保所有配对的总成本等于 0。我们的方法利用该算法来搜索在采样图像和噪声之间的有限分布的最优输运。
3. 动机
正如 Karras 等人 [2022年] 所强调的，ODE 轨迹的曲率在扩散模型的性能中起着关键作用。图 2 展示了在 EDM 中 ODE 轨迹的分析。在低噪声水平下，ODE 轨迹表现出较低的曲率，导致使用单个欧拉步长估计的图像与目标非常接近，偏差很小。这种对齐表明了 ODE 轨迹的这些区域具有更下降（downward）的曲率。
相反，在高噪声水平的片段中，如图 3-(b) 所示，模型的估计显示出模糊的图像，表明信息熵升高。在这些高噪声区域，初始估计与采样图像相距甚远，需要模型的迭代细化。这种现象表明了 ODE 过程中曲率的增加，如图 2 所示。
由于模型倾向于根据噪声数据估计后验均值 E[x0|xt, t]，在高噪声水平下，模型学会对所有可能的随机噪声扰动进行平均，导致在训练过程中信息熵增加。为了有效解决这个问题，我们提出了一个在第 4 节详细介绍的训练技术，该技术使模型能够学习更低曲率的 ODE 轨迹。从流匹配（第 2.3 节）的概念中汲取灵感，我们的目标是在训练过程中减少信息熵，特别是在高噪声水平下。
4. 近似最优输运训练
本节介绍了一种名为近似最优输运（Approximated Optimal Transport）的扩散模型的新型训练技术。我们在算法 1 中总结了详细的训练过程。
4.1 使用 AOT 进行训练
在标准扩散模型的训练中，通常通过使用与目标图像无关的随机噪声来合成噪声图像。这些模型通常是按照方程 4 进行传统训练的，其中目标图像表示为 y，噪声 ϵ 是独立采样的。这种方法构成了许多现有扩散模型的基础，并且是该领域中一个成熟的方法。
我们提出使用我们的 “近似最优输运”（AOT）技术修改这个标准过程。AOT 不依赖于随机采样的噪声，而是选择与目标图像更接近的高斯噪声 ϵ'。这种选择噪声的改变增强了目标图像和噪声之间的连接，导致更有效的学习过程，从而降低了信息熵。通过将传统的随机噪声采样替换为这种量身定制的方法，AOT 标志着从传统扩散模型实践中显着转变。
4.2 AOT 实现
AOT 的计算过程。我们的目标是确定与扩散模型训练过程中的最优输运相对应的特定 ϵ'。为了确定最优输运路径，我们旨在计算两个分布之间的 2-Wasserstein 距离，如方程 7 所述。鉴于真实图像数据集的分布是不可跟踪的，计算目标图像和噪声之间的最优输运解决方案变得具有挑战性。
为了解决这一挑战，我们在 batch 级别上近似计算两个分布之间的最优输运。我们不是为两个分布确定最优输运，而是在每次迭代中识别具有最小成本函数总和的配对。在每次迭代中，我们从数据集中抽取 N 个目标图像 y，与 N 个随机噪声配对，其中 N 表示搜索 AOT 的配对数量。我们还在每次迭代中抽取 N 个随机噪声，并仅使用这两组样本识别输运函数。
在寻找最优输运时，目标是找到一种最小化所有元素的成本函数总和（方程 6）的输运函数。在我们的情况下，只需要为 N 对元素确定这个输运是必要的。我们寻找最小化成本的索引函数 π，以下表达式概括了这一追求：
其中，π(i) 表示将每个 i 分配给唯一索引的函数，而 ϵi 和 yi 表示第 i 个选定的噪声和目标图像。为了解决这个目标，我们将其构建为分配问题，并使用在第 2.4 节中提到的 Hungarian 算法 [Kuhn，1955年] 进行求解。
首先，我们构造一个距离矩阵，记为 W，在其中 W_ij 表示第 i 个目标图像和第 j 个噪声之间的成本值。在我们的过程中，我们使用 L2 距离作为成本函数。随后，我们确定使用 Hungarian 算法匹配的最佳配对的索引。然后，我们利用这些确定的配对来合成用于训练扩散模型的噪声图像。
在每次迭代中，我们从采样噪声中选择选定的噪声 ϵ'，以确保 x_σ 分布与原始 EDM 相对应。我们指出，选择不是来自采样噪声的噪声可能会影响所选噪声的分布与高斯分布的对齐。
有条件生成的 AOT。当训练具有标记数据的模型，例如带有类别条件生成的 CIFAR-10 数据集[Krizhevsky等，2009年] 时，我们为共享相同标签的每组图像确定选定的噪声 ϵ'。我们将 N 个目标图像分成 C 个共享相同标签的集合，其中 C 表示标签类别的数量。然后，我们对 N 个噪声进行采样，并将它们分成 C 个集合，确保每个集合中的目标图像和噪声的数量相同。我们构造类似于无条件生成的 C 个距离矩阵，并对每个矩阵应用 Hungarian 算法。
我们确定每个标签对应的配对，以消除标签和噪声之间的依赖关系。如果存在相关性，则训练和采样过程可能会有所不同。我们随机采样标签和高斯噪声来合成图像，但无法考虑它们之间的相关性。通过对其进行额外适应，以确保训练目标和初始采样条件之间的一致性，我们消除了这些相关性。我们在补充材料中提供了详细的算法。
用于 GPU 内存限制的训练技术。在我们的设置中，由于 GPU 内存大小和固定 batch 大小的限制，每个 GPU 的小 batch 数，记为 B，比 N 要小。因此，在我们的训练迭代过程中，在采样训练小 batch 之前，我们在每个 GPU 上选择 N 个 AOT 配对。随后，我们从预先选择的 N 对中对每组 B 个目标图像进行采样，因为这些配对存储在临时数据加载器中。
5. 实验
6. 相关工作
考虑到 Song 等人 [2021] 使用 SDE 或 ODE 对扩散采样过程进行建模的事实，大量的研究已经致力于探索与截断误差相关的权衡。
一些研究聚焦于以模型为中心和以训练为导向的策略，以减轻这些挑战。
- PFGM [Xu等，2022] 通过将电场纳入 ODE 中，取代传统的扩散过程，进行了创新。这种方法显示出在处理 ODE 过程中增强的鲁棒性。
- 此外，PFGM++ [Xu等，2023a] 建立了 PFGM 和 EDM 之间的线性关系，从而显著改善了模型性能。
另一方面的研究探讨了使用预训练模型来增强扩散模型，主要是通过高效的采样器。
- PNDM [Liu等，2022] 建议用 Runge-Kutta 和线性多步方法 [Sauer，2011年] 等高阶数值积分器替换欧拉方法。
- 类似地，DEIS [Zhang和Chen，2023年] 主张将指数积分器 [Hochbruck和Ostermann，2010] 应用为一种高效的采样方法。
（2021|ICLR，扩散先验，VE-SDE，逼真和忠实的权衡）SDEdit：使用随机微分方程引导图像合成和编辑
（2024|ICLR，DDBM，基于分数的 ODE 和 SDE，Doob 的 h 变换）去噪扩散桥模型
（2023，InstaFlow & 整流流 & 回流 & 蒸馏）InstaFlow：一步就足以实现基于扩散的高质量文本到图像生成
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
