Adversarial Flow Models

1. Il Problema

Il lavoro affronta le limitazioni fondamentali di due grandi famiglie di modelli generativi: i GAN (Generative Adversarial Networks) e i modelli basati su Flow Matching/Consistency.

• Limiti dei GAN: Sebbene supportino la generazione in un singolo passo (one-step), l'obiettivo avversariale da solo non definisce una mappa di trasporto unica tra la distribuzione del rumore e quella dei dati. Questo porta a un'instabilità nell'addestramento, poiché il generatore può "derivare" (drift) verso mappe di trasporto arbitrarie a seconda dell'inizializzazione e della stocasticità, rendendo difficile la convergenza, specialmente con architetture standard come i Transformer.
• Limiti dei Flow Matching e Consistency Models: Questi metodi apprendono una mappa di trasporto deterministica e offrono maggiore stabilità. Tuttavia, per generare in pochi passi (o un solo passo), devono essere addestrati su tutti i timestep intermedi della flow probability per propagare la consistenza. Questo spreca capacità del modello, introduce errori di accumulo e richiede un'architettura complessa per gestire i timestep. Inoltre, spesso producono risultati leggermente sfocati se non combinati con tecniche avversariali finali.

L'obiettivo è creare un modello che unisca la stabilità e la natura deterministica dei flow models con l'efficienza e la qualità dei GAN, permettendo l'addestramento end-to-end su architetture standard senza supervisione intermedia.

2. Metodologia: Adversarial Flow Models (AFM)

Gli autori propongono una nuova classe di modelli che appartiene sia alla famiglia avversariale che a quella dei flussi.

Concetto Chiave: Mappa di Trasporto Ottimale Deterministica

A differenza dei GAN tradizionali che apprendono una mappa di trasporto arbitraria, gli AFM sono vincolati ad apprendere una mappa di trasporto ottimale deterministica (basata sul teorema di Brenier). Questo vincolo stabilizza l'addestramento avversariale eliminando l'ambiguità sulla direzione del trasporto.

Componenti Principali:

1. Obiettivo Avversariale: Utilizza un obiettivo relativistico (relativistic objective) con un discriminatore $D$ e un generatore $G$.
2. Regolarizzazione Optimal Transport (OT): Viene introdotta una perdita aggiuntiva $L_{ot}$ che minimizza il costo di trasporto quadratico $\|G(z) - z\|^2$. Questo forza il generatore a imparare una mappa specifica e stabile.
   • Il coefficiente $\lambda_{ot}$ viene decrescente durante l'addestramento: inizialmente alto per stabilizzare e guidare il modello verso la mappa ottima, poi ridotto per permettere al modello di adattarsi perfettamente alla distribuzione dei dati senza essere vincolato troppo rigidamente all'identità.
3. Gestione dei Gradienti (Gradient Normalization): Un problema critico nell'addestramento avversariale è la variazione di magnitudine dei gradienti provenienti dal discriminatore, che rende difficile bilanciare la perdita avversariale con quella OT. Gli autori introducono una tecnica di normalizzazione dei gradienti che scala i gradienti del discriminatore in base alla loro norma media (EMA), permettendo di trovare un $\lambda_{ot}$ efficace indipendentemente dalla dimensione del modello.
4. Supporto Multi-step e Any-step: Il modello può essere addestrato sia per la generazione in un singolo passo (1NFE) che in più passi, interpolando tra timesteps arbitrari lungo la flow probability.
5. Guidance (Classifier Guidance): Viene proposta una tecnica per incorporare la guida del classificatore (simile alla CFG nei diffusion models) anche per la generazione in un singolo passo, interpolando il campione generato su timesteps casuali prima di passare al classificatore temporale.

Architettura

Il modello utilizza architetture Diffusion Transformer (DiT) standard.

• Il generatore e il discriminatore sono quasi identici, con l'aggiunta di un token [CLS] nel discriminatore per la classificazione.
• Viene esplorata l'addestramento di modelli extra-deep (fino a 112 layer) tramite ripetizione di blocchi Transformer, addestrati end-to-end senza supervisione intermedia (teacher forcing).

3. Contributi Chiave

1. Unificazione Teorica: Dimostrano che combinare l'obiettivo avversariale con una regolarizzazione di trasporto ottimale crea un modello stabile che apprende una mappa deterministica, risolvendo il problema del "drift" dei GAN.
2. Efficienza di Capacità: A differenza dei metodi basati sulla consistenza, gli AFM non devono apprendere i timestep intermedi per propagare la consistenza. Questo preserva la capacità del modello per la generazione finale, permettendo a modelli più piccoli (es. B/2) di competere con modelli molto più grandi (XL/2) di altri metodi.
3. Addestramento End-to-End Profondo: Hanno dimostrato che è possibile addestrare modelli con 56 e 112 layer in un singolo passo (1NFE) senza supervisione intermedia, superando le controparti più superficiali addestrate con più passi.
4. Semplicità Architetturale: Il metodo funziona con architetture Transformer standard senza modifiche complesse, a differenza di altri GAN recenti che richiedono architetture non standard o reti di feature congelate.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet-256px in uno spazio latente (VAE).

• Prestazioni One-Step (1NFE):
  • Il modello AFM-XL/2 ha raggiunto un FID di 2.38, un nuovo stato dell'arte (SOTA) per la generazione in un singolo passo.
  • Il modello AFM-B/2 (130M parametri) ha raggiunto prestazioni comparabili a modelli Consistency XL/2 (675M parametri), dimostrando l'efficienza nella capacità del modello.
• Modelli Extra-Deep:
  • Un modello AFM-XL/2 da 112 layer ha raggiunto un FID di 1.94 con un solo passo, superando le controparti a 28 layer addestrate con 2 o 4 passi (2NFE/4NFE) a parità di parametri e computazione.
• Senza Guidance:
  • Anche senza guida (guidance-free), i modelli AFM superano i modelli Flow Matching addestrati con 250+ passi, grazie alla natura della distribuzione appresa dall'addestramento avversariale che è meno sensibile agli outlier.
• Confronto con SOTA:
  • I risultati superano significativamente metodi come MeanFlow, AlphaFlow, e GAN recenti (GAT, StyleGAN-XL) in termini di FID, specialmente nel regime a basso numero di passi.

5. Significato e Impatto

Questo lavoro rappresenta un passo significativo verso la risoluzione del compromesso tra stabilità, qualità e efficienza nella generazione di immagini.

• Stabilità dei GAN: Risolve il problema storico dell'instabilità dei GAN su architetture Transformer, rendendoli praticabili per compiti su larga scala.
• Efficienza Computazionale: Permette la generazione di alta qualità in un singolo passo (o pochi passi) senza la necessità di distillazione complessa o addestramento su tutti i timestep, riducendo il costo di inferenza.
• Scalabilità: La capacità di scalare a modelli molto profondi (112 layer) in un solo passo apre nuove direzioni di ricerca sulla relazione tra profondità del modello e qualità della generazione, suggerendo che la profondità è un fattore limitante più critico del numero di passi di inferenza.

In sintesi, gli Adversarial Flow Models offrono un framework robusto che combina i vantaggi dei due approcci dominanti, stabilendo un nuovo benchmark per la generazione di immagini rapida e di alta qualità.