COT-FM: Cluster-wise Optimal Transport Flow Matching

1. Il Problema

Il paper affronta le limitazioni attuali dei modelli Flow Matching (FM), un framework generativo che apprende un campo vettoriale deterministico per trasformare una distribuzione sorgente semplice (es. Gaussiana) in una distribuzione target complessa (i dati).

I problemi principali identificati sono:

• Traietorie Curve: I modelli FM standard spesso utilizzano accoppiamenti casuali (random coupling) o accoppiamenti ottimali a livello di batch (batch-wise optimal couplings). Queste strategie portano a campi vettoriali curvi, dove le traiettorie di trasporto si incrociano frequentemente.
• Errori di Discretizzazione: Le traiettorie curve aumentano l'errore di discretizzazione temporale durante l'integrazione dell'equazione differenziale ordinaria (ODE). Questo costringe a utilizzare un numero elevato di passaggi di campionamento (NFE - Number of Function Evaluations) per ottenere campioni di alta qualità.
• Qualità del Campionamento: Con un budget di campionamento basso (es. 1 o 2 passaggi), le traiettorie curve portano a una generazione distorta e di bassa qualità.
• Limiti delle Soluzioni Esistenti: Metodi come Rectified Flow iterativo possono appiattire le traiettorie ma rischiano il collasso del modello (model collapse) riaddestrandosi su campioni generati. Altri metodi "shortcut" (come MeanFlow) accelerano il processo ma non correggono la curvatura intrinseca del campo vettoriale, limitando il miglioramento della qualità.

2. Metodologia: COT-FM

Gli autori propongono COT-FM (Cluster-wise Optimal Transport Flow Matching), un framework "plug-and-play" che modifica il percorso di probabilità senza alterare l'architettura del modello FM sottostante. L'idea centrale è un approccio "divide et impera".

Il processo si articola in due fasi principali che si alternano:

A. Identificazione delle Distribuzioni Sorgenti per Cluster

Invece di mappare l'intera distribuzione dei dati da un'unica sorgente globale, COT-FM:

1. Clustering: Partiziona i dati target in $K$ cluster (usando etichette, descrizioni testuali o clustering non supervisionato come K-Means su feature DINO).
2. Regressione Inversa: Utilizza un modello FM pre-addestrato per invertire le traiettorie di ogni campione di dati $x_1$ fino allo spazio del rumore, ottenendo stime dei campioni sorgenti $\hat{x}_0$.
3. Stima della Sorgente: Per ogni cluster $C_k$, calcola la media ($\mu_{0,k}$) e la covarianza ($\Sigma_{0,k}$) dei campioni sorgenti inversi, definendo una distribuzione sorgente locale specifica per quel cluster: $p_{0,k} = \mathcal{N}(\mu_{0,k}, \Sigma_{0,k})$.

B. Ottimizzazione Alternata

Il framework alterna due fasi di ottimizzazione:

1. Costruzione del Campo Vettoriale Target: All'interno di ogni cluster, viene calcolato un trasporto ottimale (OT) approssimato tra la distribuzione sorgente locale $p_{0,k}$ e i dati target del cluster $C_k$. Questo crea un campo vettoriale locale più dritto e privo di incroci.
2. Fine-tuning del Modello FM: Il modello FM viene addestrato (o fine-tuned) per regressare questo nuovo campo vettoriale target, utilizzando coppie sorgente-target estratte dai trasporti ottimali calcolati a livello di cluster.

Vantaggi Chiave:

• Riduce il problema OT globale (computazionalmente costoso, $O(n^3)$) a problemi OT locali più piccoli e gestibili.
• Garantisce traiettorie più dritte perché riduce la confusione tra modalità diverse (cluster) che altrimenti si incrocerebbero.
• È agnostico rispetto al modello: funziona con qualsiasi architettura FM esistente.

3. Contributi Chiave

1. Framework Generale: Introduzione di COT-FM come metodo universale per migliorare la velocità e l'affidabilità della generazione in Flow Matching.
2. Strategia "Divide et Impera": Sostituzione dell'accoppiamento globale con accoppiamenti ottimali a livello di cluster, ottenendo distribuzioni sorgenti specifiche per ogni modalità dei dati.
3. Miglioramento senza Cambiamenti Architetturali: Il metodo non richiede modifiche alla rete neurale o ai meccanismi di input/output, rendendolo facilmente integrabile.
4. Dimostrazione Empirica: Validazione su dataset sintetici 2D, generazione di immagini (CIFAR-10, ImageNet) e manipolazione robotica (LIBERO), dimostrando che la riduzione della curvatura del flusso è la chiave per la generazione di alta qualità con pochi passaggi.

4. Risultati Sperimentali

COT-FM ha mostrato miglioramenti significativi rispetto a baseline come Rectified Flow, OT-CFM e MeanFlow:

• Dataset 2D Sintetici:
  • Riduzione drastica della distanza di Wasserstein (es. da 0.5421 a 0.1995 su "Mixture of 5-Gaussians").
  • Curvatura delle traiettorie minimizzata, indicando percorsi quasi lineari.
• Generazione di Immagini (CIFAR-10):
  • Miglioramento dell'FID (Fréchet Inception Distance) da 12.6 a 8.23 in 10 passaggi e da 4.45 a 3.97 in 50 passaggi rispetto a Rectified Flow.
  • Campionamento a 1 passaggio: FID ridotto da 378.0 a 205.0.
  • Campionamento a 2 passaggi: FID ridotto da 173 a 59.1.
• Generazione Condizionata (ImageNet 256x256):
  • Miglioramenti consistenti dell'FID su architetture SiT-B/2 e SiT-B/4, specialmente a bassi NFE.
• Manipolazione Robotica (LIBERO):
  • Con un solo passaggio di campionamento (1 NFE), COT-FM raggiunge un tasso di successo del 96.1% (Spatial) e 94.5% (Long), superando la politica FLOWER che richiede 4 passaggi per risultati simili (97.1% e 93.5%).

5. Significato e Impatto

Il lavoro di COT-FM è significativo perché risolve il compromesso fondamentale tra velocità di inferenza e qualità di generazione nei modelli generativi basati su flussi.

• Efficienza Estrema: Dimostra che è possibile ottenere generazioni di alta qualità con un numero estremamente ridotto di passaggi (anche 1 o 2), rendendo i modelli FM pratici per applicazioni in tempo reale come la robotica e il controllo.
• Teoria del Trasporto: Conferma che la curvatura delle traiettorie è il principale collo di bottiglia per la qualità nei modelli FM e che l'ottimizzazione locale (cluster-wise) è una soluzione scalabile ed efficace rispetto all'OT globale.
• Versatilità: Essendo un metodo "plug-and-play", può essere applicato a una vasta gamma di compiti generativi, dall'arte digitale al controllo robotico, senza richiedere la riprogettazione dei modelli sottostanti.

In sintesi, COT-FM rappresenta un passo avanti cruciale verso l'adozione pratica dei modelli generativi Flow Matching, rendendoli più veloci, stabili e capaci di catturare la struttura complessa dei dati reali.