LtU-ILI: An All-in-One Framework for Implicit Inference in Astrophysics and Cosmology

TL;DR

LtU-ILI addresses the challenge of performing robust implicit likelihood inference in astronomy and cosmology by delivering a unified, extensible pipeline that supports multiple neural density estimators (NPE, NLE, NRE), diverse embedding architectures, and three-stage workflows (Data, Inference, Validation). The framework unifies existing SBI/ILI codes under a common interface, adds comprehensive validation metrics (including PIT and TARP), and demonstrates broad applicability through synthetic benchmarks and real science problems such as X-ray mass estimation, halo power spectra, and GW parameter inference. Its design enables rapid hyperparameter exploration, multi-round (sequential) inference, and the integration of graph- and image-based embeddings, making it a practical tool for both exploratory analysis and production-level inference. By providing public code and benchmarks, LtU-ILI aims to standardize and accelerate trustworthy ML-based inference across astronomical datasets and survey-era science.

Abstract

This paper presents the Learning the Universe Implicit Likelihood Inference (LtU-ILI) pipeline, a codebase for rapid, user-friendly, and cutting-edge machine learning (ML) inference in astrophysics and cosmology. The pipeline includes software for implementing various neural architectures, training schemata, priors, and density estimators in a manner easily adaptable to any research workflow. It includes comprehensive validation metrics to assess posterior estimate coverage, enhancing the reliability of inferred results. Additionally, the pipeline is easily parallelizable and is designed for efficient exploration of modeling hyperparameters. To demonstrate its capabilities, we present real applications across a range of astrophysics and cosmology problems, such as: estimating galaxy cluster masses from X-ray photometry; inferring cosmology from matter power spectra and halo point clouds; characterizing progenitors in gravitational wave signals; capturing physical dust parameters from galaxy colors and luminosities; and establishing properties of semi-analytic models of galaxy formation. We also include exhaustive benchmarking and comparisons of all implemented methods as well as discussions about the challenges and pitfalls of ML inference in astronomical sciences. All code and examples are made publicly available at https://github.com/maho3/ltu-ili.

Figures (10)

• Figure 1: Demonstration of the differences between a likelihood (top-left), posterior (top-right), prior (bottom-right), and joint distribution (bottom-left) for an arbitrary one-dimensional inference problem. Here, the univariate data $x$ and parameters $\theta$ are associated by a quadratic relationship with Gaussian scatter. Given the observed data point $x_o$ (in black), our goal is to recover posterior constraints on $\theta$. In an ILI framework, we seek to capture knowledge of the joint distribution, $P(x,\theta)$, by building neural-network emulators for the likelihood (in red) or the posterior (in blue).
• Figure 2: Overview of the LtU-ILI pipeline, displaying the procedural processes when running each stage. Each of the three stages (Data, Inference, Validation) is independently configurable, meaning any setup of one stage will automatically link to the others.
• Figure 3: Example SNPE posterior inference in comparison to classical implicit (Rejection ABC) and explicit (HMC) likelihood inference for the known toy simulator in Equation \ref{['eqn:toy']}. The true value for each parameter is shown in red and contours are shown at the central $68\%$ and $95\%$ confidence intervals. Note that the prior for each $\theta_i$ is a standard normal, as $p(\theta_i) = \mathcal{N}(0,1)$.
• Figure 4: Prediction error as a function of simulation budget for various LtU-ILI training methodologies on the benchmark SLCP inference problem papamakarios2019sequential. For clarity, plotted points for different models at the same simulation budget are slightly offset horizontally. The Rejection-ABC method is shown on all subplots as a baseline. Error is defined in terms of the Classifier 2-sample Test lopez2016revisiting metric. A lower C2ST indicates a more accurate inferred posterior, with 0.5 being the optimal score. C2ST values are shown at their median and central 95% confidence interval, calculated over ten independent runs.
• Figure 5: Inference of galaxy cluster $M_{500c}$ mass from eROSITA-like mock X-ray observations, as in ho2023benchmarks. The left four subplots show random examples of single-band X-ray observables in terms of their sky-projected surface brightness and target mass (in $h^{-1}M_\odot$). The right plot is the true vs. predicted log-mass estimates on the test set, where we show the median and central 68% confidence interval of our neural posteriors. For readability, test points in the right plot have been randomly subsampled.