Joint inference of mass-maps and cosmology from weak lensing cosmic shear with diffusion models

Abstract

Upcoming Stage-IV galaxy surveys will map the large-scale structure of the Universe with unprecedented precision, requiring analysis methods that can exploit information beyond traditional two-point statistics. Field-level inference offers a principled path forward by working directly with the observed fields, capturing non-Gaussian signatures that summary statistics discard. However, existing approaches typically address either cosmological parameter estimation or field reconstruction in isolation, or rely on explicit inference frameworks that require differentiable forward models and costly MCMC sampling. We present a diffusion-model-based method that performs joint inference of weak lensing convergence maps and cosmological parameters in a single, unified framework. Built on a pixel-space vision transformer, our model learns the joint posterior distribution over both the convergence field and cosmological parameters, conditioned on noisy shear observations. By operating within the implicit inference paradigm, our approach bypasses the need for a differentiable forward model, opening the door to arbitrarily complex simulators. We demonstrate our method on simulated LSST Year-10 weak lensing data generated with log-normal convergence fields in a wCDM cosmology, showing that our approach can recover accurate joint posteriors over both the mass map and cosmological parameters from a single amortized model, validated against MCMC baselines and coverage diagnostics.