COT-FM: Emparejamiento de Flujo con Transporte Óptimo por Clústeres

Resumen

Presentamos COT-FM, un marco general que redefine la trayectoria de probabilidad en Flow Matching (FM) para lograr una generación más rápida y fiable. Los modelos FM a menudo producen trayectorias curvas debido a acoplamientos aleatorios o por lotes, lo que incrementa el error de discretización y reduce la calidad de las muestras. COT-FM soluciona este problema agrupando las muestras objetivo y asignando a cada grupo una distribución de origen dedicada, obtenida mediante la inversión de modelos FM preentrenados. Esta estrategia divide y venceras produce un transporte local más preciso y campos vectoriales significativamente más rectos, todo ello sin modificar la arquitectura del modelo. Como enfoque plug-and-play, COT-FM acelera consistentemente el muestreo y mejora la calidad de generación en conjuntos de datos 2D, benchmarks de generación de imágenes y tareas de manipulación robótica.

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.

COT-FM: Emparejamiento de Flujo con Transporte Óptimo por Clústeres

COT-FM: Cluster-wise Optimal Transport Flow Matching

Resumen

Support