COT-FM: Cluster-weises Optimal-Transport-Flow-Matching

COT-FM: Cluster-wise Optimal Transport Flow Matching

March 11, 2026

Autoren: Chiensheng Chiang, Kuan-Hsun Tu, Jia-Wei Liao, Cheng-Fu Chou, Tsung-Wei Ke

cs.AI

Zusammenfassung

Wir stellen COT-FM vor, ein allgemeines Framework, das die Wahrscheinlichkeitspfade im Flow Matching (FM) neu gestaltet, um eine schnellere und zuverlässigere Generierung zu erreichen. FM-Modelle erzeugen oft gekrümmte Trajektorien aufgrund von zufälligen oder batchweisen Kopplungen, was den Diskretisierungsfehler erhöht und die Probenqualität verringert. COT-FM behebt dies, indem Zielproben geclustert und jedem Cluster eine dedizierte Quellverteilung zugewiesen wird, die durch Umkehrung vortrainierter FM-Modelle gewonnen wird. Diese Divide-and-Conquer-Strategie führt zu einer genaueren lokalen Transportabbildung und signifikant geraderen Vektorfeldern, alles ohne Änderung der Modellarchitektur. Als Plug-and-Play-Ansatz beschleunigt COT-FM durchgängig die Erzeugung von Stichproben und verbessert die Generierungsqualität in 2D-Datensätzen, Bildgenerierungs-Benchmarks und Robotermanipulationsaufgaben.

English

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.