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

TL;DR

This work addresses efficient posterior estimation in high-dimensional inverse problems where ground-truth posterior samples are unavailable. It proposes Likelihood-Weighted Normalizing Flows, training a normalizing flow with likelihood-based importance weights to directly map a simple base density to the target posterior. A key finding is that the base distribution's topology must align with the posterior multimodality; using Gaussian Mixture bases improves reconstruction fidelity and avoids spurious mode-bridges. Quantitative evaluations on 2D and 3D benchmarks using KL divergence and Wasserstein metrics support the method's effectiveness and highlight practical considerations for base-mode initialization in multi-modal posteriors.

Abstract

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.

Figures (8)

• Figure 1: The workflow, including training the NF and then sampling from the base to pass it through the flow
• Figure 2: Ground truth densities for the 2D benchmark tasks. From left to right: Single-mode, Two-mode, and Three-mode posteriors constructed from Gaussian mixtures.
• Figure 3: Modelled posterior distributions $q_\phi(\theta)$ generated by the Normalizing Flow. While the single-mode posterior (Left) is captured accurately, the multi-modal cases (Center, Right) exhibit spurious "bridges" connecting the modes, a result of the topological mismatch between the unimodal base and multi-modal target.
• Figure 4: Isosurface visualization of the ground truth 3-dimensional posterior distribution, displaying three distinct modes in the parameter space.
• Figure 5: 2D Marginal projections of the modelled 3D posterior. The flow captures the location of the modes, but exhibits connectivity artifacts due to the unimodal base distribution.