An Optimal Transport-Based Generative Model for Bayesian Posterior Sampling

Probleemstelling

In de moderne Bayesiaanse statistiek is het uitvoeren van geldige statistische inferentie over modelparameters vaak computatief moeilijk. Het centrale obstakel is de berekening van de marginaal data dichtheid (de normalisatieconstante) in de Bayesiaanse stelling:
$$\pi_n(\theta) = \frac{\pi(\theta) p(Z^{(n)} | \theta)}{\int_\Theta \pi(\theta) p(Z^{(n)} | \theta) d\theta}$$
Deze integraal is wiskundig vaak onhandelbaar, tenzij er een geconjugeerde prior wordt gebruikt, wat de toepasbaarheid op complexe modellen beperkt.

Bestaande oplossingen hebben aanzienlijke nadelen:

1. MCMC (Markov Chain Monte Carlo): Genereert afhankelijke steekproeven en lijdt vaak aan trage mixing, vooral in hoge dimensies of bij sterk gestructureerde modellen.
2. Variational Inference (VI): Reformuleert het integratieprobleem als een optimalisatieprobleem. Hoewel dit schaalbaar is, garandeert het geen exacte steekproeven uit de doelverdeling, wat kan leiden tot vertekende schattingen en onbetrouwbare onzekerheidskwantificering.

Er is behoefte aan een methode die onafhankelijke, hoogwaardige steekproeven genereert zonder de normalisatieconstante te hoeven berekenen, en die zowel continu als discreet-mengsel ruimtes aankan.

Methodologie

De auteurs stellen een nieuwe generatieve modelbenadering voor die gebaseerd is op Optimal Transport (OT). Het doel is het leren van een deterministische afbeelding (transport map) $T^*$ die een eenvoudige referentieverdeling $\mu$ (bijv. een standaard Gaussische verdeling) transformeert naar de doel-posterior $\pi_n$.

Kernprincipes

1. Optimal Transport Beperking: In tegenstelling tot eerdere transport-map methoden die willekeurige afbeeldingen toestaan, beperkt deze methode $T^*$ tot de optimale transportkaart in de zin van Monge. Dit minimaliseert de verwachte kwadratische transportkosten:
   $$T^* = \arg \min_{T \in \mathcal{T}} \mathbb{E}_{X \sim \mu} [\|T(X) - X\|_2^2]$$
   waarbij $\mathcal{T}$ de klasse van afbeeldingen is die $\mu$ naar $\pi_n$ duwt. Deze beperking zorgt voor een unieke oplossing en een "parsimonious" (minimale) kaart.

2. Variational Formulier: Omdat het direct opleggen van de constraint $\pi_n = T_\#\mu$ analytisch onmogelijk is, minimaliseren de auteurs de Kullback-Leibler (KL) divergentie tussen de getransporteerde verdeling en de posterior:
   $$\min_{T \in \mathcal{T}} d_{KL}(T_\#\mu \parallel \pi_n)$$
   Door gebruik te maken van de invariantie van KL onder inverteerbare transformaties, wordt dit omgezet in een doelwit dat de Jacobiaan-determinant en de log-dichtheid van de posterior vereist, zonder de normalisatieconstante te hoeven kennen.

3. Structuur van de Transport Map:

   • Continue Parameters: Volgens de theorie van OT (Brenier's stelling) is de optimale kaart de gradiënt van een convexe functie ($T^* = \nabla u^*$). De auteurs benaderen $u^*$ met een familie van convexe functies die het maximum is van meerdere lokale convexe eenheden. Dit maakt de methode geschikt voor multimodale verdelingen.
   • Gemengde Discreet-Continue Parameters: Voor modellen met latente variabelen (discreet) en parameters (continu) wordt de OT kaart gefactoriseerd. De discrete component wordt bepaald door een "max"-operatie over lokale potentiaalfuncties (vergelijkbaar met een multiclass Probit-model), terwijl de continue component wordt getransformeerd via de gradiënt van de corresponderende lokale potentiaal.

4. Speciale Adaptaties:

   • Bij benadering Gaussische Posteriors: Een lineaire (affiene) transportkaart is voldoende, wat de rekenkosten drastisch verlaagt.
   • Bayesiaanse Latente Variabelen Modellen: Er wordt een mean-field benadering toegepast op de potentiaalfuncties om de exponentiële complexiteit van het discrete ruimte te reduceren, zonder de afhankelijkheid tussen variabelen volledig te negeren.

Belangrijkste Bijdragen

1. Unieke en Interpretabele Oplossing: Door de OT-structuur op te leggen, wordt de niet-uniekheid van algemene transportkaarten opgelost. De resulterende kaart is uniek en biedt een geometrisch interpreteerbaar pad tussen prior en posterior.
2. Efficiënte Steekproefneming: Eenmaal getraind, kan de kaart $T^*$ worden gebruikt om een willekeurig groot aantal onafhankelijke steekproeven te genereren, wat de correlatieproblemen van MCMC elimineert.
3. Nieuwe Inferentiële Hulpmiddelen: Omdat de OT-kaart de rangorde van quantielen behoudt (monotonie), kunnen er nieuwe tools worden ontwikkeld voor Bayesiaanse exploratieve analyse:
   • Multivariate Rangen en Quantielen: Centraal-uitwaartse quantielcontouren en tekenkrommen.
   • Bayesiaanse p-waarden: Een meetbare mate van extremiteit voor parameterwaarden in de parameter ruimte.
   • Simultane Credibele Intervallen: Efficiënte constructie van gezamenlijke onzekerheidsintervallen die afhankelijkheid tussen parameters in acht nemen.
4. Uitbreiding naar Gemengde Ruimtes: De methode is een van de eerste die OT succesvol toepast op modellen met zowel discrete als continue variabelen, een gebied waar eerdere methoden (zoals normale flows) vaak falen vanwege de inverteerbaarheidsvereiste.

Resultaten

De auteurs evalueren hun methode uitgebreid via simulaties en een real-world dataset (Gist/Yeast dataset):

• Multimodale Verdelingen: In vergelijking met Transport Monte Carlo (TMC) en Normalizing Flows (zoals Planar, MAF, NSF), presteert de OT-methode superieur bij het vangen van multimodale structuren. Normale flows hebben vaak moeite met goed gescheiden modi, terwijl TMC computatief duur wordt. De OT-methode behoudt de geometrie van de verdeling (bijv. niet-convexe vormen) en levert de laagste 2-Wasserstein-afstanden op.
• Bayesiaanse Logistische Regressie: De methode levert onzekerheidskwantificering (credibele intervallen) die zeer dicht bij die van MCMC ligt, maar met veel minder rekentijd. Het detecteert significante variabelen nauwkeuriger dan concurrenten.
• Gaussische Mengmodellen: Voor modellen met discrete latente variabelen (cluster toewijzingen) levert de methode nauwkeurige schattingen voor zowel de latente variabelen als de cluster-middelpunten, terwijl andere methoden vaak geen posteriors voor de latente variabelen kunnen genereren.
• Real-world Toepassing (Yeast Dataset): Bij het analyseren van genexpressiedata selecteert de methode dezelfde significante covariaten als MCMC, maar met een betere predictieve nauwkeurigheid dan andere transport-methoden.

Betekenis en Impact

Dit werk biedt een krachtig alternatief voor traditionele MCMC en Variational Inference methoden in de Bayesiaanse inferentie.

• Schaalbaarheid: Het biedt de snelheid van variational methods met de nauwkeurigheid van exacte steekproefneming.
• Exploratie: Het introduceert een nieuwe paradigmaverschuiving in hoe Bayesiaanse onzekerheid wordt gevisualiseerd en geanalyseerd, via OT-gebaseerde rangschikkingen en quantielen, wat dieper inzicht geeft in de geometrie van de posterior.
• Veelzijdigheid: De methode is robuust voor complexe scenario's, waaronder hoge dimensies, multimodaliteit en gemengde data-types, wat het een veelbelovende tool maakt voor moderne statistische modellering.

Kortom, de auteurs hebben een generatief model ontwikkeld dat de wiskundige elegantie van Optimal Transport combineert met praktische efficiëntie, waardoor het mogelijk wordt om complexe Bayesiaanse posteriors snel, nauwkeurig en interpreteerbaar te benaderen.