[Paper Review] JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library

J.E. Campagne, François Lanusse+8 more|arXiv (Cornell University)|Feb 10, 2023

Galaxies: Formation, Evolution, Phenomena4 citations

TL;DR

jax-cosmo is a differentiable, GPU-accelerated cosmology library built on JAX that enables end-to-end automatic differentiation of cosmological likelihoods. It accelerates Fisher matrix computations and posterior inference via gradient-based methods like Hamiltonian Monte Carlo and NUTS, achieving 5× higher effective sample size per GPU hour than traditional MCMC in a DES-Y1 3x2pt analysis.

ABSTRACT

We present jax-cosmo, a library for automatically differentiable cosmological theory calculations. It uses the JAX library, which has created a new coding ecosystem, especially in probabilistic programming. As well as batch acceleration, just-in-time compilation, and automatic optimization of code for different hardware modalities (CPU, GPU, TPU), JAX exposes an automatic differentiation (autodiff) mechanism. Thanks to autodiff, jax-cosmo gives access to the derivatives of cosmological likelihoods with respect to any of their parameters, and thus enables a range of powerful Bayesian inference algorithms, otherwise impractical in cosmology, such as Hamiltonian Monte Carlo and Variational Inference. In its initial release, jax-cosmo implements background evolution, linear and non-linear power spectra (using halofit or the Eisenstein and Hu transfer function), as well as angular power spectra with the Limber approximation for galaxy and weak lensing probes, all differentiable with respect to the cosmological parameters and their other inputs. We illustrate how autodiff can be a game-changer for common tasks involving Fisher matrix computations, or full posterior inference with gradient-based techniques. In particular, we show how Fisher matrices are now fast, exact, no longer require any fine tuning, and are themselves differentiable. Finally, using a Dark Energy Survey Year 1 3x2pt analysis as a benchmark, we demonstrate how jax-cosmo can be combined with Probabilistic Programming Languages to perform posterior inference with state-of-the-art algorithms including a No U-Turn Sampler, Automatic Differentiation Variational Inference,and Neural Transport HMC. We further demonstrate that Normalizing Flows using Neural Transport are a promising methodology for model validation in the early stages of analysis.

Motivation & Objective

• To develop a differentiable cosmology library that enables automatic differentiation of cosmological likelihoods with respect to cosmological and systematic parameters.
• To accelerate cosmological inference using JAX's JIT compilation, GPU acceleration, and batch processing for high-dimensional parameter spaces.
• To enable robust, fine-tuning-free Fisher matrix computation and gradient-based posterior inference, such as Hamiltonian Monte Carlo and variational inference.
• To support survey optimization via differentiable Fisher matrices, allowing gradient descent-based Figure of Merit maximization.
• To integrate with probabilistic programming languages like NumPyro for scalable, high-performance posterior sampling in cosmological analyses.

Proposed method

• Leverages JAX's automatic differentiation (autodiff) to compute exact gradients of cosmological observables with respect to cosmological parameters and systematics.
• Implements differentiable versions of key cosmological components: background evolution, linear and non-linear power spectra (using Halofit or Eisenstein & Hu), and angular power spectra (Cℓ) via the Limber approximation.
• Uses JAX's JIT compilation and GPU/TPU acceleration to speed up likelihood evaluations and sampling algorithms.
• Integrates with probabilistic programming frameworks like NumPyro to enable gradient-based inference techniques such as NUTS, ADVI, and Neural Transport HMC.
• Employs vmap for batch processing of cosmological simulations and parameter sweeps across multiple GPUs or CPUs.
• Supports differentiable likelihoods for 3x2pt analyses, enabling end-to-end differentiability from theory to data comparison.

Figure 1: Comparison of the radial comoving distance between CCL and jax-cosmo . More plots are available in the companion notebook \faFileCodeO .

Experimental results

Research questions

• RQ1 Can automatic differentiation in JAX enable robust, fine-tuning-free Fisher matrix computation in cosmology?
• RQ2 How does gradient-based inference with Hamiltonian Monte Carlo and NUTS compare in efficiency to traditional MCMC in a real cosmological analysis?
• RQ3 Can differentiable likelihoods enable direct optimization of survey design via gradient descent on the Figure of Merit?
• RQ4 What is the performance gain of using JAX-based samplers over optimized classical MCMC in terms of effective sample size per unit time?
• RQ5 Can normalizing flows with neural transport improve model validation in early-stage cosmological analysis?

Key findings

• The JAX NUTS sampler achieved approximately 5 times higher effective sample size per GPU per hour compared to the well-optimized Cobaya Metropolis-Hastings sampler in a DES-Y1 3x2pt analysis.
• Fisher matrices computed with jax-cosmo are exact, stable, and no longer require manual fine-tuning, significantly reducing computational and implementation overhead.
• Fisher matrices are themselves differentiable, enabling gradient-based optimization of survey design parameters such as redshift binning or observing strategy via gradient descent.
• Gradient-based inference techniques like NUTS and ADVI produce posterior samples consistent with classical MCMC while scaling more efficiently to high-dimensional parameter spaces.
• Neural Transport HMC with normalizing flows shows promise for model validation in early-stage cosmological analysis, offering fast and accurate posterior approximations.
• The library enables end-to-end differentiability of cosmological likelihoods, supporting complex inference pipelines with minimal code changes and high performance on GPU hardware.

Figure 2: Comparison of the non-linear matter power spectrum ( halofit function) between CCL and jax-cosmo . Also shown is the linear power spectrum.