PolySwyft: sequential simulation-based nested sampling

TL;DR

PolySwyft is a new simulation-based inference method combining nested sampling and neural ratio estimation, enabling efficient and accurate Bayesian inference for complex, intractable likelihood problems across various scientific fields.

Contribution

It introduces a non-amortised framework that adaptively terminates sampling using KL-divergence, improving convergence and accuracy over existing methods like TNRE.

Findings

Recovers all modes and credible regions with fewer simulator calls.

Demonstrates effectiveness on toy problems with known posteriors.

Successfully infers cosmological parameters from CMB data.

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.