Benchmarking field-level cosmological inference from galaxy redshift surveys

TL;DR

This work provides a standardized benchmark for field-level cosmological inference from galaxy redshift surveys using a fast, differentiable forward model. It systematically compares canonical and microcanonical Langevin samplers (HMC/NUTS/NUTSwG, MAMS, MCLMC) and shows that careful latent-variable preconditioning, especially Fourier-space Kaiser conditioning, yields substantial efficiency gains. The unadjusted microcanonical MCLMC emerges as the most scalable approach, achieving over an order-of-magnitude speedup in high-dimensional settings with tolerable bias controlled by the energy-variance threshold, and its performance is robust to model dimensionality up to $\operatorname{dim}(\delta_L) \approx 128^3$ in the tested regimes. The authors provide open-source benchmark code, discuss scaling to larger surveys, and outline necessary extensions for real data such as light-cone modeling and survey systematics, highlighting the practical impact for upcoming galaxy surveys.

Abstract

Field-level inference has emerged as a promising framework to fully harness the cosmological information encoded in next-generation galaxy surveys. It involves performing Bayesian inference to jointly estimate the cosmological parameters and the initial conditions of the cosmic field, directly from the observed galaxy density field. Yet, the scalability and efficiency of sampling algorithms for field-level inference of large-scale surveys remain unclear. To address this, we introduce a standardized benchmark using a fast and differentiable simulator for the galaxy density field based on $\texttt{JaxPM}$. We evaluate a range of sampling methods, including standard Hamiltonian Monte Carlo (HMC), No-U-Turn Sampler (NUTS) without and within a Gibbs scheme, and both adjusted and unadjusted microcanonical samplers (MAMS and MCLMC). These methods are compared based on their efficiency, in particular the number of model evaluations required per effective posterior sample. Our findings emphasize the importance of carefully preconditioning latent variables and demonstrate the significant advantage of (unadjusted) MCLMC for scaling to $\geq 10^6$-dimensional problems. We find that MCLMC outperforms adjusted samplers by over an order-of-magnitude, with a mild scaling with the dimension of our inference problem. This benchmark, along with the associated publicly available code, is intended to serve as a basis for the evaluation of future field-level sampling strategies. The code is readily open-sourced at https://github.com/hsimonfroy/benchmark-field-level

Figures (10)

• Figure 1: Two approaches to inferring cosmology $\Omega$ from the galaxy density field $\delta_g$, that depends on $\Omega$, initial field $\delta_L$, and nuisance parameters $b$ via a cosmological model (left panel). Top panel: inference based on summary statistics $s(\delta_g)$ marginalizes over $\delta_L$ and $\delta_g$ (e.g. by fitting a Gaussian model or a Neural Density Estimator on simulations) to obtain marginal model $\mathrm{p}(\Omega)\,\mathrm{p}(b)\,\mathrm{p}(s \mid, \Omega, b)$, then conditions on $s$ (e.g. by MCMC sampling) to obtain final posterior $\mathrm{p}(\Omega, b\mid s)$. Bottom panel: field-level inference conditions on $\delta_g$ (e.g. by MCMC sampling) to obtain full posterior $\mathrm{p}(\Omega, b, \delta_L \mid \delta_g)$, then marginalizes over $\delta_L$ (by discarding $\delta_L$ samples) to obtain final posterior $\mathrm{p}(\Omega, b \mid \delta_g)$. In this Bayesian network formalism, white (resp. grey, diamond) nodes denote latent (resp. observed, deterministic) variables, and directed (resp. undirected) edges denote computed (resp. uncomputed) dependencies.
• Figure 2: Posterior of the initial linear matter field $\delta_L$ (top row) and the observable galaxy density field $\delta_g$ (bottom row), obtained by field-level inference from the observed galaxy density field $\delta_g^\mathrm{true}$. Panels from left to right display respectively the voxel true values, posterior means, and posterior standard deviations, for a thin slice of these fields. Rightmost panels show the posterior median and $95\%$-credible intervals of (top) the transfer function $T(\delta_i, \delta_j) = \sqrt{P_{ii} / P_{jj}}$ and (bottom) the spectra correlation $R(\delta_i, \delta_j) = P_{ij} / \sqrt{P_{ii} P_{jj}}$ (with $P_{ij}(k) = \braket{\delta_i(k) \delta_j(k)^*}$) between the sampled field and the true field associated with the observation. Both the linear matter and galaxy density power spectra are well recovered, while the fields themselves are well reconstructed at large scale (spectra correlation close to 1).
• Figure 3: State of the chains at the initialization and at the end of the warmup phases, showed via the median and the $95\%$-credible region of the transfer function (left) and spectra correlation (right) between the sampled linear field and the true linear field (as defined in figure \ref{['fig:post_fields']}). The initialization based on the Kaiser posterior approximation (turquoise) provides a relevant reconstruction of the linear field $\delta_L$ at large scales (spectra correlation close to 1). From there, the subsequent warmup phases achieve convergence, first (orange) towards a fiducial power spectrum (assuming a fixed fiducial cosmology $\Omega^\mathrm{fid}$ and bias parameters $b^\mathrm{fid}$), then (indigo) towards the true power spectrum (jointly inferring $\Omega, b$ and $\delta_L$).
• Figure 4: $68\%$ and $95\%$-credible regions of the cosmological and bias parameters posterior $\Omega, b \mid \delta_g$, obtained with the five considered samplers (NUTSwG in turquoise, NUTS in orange, HMC in indigo, MAMS in magenta, MCLMC in lime). The dashed lines denote the parameter true values used to generate the (synthetic) observation.
• Figure 5: Number of model evaluations per effective sample for each sampler and each set of parameters (cosmology, galaxy bias and linear field).