Riepilogo Tecnico: Modellazione Generativa in Un Passo tramite Flussi di Gradiente di Wasserstein

Enunciato del Problema

I modelli generativi moderni, in particolare i modelli di diffusione e i metodi basati su flussi, hanno raggiunto una qualità del campione all'avanguardia, ma soffrono di elevati costi computazionali e latenza a causa della loro dipendenza dal campionamento iterativo multi-step. Sebbene esistano metodi di distillazione in uno o pochi passi, spesso ereditano segnali di addestramento da percorsi originariamente progettati per la generazione iterativa o si affidano a dinamiche euristico (ad es. Modelli Drifting) che mancano di garanzie di convergenza principiate e sono soggetti al collasso delle modalità. La sfida centrale è progettare una dinamica di addestramento che permetta a un generatore di apprendere una mappatura diretta, in un singolo passo, da una distribuzione di riferimento a una distribuzione di dati target, garantendo al contempo una convergenza stabile e una copertura globale delle modalità.

Metodologia: W-Flow

Gli autori propongono W-Flow, un framework che disaccoppia la progettazione delle dinamiche di addestramento dalla parametrizzazione del generatore. L'idea fondamentale è prescrivere l'evoluzione della distribuzione del modello verso la distribuzione target utilizzando i Flussi di Gradiente di Wasserstein (WGF) e quindi addestrare una rete neurale statica per comprimere questa evoluzione multi-step in un singolo passo di inferenza.

1. Dinamiche di Addestramento tramite Flusso di Gradiente di Wasserstein

Invece di apprendere direttamente una sequenza di passi di denoising o campi di velocità, W-Flow definisce la traiettoria di addestramento come un'evoluzione a tempo continuo nello spazio delle misure di probabilità.

• Funzionale di Energia: L'evoluzione è guidata dalla discesa più ripida di un funzionale di energia $F(q) = D(q \| p)$, dove $q$ è la distribuzione corrente del modello e $p$ è la distribuzione dei dati target.
• Campo di Velocità: Il campo di velocità a tempo continuo $V_t$ è derivato dalla prima variazione del funzionale di energia: $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$.
• Discretizzazione: Il flusso continuo è discretizzato utilizzando uno schema di Eulero esplicito. Ad ogni passo di addestramento $k$, la distribuzione del modello $q^{(k)}$ viene aggiornata tramite una mappa di trasporto locale $g^{(k)}(x) = x + \eta V^{(k)}(x)$, dove $\eta$ è la dimensione del passo.

2. Istanza con Divergenza di Sinkhorn

Per garantire che il campo di velocità sia computazionalmente trattabile e geometricamente ben comportato in dimensioni elevate, gli autori istanziano il funzionale di energia utilizzando la divergenza di Sinkhorn ($S_\epsilon$), una metrica di trasporto ottimo (OT) regolarizzata entropicamente.

• Formulazione del Campo di Velocità: Il campo di velocità indotto è dato dalla differenza tra due proiezioni baricentriche:
  $$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$$
  Qui, $T^\epsilon_{q,p}$ trasporta massa dalla distribuzione generata $q$ al target $p$, mentre $T^\epsilon_{q,q}$ rappresenta un termine di auto-trasporto che impone la coerenza con la distribuzione corrente.
• Coordinamento Globale: A differenza delle interazioni locali tramite kernel (ad es. in MMD o Modelli Drifting), il campo di velocità di Sinkhorn è derivato da un piano OT globale, garantendo che i movimenti delle particelle siano coordinati in tutto il batch per prevenire il collasso delle modalità.
• Stimatore a Due Batch: Per evitare artefatti di auto-matching in cui le particelle corrispondono a se stesse con costo zero, gli autori adottano una strategia a due batch. Calcolano il termine di auto-trasporto $T^\epsilon_{q,q}$ utilizzando un secondo batch indipendente di campioni dalla distribuzione corrente, piuttosto che mascherare la diagonale della matrice dei costi di un singolo batch.

3. Addestramento del Generatore

Il generatore $f_\theta$ è addestrato per imitare gli aggiornamenti discreti delle particelle definiti dal WGF.

• Obiettivo: La funzione di perdita minimizza la distanza quadrata tra l'output del generatore e la posizione della particella "target" derivata dall'aggiornamento di Sinkhorn:
  $$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$$
  dove $\text{sg}(\cdot)$ denota l'operatore stop-gradient, trattando la posizione aggiornata della particella come un target di regressione fisso.
• Guida Senza Classificatore (CFG): Il framework supporta la CFG modificando direttamente il campo di velocità. Invece di mescolare linearmente le distribuzioni, gli autori propongono la guida di velocità, che aggiunge un termine proporzionale alla differenza tra i campi di velocità condizionati e incondizionati. Questo approccio si è dimostrato in grado di recuperare la distribuzione target standard inclinata esponenzialmente quando si utilizza la divergenza KL e fornisce un'estensione geometricamente coerente per la divergenza di Sinkhorn.

Contributi Chiave

1. Framework Principiato in Un Passo: W-Flow stabilisce un nuovo paradigma per la modellazione generativa in un passo fondando le dinamiche di addestramento sui flussi di gradiente di Wasserstein, offrendo un'interpretazione variazionale chiara assente negli approcci euristici come i Modelli Drifting.
2. Dinamiche Basate su Sinkhorn: L'istanza di WGF con divergenza di Sinkhorn fornisce una regola di aggiornamento efficiente basata sul trasporto ottimo che cattura le discrepanze distributive globali e migliora la copertura delle modalità rispetto ai metodi basati su kernel o alle interazioni locali.
3. Convergenza Teorica: Gli autori dimostrano che le dinamiche delle particelle a campione finito di W-Flow convergono alle dinamiche distributive a tempo continuo sotto appropriate ipotesi di regolarità, fornendo una giustificazione teorica per il processo di apprendimento.
4. Stimatore a Due Batch: Viene introdotta una nuova strategia a due batch per stimare i termini di auto-trasporto, che elimina il bias di auto-matching intrinseco negli stimatori a singolo batch e nelle euristiche di mascheramento della diagonale.

Risultati Sperimentali

Gli autori valutano W-Flow su ImageNet 256×256 e FFHQ.

• Prestazioni all'Avanguardia: W-Flow raggiunge un nuovo FID all'avanguardia di 1.29 (scala XL) e 1.35 (scala L) per la generazione condizionata alla classe in un passo su ImageNet, superando significativamente i precedenti metodi in un passo (ad es. Modelli Drifting, MeanFlow) e battendo molti modelli di diffusione multi-step (ad es. LightningDiT, SiT) nonostante utilizzi una sola valutazione della funzione (NFE).
• Efficienza di Campionamento: Rispetto ai modelli di diffusione multi-step con punteggi FID simili, W-Flow produce un campionamento circa 100 volte più veloce (ad es. 77.88 immagini/sec contro 0.93 immagini/sec per SiT-XL/2 su una singola GPU H100).
• Copertura delle Modalità: Negli esperimenti con distribuzioni sbilanciate (ad es. 95% volti di anziani, 5% volti di bambini), W-Flow cattura con successo le modalità minoritarie, mentre i Modelli Drifting mostrano un collasso delle modalità, ignorando completamente la classe minoritaria.
• Velocità di Convergenza: W-Flow converge sostanzialmente più velocemente durante l'addestramento, raggiungendo le prestazioni di un Modello Drifting addestrato per 1280 epoche in sole 384 epoche.
• Trasferimento di Dominio: Il metodo dimostra una forte capacità nei compiti di trasferimento di dominio (ad es. traduzione dell'età su FFHQ), producendo traiettorie di trasporto più brevi e coerenti rispetto alle baseline euristico.

Significato e Affermazioni

Il paper afferma che i flussi di gradiente di Wasserstein forniscono una base principiata ed efficace per la modellazione generativa rapida e ad alta fedeltà. Separando la progettazione delle dinamiche di addestramento dalla parametrizzazione del generatore, W-Flow dimostra che un generatore in un passo può essere addestrato per raggiungere prestazioni comparabili a complessi processi iterativi multi-step. Gli autori sottolineano che il coordinamento globale fornito dalle dinamiche di trasporto ottimo è fondamentale per stabilizzare l'addestramento e garantire una copertura completa delle modalità, affrontando le limitazioni critiche degli approcci esistenti in un passo. I risultati suggeriscono che W-Flow può servire come alternativa scalabile per applicazioni che richiedono generazione ad alta qualità e bassa latenza, come strumenti creativi in tempo reale e simulazioni.