The Amazing Stability of Flow Matching

Titolo: L'Incredibile Stabilità del Flow Matching

1. Problema e Contesto

I modelli generativi profondi, in particolare i modelli di diffusione, hanno ottenuto risultati straordinari nella generazione di campioni di alta qualità. Tuttavia, l'addestramento di questi modelli è estremamente oneroso in termini di risorse computazionali e quantità di dati necessari.
Il Flow Matching (FM) si è proposto come un'alternativa efficiente e semplice ai modelli di diffusione, offrendo prestazioni competitive. Nonostante ciò, la stabilità dei modelli FM sotto perturbazioni dei dati (come la rimozione di porzioni significative del dataset) o delle architetture non è stata ancora approfondita.
Mentre per i modelli di diffusione è stato osservato che diverse partizioni dei dati convergono verso una base geometrica simile, non era chiaro se questo fenomeno di invarianza si applicasse anche ai modelli FM, né come tale stabilità potesse essere sfruttata per ridurre i costi di addestramento o per guidare strategie di data pruning (potatura dei dati).

2. Metodologia

Gli autori hanno condotto un'analisi empirica della stabilità dei modelli FM addestrati su sottoinsiemi di dati e con diverse configurazioni architetturali.

• Setup Sperimentale:

  • Dataset: CelebA-HQ (con esperimenti di confronto su FFHQ).
  • Architettura: Utilizzo di DiT (Diffusion Transformer) adattato al Flow Matching (FM-DiT) e, per confronto, una rete U-Net. Viene utilizzato un VQ-VAE per codificare le immagini.
  • Metriche:
    • FID (Fréchet Inception Distance): Per valutare la qualità e la diversità dei campioni generati.
    • Somiglianza Cosine ArcFace: Per quantificare la stabilità semantica e la preservazione dell'identità tra immagini generate da modelli diversi partendo dallo stesso rumore iniziale ($x_0$).

• Strategie di Perturbazione e Potatura (Pruning):
  Gli autori hanno testato la stabilità attraverso diverse manipolazioni:

  1. Sottoinsiemi Disgiunti: Addestramento su due parti casuali non sovrapposte del dataset.
  2. Scambio di Dataset: Addestramento su un dataset diverso ma dello stesso dominio (FFHQ) mantenendo lo stesso spazio latente.
  3. Rimozione di Cluster: Separazione dei dati per attributi (es. genere maschile/femminile) e addestramento su un solo gruppo.
  4. Cambiamenti Architetturali: Riduzione della capacità del modello (DiT-XL $\to$ DiT-S) e cambio radicale di architettura (DiT $\to$ U-Net).
  5. Metodi di Potatura Intelligente:
     • Gradiente (Grad): Selezione basata sulla norma del gradiente (segnale di apprendimento).
     • Perdita (Loss): Selezione basata sul valore della loss.
     • Cluster-based (Clust): Clustering delle immagini nello spazio delle feature (usando CLIP) e selezione basata sulla distanza dal centroide (campioni vicini o lontani) e sulla proporzione (bilanciata o proporzionale).

3. Risultati Chiave

• Stabilità Estrema: I modelli FM mostrano una stabilità sorprendente. Anche rimuovendo il 50% del dataset, la qualità (FID) e la diversità dei campioni generati rimangono preservate.
• Invarianza delle Traiettorie: Campioni generati da modelli addestrati su dataset diversi (o ridotti), ma inizializzati con lo stesso rumore casuale ($x_0$), evolvono verso output visivamente e semanticamente molto simili.
  • La somiglianza coseno ArcFace tra coppie corrispondenti (stesso seed) rimane alta (es. 0.69 - 0.81) anche tra modelli addestrati su sottoinsiemi disgiunti, mentre coppie non correlate scendono a ~0.34-0.37.
  • Questo vale anche quando si cambia completamente il dataset di addestramento (da CelebA-HQ a FFHQ) o si riduce drasticamente la capacità del modello.
• Impatto della Potatura:
  • La rimozione casuale (Random) degrada leggermente le prestazioni.
  • La rimozione basata sulla Loss (campioni con alta loss) peggiora significativamente il FID, poiché questi campioni spesso risiedono in regioni a bassa densità e la loro rimozione altera il campo di velocità.
  • La rimozione basata sul Gradiente (Grad) ha un impatto minimo sul FID se si mantengono i campioni ad alto gradiente, mentre la rimozione inversa (mantenere solo basso gradiente) peggiora le prestazioni.
  • Riscampionamento basato sui Cluster (Clustb): Un metodo che bilancia la distribuzione tra diversi cluster semantici ha addirittura migliorato le metriche FID rispetto al modello non potato, suggerendo che l'equilibrio dei dati è più importante della quantità assoluta.
• Robustezza Architetturale: Cambiare l'architettura (da DiT a U-Net) riduce la somiglianza fine, ma preserva gli attributi globali, indicando una stabilità strutturale del campo di flusso.

4. Contributi Principali

1. Dimostrazione della Stabilità: Il paper quantifica empiricamente quanto i modelli Flow Matching siano robusti a perturbazioni estreme di dati e architettura, mostrando che le traiettorie ODE convergono verso punti finali visivamente simili.
2. Metodi di Potatura Informata: Introduzione e valutazione di tre strategie di potatura (basate su gradiente, loss e clustering) applicate ai modelli FM.
3. Scoperta del Riscampionamento Bilanciato: Dimostrazione che un approccio basato sui cluster che bilancia la distribuzione dei dati può migliorare le prestazioni di generazione, offrendo una via per ridurre i costi computazionali senza sacrificare la qualità.

5. Significato e Implicazioni

• Efficienza Computazionale: La stabilità osservata suggerisce che è possibile addestrare modelli generativi potenti su dataset più piccoli o più bilanciati, riducendo drasticamente i costi di calcolo e di raccolta dati.
• Comprensione Teorica: I risultati indicano che la generalizzazione nei modelli FM non deriva da un singolo fattore, ma da una struttura globale del campo di velocità che viene adattata solo localmente dove necessario. Anche rimuovendo interi cluster di dati, il flusso globale rimane invariato, reindirizzando solo le traiettorie che avrebbero colpito i cluster rimossi.
• Affidabilità: La capacità di mantenere traiettorie coerenti sotto perturbazioni rende i modelli FM candidati ideali per applicazioni che richiedono robustezza e affidabilità in scenari con dati limitati o variabili.

In sintesi, il lavoro rivela che i modelli Flow Matching possiedono una "stabilità incredibile", permettendo di ottenere risultati di alta qualità anche con dati ridotti o architetture semplificate, aprendo la strada a strategie di addestramento più efficienti ed economiche.