COT-FM : Appariement de Flux par Transport Optimal Clusterisé

COT-FM: Cluster-wise Optimal Transport Flow Matching

March 11, 2026

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

cs.AI

Résumé

Nous présentons COT-FM, un cadre général qui reconfigure le chemin de probabilité dans l'appariement de flux (Flow Matching, FM) pour parvenir à une génération plus rapide et plus fiable. Les modèles FM produisent souvent des trajectoires courbes en raison de couplages aléatoires ou par lots, ce qui accroît l'erreur de discrétisation et réduit la qualité des échantillons. COT-FM résout ce problème en regroupant les échantillons cibles et en attribuant à chaque cluster une distribution source dédiée, obtenue en inversant des modèles FM pré-entraînés. Cette stratégie « diviser pour régner » produit un transport local plus précis et des champs vectoriels significativement plus rectilignes, le tout sans modifier l'architecture du modèle. En tant qu'approche prête à l'emploi, COT-FM accélère systématiquement l'échantillonnage et améliore la qualité de génération sur des ensembles de données 2D, des benchmarks de génération d'images et des tâches de manipulation robotique.

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.