Riepilogo Tecnico: Distillazione della Mappa di Transizione Stocastica per Inferenza Probabilistica Rapida

Enunciato del Problema

I modelli di diffusione hanno stabilito prestazioni all'avanguardia in termini di qualità di generazione, diversità e copertura della distribuzione rispetto ai predecessori come GAN e VAE. Tuttavia, queste prestazioni sono spesso accompagnate da costi di inferenza elevati a causa della necessità di integrare equazioni differenziali ordinarie (ODE) o equazioni differenziali stocastiche (SDE) a tempo continuo su molti passaggi.

Sebbene esistano tecniche di accelerazione esistenti come la distillazione (ad esempio, Modelli di Coerenza) e approcci basati sul Trasporto Ottimale, esse spesso affrontano un compromesso tra qualità/diversità di generazione e complessità di inferenza. Un limite critico dei metodi di distillazione attuali è la loro dipendenza dall'inferenza deterministica (ad esempio, tramite l'ODE del flusso di probabilità). Dato un campione di rumore fisso, l'output è unico. Ciò ne limita l'utilità in compiti a valle che richiedono intrinsecamente inferenza stocastica, come il campionamento posteriore di diffusione per problemi inversi, il fine-tuning basato sull'energia o il miglioramento dell'allineamento con le query degli utenti. Inoltre, molti framework di distillazione esistenti richiedono un modello insegnante preaddestrato, un'ottimizzazione a due livelli o la simulazione e la memorizzazione nella cache di traiettorie, il che complica l'addestramento e la scalabilità.

Metodologia: Distillazione della Mappa di Transizione Stocastica (STMD)

Gli autori propongono la Distillazione della Mappa di Transizione Stocastica (STMD), un framework senza insegnante progettato per accelerare l'inferenza dei modelli di diffusione preservando la struttura di transizione stocastica del processo sottostante.

Concetto Principale

A differenza dei modelli di diffusione basati su punteggi che tipicamente modellano la media della distribuzione posteriore, la STMD distilla la mappa di transizione completa associata all'SDE di campionamento. Il metodo impara un generatore condizionale che prevede la probabilità di transizione $p(x_{t'} | x_t)$ per $t' < t$ in uno o pochi passaggi, evitando la costosa integrazione numerica dell'SDE inverso.

Formulazione Tecnica

1. Flusso Condizionale Medio (MF): Il framework parametrizza le transizioni SDE utilizzando un modello di Flusso Medio Condizionale. Invece di apprendere una mappa di flusso deterministica, apprende la velocità media lungo la traiettoria del processo di diffusione.
2. Parametrizzazione x0: Per modellare la transizione $p(x_{t'} | x_t)$, gli autori utilizzano una parametrizzazione $x_0$. Esprimono la transizione come un integrale sulla posteriore $p(x_0 | x_t)$ e sulla misura di ponte nota $p(x_{t'} | x_0, x_t)$. Poiché la misura di ponte per le SDE a varianza conservativa (VP) è disponibile in forma chiusa, il problema si riduce all'apprendimento di un modello condizionale $p_\theta(x_0 | x_t)$.
3. Obiettivo di Addestramento: Il modello $u_\theta$ è addestrato per approssimare la velocità media condizionale. Gli autori definiscono una funzione di perdita Flusso Medio Condizionale (CMF):
   $$L_{dCMF} = \mathbb{E} \left[ \left\| u_\theta(z_s, r, s, x_t, t) - \text{sg}\left( (z_1 - x_0) - (s-r) \frac{d}{ds}u_\theta(z_s, r, s, x_t, t) \right) \right\|^2 \right]$$
   dove $z_s$ è un interpolante a linea retta tra i dati $x_0$ e un campione a priori $z_1$, condizionato allo stato di diffusione $x_t$ e al tempo $t$. L'operatore $\text{sg}$ indica il stop-gradient.
4. Addestramento Senza Insegnante: Il processo di addestramento non richiede un modello insegnante preaddestrato. Si basa su campioni dalla distribuzione congiunta $p(x_0, x_t)$, che possono essere generati campionando $x_0$ dalla distribuzione dei dati e applicando il kernel SDE in avanti (che è noto in forma chiusa).

Inferenza

L'inferenza viene eseguita integrando il Flusso Medio Condizionale appreso all'indietro nel tempo. Il processo comporta:

1. Campionamento di un rumore iniziale $x_1$ (o $z_1$).
2. Applicazione iterativa del campo di velocità appreso $u_\theta$ per passare dal tempo $s$ a $s-\Delta s$.
3. A ogni passo temporale di diffusione $t$, campionamento dalla distribuzione di ponte $p(x_{t-\Delta t} | x_0, x_t)$ per aggiornare lo stato, simulando efficacemente la transizione stocastica in grandi passaggi.

Contributi Chiave

1. Framework Innovativo: La proposta di STMD, un framework di distillazione basato su Flussi Medi Condizionali che consente l'apprendimento di mappe di transizione di SDE di diffusione ad alta dimensionalità utilizzando un generatore a uno o pochi passaggi tramite un obiettivo di regressione, senza alcun insegnante preaddestrato.
2. Convergenza Teorica: La derivazione dei limiti di convergenza per l'algoritmo del Flusso Medio nella distanza di Wasserstein-2. A conoscenza degli autori, questa è la prima analisi di convergenza dell'algoritmo del Flusso Medio in questa metrica. Estendono ciò per dimostrare che la loro mappa di transizione condizionale appresa converge al nucleo di transizione dell'SDE inverso sottostante.
3. Validazione Empirica: Dimostrazione del metodo su compiti di generazione e completamento di immagini sui dataset MNIST, CIFAR-10 e CelebA, mostrando prestazioni competitive rispetto alle linee di base pur mantenendo la stocasticità.

Risultati Sperimentali

Gli autori hanno valutato STMD su tre dataset:

• MNIST: Utilizzando un backbone U-Net, il modello è stato confrontato con l'algoritmo originale del Flusso Medio e un Modello Probabilistico di Diffusione Denoising (DDPM). Gli autori hanno utilizzato la Distanza di Fréchet (FD) nello spazio latente di un classificatore (poiché FID non è affidabile per MNIST). I risultati hanno mostrato che STMD performa in modo comparabile alle linee di base supportando al contempo il campionamento stocastico.
• CIFAR-10: Sono stati condotti esperimenti simili utilizzando un U-Net. Il modello ha ottenuto punteggi competitivi di Distanza di Inception di Fréchet (FID) rispetto a DDPM e alle linee di base del Flusso Medio per un numero fisso di iterazioni di addestramento, mantenendo la piena stocasticità del campionatore di diffusione.
• CelebA: Gli esperimenti hanno utilizzato un'architettura DiT latente. Il modello ha generato campioni incondizionati con un FID di 8,28 (utilizzando 4 passaggi di inferenza e 2 passaggi del Flusso Medio) ed ha eseguito con successo compiti di completamento di immagini.

In tutti i casi, i modelli sono stati addestrati da zero con gli stessi backbones di rete e conteggi di iterazioni delle linee di base per garantire un confronto equo.

Significato e Affermazioni

Il documento afferma che STMD colma il divario tra la distillazione deterministica rapida e la necessità di inferenza stocastica nei compiti a valle. Distillando la mappa di transizione invece della sola mappa di flusso, STMD preserva la struttura stocastica del processo di diffusione.

Gli autori sottolineano che il loro approccio è senza insegnante, eliminando la necessità di ottimizzazione a due livelli o memorizzazione nella cache di traiettorie, rendendo il processo di addestramento più efficiente e scalabile. Il contributo teorico di stabilire limiti di convergenza di Wasserstein-2 fornisce una base rigorosa per l'approccio del Flusso Medio, colmando un divario tra garanzie teoriche e applicazione pratica.

Il lavoro si posiziona come una soluzione per compiti che richiedono inferenza probabilistica, come problemi inversi e campionamento posteriore, dove i metodi di distillazione deterministica risultano insufficienti. Gli autori rimangono modesti riguardo alle prestazioni attuali, notando che il numero limitato di iterazioni di addestramento impedisce risultati all'avanguardia, ma lo studio comparativo valida l'efficacia del metodo rispetto alle linee di base esistenti.