Source-linked AI summary
Neural Posterior Estimation for Tomographic Weak Lensing Mass Mapping
Tim White, Shreyas Chandrashekaran, Camille Avestruz, Jeffrey Regier, the LSST Dark Energy Science Collaboration
TL;DR
Image-based weak-lensing inference must preserve field information and quantify uncertainty while avoiding the repeated likelihood evaluations used by Bayesian image methods. This paper uses amortized NPE to map multiband images to distributions over tomographic lensing maps, finding well-calibrated estimates on held-out DC2 simulations.
Problem
Image-based weak-lensing inference must preserve field information and quantify uncertainty efficiently, while Bayesian image methods require repeated likelihood evaluations over latent galaxy and instrumental variables.
Source: §1 · pp. 1–2Method
NPE trains a neural network to map multiband coadded images to variational distributions over tomographic shear and convergence maps, amortizing inference into a single forward pass.
Source: §1 · p. 2Results
On a held-out DC2 subregion, NPE produced well-calibrated variational distributions over shear and convergence maps in four redshift bins, with positively correlated means and approximately nominal marginal credible-interval coverage.
Source: §6 · p. 9Takeaways & Limitations
After training, the amortized estimator can produce rapid posterior estimates, processing each image in a single forward pass.
Source: §1 · p. 2The evaluation did not test distribution shift; discrepancies between the DC2 simulator and real images could make NPE’s posterior approximations systematically biased or miscalibrated.
Source: §6.1 · p. 9
Abstract
from arXiv · showhide
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.
1. INTRODUCTION
Weak-lensing inference must recover shear and convergence while preserving field information and quantifying uncertainty, but current multistage and likelihood-based approaches are costly. The paper proposes amortized neural posterior estimation (NPE) directly from images and evaluates it on DC2 simulations.
1. INTRODUCTION
- Field-level inference can yield tighter cosmological constraints, while compression to galaxy catalogs and summary statistics typically loses information about the shear and convergence fields.Existing pipelines also require careful propagation of statistical uncertainties across stages.Source: §1 · pp. 1–2
- Repeated likelihood evaluation over thousands or tens of thousands of nuisance variables makes Bayesian image-based methods computationally challenging for modern surveys.The nuisance variables include galaxy properties and spatially varying backgrounds and point spread functions.Source: §1 · p. 2
- NPE maps multiband images directly to variational distributions over tomographic shear and convergence maps, without first compressing pixels into galaxy catalogs.The method is likelihood-free and amortized, estimating maps in a single forward pass after training.Source: §1 · p. 2
- The study tests NPE on DC2 simulations by training on one contiguous subregion and evaluating accuracy and calibration on another.The authors describe the work as an investigation of NPE's feasibility for field-level weak-lensing inference.Source: §1 · p. 2
- The authors propose using sampled lensing maps in a subsequent simulation-based inference procedure to approximate a posterior over ΛCDM parameters.Connected with NPE, the two procedures would form an end-to-end probabilistic pipeline from images to cosmological parameters.Source: §1 · p. 2
2. THE LSST DESC DC2 SIMULATED SKY SURVEY
The paper treats the DC2 simulator's coadds, catalogs, and lensing maps as a realization of an implicit probabilistic model. It describes how image inputs and approximate tomographic targets are constructed, then represents the maps, nuisance variables, and image generation as a probabilistic chain.
2.1. Coadded images and lensing maps
- DC2's cosmoDC2 galaxies are lensed using shear and convergence maps constructed from Outer Rim lightcone particles, with map values interpolated to each galaxy's redshift and sky position.The catalog is then augmented with properties matched from the Galacticus library.Source: §2.1 · p. 3
- The imSim and GalSim pipeline renders six-band raw images, which the LSST Data Release Production pipeline processes and coadds; the analysis uses coadds for their high signal-to-noise ratio and to avoid combining overlapping-image inferences.DC2 coadds cover approximately five square degrees, compared with 300 square degrees for the raw images.Source: §2.1 · p. 3
- Each patch-level coadd is partitioned into four 6×2048×2048-pixel images because a full patch is too large to store during network training.The paper uses “coadds” or “images” to refer to these partitions unless otherwise specified.Source: §2.1 · p. 3
- Approximate training targets are formed by binning per-galaxy shear and convergence into four redshift bins and averaging each quantity within 256×256-pixel tiles.The chosen tile size aims for approximately uniform fields within tiles and enough galaxies per tile to keep sampling noise small; the resulting maps are 4×8×8.Source: §2.1 · p. 3, Figure 1 · p. 4
2.2. An implicit probabilistic model
- The target maps z comprise tomographic shear and convergence for a single image, with dimension 3×B×H×W.The simulator's maps, catalogs, and coadds are treated as a sample from this implicit generative model.Source: §2.2 · pp. 3–4
- The simulator defines an implicit prior over the tomographic maps, followed by nuisance variables conditional on those maps and an image generated from both.The nuisance variables include source properties, spatially varying background and PSF characteristics, and pixel-level noise; their dimension can reach tens of thousands.Source: §2.2 · p. 4
3. NEURAL POSTERIOR ESTIMATION
The paper uses neural posterior estimation to infer tomographic shear and convergence maps from images by learning a variational posterior. Its likelihood-free, amortized objective and spatially factorized variational family provide the basis for the inference method and its ResNet implementation.
3.1. An amortized, likelihood-free objective
- NPE trains a network to map images to variational posteriors over lensing maps, amortizing inference without evaluating the likelihood and marginalizing over nuisance variables.The objective yields statistically consistent estimates when the variational family is sufficiently flexible.Source: §3.1 · p. 5
3.2. A spatially factorized variational family
- The variational family factorizes across redshift bins, tiles, and lensing variables, although the exact posterior is expected to contain nonzero auto- and cross-correlations.The network’s receptive field can still carry large-scale structure in per-tile means; fitted marginals coincide with true marginals when their distributional form is sufficiently expressive.Source: §3.2 · pp. 5–6
- Alternative parameterizations based on shear, convergence, or the lensing potential would enforce their physical relationship, but require additional recovery steps that may introduce errors.Potential boundary effects from Kaiser-Squires inversion and inflated posterior variance from differentiating the potential were not empirically investigated.Source: §3.2 · p. 6
3.3. Network architecture
- The network is a ResNet comparable in depth to ResNet-34, downsampling images from 2048×2048 to 8×8 before predicting variational parameters.The selected architecture has 14 residual blocks, maximum channel depth 1024, group normalization, and SiLU activations.Source: §3.3 · p. 6
- Providing median PSF feature maps yields only a marginal improvement in average validation loss, while images with varied PSFs are nearly as beneficial as supplying fiducial PSF parameters.The network receives medians of four catalog PSF parameters as feature maps concatenated with each image.Source: §3.3 · p. 6
4. TRAINING AND EVALUATION
The network is trained and evaluated on spatially distinct regions of the DC2 survey. Evaluation assesses calibration, point-estimate accuracy, and whether posterior samples recover the spatial structure of the underlying lensing fields.
4.1. Training the neural network
- Training, validation, and test sets contain 313, 39, and 40 spatially disjoint images, respectively, with each set covering a contiguous region.Training takes approximately 16 hours on one NVIDIA A40 GPU; random 90-degree rotations and vertical flips augment batches.Source: §4.1 · pp. 6–7
4.2. Evaluating probabilistic shear and convergence inference
- Posterior-mean shear and convergence maps are evaluated as point estimates using RMSE and Pearson correlation across test-image tiles.These metrics are reductive because they do not assess variational-distribution calibration, which matters for downstream cosmological inference.Source: §4.2 · p. 7
- Calibration is assessed by comparing empirical coverage of marginal highest-density credible intervals with their nominal coverage for shear and convergence.Well-calibrated marginals should have empirical coverage close to nominal coverage.Source: §4.2 · p. 7
- Posterior samples are used to compare tomographic convergence auto- and cross-power spectra with the DC2 maps and to measure scale-dependent cross-correlation.The cross-correlation coefficient identifies scales where inferred maps accurately reconstruct the truth.Source: §4.2 · p. 7
5. RESULTS
NPE’s inferred tomographic maps capture key convergence structure, and its marginal variational distributions are well-calibrated overall. Results also show limits in map boundaries, higher-multipole spectra, and low-redshift signal discrimination.
5.1. Posterior mass maps
- Posterior means recover most convergence overdensities and underdensities, while sampled spectra track ground truth at low but diverge at high multipoles.Peripheral tiles have larger posterior standard deviations because the network has less surrounding-galaxy context; larger training patches could reduce boundary counts but require substantial computing resources.Source: §5.1 · pp. 7–8, Figure 3 · p. 8
5.2. Accuracy of posterior means
- RMSE is smallest in the lowest redshift bin and largest in the highest, while Pearson correlations are smallest at low redshift and higher for convergence than shear.The authors suggest low-redshift shape noise is difficult to distinguish from the lensing signal, so the network infers shear and convergence close to zero; γ1 and γ2 metrics are similar in each bin.Source: §5.2 · pp. 8–9, Table 1 · p. 9
5.3. Posterior calibration
- The 90% credible intervals do not align perfectly with the 45-degree line, and intervals are much wider in higher-redshift bins.Within each redshift bin, interval widths are approximately uniform across the range of actual shear and convergence values.Source: §5.3 · p. 9
- Empirical coverage stays close to nominal across lensing variables and redshift bins, with intervals slightly wider than nominal at most levels.Approximately 90% of the 90% intervals cover ground truth in each variable and bin.Source: §5.3 · p. 9
6. DISCUSSION
NPE produces promising, calibrated shear and convergence distributions on held-out DC2 images, while simulator mismatch and mean-field factorization constrain how the method can be extended. The discussion proposes a hierarchical route from images to cosmological parameters.
6. DISCUSSION
- On held-out DC2 data, NPE yields well-calibrated distributions whose means correlate positively with ground truth and whose credible intervals cover at approximately nominal rates.This is the paper’s first investigation of NPE for field-level weak-lensing inference.Source: §6 · p. 9
6.1. Remedies for simulator misspecification
- Because training and evaluation both used DC2-generated images, real-survey simulator mismatch could bias or miscalibrate posteriors, and sensitivity assessment is needed before Stage IV deployment.The authors identify discrepancies at any stage of the DC2 generative model as a possible source of misspecification.Source: §6.1 · p. 9
6.2. A more flexible variational family
- Mean-field factors are calibrated and recover long-range dependencies but miss tile-level auto- and cross-correlation; flows could capture those correlations, though boundary effects would persist.An autoregressive tiling scheme is proposed as a possible remedy for boundary effects.Source: §6.2 · p. 10
6.3. A path to posterior inference of cosmological parameters
- A proposed hierarchy first infers convergence maps from images, then infers cosmological parameters from maps, assuming convergence captures most image information, especially for Ωm and σ8.The parameter stage could use NPE or another neural simulation-based inference algorithm, followed by ancestral sampling.Source: §6.3 · p. 11
- The decomposition could support end-to-end probabilistic weak-lensing inference at Stage IV survey scale, while this work establishes feasibility only for the images-to-maps stage.The authors also propose extensions needed for trustworthy inference on real survey data.Source: §6.3 · p. 11
A. ADDITIONAL FIGURES
The additional figures show DC2 galaxy-property distributions and test-set visualizations of inferred shear and convergence, including correlations, power spectra, and uncertainty maps.
Figure A1
- Figure A1 shows marginal histograms and pairwise scatter plots for redshift, r-band magnitude, shear, and convergence in 50,000 DC2 galaxies.The galaxies lie in tracts 3828 or 3829.Source: Figure A1 · p. 14
Figure A2
- Figure A2 compares actual tomographic shear maps with posterior means over the 0.6 deg² test set, with right ascension and declination as map axes.The accompanying description notes that shear posterior samples were omitted because they were poor quality.Source: Figure A2 · p. 14, A. ADDITIONAL FIGURES · p. 14
Figure A3
- Figure A3 plots the actual-versus-posterior-median convergence cross-correlation coefficient against multipole in each of four redshift bins.Lines are medians over 200 posterior samples, with shaded interquartile bands.Source: Figure A3 · p. 15
Figure A4
- Figure A4 separates convergence auto-power spectra on the diagonal from cross-power spectra off the diagonal, comparing actual maps with posterior-sample summaries.The actual spectra are dashed black curves; green curves show posterior-sample medians and shaded bands show interquartile ranges.Source: Figure A4 · p. 15, A. ADDITIONAL FIGURES · p. 15
Figure A5
- Figure A5 displays tomographic standard-deviation maps for the 0.6 deg² test set.These maps depict marginal standard deviations.Source: Figure A5 · p. 16
B. NEURAL NETWORK ARCHITECTURE
The network maps PSF-augmented images to variational parameters through a residual-block backbone built from convolutional, normalization, activation, and skip-connection layers.
Network architecture
- The network fη maps a PSF-augmented image x to variational parameters ϕ through individual layers and multi-layer residual blocks.Figure B1 diagrams the inputs, outputs, and layer structure.Source: B. NEURAL NETWORK ARCHITECTURE · p. 16, Figure B1 · p. 17
Residual-block backbone
- Its backbone has 14 residual blocks, and batch size one accommodates GPU memory constraints.The blocks transform intermediate representations using input and output channel counts and a stride.Source: B. NEURAL NETWORK ARCHITECTURE · p. 16
Residual-block structure
- Each block uses two 3x3 convolutional layers with group normalization and SiLU activations, plus a skip connection containing a 1x1 convolution and group normalization.The activation citation is Y. Wu & K. He 2018.Source: B. NEURAL NETWORK ARCHITECTURE · p. 16