A Fast Generative Framework for High-dimensional Posterior Sampling: Application to CMB Delensing

Hadi Sotoudeh; Pablo Lemos; Laurence Perreault-Levasseur

astro-ph.IMMar 4, 2026

A Fast Generative Framework for High-dimensional Posterior Sampling: Application to CMB Delensing

Authors: Hadi Sotoudeh, Pablo Lemos, Laurence Perreault-Levasseur

Organizations: Department of Physics, Université de Montréal, Montréal, Canada · Mila - Quebec AI Institute, Montréal, Canada · Ciela, Montreal Institute for Astrophysics and Machine Learning, Montréal, Canada · Center for Computational Astrophysics, Flatiron Institute, New York, USA

Abstract

We introduce a deep generative framework for high-dimensional Bayesian inference that enables efficient posterior sampling. As telescopes and simulations rapidly expand the volume and resolution of astrophysical data, fast simulation-based inference methods are increasingly needed to extract scientific insights. While diffusion-based approaches offer high-quality generative capabilities, they are hindered by slow sampling speeds. Our method performs posterior sampling an order of magnitude faster than a diffusion baseline. Applied to the problem of CMB delensing, it successfully recovers the unlensed CMB power spectrum from simulated observations. The model also remains robust to shifts in cosmological parameters, demonstrating its potential for out-of-distribution generalization and application to observational cosmological data.

Figures & tables

Panel of Figure 1 — Figure 1: GRF rotation results: (a) Comparing empirical and theoretical moments for a test example. The pixel-wise empirical means and standard deviations are estimated using 1000 posterior samples. (b) TARP coverage test results.

Panel of Figure 2 — Figure 2: CMB delensing results: (a) Comparison of the lensed (input), unlensed (target), and delensed (obtained from model outputs) TT power spectra for six test examples. The delensed spectrum is represented by mean and uncertainty regions computed from 1000 posterior samples. Relative differences $\Delta\mathcal{D}_{\ell}/\mathcal{D}_{\ell}$ are computed with respect to the mean delensed spectrum. (b) TARP coverage test results. (c) Out-of-distribution performance for different $\Omega_{m}$ values, with all other cosmological parameters fixed. Each column corresponds to an $\Omega_{m}$ value, which differs from the training (base) value by a factor of the Planck satellite measurement error $\sigma_{\Omega_{m}}$ . Each row corresponds to a unique noise realization used to generate the lensed and unlensed maps.

Model	Sampling Time (s)
Our Framework	$0.31\pm 0.05$
Diffusion Baseline-100	$12.85\pm 0.09$
Diffusion Baseline-1000	$125.6\pm 0.4$

Table 1: Comparison of the time required to generate 50 posterior samples for a batch of 4 test examples on an NVIDIA H100 GPU. Uncertainties are estimated from 10 runs. The numbers next to the diffusion baselines indicate the number of denoising steps.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Data generation process for the GRF rotation experiment.

Figure 4: Data generation process for the CMB delensing experiment.

Parameter		Value
$H_{0}$	Hubble parameter ( $\text{km }\text{s}^{-1}\text{Mpc}^{-1}$ )	67.5
$\Omega_{m}$	Matter density parameter	0.27
$\Omega_{k}$	Curvature density parameter	0
$\tau$	Thomson scattering optical depth	0.06
$\ln{(10^{10}A_{s})}$	Primordial curvature perturbations	2.996
$n_{s}$	Scalar spectrum power-law index	0.965

Appendix

Table 2: Cosmological parameters for the simulated CMB maps.

Model		Input Channels	Intermediate Channels +
GRF Rotation	Mean Network	1	[32, 64, 128, 128, 128]
GRF Rotation	Dispersion Network	2 ∗	[8, 16, 32, 32, 32]
CMB Delensing	Mean Network	1	[64, 128, 256, 256, 256]
CMB Delensing	Dispersion Network	2 ∗	[16, 32, 64, 64, 128]

Appendix

Table 3: Model architecture.

Model		# Parameters	Dataset Size	Epochs	Training Time
GRF Rotation	Mean Network	7.5 M	$2^{16}$	400	25 hours
GRF Rotation	Dispersion Network	780 k	$2^{16}$	50	13 hours
CMB Delensing	Mean Network	27.8 M	$2^{13}$	400	87 hours
CMB Delensing	Dispersion Network	3.9 M	$2^{13}$	600	42 hours

Appendix

Table 4: Model size and training time.

Posterior Sampling

Posterior sampling generates draws from Bayesian posteriors to quantify uncertainty and support inference, prediction, and decision-making.

See all 69 papers

Strong Gravitational Lensing Posterior Sampling in Pixel-Space Using Diffusion Models and Recurrent Inference Machines
Jul 21, 2026Guillaume Payeur, Laurence Perreault-Levasseur, Gabriel Missael Barco +1Pixel-Space Diffusion Models Diffusion Posterior Sampling
Department of Physics, Universit´e de Montr´eal, Montreal QC, Canada · 2Mila - Quebec AI Institute, Montreal QC, Canada · 3Ciela Institute, Institute for Astrophysics and Machine Learning, Montreal QC, Canada +3
Learning the Universe: Posterior Reliability of Neural Generative Models in High-Dimensional Field-Level Inference of Cosmic Initial Conditions
Jun 8, 2026Ludvig Doeser, Jens JascheSimulation-Based Inference Normalizing Flows
The Oskar Klein Centre, Department of Physics, Stockholm University, AlbaNova University Centre, SE 106 91 Stockholm, Sweden · Center for Computational Astrophysics, Flatiron Institute, 162 5th Avenue, New York, NY 10010, USA
Deep unfolding of MCMC kernels: scalable, modular & explainable GANs for high-dimensional posterior sampling
Feb 24, 2026Jonathan Spence, Tobías I. Liaudat, Konstantinos Zygalakis +1Posterior Sampling Generative Adversarial Networks
Maxwell Institute for Mathematical Sciences, Bayes Centre, 47 Potterrow, Edinburgh, Scotland, UK · School of Mathematical and Computer Sciences, Heriot-Watt University, Edinburgh, Scotland, UK · IRFU, CEA, Université Paris-Saclay, F-91191 Gif-sur-Yvette, France +1

Jul 21, 2026J/K move · Enter open · S save

Strong Gravitational Lensing Posterior Sampling in Pixel-Space Using Diffusion Models and Recurrent Inference Machines

Guillaume Payeur, Laurence Perreault-Levasseur, Gabriel Missael Barco, Yashar Hezaveh

Department of Physics, Universit´e de Montr´eal, Montreal QC, Canada · 2Mila - Quebec AI Institute, Montreal QC, Canada · 3Ciela Institute, Institute for Astrophysics and Machine Learning, Montreal QC, Canada+3

Pixel-Space Diffusion Models Diffusion Posterior Sampling

Open paper AskPDF arXiv

Modeling galaxy-galaxy strong gravitational lenses to infer the brightness of the source galaxy and the mass distribution of the foreground galaxy is computationally challenging, particularly for high-resolution, high signal-to-noise ratio observations. In this regime, high-dimensional representations of both the source and the foreground mass distribution are necessary to model the data down to the noise level. This inference problem has been challenging for both traditional and machine learning-based methods because of its high dimensionality and its non-linearity in the foreground mass distribution. We present a method to generate joint posterior samples of the source galaxy and foreground mass distribution as pixelated images conditioned on observations. The method combines diffusion-based generative modeling and recurrent inference machines. It can model realistic gravitational lensing simulations with background and foreground galaxies drawn from cosmological hydrodynamical simulations down to the noise level.

Abstract

Posterior Sampling

Explore similar work

Strong Gravitational Lensing Posterior Sampling in Pixel-Space Using Diffusion Models and Recurrent Inference Machines

Learning the Universe: Posterior Reliability of Neural Generative Models in High-Dimensional Field-Level Inference of Cosmic Initial Conditions

Deep unfolding of MCMC kernels: scalable, modular & explainable GANs for high-dimensional posterior sampling