Sintesi Tecnica: Modellazione Generativa tramite Drifting (Deriva)

Problema

La modellazione generativa è fondamentalmente il compito di apprendere una mappatura $f$ tale che la distribuzione di pushforward di una distribuzione a priori $p_{prior}$ (ad esempio, rumore Gaussiano) corrisponda alla distribuzione dei dati $p_{data}$, denotato come $f\#p_{prior} \approx p_{data}$. Sebbene i paradigmi prevalenti come i Modelli di Diffusione e il Flow Matching raggiungano questo obiettivo raffinando iterativamente i campioni durante il tempo di inferenza (decomponendo la complessa mappatura di pushforward in una catena di trasformazioni fattibili), questo approccio richiede molteplici valutazioni della rete neurale (NFE) per campione, portando a elevati costi computazionali.

I tentativi esistenti di creare generatori a singolo step si affidano spesso alla distillazione di modelli pre-addestrati multi-step o all'approssimazione delle traiettorie di equazioni differenziali stocastiche (SDE) o ordinarie (ODE) durante l'addestramento. Gli autori propongono un nuovo paradigma che sposta fondamentalmente il processo iterativo dal tempo di inferenza al tempo di addestramento, consentendo una generazione di alta qualità con una singola valutazione della funzione neurale (1-NFE) senza fare affidamento su formulazioni SDE/ODE.

Metodologia: Modelli di Drifting (Modelli di Deriva)

Il contributo principale è il Modello di Drifting, che formula la modellazione generativa come l'evoluzione della distribuzione di pushforward durante il processo di addestramento stesso.

1. Evoluzione durante l'Addestramento

Invece di vedere l'addestramento semplicemente come l'aggiornamento dei pesi per minimizzare una perdita statica, gli autori vedono la sequenza di modelli $\{f_i\}$ prodotti dall'ottimizzazione iterativa (ad esempio, SGD) come una sequenza di distribuzioni di pushforward $\{q_i\}$. Ad ogni iterazione, un campione $x_i = f_i(\epsilon)$ viene implicitamente "trascinato" (drifted) verso $x_{i+1} = f_{i+1}(\epsilon)$ a causa degli aggiornamenti dei parametri. L'obiettivo è definire un meccanismo in cui questa deriva guidi naturalmente la distribuzione $q$ verso $p_{data}$.

2. Il Campo di Drifting (Campo di Deriva)

Per governare questa evoluzione, gli autori introducono un campo di drifting $V_{p,q}(x)$, una funzione che calcola un vettore di spostamento $\Delta x$ per un campione $x$ basandosi sulla distribuzione generata corrente $q$ e sulla distribuzione target $p$.

• Condizione di Equilibrio: Il campo è progettato in modo che $V_{p,q}(x) = 0$ se e solo se $p = q$. Quando le distribuzioni corrispondono, i campioni smettono di derivare, raggiungendo un equilibrio.
• Anti-Simmetria: Il campo soddisfa la proprietà $V_{p,q}(x) = -V_{q,p}(x)$. Questa proprietà assicura che se le distribuzioni vengono scambiate, la direzione della deriva si inverte, garantendo che $V=0$ quando $p=q$.

3. Istanziazione tramite Attrazione e Repulsione

Il campo di drifting specifico utilizzato nel paper si ispira all'algoritmo di mean-shift ed è costruito utilizzando campioni positivi da $p$ e campioni negativi da $q$:
$$V_{p,q}(x) = V^+_p(x) - V^-_q(x)$$
Dove:

• $V^+_p(x)$ è una forza attrattiva che tira $x$ verso la distribuzione dei dati $p$.
• $V^-_q(x)$ è una forza repulsiva che spinge $x$ lontano dalla distribuzione generata $q$.
• Queste forze sono calcolate utilizzando una funzione kernel $k(x, y)$ (specificamente un kernel esponenziale basato sulla distanza $\ell_2$) e normalizzate tramite operazioni di softmax per garantire la stabilità.

4. Obiettivo di Addestramento

L'obiettivo di addestramento è derivato dalla relazione di punto fisso all'equilibrio: $f(\epsilon) = f(\epsilon) + V_{p,q}(f(\epsilon))$. La funzione di perdita minimizza la distanza tra la predizione corrente e un "target congelato" (la predizione più la deriva):
$$L = \mathbb{E}_{\epsilon} \left[ \| f_\theta(\epsilon) - \text{stopgrad}(f_\theta(\epsilon) + V_{p,q}(f_\theta(\epsilon))) \|^2 \right]$$
Questa formulazione permette all'ottimizzatore di far evolvere la distribuzione muovendo i campioni verso le loro versioni derivate. Fondamentalmente, il gradiente non si propaga all'indietro attraverso la distribuzione $q$ stessa, evitando le complessità del differenziare attraverso le distribuzioni.

5. Spazio delle Caratteristiche e Guida

• Spazio delle Caratteristiche: Per gestire dati ad alta dimensionalità, la perdita di drifting è calcolata in uno spazio di caratteristiche utilizzando un encoder auto-supervisionato pre-addestrato (ad esempio, un encoder stile ResNet-MAE). Ciò assicura che campioni semanticamente simili rimangano vicini, fornendo segnali di gradiente più ricchi.
• Classifier-Free Guidance (CFG): Il metodo supporta naturalmente il CFG mescolando campioni di dati non condizionati nella distribuzione dei campioni negativi durante l'addestramento, permettendo una generazione controllabile all'inferenza senza necessità di ri-addestramento.

Contributi Chiave

1. Nuovo Paradigma: Introduce i "Modelli di Drifting", un framework in cui la distribuzione di pushforward evolve durante il tempo di addestramento, eliminando la necessità di inferenza iterativa.
2. Campo di Drifting: Propone un nuovo meccanismo di perdita basato su un campo di drifting che raggiunge l'equilibrio quando le distribuzioni generate e dei dati corrispondono, governato dall'attrazione verso i dati e dalla repulsione dalla generazione corrente.
3. Generazione One-Step allo Stato dell'Arte: Dimostra che una rete a singolo passaggio può raggiungere prestazioni competitive con, o superiori a, modelli di diffusione e flow-based multi-step.
4. Robustezza al Mode Collapse: Esperimenti su scala ridotta suggeriscono che il metodo è robusto al mode collapse; se la distribuzione generata collassa su un singolo modo, l'attrazione da parte degli altri modi nella distribuzione dei dati continua a guidare la deriva, spingendo la distribuzione a evolversi.

Risultati Sperimentali

Gli autori valutano i Modelli di Drifting su ImageNet 256×256 sia nello spazio latente che in quello dei pixel.

• Spazio Latente: Utilizzando un tokenizer SD-VAE standard, il modello raggiunge un 1-NFE FID di 1.54 (con un modello L/2 più grande). Questo supera tutti i precedenti metodi one-step (ad esempio, iMeanFlow, AdvFlow) ed è competitivo con i modelli di diffusione multi-step (ad esempio, SiT-XL/2 con 250 step raggiunge 2.06 FID).
• Spazio dei Pixel: Senza un tokenizer latente, il modello raggiunge un 1-NFE FID di 1.61, superando sostanzialmente i precedenti metodi one-step nello spazio dei pixel (ad esempio, StyleGAN-XL con 2.30 FID) e competendo con i modelli multi-step nello spazio dei pixel.
• Controllo Robotico: Nel contesto della Diffusion Policy, il modello di Drifting a singolo step eguaglia o supera le prestazioni delle policy di diffusione multi-step (100 NFE) in vari compiti di manipolazione robotica.

Significato e Rivendicazioni

Il lavoro afferma di aprire nuove opportunità per una modellazione generativa di alta qualità ed efficiente. Riformulando l'addestramento iterativo di reti neurali come un meccanismo di evoluzione della distribuzione, il lavoro sfida la necessità di formulazioni SDE/ODE per i compiti generativi.

Gli autori sottolineano che il loro metodo:

• Ammette naturalmente l'inferenza a singolo step (1-NFE), riducendo significativamente il costo computazionale al momento dell'implementazione.
• Non si affida alla distillazione da modelli multi-step pre-addestrati, ma si addestra da zero.
• Offre un'alternativa concettualmente distinta alla diffusione e al flow matching, concentrandosi sulla regola di aggiornamento ($x_{i+1} = x_i + \Delta x_i$) durante l'addestramento piuttosto che durante l'inferenza.

Gli autori rimangono modesti riguardo alle garanzie teoriche, notando che, sebbene $p=q \implies V=0$ sia dimostrato, il inverso ($V \to 0 \implies p \to q$) è verificato empiricamente e sotto specifiche ipotesi di non degenerazione, ma non è garantito per campi vettoriali arbitrari. Considerano il loro lavoro come un punto di partenza per esplorare altri meccanismi di evoluzione della distribuzione.