ICML 2025 Spotlight | 新理论框架解锁流匹配模型的引导生成

在解决离线强化学习、图片逆问题等任务中，对生成模型的能量引导（energy guidance）是一种可控的生成方法，它构造灵活，适用于各种任务，且允许无额外训练条件生成模型。同时流匹配（flow matching）框架作为一种生成模型，近期在分子生成、图片生成等领域中已经展现出巨大潜力。

然而，作为比扩散模型更一般的框架，流匹配允许从几乎任意的源分布以及耦合分布中生成样本。这在使得它更灵活的同时，也使得能量引导的实现与扩散模型有根本不同且更加复杂。因此，对于流匹配来说，如何得到具有理论保证的能量引导算法仍然是一个挑战。

针对这一问题，作者从理论上推导得到全新能量引导理论框架，并进一步提出多样的实际能量引导算法，可以根据任务特性进行灵活选择。本工作的主要贡献如下：

• 本工作首次提出了流匹配能量引导理论框架。

• 在本框架指导下，本工作提出三大类无需训练的实用流匹配能量引导算法，并可将经典扩散模型能量引导算法包含为特例。

• 本工作给出了各个流匹配能量引导算法性能的理论分析和实验比较，为实际应用提供指导。

• 论文标题：On the Guidance of Flow Matching

• 论文链接：https://arxiv.org/abs/2502.02150

• 项目地址：https://github.com/AI4Science-WestlakeU/flow_guidance

目前，本工作已被接受为 ICML 2025 spotlight poster，代码已经开源。

研究背景

已有的能量引导算法集中于扩散模型，但是流匹配模型和扩散模型相比有本质上的差别，使得它们的能量引导算法不能直接通用。简而言之，扩散模型可以被看作是流匹配模型在这些假设下的特例：源分布是高斯分布、源分布和生成分布之间没有耦合、条件速度场满足特定的线性形式。

在这些假设下，扩散模型的向量场可以和得分函数（score function）关联起来，从而能量引导向量场可以被大大简化，成为能量函数对数期望的梯度形式。在没有这些假设时，能量引导向量场则需要几乎完全重新推导。

目前虽然已经有一些工作对流匹配模型进行能量引导，但是这些流匹配模型仍然采用了高斯源分布等三个假设，所以本质上仍然是扩散模型（仅有条件向量场的系数中有细微不同）。因此，一个具有一般性的流匹配能量引导理论框架是必要的。

方法概述

首先，作者从流匹配模型基础定义出发，推导了一般的流匹配能量引导向量场。具体而言，将叠加了能量引导后的总向量场与原向量场相减，

蒙特卡洛估计

利用这一方法，在样本数不限的情况下可以计算精确的能量引导向量场。

梯度近似

也就是得到了扩散模型引导向量场中常见的「能量函数的梯度」的形式。注意到梯度前面的项和能量函数无关，可以进一步通过设置成超参数来近似，或者在一些特殊情况的流匹配模型中，可以被进一步简化。

例如，通过采用源分布是高斯分布、源分布和生成分布之间没有耦合、条件速度场满足特定的线性形式的假设（即和扩散模型相同），可以简化为经典的扩散后验采样（Diffusion Posterior Sampling, DPS）算法。

值得注意的是，虽然在扩散模型的特例中，最终形式和 DPS 相同，但是推导方式截然不同。DPS 基于扩散能量引导框架，利用 Jensen 不等式来消除不可计算的期望，但这里基于流匹配能量引导框架，则是使用泰勒展开来简化这一期望的计算。

高斯近似

实验结果

这些流匹配引导任务和扩散模型显著不同，因此针对扩散模型的精确能量引导方法（左三列，对比能量引导 CEG）完全失败。同时基于蒙特卡洛采样的引导算法取得了最接近真实（ground truth）分布的结果，佐证了它是渐进精确的和流匹配引导框架的正确性。

此外，为了从实验上比较各个引导算法优劣，作者还在离线强化学习（offline RL）和图片线性逆问题任务中测试了各个引导生成算法的效果，结果如表所示。

总体来说，在离线强化学习任务中，蒙特卡洛采样引导有最佳性能。这可能由于离线强化学习任务中需要同一个引导算法在不同时间步的条件下都产生稳定的引导采样样本，因此理论保证的能量引导算法具有最佳性能；而图片逆问题中，针对此逆问题形式设计的高斯近似引导和 GDM 有最佳性能，而蒙特卡洛采样引导由于问题维度较高不能产生合理的引导向量场。

结论

本工作针对流匹配模型中能量引导算法的空白，提出了一种新的能量引导的理论框架，并且提出几类各有优劣的实用引导算法，适用于一般的流匹配模型。此外，通过理论分析和实验对各个引导算法进行了比较，提供了实际应用指导。本工作希望为流匹配引导采样和为生成模型的进一步应用提供理论基础。

文章来自于微信公众号“机器之心”。