Joint inference of weak lensing convergence map and cosmology with diffusion models
Benjamin Remy, Chihway Chang, Rebecca Willett
astro-ph.CO, astro-ph.IM
Submitted: 2026-06-30
Comments: 12 pages, 6 figures
Code: https://github.com/b-remy/jade
License: http://creativecommons.org/licenses/by/4.0/
The gist: We present a method for joint inference of cosmological parameters and convergence maps from weak lensing observations, targeting the full posterior conditioned on the observed shear field.
Terminology
Abstract
We present a method for joint inference of cosmological parameters and convergence maps from weak lensing observations, targeting the full posterior conditioned on the observed shear field. Our approach uses implicit inference with diffusion models, learning the joint distribution from simulations, without the need to have an explicit and differentiable forward model for gradient-based MCMC sampling. We introduce a transformer-based architecture that operates in pixel space and treats cosmological parameters as additional tokens in a unified sequence, enabling efficient multimodal processing within a single network. At inference time, the trained model generates posterior samples of joint convergence maps and cosmological parameters conditioned on observed noisy shear fields. We demonstrate the method on simulated weak lensing data generated from log-normal fields in a wcdm cosmology. The model accurately reconstructs convergence maps and recovers cosmological posteriors that agree with traditional MCMC, while remaining well calibrated across the prior, with a MIRA calibration score of 0.635 plus or minus 0.017 on the joint posterior (where 0.667 is optimal). The inferred fields reproduce the correct two-point statistics as well as non-Gaussian statistics such as the one-point distribution. This work establishes diffusion-based implicit inference as a viable route toward full field-level cosmological analyses, paving the way for applications to more realistic, non-differentiable simulators.
Sources
- Posterior samples of source galaxies in strong gravitational lenses with score-based priors
- The DESI Experiment Part I: Science,Targeting, and Survey Design
- Constraining neutrino masses with weak-lensing multiscale peak counts
- Stochastic Interpolants: A Unifying Framework for Flows and Diffusions
- The Dark Energy Camera All Data Everywhere cosmic shear project V: Constraints on cosmology and astrophysics from 270 million galaxies across 13,000 deg$^2$ of the sky
- Field-Level Inference from Galaxies: BAO Reconstruction
- A Conceptual Introduction to Hamiltonian Monte Carlo
- Probing primordial non-Gaussianity by reconstructing the initial conditions
- An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
- All-in-one simulation-based inference
- Denoising Diffusion Probabilistic Models
- The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo
- Adam: A Method for Stochastic Optimization
- Euclid Definition Study Report
- Systematic-free inference of the cosmic matter density field from SDSS3-BOSS data
- Back to Basics: Let Denoising Generative Models Denoise
- Flow Matching for Generative Modeling
- Towards Practical Field-Level Inference for Weak Lensing
- AION-1: Omnimodal Foundation Model for Astronomical Sciences
- Field-level inference of cosmic shear with intrinsic alignments and baryons
Related papers
- Angular clustering and bias of photometric quasars in the Kilo-Degree Survey Data Release 4
- A Novel kinetic Sunyaev-Zel'dovich Estimator for Electron-Electron Correlations
- Magnetic fields at the dawn of structure formation I. The CARLA J1510+5958 proto-cluster
- Dark Energy Survey Year 6 Results: Weak Lensing and Galaxy Clustering Cosmological Analysis Framework
- Exploring the Impact of Systematic Bias in Type Ia Supernova Cosmology Across Diverse Dark Energy Parametrizations
- Non-Gaussian Galaxy Stochasticity and the Noise-Field Formulation