nbi: the Astronomer's Package for Neural Posterior Estimation

TL;DR

The paper introduces nbi, an open-source software framework that enhances neural posterior estimation for astronomy by providing built-in featurizers and an asymptotically exact inference algorithm, facilitating routine application to light curves and spectra.

Contribution

It presents a new framework and software, nbi, with built-in featurizers and a modified SNPE-IS algorithm for more accurate and user-friendly neural posterior estimation in astronomy.

Findings

nbi supports off-the-shelf application to astronomical data

Built-in featurizers improve handling of sequential data

SNPE-IS achieves asymptotically exact inference

Abstract

Despite the promise of Neural Posterior Estimation (NPE) methods in astronomy, the adaptation of NPE into the routine inference workflow has been slow. We identify three critical issues: the need for custom featurizer networks tailored to the observed data, the inference inexactness, and the under-specification of physical forward models. To address the first two issues, we introduce a new framework and open-source software nbi (Neural Bayesian Inference), which supports both amortized and sequential NPE. First, nbi provides built-in "featurizer" networks with demonstrated efficacy on sequential data, such as light curve and spectra, thus obviating the need for this customization on the user end. Second, we introduce a modified algorithm SNPE-IS, which facilities asymptotically exact inference by using the surrogate posterior under NPE only as a proposal distribution for importance sampling. These features allow nbi to be applied off-the-shelf to astronomical inference problems involving light curves and spectra. We discuss how nbi may serve as an effective alternative to existing methods such as Nested Sampling. Our package is at https://github.com/kmzzhang/nbi.