Learned harmonic mean estimation of the marginal likelihood for multimodal posteriors with flow matching

Technische Samenvatting: Geschatte Harmonisch Gemiddelde met Flow Matching voor Multimodale Posteriorverdelingen

Probleemstelling
De marginale waarschijnlijkheid (of de Bayesiaanse bewijslast), aangeduid als $z = p(d|M)$, is een fundamentele grootheid voor Bayesiaanse modelvergelijking, wat een kwantitatieve beoordeling via de Bayes-factor mogelijk maakt. Het berekenen van deze grootheid vereist echter het evalueren van een meerdimensionale integraal over de gehele parameterruimte, wat computationeel onhaalbaar is voor complexe, hoogdimensionale modellen. Hoewel Markov Chain Monte Carlo (MCMC) en Hamiltonian Monte Carlo (HMC) methoden uitblinken in het genereren van posterior-stalen in hoge dimensies, bieden zij niet inherent de marginale waarschijnlijkheid. Omgekeerd biedt nested sampling de bewijslast, maar beperkt het de bemonsteringsstrategie tot een specifieke geneste hiërarchie.

De "learned harmonic mean estimator" werd eerder geïntroduceerd om de marginale waarschijnlijkheid te schatten met behulp van enkel posterior-stalen, ongeacht de gebruikte bemonsteringsmethode. Deze aanpak rust op het leren van een interne doel-densiteit $\phi(\theta)$ die binnen de posterior ligt, maar dunnere staarten heeft. Eerdere implementaties met discrete normalizing flows of specifieke densiteitsestimators hadden echter moeite met sterk multimodale posteriorverdelingen. Deze architecturen slaagden er vaak niet in om complexe topologieën te vangen zonder heuristische modificaties (bijv. het handmatig construeren van Gaussian mixtures voor de basisverdeling) of leden onder het toekennen van een niet-verwaarloosbare waarschijnlijkheidsmassa aan densiteitsbruggen tussen modi, wat leidde tot een oneindige variantie in de estimator.

Methodologie
Dit werk adresseert de beperkingen van eerdere densiteitsestimators door Flow Matching te integreren in het learned harmonic mean-raamwerk. De methodologie verloopt als volgt:

1. Learned Harmonic Mean Estimator: De estimator benadert de reciproke marginale waarschijnlijkheid $\rho = z^{-1}$ met behulp van posterior-stalen $\theta_i \sim p(\theta|d)$ en een geleerde doel-densiteit $\phi(\theta)$:
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   Voor een eindige variantie van de estimator moet $\phi(\theta)$ geconcentreerd zijn binnen de gebieden met een hoge waarschijnlijkheid van de posterior en sneller afnemen in de staarten.

2. Flow Matching Architectuur: In plaats van discrete normalizing flows gebruikt de auteur Continuous Normalizing Flows (CNFs) getraind via flow matching.

   • Definitie: Een CNF parametriseert de transformatie tussen een eenvoudige basisverdeling (standaard Gaussische verdeling) en de doel-posterior via een tijdsafhankelijke snelheidveld $v_t(\theta)$.
   • Trainingsdoel: Het netwerk wordt getraind met een simulation-free regressiedoel. Een conditionele probabiliteitsroute wordt gedefinieerd als een rechte lijn tussen een ruis-monster $\theta_0$ en een posterior-monster $\theta_1$: $\theta_t = (1-t)\theta_0 + t\theta_1$. Het netwerk minimaliseert de gemiddelde kwadratische fout tussen het voorspelde snelheidveld en het ware conditionele snelheidveld $(\theta_1 - \theta_0)$.
   • Voordelen: Deze aanpak maakt architecturaal onbeperkte, hoog-capacitieve neurale netwerken (Multi-Layer Perceptrons met Swish-activatie) mogelijk en vermijdt de dure ODE-simulaties die vereist waren in eerdere CNF-methoden tijdens de training.

3. Temperatuur Schaling: Om te waarborgen dat de doel-densiteit $\phi(\theta)$ dunnere staarten heeft dan de posterior, wordt een temperatuurparameter $T \in (0, 1)$ toegepast om de variantie van de flow-basisverdeling te schalen. Dit concentreert de flow-verdeling, waardoor aan de containment-eis van de harmonic mean estimator wordt voldaan.

Belangrijkste Bijdragen

• Nieuwe Architectuur: De introductie van flow matching-gebaseerde continuous normalizing flows als de interne densiteitsestimator voor de learned harmonic mean methode.
• Omgaan met Multimodaliteit: Het aantonen dat deze architectuur complexe, sterk multimodale posterior-structuren accuraat kan vangen zonder de noodzaak van heuristische modificaties aan de basisverdeling of fijnafstemming.
• Schaalbaarheid: Het bieden van een methode die agnostisch is aan de bemonsteringsstrategie (compatibel met MCMC, HMC, NUTS, SBI, etc.) en in staat is te opereren in matig- tot hoogdimensionale ruimtes.

Experimentele Resultaten
De auteurs hebben de aanpak gevalideerd op twee uitdagende benchmarks:

1. Rastrigin Functie (2 Dimensies):

   • Een standaard benchmark met veel lokale modi.
   • Het flow matching model (6 lagen, 256 verborgen eenheden) ving alle modi succesvol met een temperatuur $T=0.98$.
   • Resultaten van 100 experimentele runs toonden aan dat de geschatte marginale waarschijnlijkheidswaarden accuraat waren en nauw aansloten bij de grondwaarheid verkregen via numerieke integratie. De foutmarges waren betrouwbaar.

2. 20-Dimensionale Gaussian Mixture:

   • Een mengeling van vijf Gaussische componenten in 20 dimensies met niet-overlappende modi en geïntroduceerde correlaties.
   • Het model (10 lagen, 256 verborgen eenheden) werd getraind op 40.000 stalen.
   • De geleerde doelverdeling bleek correct geconcentreerd binnen de posterior, waarbij de uitdagende topologie werd gevangen.
   • De geschatte marginale waarschijnlijkheden over 100 runs stemden overeen met de analytische grondwaarheid, wat de nauwkeurigheid en robuuste foutschatting van de methode bevestigt, zelfs in dit complexe, matig-dimensionale scenario.

Betekenis en Claims
Het artikel claimt dat door gebruik te maken van de expressiviteit van flow matching-getrainde CNFs, de learned harmonic mean estimator nu robuust de bewijslast voor zeer complexe, multimodale posteriors kan schatten. De primaire betekenis ligt in het wegnemen van de afhankelijkheid van heuristische aanpassingen (zoals handgekozen Gaussian mixtures) die voorheen noodzakelijk waren om met multimodaliteit om te gaan. De auteurs stellen dat deze aanpak een accurate berekening van de bewijslast mogelijk maakt met stalen van elke versnelde of schaalbare bemonsteringsmethode, wat Bayesiaanse modelvergelijking in datarijke velden zoals astrofysica faciliteert zonder de computationele knelpunten van traditionele nested sampling of de instabiliteit van de oorspronkelijke harmonic mean estimator.

Toekomstig werk, zoals door de auteurs opgemerkt, omvat het toepassen van deze methodologie op real-world problemen zoals gravitatiegolf-astrofysica en het verkennen van verdere architecturale verfijningen voor zeer hoogdimensionale taken zoals beeldvorming en veld-niveau kosmologische inferentie.