CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators

Spurio Mancini et al. (2022) (SM22 hereafter), the exploration of high-dimensional parameter spaces in particular, developed CosmoPower, a suite of - O(100) parameters and higher - necessary to accurately neural network emulators of cosmological power spectra model the physical signals and their several systematic that replaces the computation of these quantities traditionally contaminants. Sampling the posterior distribution performed with Einstein-Boltzmann solvers such in these high-dimensional spaces represents a significant as the Code for Anisotropies in the Microwave Background computational challenge for Markov Chain Monte Carlo (CAMB, Lewis & Challinor 2011) or the Cosmic (MCMC) algorithms (Roberts et al. 1997; Katafygiotis Linear Anisotropy Solving System (CLASS, Blas et al. & Zuev 2008; Liu 2009), which are traditionally used 2011). In SM22 the authors show how Bayesian inference in cosmological analyses (Lewis & Bridle 2002; Audren of cosmological parameters can be accelerated by several et al. 2013; Brinckmann & Lesgourgues 2019; Torrado & orders of magnitude using CosmoPower; the speed-up becomes Lewis 2021). Gradient-based inference methods, such as particularly relevant when the emulators are employed Hamiltonian Monte Carlo (HMC, Duane et al. 1987; Neal within an inference pipeline that can be run on 1996) and variational inference (VI, Hoffman et al. 2013; graphics processing units (GPUs). Blei et al. 2017), manage to concentrate the sampling An additional advantage in using machine learning in regions of high posterior mass, even in large parameter emulators is that they efficiently provide accurate derivatives spaces, provided one has efficient access to accurate with respect to their input parameters. This is derivatives of the likelihood function with respect to made possible by the automatic differentiation features the model parameters (Brooks et al. 2011; Neal 2011; implemented in the libraries routinely used to build these Zhang et al. 2017; Betancourt 2017).