24–28 Aug 2026
Kirchhoff Institute for Physics (KIP)
Europe/Berlin timezone

Accelerating Cosmological Inference: A Re-usable Library of JAX-based Machine Learning Emulators for Next-Generation Surveys

26 Aug 2026, 14:50
8m
HS1

HS1

Inference & Uncertainty 🔀 Inference & Uncertainty

Speaker

Dily Duan Yi Ong (University of Cambridge)

Description

The era of precision cosmology has revealed persistent tensions between independent measurements of fundamental parameters within the concordance model, most notably the Hubble constant ($H_0$), the clustering amplitude ($\sigma_8$), and spatial curvature ($\Omega_K$). The study of these discrepancies between independent datasets, which are theoretically predicted to agree, is known as tension quantification. Upcoming astronomical surveys, such as the Euclid mission and the Vera C. Rubin Observatory, will deliver data of unprecedented precision and volume, providing new insights into these tensions. However, conducting joint-constraint analyses with existing legacy datasets will require marginalising over tens of instrumental nuisance parameters, rendering traditional sampling methods computationally prohibitive.

We approach this problem by producing a comprehensive, re-usable library of machine learning emulators trained on a massive grid of nested sampling chains across 8 cosmological models and 12 astronomical surveys. Trained using DiRAC GPUs (DP456), this framework is released as part of the pip-installable package $\texttt{unimpeded}$ [2511.04661, 2511.05470] (\url{https://github.com/handley-lab/unimpeded}). Serving as a machine-learning-enhanced analogue to the Planck Legacy Archive (PLA), it enables rapid parameter estimation, cosmological model comparison and tension quantification.

The emulators are implemented with normalising flows using the JAX-native density estimation package $\texttt{margarine}$ [2205.12841]. By learning the marginal posterior directly from these pre-existing nested sampling chains, the normalising flows act as ultra-fast ``nuisance-free likelihoods'' and informative true priors. This eliminates the need to re-sample legacy instrumental systematics (e.g. calibration, foregrounds) or rely on standard Gaussian approximations. The combination of rigorous sampling and density estimation reproduces the exact posterior distributions of a full nuisance-marginalised run, but many orders of magnitude faster. Furthermore, hyperparameter tuning for these normalising flows has been systematically explored with different combinations of network architecture, learning scheduling and activation functions for optimal performance.

We believe this work represents a significant step forward in cosmological data analysis. By collapsing multi-survey joint analyses down to just the cosmological subspace, it provides a versatile, efficient and equitable platform to address current observational tensions and advance our understanding of the Universe.

Author

Dily Duan Yi Ong (University of Cambridge)

Co-authors

Dr Harry Bevins (University of Cambridge) Dr Will Handley (University of Cambridge)

Presentation materials