COT-FM: Cluster-wise Optimal Transport Flow Matching

2026-03-11

AI Summary

Flow Matching methods often suffer from curved generation trajectories due to stochasticity or batch coupling, which amplifies discretization errors and degrades sample quality. This work proposes a plug-and-play, cluster-level optimal transport strategy that reshapes probability paths in a divide-and-conquer manner: by clustering target samples and assigning each cluster a dedicated source distribution generated via the inverse of a pretrained Flow Matching model. This approach requires no modification to the underlying model architecture yet yields straighter vector fields and more precise local transport maps. Experiments demonstrate consistent improvements in both sampling speed and generation quality across 2D datasets, image synthesis, and robotic manipulation tasks.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Planning, Routing, and Scheduling: Planning with Language Models
• Intelligent Robots: Motion and Path Planning

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

We introduce COT-FM, a general framework that reshapes the probability path in Flow Matching (FM) to achieve faster and more reliable generation. FM models often produce curved trajectories due to random or batchwise couplings, which increase discretization error and reduce sample quality. COT-FM fixes this by clustering target samples and assigning each cluster a dedicated source distribution obtained by reversing pretrained FM models. This divide-and-conquer strategy yields more accurate local transport and significantly straighter vector fields, all without changing the model architecture. As a plug-and-play approach, COT-FM consistently accelerates sampling and improves generation quality across 2D datasets, image generation benchmarks, and robotic manipulation tasks.

Problem

• Flow Matching
• probability path
• discretization error
• sample quality
• vector fields

Innovation

• Flow Matching
• Optimal Transport
• Clustering
• Vector Field Straightening
• Plug-and-Play Generation