COT-FM: Allineamento del Flusso con Trasporto Ottimale a Cluster

COT-FM: Cluster-wise Optimal Transport Flow Matching

March 11, 2026

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

cs.AI

Abstract

Introduciamo COT-FM, un framework generale che rimodella il percorso di probabilità nel Flow Matching (FM) per ottenere una generazione più rapida e affidabile. I modelli FM spesso producono traiettorie curve a causa di accoppiamenti casuali o per batch, che aumentano l'errore di discretizzazione e riducono la qualità del campione. COT-FM risolve questo problema raggruppando i campioni target e assegnando a ciascun cluster una distribuzione sorgente dedicata, ottenuta invertendo modelli FM pre-addestrati. Questa strategia divide et impera produce un trasporto locale più accurato e campi vettoriali significativamente più rettilinei, il tutto senza modificare l'architettura del modello. Come approccio plug-and-play, COT-FM accelera costantemente il campionamento e migliora la qualità della generazione su dataset 2D, benchmark di generazione di immagini e compiti di manipolazione robotica.

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.