Beyond Optimal Transport: Model-Aligned Coupling for Flow Matching

📄 arXiv: 2505.23346v1 📥 PDF

作者: Yexiong Lin, Yu Yao, Tongliang Liu

分类: cs.CV

发布日期: 2025-05-29

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

提出模型对齐耦合(MAC)方法，提升Flow Matching生成质量与效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: Flow Matching 生成模型 最优传输 模型对齐 向量场学习

📋 核心要点

  1. 传统Flow Matching方法因随机耦合导致路径交叉，模型需学习非直线轨迹，生成效率低。
  2. 论文提出模型对齐耦合(MAC)，通过几何距离和模型预测误差对齐来优化耦合策略。
  3. 实验表明，MAC在少步生成设置中，显著提升了生成质量和效率，优于现有方法。

📝 摘要（中文）

Flow Matching (FM) 是一种有效的框架，用于训练模型学习将样本从源分布传输到目标分布的向量场。早期的FM方法使用随机耦合进行训练，这通常导致路径交叉，并使模型学习非直线轨迹，需要多次积分步骤才能生成高质量的样本。为了解决这个问题，最近的方法采用最优传输 (OT) 通过最小化几何距离来构建耦合，这有助于减少路径交叉。然而，我们观察到这种基于几何的耦合不一定与模型偏好的轨迹对齐，使得学习由这些耦合引起的向量场变得困难，从而阻止模型学习直线轨迹。受此启发，我们提出模型对齐耦合 (MAC)，这是一种有效的方法，不仅基于几何距离，还基于与模型基于其预测误差的偏好传输方向的对齐来匹配训练耦合。为了避免耗时的匹配过程，MAC 提出选择具有最低误差的前 k 个耦合分数用于训练。大量的实验表明，与现有方法相比，MAC 显着提高了少步设置中的生成质量和效率。

🔬 方法详解

问题定义：Flow Matching旨在学习一个向量场，将样本从源分布映射到目标分布。现有方法，如使用随机耦合或基于几何距离的最优传输(OT)构建耦合，存在路径交叉或与模型偏好轨迹不一致的问题，导致模型学习非直线轨迹，需要更多积分步骤才能生成高质量样本，计算效率较低。

核心思路：论文的核心思路是提出一种模型对齐耦合(MAC)方法，该方法在构建耦合时，不仅考虑样本之间的几何距离，还考虑模型自身的预测误差，即模型更倾向于学习哪些方向的向量场。通过将耦合与模型的偏好传输方向对齐，可以引导模型学习更直的轨迹，从而减少所需的积分步骤，提高生成效率。

技术框架：MAC方法的整体框架包括以下几个步骤：1) 训练一个初始的Flow Matching模型；2) 基于该模型，计算源分布和目标分布样本之间的预测误差；3) 根据几何距离和预测误差，构建模型对齐的耦合；4) 使用构建的耦合训练Flow Matching模型。

关键创新：MAC的关键创新在于提出了模型对齐的耦合策略。与传统的随机耦合或基于几何距离的耦合相比，MAC能够更好地利用模型自身的知识，引导模型学习更符合其偏好的向量场。这种策略能够有效地减少路径交叉，提高生成效率。

关键设计：为了避免耗时的匹配过程，MAC并没有对所有样本进行耦合，而是选择具有最低预测误差的前k个耦合分数用于训练。这个k是一个重要的超参数，需要根据具体任务进行调整。此外，损失函数的设计也至关重要，需要平衡几何距离和预测误差之间的权重，以获得最佳的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，与现有的Flow Matching方法相比，MAC在少步生成设置中显著提高了生成质量和效率。具体而言，MAC在图像生成任务上，使用更少的积分步骤，就能达到与现有方法相当甚至更好的生成效果。这表明MAC能够有效地引导模型学习更直的轨迹，从而减少计算量，提高生成速度。

🎯 应用场景

该研究成果可应用于图像生成、语音合成、分子设计等领域，尤其是在需要快速生成高质量样本的场景下，具有重要的实际价值。未来，该方法可以进一步扩展到其他生成模型中，并与其他技术相结合，以实现更高效、更灵活的生成任务。

📄 摘要（原文）

Flow Matching (FM) is an effective framework for training a model to learn a vector field that transports samples from a source distribution to a target distribution. To train the model, early FM methods use random couplings, which often result in crossing paths and lead the model to learn non-straight trajectories that require many integration steps to generate high-quality samples. To address this, recent methods adopt Optimal Transport (OT) to construct couplings by minimizing geometric distances, which helps reduce path crossings. However, we observe that such geometry-based couplings do not necessarily align with the model's preferred trajectories, making it difficult to learn the vector field induced by these couplings, which prevents the model from learning straight trajectories. Motivated by this, we propose Model-Aligned Coupling (MAC), an effective method that matches training couplings based not only on geometric distance but also on alignment with the model's preferred transport directions based on its prediction error. To avoid the time-costly match process, MAC proposes to select the top-$k$ fraction of couplings with the lowest error for training. Extensive experiments show that MAC significantly improves generation quality and efficiency in few-step settings compared to existing methods. Project page: https://yexionglin.github.io/mac