Riepilogo Tecnico: Drift Flow Matching (DFM)

Enunciato del Problema

I modelli generativi iterativi, come Flow Matching e i modelli Diffusion, hanno dimostrato forti capacità di "scaling a tempo di test", dove l'aumento del numero di passi di inferenza migliora la qualità della generazione, la controllabilità e la copertura della distribuzione. Tuttavia, questi metodi richiedono molteplici passi di inferenza, il che può essere computazionalmente costoso. Al contrario, i Modelli Drift offrono una generazione efficiente in un solo passo apprendendo una mappa di spinta diretta da una distribuzione sorgente a una distribuzione target. Eppure, il loro paradigma di generazione diretta manca di un meccanismo naturale per migliorare i risultati aumentando i passi di inferenza a tempo di test, limitando la loro flessibilità e adattabilità a diversi requisiti di qualità-efficienza.

Metodologia: Drift Flow Matching (DFM)

Gli autori propongono Drift Flow Matching (DFM), un framework che unifica l'efficienza dei Modelli Drift in un solo passo con la scalabilità dei metodi Flow Matching multi-passo. DFM raggiunge questo obiettivo apprendendo un modello di trasporto a due tempi in grado di spostare campioni da un qualsiasi tempo corrente $t$ a un qualsiasi tempo futuro $r > t$ lungo un percorso di probabilità Flow Matching.

Meccanismi Fondamentali

1. Costruzione dei Margini a Due Tempi:
   A differenza del Flow Matching standard, che esegue una regressione su un obiettivo di velocità condizionale puntuale, DFM opera su coppie di distribuzioni marginali $(p_t, p_r)$. Definisce un percorso condizionale $X^Z_t = \alpha(t)X^Z_0 + \beta(t)X^Z_1$ e $X^Z_r = \alpha(r)X^Z_0 + \beta(r)X^Z_1$ basato su un accoppiamento indipendente delle distribuzioni sorgente ($X_0$) e target ($X_1$). Questa costruzione induce distribuzioni marginali $p_t$ e $p_r$ a passi temporali arbitrari.

2. Parametrizzazione della Velocità Media:
   DFM parametrizza un campo di velocità media $u_\theta(x_t, t, r)$ per trasportare la distribuzione corrente $p_t$ alla distribuzione futura $p_r$. La mappa di trasporto è definita come:
   $$T^\theta_{t,r}(x_t) = x_t + (r - t)u_\theta(x_t, t, r)$$
   Questa formulazione permette al modello di gestire passi di dimensione arbitraria $(r-t)$. Nel limite in cui $r \to t$, la velocità media recupera la velocità marginale istantanea del Flow Matching standard.

3. Supervisione Basata sulla Deriva:
   Invece di eseguire una regressione su un vettore di velocità, DFM impiega un obiettivo basato sulla deriva per trasportare la distribuzione prevista $q^\theta_{t,r}$ (generata da $T^\theta_{t,r}$) verso la vera margine target $p_r$.

   • Campo di Deriva: Un campo di deriva $V_{q^\theta_{t,r}, p_r}$ è costruito utilizzando un'attrazione pesata da kernel verso i campioni target ($V^+_{p_r}$) e un termine di auto-correzione repulsivo lontano dalla distribuzione corrente del modello ($V^-_{q^\theta_{t,r}}$).
   • Funzione di Perdita: L'obiettivo di addestramento minimizza la magnitudine del campo di deriva utilizzando una formulazione con stop-gradient:
     $$L_{DFM}(\theta) = \frac{1}{2} \mathbb{E}_{(t,r), X_t \sim p_t} \left[ \| \hat{X}^\theta_r - \text{sg}(\hat{X}^\theta_r + V_{q^\theta_{t,r}, p_r}(\hat{X}^\theta_r)) \|^2 \right]$$
   • Stima Raggruppata: Crucialmente, il campo di deriva e la perdita devono essere calcolati raggruppando i dati per ogni coppia di tempi campionata $(t, r)$. Mescolare campioni da coppie temporali diverse corromperebbe la stima della deriva perché le marginali target $p_r$ differiscono per ogni coppia.

Contributi Chiave

• Framework Unificato: DFM colma il divario tra i Modelli Drift efficienti in un solo passo e i metodi iterativi Flow Matching, creando un singolo paradigma che supporta sia la generazione diretta che il raffinamento iterativo.
• Scaling a Tempo di Test: Il framework sblocca capacità di scaling a tempo di test per i metodi basati sulla deriva. Gli utenti possono bilanciare la velocità di campionamento con la qualità della generazione regolando il numero di passi di inferenza (NFE) senza riaddestramento.
• Principio di Apprendimento Innovativo: Il documento introduce un principio di apprendimento che sposta la regressione della velocità puntuale (standard nel Flow Matching) verso il trasporto di distribuzioni marginali tramite supervisione della deriva. Questo evita l'instabilità spesso associata ai metodi di auto-distillazione utilizzati in altri approcci di generazione in un solo passo.
• Garanzie Teoriche: Gli autori forniscono un'analisi teorica che mostra come DFM decomponga la generazione agli estremi in problemi di trasporto marginale a corto raggio controllati dalla distanza di Wasserstein ($W_2$) e che il metodo recuperi il Flow Matching standard nel limite del passo infinitesimale.

Risultati Sperimentali

Numerosi esperimenti su dati sintetici, MNIST, FFHQ, ImageNet-1k e compiti di controllo robotico dimostrano l'efficacia di DFM:

• Dati Sintetici: Le visualizzazioni mostrano che DFM può generare campioni di alta qualità in un singolo passo, migliorando ulteriormente la copertura della distribuzione e la qualità dei campioni all'aumentare di NFE.
• Generazione Condizionata (MNIST & FFHQ): DFM ottiene prestazioni competitive con l'inferenza in un solo passo (allineandosi ai baseline dei Modelli Drift) e migliora costantemente le metriche (EMD, FID) all'aumentare di NFE a 2, 5 o 10.
• Sintesi di Immagini su Larga Scala (ImageNet-1k): Sul dataset ImageNet-256x256, DFM raggiunge un FID competitivo di 1.52 con NFE=1. Aumentare NFE a 10 riduce ulteriormente il FID a 1.31, dimostrando uno scaling efficace a tempo di test comparabile ai metodi iterativi Flow Matching.
• Controllo Robotico: Nei compiti di manipolazione robotica (ad es. Lift, Can, PushT), DFM preserva la capacità di inferenza in un solo passo della Policy Drift, mostrando chiari guadagni di prestazioni con l'aumento dei passi di inferenza, superando sia la Policy Drift in un solo passo che i baseline della Policy Diffusion multi-passo in compiti specifici.

Significato e Affermazioni

Il documento afferma che DFM stabilisce un paradigma generativo più flessibile in cui la generazione diretta e il raffinamento iterativo non sono più scelte progettuali mutualmente esclusive, ma possono essere considerate all'interno di un unico framework. Collegando la modellazione generativa basata sulla deriva con la generazione iterativa basata sul flusso, DFM permette ai processi di generazione di adattare il calcolo del campionamento a diversi requisiti di qualità-efficienza. Gli autori sottolineano che questo approccio fornisce un modo innovativo per ottenere una forte generazione in un solo passo, mantenendo al contempo i benefici di scalabilità dei metodi iterativi, affrontando così una limitazione chiave dei precedenti Modelli Drift.