Latent Stochastic Interpolants

1. Il Problema

I modelli di diffusione e gli Stochastic Interpolants (SI) hanno dimostrato un grande successo nella generazione di dati complessi, permettendo di trasformare una distribuzione semplice (prior) in una distribuzione complessa (dati) tramite processi stocastici continui. Tuttavia, l'uso degli SI presenta due limitazioni fondamentali quando si applicano a modelli a variabili latenti:

1. Mancanza di ottimizzazione congiunta: Gli SI tradizionali richiedono l'accesso diretto ai campioni delle due distribuzioni (prior e target) nello spazio delle osservazioni. Questo impedisce il loro utilizzo in modelli a variabili latenti dove l'encoder, il decoder e il generatore devono essere ottimizzati congiuntamente in uno spazio latente non osservato.
2. Inefficienza computazionale: Applicare gli SI direttamente nello spazio delle osservazioni ad alta dimensionalità (es. immagini ad alta risoluzione) è computazionalmente costoso.
3. Prior rigide: I modelli di diffusione standard spesso sono vincolati a prior semplici (es. Gaussiana isotropa), limitando la flessibilità della generazione.

Il lavoro mira a colmare il divario permettendo l'apprendimento congiunto di un encoder, un decoder e un modello generativo SI all'interno di uno spazio latente, mantenendo la flessibilità degli SI e riducendo i costi computazionali.

2. Metodologia: Latent Stochastic Interpolants (LSI)

Gli autori propongono Latent Stochastic Interpolants (LSI), un framework che estende gli SI allo spazio latente attraverso un approccio basato su un Lower Bound della Verosimiglianza (ELBO) derivato in tempo continuo.

A. Formulazione Teorica

Il metodo si basa su due pilastri teorici:

1. ELBO per variabili latenti dinamiche: Utilizzando il teorema di Girsanov, gli autori derivano un ELBO per processi stocastici definiti da Equazioni Differenziali Stocastiche (SDE). L'obiettivo massimizza la verosimiglianza dei dati osservati penalizzando la divergenza tra il percorso del modello e quello approssimato (variational posterior).
2. Ponte di Diffusione (Diffusion Bridge): Per evitare la simulazione costosa delle SDE durante l'addestramento, viene costruita un'approssimazione variazionale che permette il campionamento diretto delle variabili latenti $z_t$ a qualsiasi tempo $t$, senza integrare l'SDE passo dopo passo.

B. Costruzione dell'Interpolante Latente

Invece di interpolare direttamente tra i dati osservati, LSI interpola tra:

• Una distribuzione prior $p_0(z_0)$ nello spazio latente.
• La distribuzione aggregata approssimata $p_\theta(z_1) = \int p_\theta(z_1|x)dx$ definita dall'encoder.

L'interpolante $z_t$ è parametrizzato come:
$$z_t = \eta_t \epsilon + \kappa_t z_1 + \nu_t z_0$$
dove $\epsilon \sim \mathcal{N}(0, I)$, e $z_1$ è l'encoding stocastico dell'immagine $x$. Questa formulazione permette di campionare $z_t$ direttamente dato $z_0$ e $z_1$, rendendo l'addestramento senza simulazione (simulation-free).

C. Funzione di Perdita (Loss Function)

L'obiettivo di addestramento è un ELBO che combina:

1. Termine di Ricostruzione: $-\ln p_\theta(x_1|z_1)$, che garantisce che il decoder possa ricostruire l'immagine dal latente.
2. Termine di Dinamica (Flow Matching): Una perdita che allinea il campo di velocità appreso $h_\theta$ con il campo teorico derivato dal ponte di diffusione.

La perdita totale è:
$$\mathcal{L} = \mathbb{E} \left[ -\ln p_\theta(x_1|z_1) + \frac{\beta_t}{2} \| \text{target} - h_\theta(z_t, t) \|^2 \right]$$
Dove $\beta_t$ è un peso che bilancia ricostruzione e generazione. Gli autori introducono una parametrizzazione chiamata InterpFlow, che si è dimostrata più stabile e performante rispetto ad altre varianti (come NoisePred o Denoising).

D. Campionamento e Prior Flessibili

• Campionamento: Il modello supporta sia campionamento deterministico (tramite ODE di flusso di probabilità) che stocastico. È possibile utilizzare la Classifier-Free Guidance (CFG) per migliorare la qualità dei campioni condizionati.
• Prior Arbitrarie: A differenza dei modelli di diffusione standard, LSI non è vincolato a prior Gaussiane. Gli esperimenti mostrano che il framework funziona efficacemente con prior Uniformi, Laplaciane o Mixture di Gaussiane.

3. Contributi Chiave

1. LSI Framework: Un nuovo metodo per l'addestramento end-to-end di modelli generativi a variabili latenti con dinamiche temporali continue, dove encoder, decoder e generatore sono ottimizzati congiuntamente.
2. Prospettiva Unificante: Integrazione della formulazione continua degli SI all'interno dei modelli a variabili latenti, sfruttando i processi stocastici continui per derivare un obiettivo di training principato.
3. Obiettivo ELBO Principato: Un nuovo ELBO che mantiene i vantaggi degli SI (addestramento senza simulazione, scelta flessibile del prior) abilitando al contempo l'ottimizzazione congiunta nello spazio latente.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul benchmark standard ImageNet a diverse risoluzioni (64x64, 128x128, 256x256).

• Prestazioni (FID): I modelli LSI raggiungono FID (Fréchet Inception Distance) comparabili a quelli dei modelli SI addestrati direttamente nello spazio delle osservazioni (pixel), nonostante operino in uno spazio latente a dimensionalità ridotta.
  • Esempio a 128x128: LSI ottiene un FID di 3.12, contro 3.46 del modello nello spazio delle osservazioni.
• Efficienza Computazionale:
  • Durante il campionamento, l'encoder non viene utilizzato e il decoder viene chiamato una sola volta. Il modello latente (più piccolo) viene eseguito ripetutamente.
  • Rispetto ai modelli nello spazio delle osservazioni, LSI offre un risparmio significativo di FLOPs (Floating Point Operations) durante il campionamento: ~30% in meno a 128x128 e ~65% in meno a 256x256 per 100 step di campionamento.
• Ottimizzazione Congiunta: L'analisi del parametro di bilanciamento $\beta$ mostra che l'addestramento congiunto è cruciale. Un $\beta$ troppo basso (che isola l'encoder) porta a prestazioni peggiori, mentre un $\beta$ ottimizzato migliora la qualità generativa (riduzione del FID del ~17%).
• Robustezza al Prior: Il modello mantiene buone prestazioni anche con prior non Gaussiane (es. Laplaciana, Uniforme), dimostrando la flessibilità del framework.

5. Significato e Impatto

Il lavoro di Latent Stochastic Interpolants rappresenta un passo avanti significativo nell'evoluzione dei modelli generativi:

• Superamento dei limiti degli SI: Risolve il problema fondamentale che impediva l'uso degli SI in modelli a variabili latenti, aprendo la strada a generative modeling più efficiente.
• Efficienza Scalabile: Dimostra che è possibile ottenere alta qualità generativa riducendo drasticamente il costo computazionale del campionamento, rendendo i modelli più accessibili per applicazioni reali.
• Flessibilità Teorica: Fornisce un quadro teorico unificato che combina i vantaggi dei modelli di diffusione, dei flussi normali (Normalizing Flows) e dei modelli variazionali, permettendo la scelta di prior complesse e dinamiche temporali flessibili.

In sintesi, LSI offre un approccio scalabile e principato per la generazione di immagini ad alta risoluzione, combinando l'efficienza degli spazi latenti con la flessibilità e la potenza teorica degli Stochastic Interpolants.