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

Technische Zusammenfassung: Schätzung des harmonischen Mittels mittels Flow Matching für multimodale Posteriori-Verteilungen

Problemstellung
Die marginale Likelihood (oder Bayes-Evidenz), bezeichnet als $z = p(d|M)$, ist eine fundamentale Größe für den Bayesschen Modellvergleich und ermöglicht eine quantitative Bewertung mittels des Bayes-Faktors. Die Berechnung dieser Größe erfordert jedoch die Auswertung eines mehrdimensionalen Integrals über den gesamten Parameterraum, was für komplexe, hochdimensionale Modelle rechnerisch prohibitiv ist. Während Markov-Chain-Monte-Carlo- (MCMC) und Hamiltonian-Monte-Carlo-Methoden (HMC) exzellent darin sind, Posteriori-Stichproben in hohen Dimensionen zu generieren, liefern sie nicht inhärent die marginale Likelihood. Umgekehrt bietet das Nested Sampling zwar die Evidenz, schränkt jedoch die Sampling-Strategie auf eine spezifische hierarchische Struktur ein.

Der „gelernte harmonische Mittelwert-Schätzer“ wurde zuvor eingeführt, um die marginale Likelihood unter Verwendung ausschließlich von Posteriori-Stichproben zu schätzen, unabhängig von der verwendeten Sampling-Methode. Dieser Ansatz stützt sich auf das Erlernen einer internen Zieldichte $\phi(\theta)$, die innerhalb der Posteriori-Verteilung liegt, jedoch dünnere Enden (Tails) aufweist. Frühere Implementierungen unter Verwendung diskreter Normalisierende Flüsse oder spezialisierter Dichteschätzer hatten jedoch Schwierigkeiten mit hochgradig multimodalen Posteriori-Verteilungen. Diese Architekturen konnten komplexe Topologien oft nicht ohne heuristische Modifikationen (z. B. die manuelle Konstruktion von Gaußschen Mischungen für die Basisverteilung) erfassen oder wiesen eine nicht vernachlässigbare Wahrscheinlichkeitsmasse in den Dichteverbindungen zwischen den Modi auf, was zu einer unendlichen Varianz des Schätzers führte.

Methodik
Diese Arbeit adressiert die Einschränkungen früherer Dichteschätzer durch die Integration von Flow Matching in das Framework des gelernten harmonischen Mittelwerts. Die Methodik verläuft wie folgt:

1. Gelernter harmonischer Mittelwert-Schätzer: Der Schätzer approximiert die reziproke marginale Likelihood $\rho = z^{-1}$ unter Verwendung von Posteriori-Stichproben $\theta_i \sim p(\theta|d)$ und einer gelernten Zieldichte $\phi(\theta)$:
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   Damit der Schätzer eine endliche Varianz besitzt, muss $\phi(\theta)$ in den Regionen hoher Wahrscheinlichkeit der Posteriori konzentriert sein und in den Enden schneller abfallen.

2. Flow-Matching-Architektur: Anstelle diskreter normalisierender Flüsse verwenden die Autoren Continuous Normalizing Flows (CNFs), die mittels Flow Matching trainiert werden.

   • Definition: Eine CNF parametrisiert die Transformation zwischen einer einfachen Basisverteilung (Standard-Gauß-Verteilung) und der Ziel-Posteriori über ein zeitabhängiges Geschwindigkeitsfeld $v_t(\theta)$.
   • Trainingsziel: Das Netzwerk wird mittels eines simulationsfreien Regressionsziels trainiert. Ein konditionaler Pfad der Wahrscheinlichkeit wird als gerade Linie zwischen einer Rauschprobe $\theta_0$ und einer Posteriori-Probe $\theta_1$ definiert: $\theta_t = (1-t)\theta_0 + t\theta_1$. Das Netzwerk minimiert den mittleren quadratischen Fehler zwischen dem vorhergesagten Geschwindigkeitsfeld und dem wahren konditionalen Geschwindigkeitsfeld $(\theta_1 - \theta_0)$.
   • Vorteile: Dieser Ansatz ermöglicht architektonisch uneingeschränkte, hochkapazitive neuronale Netze (Multi-Layer-Perceptrons mit Swish-Aktivierung) und vermeidet die teuren ODE-Simulationen, die in früheren CNF-Methoden erforderlich waren.

3. Temperatur-Skalierung: Um sicherzustellen, dass die Zieldichte $\phi(\theta)$ dünnere Enden als die Posteriori-Verteilung hat, wird ein Temperaturparameter $T \in (0, 1)$ angewendet, um die Varianz der Flow-Basisverteilung zu skalieren. Dies konzentriert die Flow-Verteilung und erfüllt die Anforderung der Eingrenzung (Containment).

Zentrale Beiträge

• Neuartige Architektur: Die Einführung von Flow-Matching-basierten Continuous Normalizing Flows als interner Dichteschätzer für die Methode des gelernten harmonischen Mittelwerts.
• Umgang mit Multimodalität: Nachweis, dass diese Architektur komplexe, hochgradig multimodale Posteriori-Strukturen präzise erfassen kann, ohne dass heuristische Modifikationen der Basisverteilung oder Feinabstimmungen erforderlich sind.
• Skalierbarkeit: Bereitstellung einer Methode, die gegenüber der Sampling-Strategie agnostisch ist (kompatibel mit MCMC, HMC, NUTS, SBI etc.) und in moderaten bis hohen Dimensionen operieren kann.

Experimentelle Ergebnisse
Die Autoren validierten den Ansatz an zwei anspruchsvollen Benchmarks:

1. Rastrigin-Funktion (2 Dimensionen):

   • Ein Standard-Benchmark mit vielen lokalen Modi.
   • Das Flow-Matching-Modell (6 Schichten, 256 Einheiten) erfasste erfolgreich alle Modi mit einer Temperatur von $T=0,98$.
   • Ergebnisse aus 100 Testläufen zeigten, dass die geschätzten Werte der marginalen Likelihood akkurat waren und eng mit der durch numerische Integration ermittelten Grundwahrheit übereinstimmten. Die Fehlerschätzungen waren zuverlässig.

2. 20-Dimensionale Gaußsche Mischung:

   • Eine Mischung aus fünf Gaußschen Komponenten in 20 Dimensionen mit nicht überlappenden Modi und eingeführten Korrelationen.
   • Das Modell (10 Schichten, 256 Einheiten) wurde mit 40.000 Stichproben trainiert.
   • Die gelernte Zielverteilung konnte gezeigt werden, dass sie korrekt innerhalb der Posteriori konzentriert ist und die anspruchsvolle Topologie erfasst.
   • Die geschätzten marginalen Likelihoods über 100 Läufe hinweg stimmten mit der analytischen Grundwahrheit überein, was die Genauigkeit der Methode und die robuste Fehlerschätzung selbst in diesem komplexen, moderat-dimensionalen Szenario bestätigt.

Bedeutung und Ansprüche
Das Paper behauptet, dass durch die Nutzung der Ausdrucksstärke von durch Flow Matching trainierten CNFs der gelernte harmonische Mittelwert-Schätzer nun in der Lage ist, die Evidenz für hochkomplexe, multimodale Posterioris robust zu schätzen. Die primäre Bedeutung liegt in der Beseitigung der Abhängigkeit von heuristischen Anpassungen (wie handverlesenen Gaußschen Mischungen), die zuvor zur Handhabung von Multimodalität notwendig waren. Die Autoren behaupten, dass dieser Ansatz eine genaue Berechnung der Evidenz unter Verwendung von Stichproben aus beliebigen beschleunigten oder skalierbaren Sampling-Methoden ermöglicht, was die Bayesische Modellvergleich in datenreichen Feldern wie der Astrophysik ohne die Rechenengpässe des traditionellen Nested Samplings oder die Instabilität des ursprünglichen harmonischen Mittelwert-Schätzers erleichtert.

Zukünftige Arbeiten, wie von den Autoren angemerkt, umfassen die Anwendung dieser Methodik auf reale Probleme wie die Gravitationswellen-Astrophysik sowie die Untersuchung weiterer architektonischer Verfeinerungen für sehr hochdimensionale Aufgaben wie die Bildgebung oder die kosmologische Inferenz auf Feldebene.