Search papers, labs, and topics across Lattice.
The paper introduces Cluster-wise Optimal Transport Flow Matching (COT-FM), a framework that improves Flow Matching by optimizing the probability path between source and target distributions. COT-FM clusters target samples and learns a dedicated source distribution for each cluster by reversing a pretrained FM model, thereby straightening the flow and reducing discretization error. Experiments across image generation, 2D datasets, and robotic manipulation show that COT-FM accelerates sampling and improves generation quality without modifying the underlying FM architecture.
Straighter flows, better generations: COT-FM carves up complex generative tasks into simpler, cluster-specific flows, leading to faster and more reliable sampling.
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.
