PolySwyft: sequential simulation-based nested sampling

TL;DR

PolySwyft tackles likelihood-free Bayesian inference for forward simulators by uniting nested sampling with neural ratio estimation in an NSNRE framework. It employs a dead-measure truncation driven by the observation and uses KL-based criteria to terminate sampling, enabling efficient, accurate recovery of multimodal posteriors while preserving Bayesian validity. Demonstrations on MVG, GMM, and a CosmoPower-driven CMB example show reduced simulator calls and robust posterior recovery compared to TNRE and standard NS baselines. The work provides open-source tooling and outlines concrete future directions, including marginalization, autoregressive models, and dynamic sampling strategies, with broad applicability in cosmology and beyond.

Abstract

We present PolySwyft, a novel, non-amortised simulation-based inference framework that unites the strengths of nested sampling (NS) and neural ratio estimation (NRE) to tackle challenging posterior distributions when the likelihood is intractable but a forward simulator is available. By nesting rounds of NRE within the exploration of NS, and employing a principled KL-divergence criterion to adaptively terminate sampling, PolySwyft achieves faster convergence on complex, multimodal targets while rigorously preserving Bayesian validity. On a suite of toy problems with analytically known posteriors of a dim(theta,D)=(5,100) multivariate Gaussian and multivariate correlated Gaussian mixture model, we demonstrate that PolySwyft recovers all modes and credible regions with fewer simulator calls than swyft's TNRE. As a real-world application, we infer cosmological parameters dim(theta,D)=(6,111) from CMB power spectra using CosmoPower. PolySwyft is released as open-source software, offering a flexible toolkit for efficient, accurate inference across the astrophysical sciences and beyond.

Figures (15)

• Figure 1: Nested Sampling Neural Ratio Estimation (NSNRE) meta-algorithm cycle. There are 5 distinct phases: 1. Sample from the prior for the initial training dataset 2. Use a forward simulator to sample joint $(\theta,D)$ 3. (Re-)Train an NRE on joint $\mathcal{J}$ and disjoint $\pi \times Z$ dataset 4. Use an NS on NRE to generate new samples $\theta$ around observation $D_{\mathrm{obs}}$. 5. Terminate if the KL-divergence criterion is fulfilled; otherwise, continue with step 2 (dashed arrow).
• Figure 2: Nested sampling algorithm with a generic Sampler.
• Figure 3: A typical dead points distribution for a parameter pair where one recursively zooms into the exponentially dense regions of dead points. The dead points have constant density in $\log X$, while the live points (larger coloured dots) have uniform density until termination. The plots were generated with code provided by hu_aeons_2023.
• Figure 4: PolySwyft algorithm
• Figure 5: A simple dynamic nested sampling mechanism that increases the number of live points $n_{\mathrm{live}}$ at the $99.9\%$ posterior contour (red) that was found using a quick initial run (blue). The x-axis is in negative $\log X$ prior volume contraction scale. Here, the initial blue run determines the posterior contours of the current ratio estimator that a second run leverages for dynamically adjusting the live points at a given contour. In principle, the live point profile can be adjusted to any profile.