[Paper Review] Field-Level Inference with Microcanonical Langevin Monte Carlo

Adrian E. Bayer, Uroš Seljak+1 more|arXiv (Cornell University)|Jul 18, 2023

Markov Chains and Monte Carlo Methods9 citations

TL;DR

The paper applies Microcanonical Langevin Monte Carlo (MCLMC) to field-level inference in cosmology, showing substantial efficiency gains over HMC for sampling high-dimensional initial conditions and cosmological parameters. MCLMC achieves higher effective sample size per gradient evaluation, with scalability that grows favorably with dimensionality across nonlinearity levels.

ABSTRACT

Field-level inference provides a means to optimally extract information from upcoming cosmological surveys, but requires efficient sampling of a high-dimensional parameter space. This work applies Microcanonical Langevin Monte Carlo (MCLMC) to sample the initial conditions of the Universe, as well as the cosmological parameters $σ_8$ and $Ω_m$, from simulations of cosmic structure. MCLMC is shown to be over an order of magnitude more efficient than traditional Hamiltonian Monte Carlo (HMC) for a $\sim 2.6 imes 10^5$ dimensional problem. Moreover, the efficiency of MCLMC compared to HMC greatly increases as the dimensionality increases, suggesting gains of many orders of magnitude for the dimensionalities required by upcoming cosmological surveys.

Motivation & Objective

• Motivate field-level inference as a means to optimally extract information from upcoming cosmological surveys.
• Demonstrate sampling of high-dimensional initial conditions and cosmological parameters using Microcanonical Langevin Monte Carlo (MCLMC).
• Compare MCLMC against Hamiltonian Monte Carlo (HMC) in terms of efficiency and scaling with dimensionality.
• Assess robustness across different levels of nonlinearity in forward modeling (ZA, 2LPT, PM).
• Provide guidance on scalability for data volumes of upcoming surveys.

Proposed method

• Review field-level inference and the target posterior for initial modes s and cosmological parameters λ.
• Describe Microcanonical Hamiltonian/Langevin Monte Carlo (MCHMC/MCLMC) dynamics and their ergodicity-enhancing momentum resampling.
• Introduce the MCLMC equations with projection P(u) and force f(z) and a Langevin diffusion term η dW.
• Use differentiable forward modeling (pmwd) to generate 3D nonlinear dark matter fields and sample across 1LPT, 2LPT, and 5-step PM.
• Precondition parameters by approximate posterior variance to improve mixing.
• Evaluate effective sample size per gradient under varying dimensionality and nonlinear forward models.

Experimental results

Research questions

• RQ1 How does MCLMC perform relative to HMC in sampling efficiency (ESS per gradient) for field-level cosmological inference?
• RQ2 How does the efficiency advantage of MCLMC scale with dimensionality and degrees of nonlinearity in the forward model?
• RQ3 Can MCLMC recover the same posterior as HMC for initial conditions and cosmological parameters (Ωm, σ8)?
• RQ4 What is the impact of forward-model nonlinearities (ZA, 2LPT, PM) on sampling performance?
• RQ5 Is MCLMC feasible for the dimensionalities expected in upcoming cosmological surveys?

Key findings

• MCLMC produces samples in close agreement with HMC for both initial modes and cosmological parameters.
• MCLMC outperforms HMC in all cases, with the largest gains at higher nonlinearity such as 2LPT (up to two orders of magnitude in ESS) and moderate gains for PM steps.
• ESS per gradient evaluation for MCLMC shows a pronounced advantage over HMC as dimensionality increases, with a factor of approximately 40–80 in the 64^3 case for modes and parameters respectively.
• MCLMC's ESS degrades more slowly with increasing dimensionality than HMC, due to the lack of a strict energy-based step-size constraint in MCLMC.
• The dimensionality-dependent efficiency gain suggests potential orders-of-magnitude improvements for the high-dimensional maps required by future surveys.