Joint inference of mass-maps and cosmology with diffusion models

Benjamin Remy

Open Skai Conference

Skai Institute, Chicago, September 2026

with Chihway Chang and Rebecca Willett

slides.com/benjaminremy/openskai26

Weak lensing for cosmology

Credit: Jessie Muir adapted by Justine Zeghal

Shear

Convergence

\gamma

\kappa

How to optimally infer cosmology from observing ?

\gamma_\text{obs}

How to reconstruct with a non-linear model?

\kappa

\theta

Weak lensing mass-mapping as an inverse problem

\theta

N-body +

ray-tracing

(e.g. TNG, Gower Street)

\kappa

\gamma^\text{obs}

p(\theta)

p(\kappa\mid \theta)

p(\gamma^{\text{obs}} \mid \kappa, \theta)

Forward model

\delta_\text{IC}

Inference (inverse problem)

p(\theta\mid \kappa, \gamma_\text{obs})

p(\kappa\mid \gamma_\text{obs})

\gamma^{\text{obs}}

Running N-body simulations, we

implicitly sample from the joint distribution

\theta, \kappa, \gamma \sim p(\theta)p(\kappa\mid \theta)p(\gamma\mid \kappa, \theta)

Implicit inference uses

simulations to learn

the implicit distributions

(posterior, likelihood, likelihood ratio)

(\theta, \kappa, \gamma)

Full field inference

Full field inference x Implicit inference (LFI)

Field reconstruction

Cosmological inference

NPE (Jeffrey et al. 2020, Lanzieri et al. 2025)

Targets

p(\theta \mid \gamma_\text{obs})

Targets

p(\kappa \mid \gamma_\text{obs}, \theta_\text{fid})

Posterior sampling with diffusion models (Remy et al. 2023)

Implicit inference for weak lensing full field

How can we combine them in a joint inference framework ?

p(\theta, \kappa \mid \gamma_\text{obs})

Text

A diffusion model learns to reverse the noising process, by learning

the prior score function

\nabla_x \log p_t(x)

U-net architecture, well suited for 2D, 3D fields

But not for inferring cosmological parameters...

Tweedy formula

\color{blue}{d_\phi(x^\prime, t)}\color{black}{, \quad x^\prime=x+n, \quad n\sim\mathcal{N}(0,\sigma_t)}

\color{blue}{d_\phi(x^\prime, t) = x^\prime + \sigma_t^2 \nabla_x \log p_t(x)}

Posterior inference with diffusion model

A diffusion model learns to reverse the noising process, by learning

the posterior score function

\nabla_x \log p_t(x\mid y)

U-net architecture, well suited for 2D, 3D fields

But not for inferring cosmological parameters...

Tweedy formula

\color{blue}{d_\phi(x^\prime, t, y)}\color{black}{, \quad x^\prime=x+n, \quad n\sim\mathcal{N}(0,\sigma_t)}

\color{blue}{d_\phi(x^\prime, t, y) = x^\prime + \sigma_t^2 \nabla_x \log p_t(x \mid y)}

Posterior inference with diffusion model

p(\theta, \kappa)

We need to model to design a denoiser architecture

to learn the score function.

\left[\theta_t, \kappa_t\right] = \left[\theta, \kappa\right] + \sigma_tn

d^\star(\theta_t, \kappa_t, \sigma_t) = \left[\theta^\star, \kappa^\star\right]

We learn the joint denoiser

Learning the joint distribution

And now have the joint score function to run a diffusion model

\nabla_{\kappa_t, \theta_t} \log p(\kappa_t, \theta_t)

p(\kappa, \theta)

\nabla_{\kappa_t, \theta_t} \log p(\kappa_t, \theta \mid \gamma_{t}^\text{obs})

p(\kappa, \theta\mid \gamma_\text{obs})

Learning the joint distribution

We need to model to design a denoiser architecture

to learn the score function.

\left[\theta_t, \kappa_t\right] = \left[\theta, \kappa\right] + \sigma_tn

We learn the joint denoiser

And now have the joint score function to run a diffusion model

d^\star(\theta_t, \kappa_t, \sigma_t, \gamma_\text{obs}) = \left[\theta^\star, \kappa^\star\right] \mid \gamma_\text{obs}

\nabla_{\kappa_t, \theta_t} \log p(\kappa_t, \theta \mid \gamma_{t}^\text{obs})

p(\kappa, \theta\mid \gamma_\text{obs})

Learning the joint distribution

We need to model to design a denoiser architecture

to learn the score function.

\left[\theta_t, \kappa_t\right] = \left[\theta, \kappa\right] + \sigma_tn

We learn the joint denoiser

And now have the joint score function to run a diffusion model

d^\star(\theta_t, \kappa_t, \sigma_t, \gamma_\text{obs}) = \left[\theta^\star, \kappa^\star\right] \mid \gamma_\text{obs}

Learning the joint distribution

p(\theta, \kappa)

Mocked weak lensing maps

with sbi_lens

128x128x5 maps, spanning

10x10 deg²

LSST Y10-like survey setting

Results

Learning the joint distribution

p(\theta, \kappa)

Results

A fast and differentiable (JAX) log-normal mass maps simulator

sbi-lens

(Lanzieri, Zeghal, et al. 2024 )

Learning the joint distribution

p(\theta, \kappa)

Results

Learning the joint distribution

p(\theta, \kappa)

Results

Learning the joint distribution

p(\theta, \kappa)

Results

Learning the joint distribution

p(\theta, \kappa)

Results

Learning the joint distribution

p(\theta, \kappa)

Results

Learning the joint distribution

p(\theta, \kappa)

Results

Cosmology marginal

p(\theta\mid \gamma)

TARP coverage test (Lemos et al. 2023)

MIRA calibration score (Sharief, Zeghal, et al. 2026)

on the joint distribution, where 2/3 is optimal

0.6428 \pm 0.0149

Results

Explicit conditioning

\nabla_{\theta, \kappa} \log p_t(\theta, \kappa \mid \gamma_\text{obs}) = \nabla_{\theta, \kappa} \log p_t(\gamma_\text{obs}\mid {\theta, \kappa}) + \color{green}{\nabla_{\theta, \kappa} \log p_t({\theta, \kappa})}

= \color{blue}{\nabla_{\kappa} \log p_t(\gamma_\text{obs}\mid {\kappa})}

+ \color{green}{\nabla_{\theta, \kappa} \log p_t({\theta, \kappa})}

\underset{\nabla_{\kappa_t} \mathbb{E}[\kappa\mid \kappa_t]^\top A^\top (\Sigma_y + A \mathbb{V}[\kappa\mid\kappa_t]A^\top)^{-1}(\gamma_\text{obs} - A \mathbb{E}[\kappa\mid\kappa_t])}{\underbrace{\color{white}{\nabla_{\kappa} \log p_t(\gamma_\text{obs}\mid {\kappa})}}}

Remy et al. (2022)

Rozet et al. (2024)

Joint prior

\gamma_\text{obs} = A\kappa + n

Linear model

Gaussian likelihood

and include

- lensing operator

- survey mask

- survey noise

Explicit conditioning