Sintesi Tecnica: Dinamiche Stocastiche Data-to-Energy

Definizione del Problema

Il problema del ponte di Schrödinger (SB) mira a trovare un sistema dinamico stocastico che colleghi due distribuzioni marginali, $p_0$ e $p_1$, minimizzando un costo di trasporto rispetto a un processo di riferimento (tipicamente un SDE di Itô). Questo problema generalizza il trasporto ottimo nell'ambito stocastico e costituisce la base per moderne tecniche di modellazione generativa come i modelli di diffusione e il flow matching.

Gli algoritmi esistenti per risolvere il problema SB, in particolare quelli basati sulla procedura di Iterative Proportional Fitting (IPF), richiedono generalmente l'accesso a campioni da entrambe le distribuzioni marginali. Questa limitazione impedisce l'applicazione dei metodi SB in scenari in cui una o entrambe le distribuzioni sono definite solo da densità non normalizzate (funzioni di energia) $p(x) = e^{-E(x)}/Z$, dove la funzione di partizione $Z$ è sconosciuta e non sono disponibili campioni diretti. Questo divario è critico per applicazioni come l'inferenza bayesiana dei posteriori negli spazi latenti di modelli generativi, dove il posteriore è spesso noto solo a meno di una costante di normalizzazione tramite una funzione di ricompensa o verosimiglianza.

Metodologia

Il documento propone il Data-to-Energy (e Energy-to-Energy) IPF, un metodo generale per modellare ponti di Schrödinger quando le distribuzioni sono fornite come densità non normalizzate. L'approccio estende l'algoritmo IPF standard integrando tecniche di apprendimento per rinforzo off-policy e campionamento di diffusione.

1. Framework IPF Generalizzato

L'algoritmo IPF standard aggiorna iterativamente i processi in tempo diretto e inverso ($\overrightarrow{P}_t$ e $\overleftarrow{P}_t$) risolvendo problemi di semi-ponte.

• Passaggio Backward (Data-to-Data): Quando sono disponibili campioni da $p_0$, il processo backward viene addestrato tramite massima verosimiglianza (passaggio IPF standard).
• Passaggio Forward (Data-to-Energy): Quando $p_1$ è data da una densità non normalizzata $e^{-E_1(x)}$, il processo forward non può essere addestrato tramite la verosimiglianza standard. Invece, gli autori introducono una perdita di log-varianza (conditional VarGrad). Questa perdita minimizza la varianza del logaritmo del rapporto tra la probabilità della traiettoria forward e la densità target, pesata dalla probabilità della traiettoria inversa:
  $$\text{LLV}(x_0, \theta) = \text{Var} \left( \log \frac{\overrightarrow{p}_\theta(\tau | x_0)}{p_1(x_1) \overleftarrow{p}_\phi(\tau | x_1)} \right)$$
  Questo obiettivo impone la proporzionalità richiesta dalla soluzione SB senza necessitare di campioni da $p_1$.

2. Strategie di Apprendimento per Rinforzo Off-Policy

Per affrontare la sfida del campionamento da densità non normalizzate complesse e ad alta dimensionalità dove l'addestramento on-policy fallisce nel coprire tutti i modi, gli autori adattano tecniche off-policy:

• Replay Buffer: Un buffer memorizza campioni finali $x_1$ dal processo forward. Durante l'addestramento, $x_0$ viene generato campionando $x_1$ dal buffer ed eseguendo una traiettoria inversa. Ciò focalizza l'esplorazione sulle regioni ad alta densità del target.
• Aggiornamenti Langevin: Il buffer viene periodicamente aggiornato utilizzando la dinamica di Langevin non aggiustata sulla densità target $p_1$ per correggere fit imperfetti.
• Riutilizzo delle Traiettorie: L'algoritmo riutilizza le traiettorie inverse per formare batch di traiettorie che condividono un punto iniziale $x_0$, migliorando l'efficienza della stima della perdita basata sulla varianza.
• Mixing: La politica di addestramento mescola campioni on-policy con campioni off-policy dal buffer per bilanciare esplorazione e sfruttamento.

3. Coefficienti di Diffusione Apprendibili

Un contributo metodologico secondario riguarda la generalizzazione dell'algoritmo IPF per apprendere non solo le funzioni di drift, ma anche i coefficienti di diffusione ($\sigma_t$) delle dinamiche. Sostituendo le varianze fisse nei passi discretizzati di Euler-Maruyama con funzioni apprendibili, il metodo mira a correggere gli errori di discretizzazione temporale, un concetto ispirato ai recenti lavori sui campionatori di diffusione.

Contributi Chiave

1. Primo Metodo SB Data-to-Energy Generale: Il documento introduce il primo algoritmo capace di inferire ponti di Schrödinger quando una o entrambe le distribuzioni sono fornite come densità non normalizzate senza accesso a campioni. Esso apprende con successo ponti stocastici con basso costo di trasporto tra dataset sintetici e densità, con prestazioni pari al data-to-data IPF che utilizza campioni reali.
2. Coefficienti di Diffusione Apprendibili: Gli autori dimostrano che gli algoritmi data-to-data IPF esistenti possono essere sostanzialmente migliorati apprendendo il coefficiente di diffusione oltre al drift. Questo generalizza le recenti scoperte per i campionatori di diffusione all'intero contesto SB, mostrando come la varianza apprendibile permetta una modellazione più accurata, specialmente con discretizzazioni temporali grossolane.
3. Campionamento del Posteriore Outsourced: Il metodo viene applicato al problema del campionamento di distribuzioni posteriori negli spazi latenti di modelli generativi (es. GAN, VAE). Modellando un ponte tra la distribuzione prior del latente e un posteriore pesato da una ricompensa, gli autori creano un metodo di traduzione immagine-immagine data-free. Questo approccio preserva il contenuto semantico (stile, sfondo) non vincolato dal target, a differenza dei comuni campionatori di diffusione.

Risultati Sperimentali

L'efficacia del metodo proposto è dimostrata attraverso diversi benchmark:

• Benchmark Sintetici 2D: Su compiti come Gaussian to Gaussian Mixture Models (GMM) e distribuzioni "Two Moons", il Data-to-Energy IPF raggiunge prestazioni comparabili al Data-to-Data IPF. L'inclusione della varianza apprendibile riduce significativamente la distanza Wasserstein-2 e la Path KL, specialmente quando il numero di passi di discretizzazione è ridotto.
• Bridging Energy-to-Energy: I risultati preliminari mostrano che l'algoritmo può collegare con successo due distribuzioni definite esclusivamente da funzioni di energia, senza accesso a campioni da nessuna delle due.
• Traduzione di Immagini (MNIST & CIFAR-10): In esperimenti di spazio latente utilizzando SN-GAN e StyleGAN, il metodo traduce campioni prior verso posteri di classe target (es. generare "cani" o "cinque").
  • Metriche: Il metodo ottiene punteggi FID inferiori rispetto al rejection sampling e alla dinamica di Langevin in diverse classi.
  • Qualitativo: L'ispezione visiva conferma che il ponte preserva le caratteristiche di stile (es. spessore del tratto, orientamento dello sfondo) pur trasformando il contenuto nella classe target, suggerendo un'efficace preservazione dell'informazione semantica non vincolata.
• Studi di Ablazione: Gli esperimenti confermano che i replay buffer, gli aggiornamenti di Langevin e il riutilizzo delle traiettorie migliorano significativamente le metriche di Path KL. Il riutilizzo delle traiettorie backward migliora la Path KL ma può impattare negativamente il mean log-reward, suggerendo un compromesso tra copertura dei modi e costo di trasporto.

Significato e Rivendicazioni

Il documento rivendica di fornire un framework generale per i ponti di Schrödinger data-to-energy, colmando una lacuna nella letteratura dove i metodi SB erano precedentemente limitati a scenari con accesso ai campioni per entrambe le marginali.

Gli autori posizionano il loro lavoro come un passo significativo verso la modellazione generativa data-free scalabile e l'inferenza dei posteriori. Sfruttando tecniche di apprendimento per rinforzo off-policy, il metodo consente di apprendere la dinamica del trasporto ottimo in spazi latenti ad alta dimensionalità dove sono disponibili solo funzioni di energia. La capacità di apprendere i coefficienti di diffusione potenzia ulteriormente la robustezza di queste dinamiche contro gli errori di discretizzazione.

Gli autori notano con modestia che, sebbene il metodo sia scalabile a compiti di generazione di immagini ad alta dimensionalità, è necessario ulteriore lavoro per affrontare il mode collapse e migliorare ulteriormente la copertura dei modi. Identificano inoltre potenziali direzioni future, come il fine-tuning della ricompensa condizionato al testo e i problemi di ponte di Schrödinger discreti, ma non rivendicano un successo immediato in queste aree inesplorate.