FGM：刷新流匹配模型单步文本到图像生成纪录！(浙大&卡内基梅隆&北大&西湖大学)

文章链接：https://arxiv.org/pdf/2410.19310

亮点直击

• 创新性的一步蒸馏框架：提出了一种名为“流生成匹配 (FGM)”的概率框架，通过一步蒸馏的方式将流匹配模型的采样过程简化为单步生成，大幅提高计算效率，保持了原始模型的生成质量。
• 对流匹配模型的蒸馏优化：FGM 弥补了流匹配模型在蒸馏方面的研究空白，实现了流匹配模型的一步采样，克服了该模型在采样过程中的计算资源限制。
• 出色的生成表现：在 CIFAR10 和大规模文本生成图像等任务上，FGM 在单步生成中取得了优异的生成表现，与多步生成模型的性能相当，并刷新了流匹配模型的一步生成记录。
• 提升流模型的实用性：FGM 不仅扩展了蒸馏技术的应用场景，还显著提升了流匹配模型的实际应用潜力，特别适合需要快速高效采样的应用场景。

总结速览

解决的问题
在AIGC领域中，流匹配模型因其坚实的理论基础和强大的大规模生成能力而取得了显著成功。然而，这些模型的采样过程对计算资源要求极高，因为它们需要多步数值常微分方程 (ODEs) 的计算。

提出的方案
本文提出了一种名为“流生成匹配 (Flow Generator Matching, FGM)”的创新方法，通过将流匹配模型的采样加速为一步生成，降低了对计算资源的需求，同时保证了模型性能。

应用的技术
FGM 模型在 CIFAR10 无条件生成基准测试上获得了优异的成绩，创造了3.08的最新 Fréchet Inception Distance (FID) 分数，超越了原有的50步流匹配模型。此外，通过 FGM 对 Stable Diffusion 3 进行了蒸馏，生成了一步文本生成图像的 MM-DiT-FGM 模型。

达到的效果
在 GenEval 基准测试上，MM-DiT-FGM 展现出卓越的生成效果，在单步生成的高效性下，生成质量可与多步模型相媲美。

流生成匹配

本节介绍流生成匹配(Flow Generator Matching, FGM)，这是一种专为流匹配模型设计的一步蒸馏设计的通用方法。首先定义问题设置和符号，然后介绍我们的匹配目标函数以及 FGM 如何最小化该目标。最后，将 FGM 与现有的流蒸馏方法进行比较，突出我们方法的经验和理论优势。

问题设置

问题表述。本文的框架基于一个预训练的流匹配模型，该模型能够准确地逼近边际向量场 。流  在噪声分布和数据分布之间建立了联系。同时，也知道条件转移 ，这意味着 。假设预训练的流匹配模型对数据分布有足够好的逼近，即  是真实的数据分布。

本文的目标是训练一个一步生成器模型 ，直接将随机噪声  转换为样本 。令  表示学生模型生成样本  的分布， 表示通过  转移的边际概率路径。

该学生边际概率路径隐式地引导了生成该路径的流向量场 ，该向量场未知且难以处理。

不可处理的目标。一步流生成匹配的目标是使学生分布  与数据分布  匹配。为此，考虑将边际向量场  与预训练的向量场  进行匹配，从而使得分布  和  能够彼此匹配。

本节定义流生成匹配的目标。基于之前的讨论，目标是最小化隐式向量场  和预训练流模型向量场  之间的期望  距离，表示为

注意到样本  依赖于参数 。在必要时，可以使用  来强调这种参数依赖关系。

显然，当且仅当所有诱导向量场一致时，，即  几乎处处对  成立。因此，这意味着  几乎处处对  成立，这表明两个分布  几乎处处对  匹配。然而，尽管最小化目标可以得到一个强大的一步生成器，但它是不可处理的，因为我们不知道  和生成器分布  之间的关系。下面介绍本文的主要贡献：一个具有理论保证的、与（4.1）等价且可处理的训练目标。

可处理的目标

本文的目标是优化参数  以最小化目标。然而，隐式向量场  是未知且不可处理的，因此无法直接最小化该目标。不过，通过对  求损失函数的梯度，可以得到

其中， 和  定义为

公式 (4.3) 中的梯度考虑了关于参数  的所有导数。

请注意，第一个梯度  可以通过停止对  的 -梯度来获得，即 。这导致了一个替代损失函数，其梯度与  一致。

然而，第二个梯度 (4.5) 包含一个不可处理的项 。对于学生生成器，我们只能从条件概率路径中有效地采样，但向量场  及其关于  的梯度未知。幸运的是，在本文中，我们提出了如下定理 4.2，使得学生向量场的 -梯度更加可处理。在此之前，我们需要首先引入一个新的流乘积恒等式（Flow Product Identity），这是我们的贡献之一，即定理 4.1。

定理 4.1 （流乘积恒等式） 设  为一个向量值函数，使用 4.1 节中的符号，在温和的条件下，该恒等式成立：

接下来，展示可以引入一个等效的可处理损失函数，该损失函数与不可处理的损失函数 (4.1) 在参数梯度上相同，如定理 4.2 所示。

定理 4.2. 如果分布  满足某些广义正则性条件，则对于所有不依赖于参数的向量值函数 ，对于所有参数 ，该等式成立：

公式 (4.8) 表明，无法处理的梯度  的期望可以与来自学生模型的可微样本的可处理期望进行互换。

恒等式 (4.8) 的直接结果是，梯度 Grad2() 与以下可处理损失函数 (4.9) 一致，该损失函数在生成器向量中对  施加了停止梯度操作 ：

将 (4.6) 和 (4.9) 结合到 (4.3) 中，得到一个等效的损失函数，用于最小化原始目标，即：

定义在 (4.6) 和 (4.9) 中的 L1(θ) 和 L2(θ)。这产生了通过最小化损失函数 (4.10) 提出的流生成匹配 (FGM) 目标。下算法 1 总结了流生成匹配的伪算法，将预训练的流匹配模型蒸馏成一个单步生成器。

与扩散蒸馏的区别 FGM 从单步扩散蒸馏中获得灵感，旨在最小化分布发散，然而，所得到的理论与单步扩散蒸馏的理论本质上是不同的。FGM 和单步扩散蒸馏之间最显著的区别在于，流匹配并不意味着对概率密度进行显式建模，如扩散模型所做。因此，分布发散的定义不能应用于流模型及其蒸馏。然而，FGM 通过直接处理流匹配目标而不是分布发散来克服这个问题。主要的见解是我们提出的显式-隐式梯度等效理论绕过了不可处理的流匹配目标，从而产生具有理论保证的强大实用算法。定理 4.2 也可能为未来对流匹配模型的研究带来新的贡献。

与其他流蒸馏方法的比较 目前有少数现有工作尝试将流模型加速到单步或少步生成模型。一致性流匹配 (CFM)是最近的一项工作，将预训练的流模型蒸馏为一个或两个步的模型。尽管 CFM 显示出不错的结果，但在理论和实际方面与我们的 FGM 不同。首先，CFM 背后的理论建立在流模型的轨迹一致性上，这直接从一致性模型中推广。相反，我们的 FGM 是以流匹配目标为出发点，试图用真值教师流来训练单步生成器的隐式流，并具有理论保证。在实际方面，在 CIFAR10 生成上，展示了训练的单步 FGM 模型在流模型中达到了 3.08 的新的 SoTA FID，超过了 CFM 最好的 2 步生成结果 5.34 的 FID。如此强大的实证表现使得 FGM 成为在标准基准上加速流匹配模型的可靠解决方案。除了简单的 CIFAR10 生成，在第 5.2 节中，我们还使用 FGM 蒸馏领先的大规模文本到图像流模型，获得了几乎没有性能下降的非常强大的单步文本到图像模型。

实验

这里进行了实验以评估 FGM 的有效性和灵活性。实验涵盖了标准评估基准、无条件 CIFAR10 图像生成和使用Stable Diffusion 3 (SD3)的大规模文本到图像生成。这些实验展示了 FGM 在构建高效单步生成器的能力，同时保持高质量样本。

单步 CIFAR10 生成

实验设置 首先在 CIFAR10 数据集上评估 FGM 的有效性，这是生成模型性能的标准测试平台。在 CIFAR10 上使用 ReFlow 目标 (3.6) 预训练流匹配模型进行条件和无条件生成。参考了用于 EDM 模型的神经网络架构。以 512 的批量大小训练条件和无条件模型 20,000 张图像，得到的内部训练流模型显示出 CIFAR10 无条件 FID 为 2.52，生成步为 300，这略微低于原始 ReFlow 模型，其 FID 为 2.58，使用 127 个生成步。然而，在下表 1 中，发现这样的略微较差模型并不会影响强单步生成器的蒸馏。

这些流模型作为流生成匹配 (FGM) 的教师模型。然后应用 FGM 从流模型中蒸馏单步生成器。通过FID评估生成图像的质量。较低的 FID 分数表示样本质量和多样性更高。

注意损失 (4.11) 和损失 (4.12) 共同构成 FGM 损失的完整参数梯度。我们发现这两个损失在使用多层感知机进行 2D 数据集生成时效果很好。在实践中，发现使用损失 (4.11) 在 CIFAR10 模型上会导致不稳定，这与 Poole et al. (2022) 的观察相似，即其雅可比项的条件数可能不适定。因此，在训练时我们不使用损失 (4.11)，而是观察到良好的性能。

用预训练流模型初始化生成器 受扩散蒸馏技术的启发，用预训练流模型初始化单步生成器。回忆流模型的训练目标 (3.6)，预训练流模型 θ 大致预测从随机噪声到数据的方向。因此，我们使用预训练流来构建我们的单步生成器。具体来说，构建单步生成器：

θ 是生成器的可学习参数，而 是预先确定的最佳时间步。

定量评估 使用 Fretchet Inception Distance（FID）评估每个模型，这是评估 CIFAR10 数据集上图像生成结果的黄金标准。上表1和表2总结了生成模型在 CIFAR10 数据集上的 FID。在无条件生成任务中，我们的教师流模型在50个生成步和100个生成步下分别获得了 3.67 和 2.93 的 FID。FGM 模型在仅使用一个生成步的情况下取得了 3.08 的 FID，显著超越了教师模型在50个生成步下的结果，幅度达到了16%。在 CIFAR10 条件生成任务中，我们的一步 FGM 模型获得了 2.58 的 FID，超越了教师流模型在100个生成步下的 2.87 FID。在 CIFAR10 生成基准上的结果展示了 FGM 的优越性能，它能够在显著幅度上超越多步教师流模型。

CIFAR-10 生成任务相对简单。进行了实验，训练大型的一步文本到图像生成器，采用来自顶级性能的基于Transformer的流模型进行蒸馏。

文本到图像生成实验：

下面展示了通过 FGM 蒸馏的一步 T2I 生成器在其他行业级模型中展现出最先进的结果。在进一步的实验中，FGM 被应用于从顶尖的流匹配模型（如Stable Diffusion 3 Medium）中训练单步文本到图像生成器。用于训练的大型数据集包含大约 200 万条记录，增强了模型生成高质量图像的能力。尽管在使用某些损失函数时训练动态存在初始不稳定性，FGM 仍然保持了与多步模型的竞争力，展示了其在处理复杂文本输入生成图像方面的鲁棒性。

总的来说，提出的 FGM 框架为生成模型的效率和有效性带来了有希望的进展，为流匹配和文本到图像生成技术的未来发展铺平了道路。

定性评估 对所提出的蒸馏方法的定性评估，以分析其性能。下图2展示了多个样本输出，比较了教师模型 Hyper-SD3和 Flash-SD3的方法。结果显示，即使在仅使用单步生成的情况下，视觉质量依然很高，尤其是在细节和颜色再现方面。特别是，一步的 MM-DiT-FGM 在每个生成图像上展现了美学上的光影效果。与现有的蒸馏方法相比，我们的模型在生成质量上达到了可比水平，但成本显著降低。这种优势使得 FGM 在需要实时交互的应用中变得切实可行。

GAN 损失的整合 显然，纯 FGM 算法在训练时并不依赖任何图像数据。近年来，许多研究表明，在蒸馏中引入 GAN 损失对于提高生成图像的高频细节是有益的。因此，也将 GAN 损失与 FGM 结合用于训练单步文本到图像模型，并发现了其益处。

在训练过程中，观察到在 FGM 效率较低的某些噪声调度间隔中，GAN 损失能够提供有效的梯度，从而改善模型输出的质量。因此，GAN 损失的一个重要优势是能够补偿 FGM 在某些噪声调度中的效率不足，从而对我们的损失函数形成补充。

结论

本文介绍了流生成匹配（FGM），一种强有力的概率性单步蒸馏方法，专为流匹配模型而设计。本文建立了 FGM 的理论基础，并验证了其在单步 CIFAR10 生成和大规模单步文本到图像生成方面的强大经验表现。

尽管 FGM 具有坚实的理论基础和强大的经验表现，但它仍然存在一些局限性。第一个局限性是目前 FGM 仍需额外的流模型来近似生成器引起的流向量。这一要求增加了蒸馏的内存开销，并在预训练的流模型和生成器的模型大小较大时可能带来挑战。其次，FGM 是一种完全不依赖图像数据的方法，这意味着在蒸馏时不需要真实的图像数据。然而，众所周知的一个论点是，持续引入高质量图像数据对于提高文本到图像生成模型的性能至关重要。希望未来的工作能够探索如何将数据整合到蒸馏过程中。

本文转自AI生成未来 ，作者：AI生成未来

原文链接:​​https://mp.weixin.qq.com/s/EhqRJzrs6UpYaxoHYFLghQ​​