Riepilogo Tecnico: Esplorazione e Sfruttamento della Stabilità nel Flow Matching Latente

1. Enunciato del Problema

L'addestramento di modelli generativi, in particolare i Modelli di Diffusione e le loro controparti basate sul Flow Matching (FM), rimane computazionalmente costoso. Questi modelli richiedono tipicamente dataset massicci, risorse computazionali estese, lunghi tempi di addestramento e, nelle configurazioni condizionali, annotazioni estese. Sebbene il Flow Matching offra una formulazione deterministica per la generazione di campioni tramite Equazioni Differenziali Ordinarie (ODE) che spesso riduce i passaggi di campionamento rispetto alle Equazioni Differenziali Stocastiche (SDE), il costo dell'addestramento di questi modelli su dataset in scala completa rimane un collo di bottiglia.

La domanda centrale affrontata da questo lavoro è se sia possibile ridurre significativamente la dimensione del dataset e la capacità del modello senza sacrificare la qualità della generazione. Nello specifico, gli autori investigano la stabilità dei modelli Latent Flow Matching (LFM) sotto perturbazioni maggiori, inclusi il sottocampionamento del dataset (pruning), modifiche architetturali e configurazioni di addestramento alterate. Essi ipotizzano che, se le traiettorie di trasporto LFM sono intrinsecamente stabili, si potrebbe sfruttare questa proprietà per derivare algoritmi di addestramento e inferenza più efficienti.

2. Metodologia

2.1. Fondamento Teorico: Stabilità nel Flow Matching

Il lavoro si basa sull'osservazione che il Rectified Flow ammette un'espressione in forma chiusa per il campo di velocità. Per un dataset finito, il campo di velocità ottimale è una somma pesata dei campioni di addestramento, dove i pesi sono determinati da una funzione softmax sulle distanze quadrate tra lo stato corrente e i campioni di addestramento scalati. Gli autori notano che questa softmax diventa "piccata" (peaked) nelle prime fasi della fase di trasporto, il che significa che un singolo campione di addestramento domina spesso il campo di velocità per una data traiettoria.

Questa intuizione teorica suggerisce che la traiettoria di trasporto è determinata in larga misura dalla presenza di specifici campioni "dominanti". Di conseguenza, la rimozione di campioni non dominanti (pruning) dovrebbe lasciare la maggior parte delle traiettorie invariate. Gli autori formalizzano questo concetto definendo una metrica di assegnazione: per una traiettoria che inizia dal rumore $x_0$, determinano a quale campione di addestramento $x_i$ essa termina. Ipotizzano che questa assegnazione rimanga invariata anche quando una grande frazione del dataset viene rimossa, a condizione che il campione dominante per quella traiettoria venga mantenuto.

2.2. Strategie di Pruning dei Dati

Per sfruttare questa stabilità, gli autori propongono tre criteri di pruning dei dati adattati dai modelli discriminativi ma riformulati per l'obiettivo FM (che utilizza percorsi di rumore condivisi e istanti temporali):

1. Punteggio basato sul Gradiente (G): I campioni sono classificati in base alla magnitudine della norma del gradiente della perdita.
2. Punteggio basato sulla Perdita (L): I campioni sono classificati in base alla magnitudine della perdita calcolata lungo percorsi di rumore condivisi.
3. Punteggio basato sul Clustering (C): I campioni sono raggruppati in uno spazio di embedding pre-addestrato (ad es. CLIP). Gli autori propongono due varianti specifiche:
   • Proporzionale ($C_p$): Seleziona campioni in proporzione alla dimensione del cluster.
   • Bilanciato ($C_b$): Seleziona un numero uguale di campioni da ciascun cluster per mitigare gli squilibri distribuzionali.
   • Basato su Kernel ($C_b^\kappa$): Seleziona campioni in modo greedy per corrispondere alla media del kernel Gaussiano del cluster nello spazio latente.

Gli autori valutano anche l'"inverso" di questi criteri (selezionando i punteggi più bassi) per testare la robustezza dell'affermazione sulla stabilità.

2.3. Inferenza da Grezzo a Fine (C2F)

Sfruttando la stabilità delle traiettorie attraverso le capacità dei modelli, gli autori propongono una strategia di inferenza in due fasi:

1. Modello Grezzo: Un'architettura leggera (ad es. DiT-S/2) viene addestrata su un sottoinsieme pruning dei dati. Gestisce la prima fase della traiettoria ($t \in [0, t_0)$).
2. Modello Fine: Un'architettura ad alta capacità (ad es. DiT-XL/2), pre-addestrata sull'intero dataset, gestisce la seconda fase ($t \in [t_0, 1]$).

Per garantire una transizione fluida nel punto di "giuntura" $t_0$, il Modello Grezzo viene affinato con una funzione di perdita di continuità. Questa perdita minimizza la differenza tra i campi di velocità del Modello Grezzo e del Modello Fine in $t_0$. La traiettoria del Modello Fine viene utilizzata per generare lo stato target $x_{t_0}$ tramite inversione ODE (integrazione all'indietro), assicurando che il Modello Grezzo impari ad allinearsi al percorso del Modello Fine.

3. Contributi Chiave

1. Dimostrazione Empirica della Stabilità LFM: Gli autori mostrano che modelli LFM addestrati indipendentemente convergono verso traiettorie di trasporto altamente coerenti sotto sottocampionamento del dataset, modifiche architetturali e diverse inizializzazioni casuali. Questa stabilità è legata alle proprietà teoriche dell'obiettivo FM.
2. Algoritmi di Pruning Pratici: Introducono tre criteri di pruning (Gradiente, Perdita, Clustering) adatti per FM. Dimostrano che un pruning aggressivo (ad es. rimozione del 50-95% dei dati) non degrada le prestazioni e può persino migliorare la velocità di convergenza e la qualità della generazione sotto budget computazionali limitati.
3. Accelerazione dell'Inferenza da Grezzo a Fine: Propongono una strategia di generazione a due modelli che suddivide la traiettoria. Questo approccio ottiene un'accelerazione significativa dell'inferenza (circa 2.15$\times$) mantenendo la qualità della generazione, a condizione che l'ipotesi di stabilità valga attraverso le architetture dei modelli.
4. Mitigazione del Bias tramite Pruning Bilanciato: Gli autori mostrano che il pruning basato su clustering bilanciato ($C_b$) può mitigare i bias nei dati generati (ad es. genere, tono della pelle) senza richiedere un bilanciamento esplicito consapevole delle etichette, equalizzando la rappresentazione dei cluster semantici nel set di addestramento.

4. Risultati Sperimentali

I metodi sono stati valutati su CelebA-HQ, FFHQ e ImageNet-1K.

• Stabilità sotto Pruning: Su CelebA-HQ, i modelli addestrati su sottoinsiemi disgiunti (prunati al 50%) hanno prodotto immagini visivamente simili quando avviati dallo stesso seme di rumore $x_0$. La similarità coseno ArcFace tra modelli prunati e non prunati è variata da 0.79 a 0.83, significativamente più alta rispetto alle coppie casuali (0.37).
• Metriche di Prestazione:
  • CelebA-HQ: Il pruning del 50% dei dati utilizzando il clustering bilanciato ($C_b$) ha leggermente migliorato la Fréchet Inception Distance (FID) da 24.24 (non prunato) a 22.80.
  • ImageNet: Il pruning del 75% dei dati ($pr=0.75$) ha portato a una convergenza più rapida e ha mantenuto o leggermente migliorato l'FID rispetto all'intero dataset. Anche con un pruning del 95%, il modello ha converguto più rapidamente inizialmente, sebbene le prestazioni si siano degradate dopo un certo numero di iterazioni.
• Robustezza alle Perturbazioni:
  • Architettura: L'aumento della capacità del modello (da DiT-S/2 a DiT-XL/2) ha comportato cambiamenti minimi nella traiettoria (similarità $s=0.81$). Il passaggio a un'architettura U-Net ha ridotto la similarità ($s=0.55$) ma ha preservato gli attributi grezzi.
  • Spazio Latente: La stabilità è sensibile ai cambiamenti nella parametrizzazione dello spazio latente (ad es. cambiare il seme VAE o applicare trasformazioni invertibili), il che ha rotto la stabilità.
  • Obiettivo: Il passaggio dal Flow Matching alla diffusione basata sul punteggio ha alterato completamente le traiettorie, indicando che la stabilità è specifica per l'obiettivo FM.
• Efficienza da Grezzo a Fine: Su CelebA-HQ, l'approccio C2F con $t_0=0.7$ ha ridotto il tempo di inferenza da 93.95 ms/immagine (solo Fine) a 43.53 ms/immagine (circa 2.15$\times$ di accelerazione) mantenendo l'FID.
• Equità: Il pruning bilanciato ($C_b$) ha ridotto la divergenza KL verso una distribuzione uniforme attraverso gli attributi (genere, età, tono della pelle) più efficacemente rispetto ai modelli non prunati, nonostante fosse agnostico rispetto alle etichette.

5. Significato e Affermazioni

Gli autori affermano che la stabilità osservata è una proprietà caratteristica del paradigma Rectified Flow Matching sotto rappresentazioni latenti fisse, piuttosto che un artefatto di specifici algoritmi di addestramento.

• Efficienza: Il significato principale è la capacità di ridurre il tempo di addestramento e i requisiti di dati senza sacrificare la qualità. Gli autori dimostrano che per dataset grandi come ImageNet, il pruning può accelerare la convergenza, e per dataset più piccoli come CelebA-HQ, può migliorare la qualità della generazione.
• Accelerazione dell'Inferenza: L'approccio C2F offre un modo principiato per ridurre i costi di inferenza scaricando le prime fasi, grezze, della generazione su modelli leggeri, affidandosi alla stabilità della traiettoria per garantire che il modello ad alta capacità debba solo rifinire i dettagli finali.
• Controllo del Bias: Il lavoro suggerisce che il pruning del dataset, specificamente tramite clustering bilanciato, può servire come meccanismo per controllare e mitigare gli squilibri distribuzionali nei dati generati, offrendo una strada verso modelli generativi più equi senza un complesso riponderaggio consapevole delle etichette.

Il documento conclude che, sebbene la stabilità sia robusta ai cambiamenti nei dati e nella capacità del modello, è sensibile ai cambiamenti nella struttura dello spazio latente e nell'obiettivo generativo stesso. Questa intuizione permette ai praticanti di progettare pipeline più efficienti dal punto di vista delle risorse mantenendo una generazione ad alta fedeltà.