One-Sided Quantile Coupling for Flow Matching

Sintesi Tecnica: Accoppiamento a Quantili Unilaterale per il Flow Matching

Definizione del Problema
Il Flow Matching (FM) addestra modelli generativi a tempo continuo attraverso la regressione del campo di velocità di un percorso di probabilità che connette una distribuzione sorgente semplice (tipicamente una Gaussiana standard) a una distribuzione di dati target. La qualità dell'accoppiamento — ovvero l'associazione tra campioni sorgente e target — influenza criticamente la geometria delle traiettorie di addestramento. Il FM standard utilizza un accoppiamento casuale indipendente, il che spesso risulta in percorsi lunghi e conflittuali che aumentano la varianza del campo di velocità target e ostacolano l'efficienza dell'addestramento.

Seborems l'accoppiamento strutturato ispirato al Trasporto Ottimale (OT) può raddrizzare questi percorsi e migliorare la qualità della generazione, i metodi esistenti (ad es., OT-CFM) si affidano alla risoluzione di un problema di assegnazione mini-batch. Tali approcci richiedono la costruzione di matrici di costo a coppie e la risoluzione di problemi di assegnazione (ad es, tramite l'algoritmo di Hungarian o le iterazioni di Sinkhorn), portando a costi computazionali che scalano almeno quadraticamente (o cubicamente per soluzioni esatte) rispetto alla dimensione del batch. Ciò limita la scalabilità e può introdurre instabilità. Il presente articolo affronta la questione se i benefici geometrici dell'accoppiamento strutturato possano essere ottenuti senza accoppiare due batch pre-campionati, evitando così l'alto costo del trasporto tra batch.

Metodologia: Quantile Coupling Flow Matching (QC-FM)
Gli autori propongono il Quantile Coupling Flow Matching (QC-FM), uno schema di accoppiamento unilaterale leggero. A differenza dei metodi tradizionali che campionano sia i batch di rumore che quelli di dati per poi ri-accoppiarli, il QC-FM campiona solo il batch di dati e costruisce direttamente i campioni sorgente accoppiati.

1. Costruzione Unilaterale:

   • Dato un batch di dati $\mathcal{B}_x$, il metodo proietta i punti dati su $k$ direzioni ortogonali casuali (fette o slices) definite da una matrice $U$.
   • Per ogni direzione, i punti dati vengono classificati all'interno del batch. Questi ranghi empirici vengono mappati sui quantili gaussiani utilizzando la funzione CDF inversa della normale standard.
   • Il codice latente sorgente $e^{(i)}$ è costruito per soddisfare tali vincoli di quantile proiettati. Nello specifico, la proiezione della sorgente sulle direzioni delle fette corrisponde ai quantili gaussiani derivati dai ranghi dei dati.
   • Le dimensioni rimanenti (il complemento ortogonale dello spazio sottostante delle fette) sono riempite campionando da una distribuzione Gaussiana condizionale. Ciò assicura che la sorgente rimanga stocastica nelle direzioni non vincolate.

2. Estensioni Ibride Gaussian-Remainder:
   Poiché l'intero accoppiamento QC (dove tutti i punti sorgente sono costruiti tramite ranghi) devia dal prior Gaussiano i.i.d. standard per $k \ge 2$ a causa delle dipendenze tra i codici delle fette, gli autori introducono due strategie ibride per mantenere i benefici del QC preservando al contempo il prior:

   • QC-FM-Mixture: Applica la costruzione QC a un sottoinsieme "ancora" del batch (ad es., il 20%) e accoppia i restanti campioni sorgente "residui" con i restanti campioni dati utilizzando un accoppiamento casuale indipendente (l'approccio Baseline). Questo crea un mix di accoppiamenti strutturati e non strutturati.
   • QC-FM-Adjacency: Applica il QC al sottoinsieme ancora, ma accoppia i campioni sorgente residui in base alla geometria locale indotta dalle ancore. I punti dati residui vengono assegnati alla loro ancora più vicina (nello spazio strutturato) e i campioni sorgente residui vengono accoppiati all'interno di questi gruppi. Ciò propaga il bias strutturato senza duplicare i target.

Contributi Chiave

• Algoritmo QC-FM: Un nuovo accoppiamento strutturato, locale al batch, che impone un trasporto monodimensionale monotono lungo direzioni casuali. Elimina la necessità di matrici di costo a coppie e risolutori di assegnazione, riducendo il costo computazionale dominante a $O(kB \log B)$ (ordinamento) anziché $O(B^2)$ o $O(B^3)$.
• Analisi Teorica:
  • Riduzione della Varianza: Gli autori dimostrano che, per un frame fisso, l'accoppiamento QC elimina la varianza di regressione irriducibile lungo ogni fetta selezionata, rendendo il flusso ideale esattamente rettilineo in quelle direzioni.
  • Preservazione del Prior: Quantificano che per $k=1$, la marginale della sorgente è esattamente gaussiana. Per $k \ge 2$, la deviazione dal prior Gaussiano è confinata alla copula dei codici delle fette e il costo di trasporto di tale deviazione è limitato.
• Schemi di Addestramento Ibridi: L'introduzione delle varianti Mixture e Adjacency che bilanciano l'iniezione di bias geometrico con la preservazione del segnale Gaussiano esplicito della Baseline.

Risultati Sperimentali
Il metodo è stato valutato su CIFAR-10, CelebA-64, FFHQ-64 e ImageNet-64 sotto budget di addestramento e architetture equivalenti.

• Prestazioni: QC-FM-Mixture ha superato costantemente la Baseline dell'accoppiamento indipendente e l'OT-CFM esatto con l'algoritmo di Hungarian su tutti e quattro i dataset. I miglioramenti rispetto alla Baseline sono variati dal 6,3% (CIFAR-10) al 12,9% (FFHQ) in termini di FID-50K.
• Efficienza: In termini di costo computazionale, QC-FM è significativamente più veloce dei metodi basati su OT. Con una dimensione del batch di 2048, QC-FM-Mixture è stato oltre 800 volte più veloce del risolutore Hungarian OT e 35 volte più veloce del risolutore Sinkhorn.
• Diagnostica su Toy Datasets: Sui dataset sintetici (Checkerboard e 8-Gaussians), le varianti QC-FM hanno prodotto traiettorie apprese quasi rettilinee (Rapporto di Lunghezza del Percorso $\approx 1$) e una varianza della velocità target significativamente inferiore rispetto alla Baseline, confermando i benefici geometrici dell'accoppiamento.

Significato e Rivendicazioni
L'articolo sostiene che preservare la struttura dei ranghi proiettati è un modo semplice, scalabile ed efficace per iniettare un utile bias geometrico negli accoppiamenti di Flow Matching senza risolvere un problema di trasporto mini-batch.

• Scalabilità: QC-FM offre una via di mezzo attraente tra la Baseline non strutturata e i costosi metodi di trasporto tra batch. È computazionalmente leggero e controllabile tramite il numero di proiezioni ($k$).
• Praticità: Il metodo non richiede componenti appresi per l'inferenza; la generazione parte comunque dalla Gaussiana standard. L'accoppiamento viene costruito "on-the-fly" durante l'addestramento.
• Limitazioni: Gli autori notano con modestia che QC-FM è un surrogato locale al batch per l'accoppiamento strutturato, non un'approssimazione del Trasporto Ottimale globale. I quantili sono stimati all'interno dei mini-batch e la teoria controlla la varianza per singola fetta piuttosto che garantire la rettilineità globale delle traiettorie o specifici miglioramenti del FID. La selezione del numero di fette ($k$) e del rapporto delle ancore ($p$) dipende attualmente da euristiche e dalla calibrazione specifica per il dataset.

In conclusione, il lavoro dimostra che l'accoppiamento a quantili unilaterale può migliorare la qualità del campionamento e l'efficienza dell'addestramento nel Flow Matching strutturando l'accoppiamento rumore-dati attraverso la preservazione dei ranghi, offrendo un'alternativa pratica ai metodi di trasporto ottimale computazionalmente proibitivi.