COT-FM: Cluster-wise Optimal Transport Flow Matching
This work addresses a specific bottleneck in generative modeling for researchers and practitioners, offering a plug-and-play solution to enhance Flow Matching methods, though it is incremental as it builds on existing FM frameworks without altering model architecture.
The paper tackled the problem of curved trajectories in Flow Matching models, which increase discretization error and reduce sample quality, by introducing COT-FM, a framework that reshapes probability paths using clustering and dedicated source distributions, resulting in faster sampling and improved generation quality across various datasets and tasks.
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.