CosmoPower-JAX: Differentiable Cosmology Emulator

• CosmoPower-JAX is a fully differentiable, GPU-accelerated emulator framework for cosmological power spectra that employs PCA-based compression and neural emulation for fast parameter inference.
• It integrates principal component analysis with a multi-layer perceptron, trained on extensive simulations, achieving sub-percent accuracy compared to traditional Boltzmann solvers.
• The framework leverages JAX’s jit compilation, vectorized batch processing, and automatic differentiation to support ultra-fast, high-dimensional gradient-based sampling.

CosmoPower-JAX is a fully differentiable, GPU-accelerated emulator framework for cosmological power spectra, implemented in JAX and designed to facilitate ultra-fast gradient-based parameter inference for current and next-generation cosmological surveys. It extends the CosmoPower methodology to a pure JAX environment, combining principal-component compression and neural emulation with just-in-time compilation, vectorized batch prediction, and automatic differentiation to yield rapid and scalable predictions of $C_\ell$ (CMB angular power spectra) and $P(k)$ (3D matter power spectra). This enables precise, high-dimensional Bayesian inference with efficiency and scalability suitable for analyses involving up to hundreds of free parameters while achieving sub-percent accuracy with respect to traditional Boltzmann solvers such as CLASS or CAMB (Piras et al., 2023, Bolliet et al., 10 Jul 2025).

1. Theoretical Framework and Neural Emulation

CosmoPower-JAX is built upon a two-stage emulation strategy. The first stage compresses output from Boltzmann solvers (CLASS) into a reduced set of basis weights using principal-component analysis (PCA) or singular-value decomposition (SVD). The training data comprise calculated spectra across a grid of cosmologies $\{\theta_n\}$, producing matrices for $P_\text{CLASS}(k,z;\theta_n)$ and $C_\ell^\text{CLASS}(\theta_n)$. Decomposition yields basis sets $\{\phi_i(k,z)\}$ and $\{\psi_j(\ell)\}$ and corresponding mode amplitudes $\alpha_i(\theta_n)$, $\beta_j(\theta_n)$ such that: $$P_\text{CLASS}(k,z;\theta_n) \simeq \sum_{i=1}^{N_\text{modes}} \alpha_i(\theta_n) \phi_i(k,z),\quad C_\ell^\text{CLASS}(\theta_n) \simeq \sum_{j=1}^{M_\text{modes}} \beta_j(\theta_n) \psi_j(\ell).$$ In the second stage, a multi-layer perceptron (MLP) or Gaussian process (GP) regresses these mode weights as smooth functions of cosmological parameters. For new parameter vectors $P(k)$0, the emulator predicts new weights and reconstructs the spectra: $P(k)$1 or directly $P(k)$2, where $P(k)$3 is the smooth interpolation and $P(k)$4 a residual noise GP (Bolliet et al., 10 Jul 2025).

The neural network architecture in CosmoPower-JAX consists of four fully connected layers of 512 units each, employing a “swish-like” activation function. The network is trained with mean squared error loss over large datasets: $P(k)$5 matter-power spectra and $P(k)$6 CMB spectra. Overfitting is mitigated by early stopping on a validation set, and the Adam optimizer is used with decaying learning rates (Piras et al., 2023).

2. Training Data Construction and Compression

Training cosmologies $P(k)$7 are sampled via Latin-Hypercube or Sobol sequences covering the target region of parameter space, typically for Stage IV survey forecasts. Parameters include $P(k)$8. For each sampled $P(k)$9, the full CLASS Boltzmann solver provides high-accuracy spectra (e.g., $\{\theta_n\}$0 up to $\{\theta_n\}$1, $\{\theta_n\}$2 on a 2D $\{\theta_n\}$3 grid). The resultant output matrices are subjected to PCA/SVD, with 10–20 dominant modes retained—accounting for $\{\theta_n\}$4 of the variance—which compresses the emulation to interpolation of a small number ($\{\theta_n\}$5) of amplitude functions. This compression increases smoothness in parameter dependency and dramatically reduces the computational complexity of the regression problem (Bolliet et al., 10 Jul 2025).

3. JAX Implementation and Computational Techniques

CosmoPower-JAX leverages JAX features to achieve high throughput and differentiability:

• jit compilation: The main emulation and likelihood routines are decorated with @jax.jit, compiling Python code into efficient XLA kernels.
• vmap and batch evaluation: jax.vmap enables batch processing of parameter sets, enabling predictions across tens to hundreds of cosmologies in a single kernel call, without explicit loops.
• automatic differentiation: The entire emulator is written as a pure JAX function, allowing gradients with respect to any $\{\theta_n\}$6 via jax.grad or jax.jacobian, which is critical for Hamiltonian Monte Carlo (HMC) and Fisher forecasting.
• multi-GPU scaling: Combining JAX's pmap and NumPyro's chain vectorization, CosmoPower-JAX distributes MCMC chains across multiple GPUs achieving near-perfect scaling (Piras et al., 2023).

The constant pre-computed basis vectors $\{\theta_n\}$7 and $\{\theta_n\}$8 ensure that the runtime cost per parameter point is dominated by small dense-matrix multiplications and lightweight MLP passes, yielding exceptional evaluation speeds.

4. Performance Benchmarks and Accuracy

CosmoPower-JAX delivers significant speed-outs versus traditional Boltzmann solvers:

• P(k): On a laptop CPU, $\{\theta_n\}$9 at 50 $P_\text{CLASS}(k,z;\theta_n)$0-points and 10 redshifts evaluates in $P_\text{CLASS}(k,z;\theta_n)$10.4 ms (jitted), compared to $P_\text{CLASS}(k,z;\theta_n)$21 ms for Fortran/GPU CosmoPower and $P_\text{CLASS}(k,z;\theta_n)$31000 ms for CLASS halofit (Bolliet et al., 10 Jul 2025).
• $P_\text{CLASS}(k,z;\theta_n)$4: $P_\text{CLASS}(k,z;\theta_n)$5 evaluations ($P_\text{CLASS}(k,z;\theta_n)$6) require $P_\text{CLASS}(k,z;\theta_n)$70.8 ms with CosmoPower-JAX, $P_\text{CLASS}(k,z;\theta_n)$80.03 s for CLASS_SZ fast mode, and $P_\text{CLASS}(k,z;\theta_n)$9 s for full CLASS.
• Batch throughput: For 64 cosmologies in parallel, per-point evaluation is $C_\ell^\text{CLASS}(\theta_n)$0100 $C_\ell^\text{CLASS}(\theta_n)$1s/cosmology for $C_\ell^\text{CLASS}(\theta_n)$2 and 200 $C_\ell^\text{CLASS}(\theta_n)$3s/cosmology for $C_\ell^\text{CLASS}(\theta_n)$4.
• Total speedup: $C_\ell^\text{CLASS}(\theta_n)$5 over CLASS for $C_\ell^\text{CLASS}(\theta_n)$6; $C_\ell^\text{CLASS}(\theta_n)$7 for $C_\ell^\text{CLASS}(\theta_n)$8.

Accuracy metrics include:

• Maximum emulation error: $C_\ell^\text{CLASS}(\theta_n)$9 for $\{\phi_i(k,z)\}$0 in $\{\phi_i(k,z)\}$1 h/Mpc, $\{\phi_i(k,z)\}$2; $\{\phi_i(k,z)\}$3 for $\{\phi_i(k,z)\}$4 in $\{\phi_i(k,z)\}$5 (TT/TE/EE/lensing).
• Derivative fidelity: JAX autodiff gradients agree with classic finite-difference derivatives at sub-percent differences (Piras et al., 2023, Bolliet et al., 10 Jul 2025).
• Posterior reliability: KL divergence between emulator-generated and CAMB-based posterior marginals $\{\phi_i(k,z)\}$6; Fisher ellipses agree at $\{\phi_i(k,z)\}$7.

5. Inference, Gradient-Based Sampling, and Scalability

Gradient-based samplers are natively supported, leveraging the differentiable emulator:

• Hamiltonian Monte Carlo (HMC) and No-U-Turn Sampler (NUTS) pipelines execute on CPU or GPU, with HMC trajectory parameters and parameter reparametrization optimized for efficiency and geometry.
• Example scalability: Stage IV cosmic shear (37 parameters) with NUTS + CosmoPower-JAX on three A100 GPUs completes in $\{\phi_i(k,z)\}$81 day (compared to $\{\phi_i(k,z)\}$94 months for PolyChord + CAMB + CCL on CPU); a triple survey 3$\{\psi_j(\ell)\}$02pt (157 parameters) analysis runs in $\{\psi_j(\ell)\}$13 days (vs. $\{\psi_j(\ell)\}$26 years for CAMB-based pipeline) (Piras et al., 2023).
• Vectorized MCMC: Chain parallelism across multiple GPUs via pmap enables O(100–1000) effective chains, improving convergence and statistical power.
• Integrability: Any JAX-based likelihood or observable (e.g., jax-cosmo) can be employed within the framework, enabling fast forward modeling and uncertainty quantification for non-Gaussian or field-level likelihoods.

6. Limitations and Domain of Validity

• Training domain boundaries: The emulator is strictly valid within the convex hull of training parameters; extrapolation leads to uncontrolled error growth. Retraining is needed for new physical extensions (e.g., massive neutrino, $\{\psi_j(\ell)\}$3CDM, varied $\{\psi_j(\ell)\}$4) and for cross-model combinations.
• Resolution limits: Small-scale or rapid-oscillation spectral features (from e.g. reionization or neutrino-induced phase shifts) may be inadequately captured in a low-mode PCA truncation, necessitating careful mode-count selection for explorations beyond $\{\psi_j(\ell)\}$5 regions (Bolliet et al., 10 Jul 2025).
• High-resolution regime: Ultra-high $\{\psi_j(\ell)\}$6 ($\{\psi_j(\ell)\}$7) or $\{\psi_j(\ell)\}$8 ($\{\psi_j(\ell)\}$9/Mpc) are beyond validated emulator range and require fallback to full Boltzmann evaluation or specialized fitting functions.
• Sampling limitations: Efficacy of NUTS/HMC may degrade for multi-modal or highly anisotropic posteriors; emulators are at risk of introducing bias if the target likelihood explores extreme parameter ranges (Piras et al., 2023).
• Posterior exploration: Extreme extrapolation may bias results, and appropriate coverage diagnostics should be employed.

7. Practical Implications and Future Prospects

CosmoPower-JAX enables high-dimensional, fast, and accurate likelihood evaluation for next-generation cosmological analyses. By making all spectral calculations differentiable, it unlocks the possibility of field-level or non-Gaussian inference via automatic differentiation, robust Fisher forecasting, and order-of-magnitude improvements in sampling efficiency. The public availability and swift retrainability of the emulator make it straightforward to incorporate additional cosmology, baryonic, or nuisance parameters. Linear GPU scaling is achievable, with demonstrated analyses reaching $\alpha_i(\theta_n)$0–$\alpha_i(\theta_n)$1 parameter posteriors in hours to days on moderate hardware (Piras et al., 2023, Bolliet et al., 10 Jul 2025).

Planned extensions include wider adoption within jax-cosmo, expansion to cross-spectra emulation, generalized gradient samplers (e.g., Microcanonical or tempered HMC), JAX/Pyro-based variational inference, and emulation of higher-order statistics (including bispectra). The extension to hardware accelerators such as TPUs is under consideration to further expand computational accessibility and throughput.

Overall, CosmoPower-JAX establishes a computational paradigm for cosmological parameter inference—fully differentiable, massively parallel, and demonstrably capable of supporting the accuracy and complexity requirements posed by forthcoming Stage IV and next-generation survey projects.