Neural Posterior Estimation for Inferring Weak Lensing Shear
Authors: Tim White, Dingrui Tao, Camille Avestruz, Jeffrey Regier, the LSST Dark Energy Science Collaboration
Organizations: Department of Statistics, University of Michigan · Department of Physics, University of Michigan
Abstract
The prevailing approach to inferring weak gravitational lensing shear from images involves detecting galaxies, estimating their ellipticities, and calibrating these estimates to correct for image noise, selection bias, and model misspecification. Characterizing the statistical model and assumptions underlying this pipeline is challenging, which makes it difficult to propagate uncertainty through its various stages. As an alternative, we propose to infer shear using neural posterior estimation (NPE), a type of simulation-based inference. We train a deep neural network to map a simulated multiband image to a variational distribution over the underlying shear field, thereby folding galaxy detection, deblending, measurement, and calibration into a single implicit inference step. Once trained, the network accounts for all features present in the simulated images, including potential sources of bias. In experiments on simulated constant-shear images with increasingly complex observational effects, NPE produces accurate and well-calibrated posterior approximations for both shear components in the presence of blended galaxies, spatially varying point spread functions, stars, and detector artifacts. These results demonstrate that NPE can be a viable shear estimation method in settings where all anticipated features and artifacts can be simulated, a requirement that will become increasingly feasible as simulation fidelity improves in the coming decades.
Weak gravitational lensing shear and convergence trace the distribution of baryonic and dark matter across space, making them a powerful probe of cosmic structure. Inferring shear and convergence from images is a challenging inverse problem. The prevailing approach to this task estimates shear from weighted averages of galaxy ellipticities, calibrates these estimates to account for systematic biases, and transforms them to reconstruct convergence, a multistage procedure that requires substantial computational resources and meticulous handling of statistical uncertainties. As an alternative, we propose a probabilistic approach to field-level weak lensing inference in which we train a deep neural network to directly map a multiband image to a variational distribution over the underlying tomographic shear and convergence fields. This neural posterior estimation (NPE) procedure implicitly marginalizes over nuisance variables in the cosmological forward model and does not require evaluating the likelihood function. It is also amortized, so it enables rapid posterior inference for astronomical surveys once the neural network is trained. When evaluated on synthetic images from the LSST-DESC DC2 Simulated Sky Survey, NPE produces well-calibrated variational distributions for shear and convergence that are consistent with the ground truth. We describe how maps sampled from these variational distributions could be used in a subsequent simulation-based inference procedure to approximate the posterior distribution over cosmological parameters.
Tim White, Shreyas Chandrashekaran, Camille Avestruz +2
Upcoming stage-IV surveys such as Euclid and Rubin will deliver vast amounts of high-precision data, opening new opportunities to constrain cosmological models with unprecedented accuracy. A key step in this process is the reconstruction of the dark matter distribution from noisy weak-lensing shear measurements. Current deep-learning-based mass-mapping methods achieve high reconstruction accuracy, but either require retraining a model for each new observed sky region (limiting practicality) or rely on slow Markov chain Monte Carlo sampling. Efficient exploitation of future survey data therefore calls for a new method that is accurate, flexible, and fast at inference. In addition, an uncertainty quantification with coverage guarantees is essential for a reliable cosmological parameter estimation. We introduce PnPMass, a plug-and-play approach for weak-lensing mass mapping. The algorithm produces point estimates by alternating between a gradient descent step with a carefully chosen data fidelity term and a denoising step implemented with a single deep-learning model trained on simulated data corrupted by Gaussian white noise. We also propose a fast sampling-free uncertainty quantification scheme based on moment networks, with calibrated error bars obtained through conformal prediction to ensure coverage guarantees. Finally, we benchmark PnPMass against model-driven and data-driven mass-mapping techniques. PnPMass achieves a performance close to that of the currently best deep-learning methods while offering fast inference. It converges in just a few iterations, and it requires only a single training phase, regardless of the noise covariance of the observations. It therefore combines flexibility, efficiency, and reconstruction accuracy while delivering tighter error bars than existing approaches, making it well suited for upcoming weak-lensing surveys.
We perform a realistic KiDS-Legacy mock analysis with field-level neural compression and simulation-based inference using just 60 N-body simulations. The weak lensing shear field encodes substantially more cosmological information than standard two-point summary statistics such as the power spectrum. Field-level inference can fully exploit this information, but physical realism at the field-level requires very high-fidelity simulations. This poses a major challenge for simulation-based inference (SBI): accurate empirical density modelling and deep-learning-based neural compression require tens of thousands of training samples, but achieving physical realism at the field level makes each simulation extremely costly. We demonstrate that multifidelity SBI can alleviate this tension by substantially reducing the number of high-fidelity simulations needed for accurate cosmological inference. We pre-train neural inference models on realistic KiDS-Legacy-like shear mocks using fast log-normal \texttt{GLASS} simulations and fine-tune them on a small set of high-fidelity N-body simulations. We show that 60 high-fidelity simulations are sufficient to obtain informative and well-calibrated cosmological posteriors, enabling at least an order-of-magnitude reduction in simulation cost for accurate field-level inference in a realistic setting.