COT-FM: Cluster-wise Optimal Transport Flow Matching

TL;DR

COT-FM introduces a clustering-based method to improve the accuracy and speed of flow matching models, resulting in more reliable and higher-quality sample generation across various tasks.

Contribution

It proposes a novel clustering approach that refines flow matching by localizing transport, enhancing both efficiency and sample quality without altering model architecture.

Findings

Accelerates sampling across multiple datasets

Improves sample quality in image generation benchmarks

Enhances performance in robotic manipulation tasks

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.