⚛️ high-energy experiments

Amortized Inference of Multi-Modal Posteriors using Likelihood-Weighted Normalizing Flows

This paper introduces an amortized inference method using likelihood-weighted normalizing flows that overcomes the limitations of standard unimodal base distributions by initializing with a Gaussian Mixture Model, thereby accurately capturing multi-modal, non-Gaussian posteriors in high-dimensional inverse problems as demonstrated in multi-modal benchmarks and a heavy flavor physics application.

Original authors: Rajneil Baruah

Published 2026-08-04

📖 3 min read🧠 Deep dive

CC BY 4.0

Original authors: Rajneil Baruah

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery, but you don't have the crime scene photos. Instead, you only have a list of suspects (the "prior") and a magic mirror that tells you how likely each suspect is to be the culprit based on the evidence (the "likelihood"). Your goal is to figure out the true profile of the criminal. In the world of physics and finance, scientists face this exact problem every day: they need to find hidden parameters that explain the data they see. Usually, they do this by running millions of slow, computer-heavy simulations to map out the "posterior"—the final, most likely shape of the truth. But this takes forever, like trying to find a needle in a haystack by checking every single piece of hay one by one.

Recently, scientists started using a clever trick called "Normalizing Flows." Think of this as a magical shape-shifter. It starts with a simple, smooth blob of clay (a standard distribution) and learns to stretch, twist, and squish it until it perfectly matches the complex, jagged shape of the truth. The problem is, usually, you need to show the shape-shifter a bunch of finished examples of the truth to teach it how to do the job. But what if you don't have those examples? What if you only have the magic mirror and the list of suspects? This is the puzzle the paper tackles: Can we teach the shape-shifter to find the truth using only the likelihood scores, without ever seeing the final answer?

The author of this paper say, "Yes, we can!" They developed a new way to train these shape-shifters, which they call "Likelihood-Weighted Normalizing Flows." Instead of feeding the model finished answers, they feed it random guesses and tell the model, "This guess is good, so stretch the clay toward it; that guess is bad, so ignore it." They found that this method works incredibly well, but with one catch: the starting clay matters. If you start with a single, round ball of clay (a simple Gaussian), and the truth is actually two separate islands of land, the clay must stretch a bridge between them to connect the islands. This creates a fake "bridge" of probability where there shouldn't be one. However, if you start with a clay model that already has two separate lumps (a Gaussian Mixture), the shape-shifter can perfectly match the two islands without building a fake bridge.

To prove this, the researchers tested their method on some tricky math puzzles and then on a real, complex problem in particle physics: figuring out the secrets of a particle called the B meson. This particle behaves in a way that creates two distinct possibilities for its properties, and the "truth" isn't a perfect 50/50 split between them; one possibility is actually about 81% likely, while the other is only 19%. When they used their new method, it correctly identified both possibilities and got the 81/19 split almost exactly right, matching the results of the slow, traditional computer methods but doing it in a fraction of the time. The study shows that while a simple starting shape can work if the truth is very lopsided, matching the "topology" (the number of separate pieces) of your starting clay to the number of pieces in the truth is the secret to getting the most accurate map. This means scientists can now solve these complex, multi-part mysteries much faster, without needing to wait weeks for the computer to catch up.

Technical Summary: Amortized Inference of Multi-Modal Posteriors using Likelihood-Weighted Normalizing Flows

Problem Statement

Scientific inquiry across high-energy physics, cosmology, and finance often relies on solving inverse problems: inferring the posterior distribution of theoretical parameters given observational data. Traditional methods like Markov Chain Monte Carlo (MCMC) and Nested Sampling are statistically robust but suffer from the "curse of dimensionality," often requiring weeks or months to converge in high-dimensional spaces with expensive likelihood evaluations.

While Machine Learning (ML) offers alternatives like Normalizing Flows (NFs) for efficient sampling, standard NF training relies on Maximum Likelihood Estimation (MLE), which requires a dataset of samples drawn from the target posterior. In many scientific contexts, such as Simulation-Based Inference (SBI), these posterior samples are unavailable; researchers typically possess only a prior distribution and a "black-box" simulator to evaluate likelihoods. Training NFs on prior samples alone causes the model to trivially reconstruct the prior, failing to capture likelihood-informed structures. Furthermore, capturing complex, multi-modal, and non-Gaussian posterior topologies remains a significant challenge for standard flow architectures.

Methodology

The paper proposes a novel framework for Amortized Posterior Estimation using Likelihood-Weighted Normalizing Flows (LW-NF). The core innovation is training a Normalizing Flow directly on samples drawn from the prior, weighted by their likelihood values, rather than on posterior samples.

Theoretical Framework

The method reframes posterior inference as a density estimation problem. By minimizing the Kullback-Leibler (KL) divergence between the true posterior $p(\theta|D)$ and the flow approximation $q_\phi(\theta)$ , the author derives a loss function equivalent to a likelihood-weighted negative log-likelihood:
$\mathcal{L}(\phi) = -\frac{1}{N} \sum_{i=1}^N L(\theta_i) \log q_\phi(\theta_i)$
where $\theta_i \sim p(\theta)$ (the prior) and $L(\theta_i) = p(D|\theta_i)$ is the likelihood. This approach allows the flow to learn the transformation from a simple base distribution to the complex posterior without ever observing a "true" posterior sample.

Topological Constraints and Base Distributions

A critical theoretical insight of the work is the impact of the base distribution's topology on the model's ability to represent multi-modal posteriors. Since Normalizing Flows are diffeomorphisms (bijective mappings), they preserve the connectivity of the base distribution.

Unimodal Base: If a standard unimodal Gaussian is used as the base for a multi-modal target, the flow is mathematically forced to create "spurious probability bridges" connecting the disconnected modes to maintain topological continuity.
Multi-modal Base: The author demonstrates that initializing the flow with a Gaussian Mixture Model (GMM) where the number of components matches the cardinality of the target modes significantly improves reconstruction fidelity, eliminating these artificial bridges.

Key Contributions

Likelihood-Weighted Training Objective: The paper establishes a method to train Normalizing Flows using only prior samples and pointwise likelihood evaluations, bypassing the need for expensive posterior sampling or large simulation budgets typical of Neural Posterior Estimation (NPE).
Topological Analysis: The study provides a systematic analysis of how the modality of the base distribution affects the reconstruction of multi-modal posteriors. It proves that matching the cardinality of the base distribution to the target modes is essential for high-fidelity reconstruction, particularly in avoiding topological artifacts.
Physics Application: The method is applied to a complex, real-world problem in heavy flavor physics: the extraction of Wolfenstein parameters ( $\bar{\rho}, \bar{\eta}, |V_{cb}|, g$ ) from the CP asymmetry in $B^0 \to J/\psi K^0$ . This problem is characterized as simultaneously multi-modal, non-Gaussian, and asymmetric in mode weights.

Results

The efficacy of the method was validated on synthetic 2D and 3D benchmarks and the 4D physics problem.

Synthetic Benchmarks:
- On multi-modal tasks, flows with unimodal (Gaussian) bases exhibited significant Wasserstein distance errors and visual "bridges" between modes.
- Flows initialized with GMM bases matching the target mode count achieved reconstruction fidelity at the statistical noise floor, with metrics (KL divergence, Wasserstein distance) indistinguishable from the ground truth.
- The method performed better on Gaussian posteriors than non-Gaussian ones, though the trend of matching base cardinality held true for both.
Physics Application ( $B^0 \to J/\psi K^0$ ):
- The true posterior is bimodal with an asymmetric weight ratio of approximately 81% to 19% (due to the Jacobian of the prior transformation), not the standard 50/50 split.
- MCMC Reference: A well-converged MCMC run established the reference posterior with a mode fraction $f_1 \approx 0.809$ .
- LW-NF Performance: Both the single-Gaussian base ( $k=1$ ) and the two-component GMM base ( $k=2$ ) successfully recovered the two modes with the correct asymmetric weights ( $f_1 \approx 0.824$ and $0.819$, respectively).
- Metrics: The flows achieved mean marginal Jensen-Shannon divergences of $\sim 10^{-3}$ and Wasserstein distances of $\sim 7 \times 10^{-3}$ relative to the MCMC reference.
- Base Comparison: While the GMM-2 base was marginally more accurate on minority-mode localization, the single-Gaussian base was deemed "adequate" for this specific problem because the strong weight asymmetry ( $81/19$ ) allowed the unimodal base to grow a "modest tail" for the secondary mode rather than needing to split mass evenly.

Significance and Claims

The paper claims that the LW-NF approach offers a powerful amortized inference framework. Once trained, the flow generates posterior samples via a single forward pass, enabling near-instantaneous posterior evaluation for repeated analyses (e.g., global CKM fits), a task where MCMC would require re-equilibrating chains for every new dataset.

The author emphasizes that the method's ability to reproduce the physically correct asymmetric mode weights (81/19) rather than a symmetric split serves as a rigorous validation of its fidelity. The work suggests that while a single Gaussian base can suffice for highly asymmetric posteriors, topological alignment (matching base and target modality) is a critical design principle for robust inference in general multi-modal scenarios. The paper concludes that this approach occupies a "sweet spot" between likelihood-accessible paradigms and simulation-based inference, providing a faster tool for resolving multi-modal structures in high-dimensional scientific problems.