Sintesi Tecnica: FlowConsist

Problema

I modelli di flusso veloci, inclusi i Consistency Models e MeanFlow, mirano ad accelerare il campionamento generativo apprendendo a predire direttamente gli integrali di percorso delle Equazioni Differenziali Ordinarie (ODE), consentendo la generazione in un singolo step o in pochi step. Tuttavia, gli autori sostengono che gli attuali paradigmi di addestramento soffrono di due problemi fondamentali che degradano le prestazioni:

1. Deriva Sistematica della Traiettoria (Systematic Trajectory Drift): Gli attuali metodi costruiscono velocità condizionali accoppiando casualmente campioni di rumore ($\epsilon$) e di dati ($x$). Sebbene il valore atteso di queste velocità condizionali sia uguale alla velocità marginale, le singole traiettorie condizionali intersecano molteplici traiettorie marginali distinte. L'addestramento su queste velocità condizionali causa una deriva dell'obiettivo di ottimizzazione rispetto al vero percorso ODE marginale singolare, portando a traiettorie inconsistenti.
2. Accumulo di Errore (Error Accumulation): Gli errori di approssimazione introdotti da un modello si accumulano su ampi intervalli temporali. Nei modelli di flusso veloce che si affidano a velocità auto-predette per l'addestramento, questi errori si sommano, causando deviazioni significative tra i campioni generati e la distribuzione dei dati reali, in particolare nella generazione a singolo step (1-NFE).

Metodologia

Gli autori propongono FlowConsist, un framework di addestramento progettato per imporre la coerenza della traiettoria e correggere l'accumulo di errore senza richiedere un modello insegnante pre-addestrato.

1. Imporre la Coerenza della Traiettoria (Sostituire le Velocità Condizionali)

L'intuizione teorica centrale è che fare affidamento sulle velocità condizionali ($v_t = \epsilon - x$) introduce un termine di varianza ($L_{var}$) nella funzione di perdita, che costringe il modello a sopprimere i gradienti in direzioni ad alta varianza e impedisce un adattamento fedele delle traiettorie curve.

• Spostamento Teorico: Invece di utilizzare la velocità condizionale accoppiata casualmente come surrogato della velocità marginale, FlowConsist sostituisce la velocità condizionale nell'obiettivo di coerenza con la velocità marginale predetta dal modello stesso ($u_\theta$).
• Implementazione: Gli autori derivano un nuovo obiettivo di addestramento (Eq. 7) che combina un termine standard di Flow Matching con un termine di regolarizzazione della coerenza della traiettoria.
  • L'obiettivo minimizza la differenza tra la velocità media predetta e un target derivato dalla predizione della velocità marginale del modello stesso.
  • Fondamentalmente, la velocità condizionale viene utilizzata solo in un termine che converge alla velocità marginale in valore atteso una volta addestrato il modello, eliminando efficacementmente la deriva sistematica durante il processo di ottimizzazione.
• Classifier-Free Guidance (CFG): Il metodo incorpora il CFG trattando la scala di guidance $\omega$ come un input secondario, permettendo al modello di apprendere una rappresentazione unificata che può essere regolata durante l'inferenza.

2. Rettifica della Traiettoria (Affrontare l'Accumulo di Errore)

Per affrontare la divergenza monotona causata dagli errori di approssimazione cumulativi sugli step temporali, gli autori introducono una strategia di rettifica della traiettoria.

• Allineamento della Velocità Marginale: A differenza dei metodi precedenti che correggono solo il rumore iniziale, FlowConsist allinea la distribuzione marginale dei campioni generati dal modello con la distribuzione dei dati reali in ogni step temporale lungo la traiettoria.
• Meccanismo Senza Insegnante (Teacher-Free): Invece di fare affidamento su un modello insegnante pre-addestrato (come in DMD) per stimare le velocità marginali reali, FlowConsist utilizza un modello ausiliario ($G_\psi$) addestrato per stimare la velocità marginale dei campioni generati dal modello principale ($F_\theta$).
• Processo di Correzione: Il modello principale viene addestrato per allineare la sua velocità marginale predetta con la velocità stimata dal modello ausiliario. Questo "tira" efficacemente le predizioni del modello verso la distribuzione dei dati reali, rettificando gli errori cumulativi che altrimenti si accumulerebbero in una mappatura a singolo step.

Contributi Chiave

1. Identificazione della Deriva della Traiettoria: Il paper fornisce una prova teorica (Teorema 1 & 2) dimostrando che le velocità condizionali nel Flow Matching possiedono intrinsecamente una covarianza non nulla, portando a una inevitabile deriva della traiettoria quando utilizzate come stime di velocità istantanea nei modelli di flusso veloce.
2. Nuovo Obiettivo di Addestramento: Gli autori propongono un'alternativa rigorosa (Eq. 7) che sostituisce le velocità condizionali con velocità marginali auto-predette, garantendo che l'ottimizzazione rimanga ancorata a una traiettoria marginale coerente.
3. Strategia di Rettifica della Traiettoria: Viene introdotto un metodo senza insegnante (Eq. 11) per allineare le distribuzioni marginali ad ogni step temporale, mitigando l'accumulo di errori di approssimazione.
4. Prestazioni allo Stato dell'Arte: Il metodo raggiunge un nuovo stato dell'arte su ImageNet 256×256 con un Fréchet Inception Distance (FID) di 1.52 utilizzando solo 1 step di campionamento (1-NFE), senza distillazione da un modello insegnante.

Risultati Sperimentali

• Dataset e Setup: Gli esperimenti sono stati condotti su ImageNet 256×256 utilizzando un'architettura DiT e uno spazio latente VAE pre-addestrato.
• Studi di Ablazione:
  • Sostituire lo standard MeanFlow loss con la proposta marginal trajectory consistency loss (Eq. 7) ha migliorato l'FID da 5.86 a 4.99.
  • L'aggiunta della marginal velocity alignment loss (Eq. 11) ha ulteriormente ridotto l'FID a 4.31.
  • Il metodo ha dimostrato robustezza attraverso diverse scale di CFG e rapporti di integrazione dei due componenti di perdita.
• Confronto: FlowConsist-XL/2 ha raggiunto un FID di 1.52 (1-NFE), superando il precedente miglior modello di flusso veloce non distillato (iMF-XL/2 a 1.72) e avvicinandosi alle prestazioni dei metodi con distillazione. Ha inoltre superato i modelli di diffusione multi-step in termini di efficienza (1 step vs 250+ step) mantenendo al contempo una qualità competitiva.

Significato

Il paper sostiene che FlowConsist offre una nuova prospettiva per la modellazione di flussi veloci affrontando le cause profonde delle limitazioni delle prestazioni esistenti: il disallineamento teorico tra velocità condizionali e marginali, e l'accumulo pratico di errori di approssimazione. Rettificando questi problemi, gli autori dimostrano che una generazione ad alta fedeltà in un singolo step è realizzabile attraverso obiettivi di addestramento raffinati, colmando il divario tra i paradigmi senza insegnante e quelli con distillazione. Il lavoro sottolinea che la coerenza della traiettoria è un requisito fondamentale affinché i modelli di flusso veloce possano eseguire con successo mappature lineari tra gli step temporali.