An Optimal Transport-Based Generative Model for Bayesian Posterior Sampling

1. Il Problema

L'inferenza bayesiana moderna si scontra con la difficoltà di campionare da distribuzioni posteriori $\pi_n(\theta)$ quando la costante di normalizzazione (la densità marginale dei dati) è intrattabile.

• Limiti degli approcci attuali:
  • MCMC (Markov Chain Monte Carlo): Sebbene asintoticamente esatto, soffre di tempi di mixing lenti, specialmente in spazi ad alta dimensionalità o con strutture complesse, generando campioni altamente correlati.
  • Inferenza Variazionale (VI): Scalabile ma fornisce solo approssimazioni della densità target, non garantendo il campionamento esatto e potenzialmente introducendo bias nelle stime e nella quantificazione dell'incertezza.
  • Mappe di Trasporto Esistenti: I metodi basati su mappe di trasporto (come i flussi normalizzanti o le mappe triangolari) spesso non impongono vincoli strutturali sufficienti per garantire l'unicità della mappa, portando a instabilità computazionale o inefficienza.

2. Metodologia

Il paper propone un nuovo approccio generativo basato sulla Trasporto Ottimo (Optimal Transport - OT) per apprendere una mappa deterministica $T^*$ che spinge una distribuzione di riferimento semplice $\mu$ (es. Gaussiana standard) verso la distribuzione target $\pi_n$.

Formulazione del Problema

L'obiettivo è trovare la mappa $T^*$ che minimizza il costo di trasporto quadratico (problema di Monge):
$$T^* = \arg \min_{T \in \mathcal{T}} \mathbb{E}_{X \sim \mu} [\|T(X) - X\|^2_2]$$
soggetto al vincolo $T\#\mu = \pi_n$ (la mappa spinge $\mu$ su $\pi_n$).

Poiché il vincolo è difficile da imporre direttamente quando $\pi_n$ è noto solo a meno di una costante, il problema viene riformulato come ottimizzazione della divergenza Kullback-Leibler (KL):
$$\min_{T \in \mathcal{T}} d_{KL}(T\#\mu \parallel \pi_n)$$
Grazie all'invarianza della KL sotto trasformazioni invertibili, questo è equivalente a massimizzare:
$$\max_{T \in \mathcal{T}} \mathbb{E}_{\mu} [\log(\pi_n \circ T) + \log |\det(J_T)|]$$
dove $J_T$ è la matrice Jacobiana.

Struttura della Mappa di Trasporto

La contribuzione chiave risiede nella definizione della classe di funzioni $\mathcal{T}$, vincolata dalla teoria del trasporto ottimo per garantire l'unicità e l'efficienza:

1. Parametri Continui:

   • Secondo il teorema di Brenier, la mappa OT ottima è il gradiente di una funzione convessa ($T^* = \nabla u^*$).
   • La classe $\mathcal{T}$ è costruita come il gradiente di una funzione convessa approssimata come il massimo di funzioni convesse locali (unità convesse):
     $$u(x) = \max_{k \in [L]} u_k(x), \quad \text{dove } u_k(x) = \sum_{m=1}^M f(x; \alpha_m, \beta_m, w_m, v_m)$$
   • Le funzioni $f$ sono unità convesse basate su integrali di funzioni univariate crescenti e limitate (es. Tanh, Softsign). L'uso dell'operatore max permette di catturare naturalmente distribuzioni multimodali ben separate.

2. Parametri Misti (Continui e Discreti):

   • Per modelli con variabili latenti discrete (es. modelli a mistura), la mappa non è più invertibile.
   • Il paper estende la teoria OT a variabili miste, mostrando che la potenziale ottima ha la forma:
     $$u^*(x^{(1)}, x^{(2)}) = \kappa \max_{k \in [K]} \{ \langle x^{(1)}, b_k \rangle + \phi_k(x^{(2)}) \}$$
     dove $x^{(1)}$ è la parte discreta (codificata tramite embedding $b_k$) e $x^{(2)}$ è la parte continua.
   • La mappa risultante seleziona dinamicamente il componente discreto $\tau$ e applica una mappa locale $\nabla \phi_\tau$ per la parte continua.

3. Approssimazioni Specifiche:

   • Posteriori quasi-Gaussiani: Si può restringere la classe a mappe affini, riducendo drasticamente il costo computazionale mantenendo l'accuratezza statistica (teorema di Bernstein-von Mises).
   • Modelli Latenti: Viene introdotta un'approssimazione "mean-field" sulla funzione potenziale per gestire spazi discreti esponenzialmente grandi (es. modelli a mistura gaussiane), decomponendo la funzione globale in una somma di funzioni locali indipendenti.

3. Contributi Chiave

• Unicità e Stabilità: A differenza dei flussi normalizzanti generici, l'imposizione della struttura OT (gradiente di potenziale convesso) garantisce l'unicità della mappa appresa, migliorando la stabilità dell'addestramento e l'interpretabilità.
• Gestione della Multimodalità: L'uso dell'operatore max sulle potenziali locali permette di modellare distribuzioni con supporti sconnessi o multimodali in modo naturale, superando i limiti delle mappe triangolari o dei flussi planari che tendono a fallire su questi casi.
• Supporto per Variabili Miste: Estensione formale della teoria OT a spazi misti (discreto-continuo), cruciale per modelli bayesiani con variabili latenti categoriali.
• Nuovi Strumenti Inferenziali: La natura della mappa OT (che preserva l'ordinamento dei ranghi) abilita la costruzione di:
  • Quantili Multivariati Center-Outward: Contorni di quantile che rispettano la geometria della distribuzione target.
  • Ranghi e Segni: Strumenti per l'analisi esplorativa bayesiana e la visualizzazione delle dipendenze.
  • P-value Bayesiani: Una misura di "estremità" di un parametro nello spazio dei parametri basata sulla distanza dal centro nella distribuzione di riferimento.

4. Risultati Sperimentali

Gli autori hanno valutato il metodo su diversi scenari simulati e su dati reali (dataset del lievito):

• Miscele Gaussiane (Multimodalità): Il metodo OT supera significativamente i flussi normalizzanti (Planar, MAF, NSF) e il Transport Monte Carlo (TMC) nella stima di distribuzioni multimodali ad alta dimensionalità, mostrando errori di Wasserstein ($W_2$) inferiori e una migliore capacità di catturare tutti i modi della distribuzione.
• Regressione Logistica Bayesiana: Il metodo produce intervalli credibili e selezioni di variabili coerenti con l'MCMC (benchmark), ma con tempi di calcolo molto ridotti. L'analisi dei ranghi center-outward rivela correlazioni tra parametri che le analisi marginali nascondono.
• Modelli a Mistura Gaussiane (Latenti): Il metodo "mean-field" proposto riesce a stimare sia le variabili latenti discrete che i parametri continui, cosa che i flussi normalizzanti standard non possono fare direttamente.
• Efficienza: Una volta appresa la mappa $T^*$, il campionamento di un numero arbitrario di campioni indipendenti è istantaneo, offrendo un vantaggio computazionale rispetto alle catene MCMC.

5. Significato e Implicazioni

Questo lavoro rappresenta un ponte significativo tra la teoria del trasporto ottimo e l'inferenza bayesiana pratica.

• Superamento del Compromesso: Offre un'alternativa che combina la velocità e l'indipendenza dei campioni tipica dei metodi generativi con la precisione e l'interpretabilità spesso riservate ai metodi basati su catene di Markov.
• Nuova Visione Statistica: Trasforma il problema del campionamento in un problema di ottimizzazione strutturata, permettendo non solo di generare campioni, ma di derivare nuove statistiche descrittive (ranghi, quantili multivariati) che arricchiscono l'analisi esplorativa bayesiana.
• Scalabilità: La capacità di gestire spazi misti e multimodali rende questo approccio promettente per modelli bayesiani complessi e moderni che erano finora difficili da analizzare in modo efficiente.

In sintesi, il paper propone un framework unificato che utilizza i vincoli geometrici del trasporto ottimo per creare un generatore di posteriori bayesiani robusto, unico e ricco di informazioni inferenziali.