Benchmarking field-level cosmological inference from galaxy redshift surveys
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
💡 Research Summary
This paper presents a systematic benchmark for field‑level Bayesian inference applied to galaxy redshift surveys, focusing on the joint estimation of cosmological parameters and the initial linear matter density field (δ_L) directly from the observed galaxy density field (δ_g). The authors build a fast, differentiable forward model using JaxPM, which incorporates Gaussian initial conditions, non‑linear large‑scale structure evolution via 2‑LPT and a Particle‑Mesh N‑body solver, redshift‑space distortions, and a second‑order Lagrangian bias expansion. Observational noise is added as Gaussian shot noise, yielding a hierarchical model with two cosmological parameters (Ω_m, σ₈), four bias parameters, and a high‑dimensional latent field whose dimensionality ranges from ~2.6 × 10⁵ to >2 × 10⁶ depending on the grid resolution.
To explore this high‑dimensional posterior, the study evaluates several Markov Chain Monte Carlo (MCMC) algorithms: standard Hamiltonian Monte Carlo (HMC), the No‑U‑Turn Sampler (NUTS) both in its vanilla form and within a Gibbs scheme, adjusted micro‑canonical Langevin Monte Carlo (MAMS), and unadjusted micro‑canonical Langevin Monte Carlo (MCLMC). A key contribution is the systematic investigation of preconditioning strategies, including parameter standardization, mass‑matrix conditioning based on prior covariances, and Fourier‑space scaling of the latent field. Performance is quantified by the number of model evaluations required per effective sample (ESS per gradient evaluation), acceptance rates, and scaling behavior with dimensionality.
Results show that without preconditioning all samplers suffer severe efficiency loss as dimensionality grows, with HMC and NUTS requiring increasingly small step sizes to maintain reasonable acceptance. Preconditioning dramatically improves performance, but the most striking finding is that unadjusted MCLMC outperforms all other methods by more than an order of magnitude in ESS per model evaluation, even at >10⁶ dimensions. MCLMC’s isokinetic Langevin dynamics preserve kinetic energy without a Metropolis correction, allowing large, stable steps and resulting in a mild scaling of efficiency with dimension (approximately O(d^0.1)). Adjusted MAMS performs better than HMC/NUTS but remains substantially slower than MCLMC. The authors also discuss bias introduced by the lack of Metropolis adjustment and demonstrate that choosing a small friction coefficient mitigates this effect.
The benchmark suite, along with fully open‑source code and synthetic data, is released publicly (https://github.com/hsimonfroy/benchmark-field-level) to serve as a reference platform for future algorithmic developments. The authors conclude that micro‑canonical, unadjusted samplers—particularly MCLMC—offer a scalable path forward for field‑level cosmological inference in upcoming massive surveys such as DESI, Euclid, and LSST, enabling near‑lossless extraction of cosmological information from the full galaxy density field.
Comments & Academic Discussion
Loading comments...
Leave a Comment