Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling

1. Il Problema

I modelli generativi basati su flussi (Flow-Based Models), in particolare il Conditional Flow Matching (CFM), offrono un'alternativa efficiente ai modelli di diffusione per la generazione di immagini e video. Tuttavia, come i modelli di diffusione, richiedono l'integrazione numerica di un'Equazione Differenziale Ordinaria (ODE) durante l'inferenza, il che comporta molteplici valutazioni della rete neurale (NFE - Number of Function Evaluations) e rende la generazione lenta rispetto ai metodi a passaggio singolo (come le GAN).

La velocità di inferenza è strettamente legata alla curvatura delle traiettorie di campionamento. Traiettorie più dritte richiedono meno passi di integrazione per raggiungere la stessa accuratezza.

• Il limite attuale: La curvatura dipende da come i dati ($x$) sono accoppiati al rumore ($z$) durante l'addestramento. L'accoppiamento ottimale è dato dal Trasporto Ottimo (Optimal Transport - OT), che minimizza il costo di trasporto tra le distribuzioni.
• Il collo di bottiglia: Calcolare il piano di trasporto ottimo globale per dataset di grandi dimensioni è computazionalmente proibitivo ($O(n^3)$). Le soluzioni attuali utilizzano l'OT su mini-batch (es. OT-CFM, BatchOT), ma questo approccio è limitato: le assegnazioni ottimali calcolate localmente su un mini-batch non garantiscono un'ottimalità globale, portando a traiettorie ancora curve e a una scarsa scalabilità su dataset grandi.

2. Metodologia: LOOM-CFM

Gli autori introducono LOOM-CFM (Looking Out Of Minibatch-CFM), un metodo iterativo progettato per estendere la portata dell'OT su mini-batch, migliorando l'approssimazione del piano di trasporto globale.

Meccanismo Principale

Invece di calcolare un accoppiamento ottimo locale e scartarlo alla fine di ogni passo di addestramento (come fanno i metodi precedenti), LOOM-CFM preserva e aggiorna iterativamente le assegnazioni rumore-dati tra i diversi mini-batch:

1. Inizializzazione: Si parte con un accoppiamento casuale o identico tra dati e rumore.
2. Iterazione: Ad ogni passo di addestramento $k$:
   • Viene campionato un mini-batch di dati $\{x_{n_j}\}$.
   • Si recuperano i campioni di rumore associati a questi dati dall'iterazione precedente $\{z_{\tau_{k-1}(n_j)}\}$.
   • Si risolve un problema di trasporto ottimo locale all'interno di questo mini-batch per trovare una permutazione $\omega_k$ che minimizzi il costo locale.
   • L'assegnazione globale viene aggiornata componendo la permutazione locale con quella precedente: $\tau_k = \omega_k \circ \tau_{k-1}$.
3. Addestramento: Il gradiente viene calcolato utilizzando le nuove assegnazioni aggiornate, che vengono salvate per le iterazioni future.

Questo processo assomiglia a un "tessuto" (da qui il nome Loom), dove le informazioni locali si propagano globalmente, permettendo al sistema di convergere verso un piano di trasporto più vicino all'ottimo globale rispetto ai metodi puramente locali.

Prevenzione dell'Overfitting (Multiple Noise Caches)

Un rischio nell'usare un set fisso di rumore è l'overfitting del modello su quelle specifiche coppie. Per mitigare questo problema senza costi computazionali aggiuntivi significativi, LOOM-CFM assegna multiple istanze di rumore (cache) a ciascun punto dati. Durante ogni passo di addestramento, viene selezionata casualmente una delle cache di rumore associate al dato. Questo aumenta artificialmente la dimensione del dataset e permette di utilizzare nuove istanze di rumore all'inferenza, migliorando la qualità del campione.

3. Contributi Chiave

1. Algoritmo Iterativo LOOM-CFM: Un nuovo algoritmo che ottimizza le assegnazioni globali dati-rumore aggiornando dinamicamente le assegnazioni locali tra i mini-batch, superando i limiti di scalabilità dell'OT su mini-batch standard.
2. Analisi di Convergenza: Viene fornita una dimostrazione teorica che l'algoritmo converge in un numero finito di passi a una soluzione stazionaria priva di cicli alternati negativi di lunghezza inferiore alla dimensione del mini-batch.
3. Tecnica di Cache Multipla: Un metodo semplice per prevenire l'overfitting assegnando più rumore per dato, che non altera la distribuzione sottostante dei dati ma migliora la generalizzazione.
4. Compatibilità: Il metodo è compatibile con tecniche avanzate come la distillazione (Reflow) e l'addestramento nello spazio latente (Latent Flow Matching), permettendo la sintesi di immagini ad alta risoluzione.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su CIFAR-10, ImageNet-32, ImageNet-64 e FFHQ-256. La metrica principale è il FID (Fréchet Inception Distance) a basso numero di NFE (passi di integrazione).

• Miglioramento Speed-Quality: LOOM-CFM supera sistematicamente i metodi basati su OT a mini-batch (come OT-CFM e BatchOT).
  • Su CIFAR-10: Riduzione del FID del 41% con 12 NFE rispetto ai metodi OT standard.
  • Su ImageNet-32: Riduzione del FID del 46% con 12 NFE.
  • Su ImageNet-64: Riduzione del FID del 54% con 12 NFE.
• Traiettorie più dritte: L'analisi della curvatura delle traiettorie conferma che LOOM-CFM produce percorsi di generazione molto più lineari, permettendo una generazione di alta qualità con pochissimi passi (es. 4-8 NFE).
• Inizializzazione per Reflow: LOOM-CFM funge da eccellente inizializzazione per l'algoritmo di Reflow, eliminando la necessità di multiple iterazioni di retraining per ottenere traiettorie dritte.
• Alta Risoluzione: Applicato allo spazio latente di un autoencoder pre-addestrato (simile a Stable Diffusion), LOOM-CFM genera immagini FFHQ 256x256 con un FID competitivo usando un ordine di grandezza in meno di passi rispetto ai metodi precedenti.

5. Significato e Impatto

LOOM-CFM rappresenta un avanzamento significativo nell'efficienza dei modelli generativi basati su flussi.

• Efficienza: Risolve il compromesso tra velocità di inferenza e qualità della generazione, rendendo i modelli CFM competitivi con le GAN in termini di velocità, mantenendo la stabilità di addestramento dei modelli di diffusione.
• Scalabilità: Supera il limite fondamentale dell'OT su mini-batch, rendendo fattibile l'uso di accoppiamenti ottimali approssimati su dataset di grandi dimensioni senza costi computazionali proibitivi.
• Semplicità: A differenza di altri approcci che richiedono modelli ausiliari o funzioni di perdita complesse, LOOM-CFM modifica solo la strategia di campionamento e accoppiamento dei dati, mantenendo l'architettura del modello e l'obiettivo di addestramento originali.

In sintesi, il lavoro dimostra che migliorare la "mappatura" tra rumore e dati attraverso un meccanismo di memoria globale delle assegnazioni locali è la chiave per accelerare drasticamente l'inferenza nei modelli di generazione moderna.