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

We present a novel technique for amortized posterior estimation using Normalizing Flows trained with likelihood-weighted importance sampling. This approach allows for the efficient inference of theoretical parameters in high-dimensional inverse problems without the need for posterior training samples. We implement the method on multi-modal benchmark tasks in 2D and 3D to check for the efficacy. A critical observation of our study is the impact of the topology of the base distributions on the modelled posteriors. We find that standard unimodal base distributions fail to capture disconnected support, resulting in spurious probability bridges between modes. We demonstrate that initializing the flow with a Gaussian Mixture Model that matches the cardinality of the target modes significantly improves reconstruction fidelity, as measured by some distance and divergence metrics.

💡 Research Summary

The paper introduces a novel amortized inference framework that leverages likelihood‑weighted importance sampling to train Normalizing Flows (NFs) without ever requiring samples from the true posterior. Traditional simulation‑based inference (SBI) methods such as Neural Posterior Estimation (NPE) depend on large batches of simulated data drawn from the posterior, which is often infeasible when the likelihood is expensive to evaluate. Instead, the authors draw a static dataset of parameters {θ_i} from a simple prior (typically uniform or a standard Gaussian) and compute the corresponding data likelihoods L(θ_i)=p(D|θ_i). These likelihood values serve as importance weights. The training objective becomes the weighted negative log‑likelihood

 L(ϕ)=−(1/N)∑_{i=1}^N L(θ_i) log q_ϕ(θ_i),

where q_ϕ(θ) is the density induced by the flow f_ϕ mapping a base latent distribution p_Z(z) to the parameter space. By expanding the KL divergence between the true posterior p(θ|D) and q_ϕ(θ), the authors show that minimizing this weighted loss is mathematically equivalent to minimizing KL(p(θ|D)‖q_ϕ(θ)). Consequently, the flow learns to approximate the posterior directly from prior samples weighted by their likelihoods, achieving amortized inference: once trained, sampling from the flow yields independent posterior draws at negligible cost.

A central theoretical insight is that a Normalizing Flow is a diffeomorphism; it preserves the topological properties of the base distribution. If the base distribution is unimodal and connected, the transformed density cannot represent a truly disconnected support. This leads to “probability bridges” that artificially connect distinct modes of a multimodal posterior. The authors empirically demonstrate this effect on synthetic 2‑D and 3‑D benchmark problems with one, two, and three Gaussian mixture modes. With a standard Gaussian base, the flow captures the locations of the modes but creates low‑density bridges between them. Quantitatively, KL divergence remains modest (<0.05) while the average Wasserstein‑1 distance grows substantially for multimodal cases, indicating a mismatch in the geometry of the reconstructed distribution.

To remedy the topological mismatch, the paper proposes initializing the base distribution as a Gaussian Mixture Model (GMM) whose number of components matches the true number of posterior modes. By providing a base with the same disconnected support, the flow can map each component to a distinct posterior mode without being forced to create bridges. Experiments with “Model‑2D‑k” and “Model‑3D‑k” (k=1,2,3) confirm that when k equals the true number of modes, the spurious connections disappear and both Wasserstein‑1 and W‑sliced KL metrics improve dramatically (e.g., Model‑2D‑3 achieves W₁≈0.08 versus 0.13 for the unimodal base). The loss values across different base configurations remain nearly identical, underscoring that the improvement stems purely from better topological alignment rather than a change in optimization difficulty.

The authors also test non‑Gaussian target posteriors (products of non‑Gaussian factors) while still using Gaussian‑mixture bases. Even in this more challenging setting, matching the number of mixture components yields the best reconstruction metrics, demonstrating that the approach is robust to deviations from Gaussianity as long as the support topology is respected.

Overall contributions:

1. Likelihood‑Weighted NF Training – A simple, gradient‑based loss that directly incorporates data likelihoods, enabling posterior approximation without posterior samples.
2. Topological Analysis – Formal discussion of how the diffeomorphic nature of flows imposes connectivity constraints, leading to identifiable failure modes for multimodal posteriors.
3. Mode‑Matching Base Distributions – Empirical evidence that initializing the base with a GMM whose component count matches the true posterior modes eliminates artificial bridges and substantially improves distributional fidelity.

The method is computationally efficient: the expensive likelihood evaluations are performed once during dataset construction, after which training proceeds like any standard NF. Once trained, the model provides amortized samples for any new observation without further simulation. This makes the approach attractive for high‑dimensional scientific inference problems where forward simulations are costly (e.g., particle physics phenomenology, cosmological parameter estimation, or complex systems modeling).