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

Technical Summary: Learned Harmonic Mean Estimation with Flow Matching for Multimodal Posteriors

Problem Statement
The marginal likelihood (or Bayesian evidence), denoted as $z = p(d|M)$, is a fundamental quantity for Bayesian model comparison, enabling quantitative assessment via the Bayes factor. However, computing this quantity requires evaluating a multidimensional integral over the entire parameter space, which is computationally prohibitive for complex, high-dimensional models. While Markov chain Monte Carlo (MCMC) and Hamiltonian Monte Carlo (HMC) methods excel at generating posterior samples in high dimensions, they do not inherently provide the marginal likelihood. Conversely, nested sampling provides the evidence but restricts the sampling strategy to a specific nested hierarchy.

The "learned harmonic mean estimator" was previously introduced to estimate the marginal likelihood using only posterior samples, regardless of the sampling method used to obtain them. This approach relies on learning an internal target density $\phi(\theta)$ that is contained within the posterior but has thinner tails. However, previous implementations using discrete normalizing flows or bespoke density estimators struggled with highly multimodal posteriors. These architectures often failed to capture complex topologies without heuristic modifications (e.g., manually constructing Gaussian mixtures for the base distribution) or suffered from assigning non-negligible probability mass to density bridges between modes, leading to infinite variance in the estimator.

Methodology
This work addresses the limitations of previous density estimators by integrating Flow Matching into the learned harmonic mean framework. The methodology proceeds as follows:

1. Learned Harmonic Mean Estimator: The estimator approximates the reciprocal marginal likelihood $\rho = z^{-1}$ using posterior samples $\theta_i \sim p(\theta|d)$ and a learned target density $\phi(\theta)$:
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   For the estimator to have finite variance, $\phi(\theta)$ must be concentrated within the high-probability regions of the posterior and decay faster in the tails.

2. Flow Matching Architecture: Instead of discrete normalizing flows, the authors employ Continuous Normalizing Flows (CNFs) trained via flow matching.

   • Definition: A CNF parametrizes the transformation between a simple base distribution (standard Gaussian) and the target posterior via a time-dependent velocity field $v_t(\theta)$.
   • Training Objective: The network is trained using a simulation-free regression objective. A conditional probability path is defined as a straight line between a noise sample $\theta_0$ and a posterior sample $\theta_1$: $\theta_t = (1-t)\theta_0 + t\theta_1$. The network minimizes the mean squared error between the predicted velocity field and the true conditional vector field $(\theta_1 - \theta_0)$.
   • Advantages: This approach allows for architecturally unrestricted, high-capacity neural networks (Multi-Layer Perceptrons with Swish activation) and avoids the expensive ODE simulations required during training in earlier CNF methods.

3. Temperature Scaling: To ensure the target density $\phi(\theta)$ has thinner tails than the posterior, a temperature parameter $T \in (0, 1)$ is applied to scale the variance of the flow's base distribution. This concentrates the flow distribution, satisfying the containment requirement for the harmonic mean estimator.

Key Contributions

• Novel Architecture: The introduction of flow matching-based continuous normalizing flows as the internal density estimator for the learned harmonic mean method.
• Handling Multimodality: Demonstrating that this architecture can accurately capture complex, highly multimodal posterior structures without the need for heuristic modifications to the base distribution or fine-tuning.
• Scalability: Providing a method that is agnostic to the sampling strategy (compatible with MCMC, HMC, NUTS, SBI, etc.) and capable of operating in moderate-to-high dimensional spaces.

Experimental Results
The authors validated the approach on two challenging benchmarks:

1. Rastrigin Function (2 Dimensions):

   • A standard benchmark with many local modes.
   • The flow matching model (6 layers, 256 hidden units) successfully captured all modes with a temperature $T=0.98$.
   • Results from 100 experimental runs showed that the estimated marginal likelihood values were accurate and closely matched the ground truth obtained via numerical integration. The error estimates were reliable.

2. 20-Dimensional Gaussian Mixture:

   • A mixture of five Gaussian components in 20 dimensions with non-overlapping modes and introduced correlations.
   • The model (10 layers, 256 hidden units) was trained on 40,000 samples.
   • The learned target distribution was shown to be properly concentrated within the posterior, capturing the challenging topology.
   • The estimated marginal likelihoods across 100 runs aligned with the analytic ground truth, confirming the method's accuracy and robust error estimation even in this complex, moderate-dimensional scenario.

Significance and Claims
The paper claims that by leveraging the expressivity of flow matching-trained CNFs, the learned harmonic mean estimator can now robustly estimate the evidence for highly complex, multimodal posteriors. The primary significance lies in removing the reliance on heuristic adjustments (such as hand-picked Gaussian mixtures) that were previously necessary to handle multimodality. The authors assert that this approach enables accurate evidence calculation using samples from any accelerated or scalable sampling method, facilitating Bayesian model comparison in data-rich fields like astrophysics without the computational bottlenecks of traditional nested sampling or the instability of the original harmonic mean estimator.

Future work, as noted by the authors, includes applying this methodology to real-world problems such as gravitational wave astrophysics and exploring further architectural refinements for very high-dimensional tasks like imaging and field-level cosmological inference.