Simplifying Flow Matching Transformations with Low-Rank Mixture Models

Sintesi Tecnica: Semplificare le Trasformazioni di Flow Matching con Modelli di Miscela a Basso Rango

Problematica
I flussi di normalizzazione (normalizing flows), in particolare i Continuous Normalizing Flows (CNF), apprendono una mappatura invertibile tra una distribuzione di dati complessa e una distribuzione di base semplice, tipicamente una normale standard. Sebbene il Flow Matching (FM) abbia permesso di addestrare i CNF in modo scalabile attraverso la regressione diretta verso il campo vettoriale dell'ODE, la scelta di una densità di base normale standard crea spesso un disallineamento topologico con la distribuzione dei dati target. Questo disallineamento costringe la trasformazione del flusso appresa a diventare inutilmente complessa per colmare il divario tra lo spazio latente semplice e il manifold dei dati complesso. Di conseguenza, tale complessità porta a una convergenza dell'addestramento più lenta, costi computazionali più elevati durante l'inferenza (richiedendo più step del solver ODE) e prestazioni generative subottimali.

Metodologia
Gli autori propongono di sostituire la densità di base normale standard con Miscele di Analizzatori di Componenti Principali Probabilistici (MPPCA). L'MPPCA è un modello di miscela gaussiana a basso rango che offre due vantaggi critici:

1. Verosimiglianza Analitica: A differenza di alcune distribuzioni di base apprese, l'MPPCA mantiene un'espressione esatta della verosimiglianza, preservando la trattabilità necessaria per la stima della densità nei flussi di normalizzazione.
2. Fitting Efficiente: I modelli MPPCA possono essere adattati rapidamente ed economicamente utilizzando l'algoritmo di Expectation-Maximization (EM), anche su dati ad alta dimensionalità, vincolando le matrici di covarianza a un basso rango (modellando le covarianze su un sottospazio a bassa dimensionalità appreso).

Il workflow proposto prevede una strategia di "warm-start":

1. Pre-addestramento: Un modello MPPCA viene adattato ai dati di addestramento tramite EM per approssimare da vicino la distribuzione dei dati target.
2. Addestramento del Flusso: Il flusso di normalizzazione continua viene quindi addestrato utilizzando l'obiettivo di Flow Matching, con l'MPPCA pre-adattato che funge da distribuzione di base iniziale ($p_0$).
3. Integrazione: Questo approccio è testato sia con gli obiettivi di Flow Matching standard di tipo Variance-Preserving (VP) che con quelli basati sul Trasporto Ottimale (OT).

Contributi Chiave
Il documento delinea tre contributi primari:

• Costruzione di una Base Espressiva: Gli autori costruiscono distribuzioni di base espressive per i flussi utilizzando modelli MPPCA e dimostrano che la procedura di fitting è computazionalmente economica rispetto all'addestramento del flusso successivo.
• Sinergia con il Trasporto Ottimale: Dimostrano che le distribuzioni di base MPPCA possono essere combinate efficacementmente con gli obiettivi di flow matching basati su OT. Questa combinazione genera campioni di qualità superiore riducendo simultaneamente il numero di step di integrazione richiesti durante l'inferenza.
• Validazione Empirica: L'approccio è valutato quantitativamente su compiti di stima della densità impegnativi (dataset tabulari UCI) e compiti generativi ad alta dimensionalità (dataset di immagini inclusi FashionMNIST, CIFAR-10 e CelebA), dimostrando guadagni costanti rispetto ai baseline.

Risultati Sperimentali
Gli autori hanno valutato il loro metodo contro i baseline che utilizzano una normale standard ($N(0, I)$) su molteplici dataset:

• Efficienza dell'Addestramento: Sui dataset tabulari (Hepmass, Miniboone, Bsds300), i flussi con basi MPPCA hanno richiesto significativamente meno epoche di addestramento per convergere. Ad esempio, sul dataset Miniboone, il VP-MPPCA ha richiesto circa 24 epoche rispetto alle circa 54 del VP-Normal. Il tempo totale di addestramento è stato ridotto nonostante l'overhead del fitting dell'MPPCA, poiché il tempo di fitting EM ha rappresentato meno del 3% del tempo totale di addestramento.
• Complessità dell'Inferenza: Il numero di valutazioni della funzione (NFE) richieste dai solver ODE a step adattivi è stato costantemente inferiore per i modelli basati su MPPCA. Ciò indica che i campi vettoriali appresi erano più semplici e meno complessi, riducendo l'onere computazionale durante la generazione.
• Qualità Generativa: Sui dataset di immagini, il modello OT-MPPCA ha ottenuto i punteggi più bassi per la Fréchet Inception Distance (FID) e il Numero di Bin Statisticamente Diversi (NDB). In particolare, l'ispezione visiva dei campioni generati (Figura 3) ha rivelato che i modelli inizializzati con MPPCA producevano campioni strutturati e significativi dopo anche solo il primo epoch di addestramento, mentre i baseline con normale standard richiedevano molte più epoche per raggiungere una struttura simile.

Significatività e Rivendicazioni
Il paper sostiene che apprendere una distribuzione di base espressiva funga da potente bias induttivo che semplifica la trasformazione del flusso. Inizializzando il flusso con una distribuzione che approssima da vicino i dati (minimizzando la divergenza KL tra la base e i dati), il modello è sollevato dall'onere di apprendere trasformazioni globali complesse. Invece, il flusso deve solo raffinare la distribuzione di base pre-strutturata.

Gli autori sottolineano che questo approccio è pratico per compiti ad alta dimensionalità perché i modelli MPPCA evitano i problemi di memoria e condizionamento associati alle miscele gaussiane a rango pieno. Il lavoro suggerisce che il "warm-starting" dell'addestramento del flusso con un modello di miscela a basso rango è una strategia percorribile per migliorare sia l'efficienza dell'addestramento che la velocità di inferenza senza sacrificare la qualità generativa. Il lavoro futuro propone di esplorare l'estensione del fitting MPPCA a dimensioni ancora più elevate e di combinare questo approccio con altri obiettivi di flow matching come l'action matching e lo Schrödinger Bridge flow matching.