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

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

Problem Statement
Modeling galaxy-galaxy strong gravitational lenses to infer the brightness distribution of a source galaxy and the mass distribution of a foreground galaxy is a computationally demanding inverse problem. This challenge is exacerbated in the regime of high-resolution, high signal-to-noise ratio (SNR) observations, where low-dimensional parametric assumptions (e.g., simple elliptical power laws) are insufficient to capture the complexity of the data. To model observations down to the noise level, high-dimensional, pixelated representations of both the source and the foreground mass are required. However, representing the lens and source as pixelated images renders the forward model non-linear with respect to the foreground mass, creating a high-dimensional, non-linear inverse problem characterized by complex, multimodal posterior distributions and severe parameter degeneracies. Traditional sampling algorithms and standard machine learning approaches often struggle to effectively sample these distributions.

Methodology: The DiRIM Framework
The authors propose Diffusion Recurrent Inference Machine (DiRIM), a Bayesian method designed to generate joint posterior samples of the source galaxy brightness ($s$) and the foreground mass distribution (represented as a convergence map, $\kappa$) conditioned on observations ($y$). The framework integrates two distinct machine learning paradigms:

1. Diffusion-Based Generative Modeling: The method utilizes score-based diffusion models to sample from complex probability distributions. The process involves a forward stochastic differential equation (SDE) that gradually adds noise to the data, and a reverse SDE that reconstructs the data from noise. To sample from the posterior $p(x|y)$, the reverse SDE requires the posterior score $\nabla_x \log p_t(x|y)$.
2. Recurrent Inference Machines (RIMs): RIMs are employed to solve the inverse problem by learning a data-driven, iterative inference algorithm. Instead of producing a single Maximum A Posteriori (MAP) estimate, the RIM framework is adapted here to iteratively refine the estimate of the posterior score.

The DiRIM Architecture:

• Input: The neural network $g_\theta$ (a U-Net architecture) takes as input the noisy model parameters at time $t$ ($x_t$), the observation $y$, the time step $t$, the likelihood gradients ($\nabla L$), and normalized lensed image residuals ($\delta y / \sigma_y$).
• Iterative Refinement: Unlike standard diffusion models that perform a single denoising step per time step, DiRIM uses the RIM framework to iteratively update its estimate of the denoised source and convergence map ($\hat{x}_0$) within each diffusion time step.
• Score Estimation: The network is trained to minimize a conditional denoising score-matching loss. The final output of the RIM iterations ($\hat{x}_0^{(M)}$) is used to compute the posterior score required for the reverse SDE:
  $$\nabla_{x_t} \log p_t(x_t|y) = \frac{\hat{x}_0^{(M)} - x_t}{\sigma^2(t)}$$
• Training Data: The model is trained on realistic mock observations generated from cosmological magneto-hydrodynamical simulations (IllustrisTNG) and analytic profiles. The source images are drawn from the SKIRT-TNG dataset (dust radiative transfer processed), and convergence maps are either smoothed from simulation halos or rendered from analytic profiles (Elliptical Power Law + multipoles + NFW subhalos).

Key Contributions

• First Pixel-Space Joint Posterior Sampling: This work presents the first method capable of sampling the joint posterior of source and foreground mass as pixelated images for strong gravitational lensing, moving beyond point estimates or low-dimensional parametric fits.
• Integration of RIMs and Diffusion: The paper introduces a novel architecture where a RIM is embedded within a conditional diffusion model to recursively refine posterior scores, addressing the non-linearity and high dimensionality of the lensing problem.
• Noise-Level Modeling: The method demonstrates the ability to model observations down to the noise level, a capability often lost when using simplified parametric models.

Results
The authors trained and evaluated DiRIM on two datasets: one using convergence maps from cosmological simulations and another using analytic profiles.

• Calibration: Tests of Accuracy with Random Points (TARP) plots indicate that the posterior samples are nearly perfectly calibrated, suggesting the estimator is unbiased.
• Goodness of Fit: The distribution of $\chi^2$ values for reconstructed images closely matches the theoretical $\chi^2$ distribution, confirming that the models capture the data down to the noise level.
• Complexity Handling: The method successfully models systems with complex mass distributions, including multiple main deflectors and subhalos. In cases where analytic parametric models failed to fit the data to the noise level due to lack of flexibility, DiRIM succeeded.
• Subhalo Detection: When applied to analytic maps containing NFW subhalos, the method detected the subhalos in over 98% of posterior samples, achieving a detection significance of $>3\sigma$.
• Degeneracy Sampling: The method effectively samples degenerate directions in parameter space, such as the correlation between source size and the radial slope of the convergence map.
• Robustness (OOD): The model demonstrated robustness to out-of-distribution (OOD) data, successfully inferring the morphology of sources and convergence maps that differed significantly from the training distribution (e.g., AI-generated "7" shaped galaxies and summed convergence maps).
• Ablation: Removing the RIM component (setting iterations $M=1$) resulted in biased inference and residuals with significant structure, confirming that the iterative refinement of scores is essential for the method's success.

Significance and Claims
The paper claims that DiRIM constitutes a "first proof of concept" regarding the ability of diffusion-based generative models combined with recurrent inference machines to solve high-dimensional, non-linear inverse problems in astrophysics. The authors emphasize that their method leverages flexible, high-dimensional representations to extract information from data that is inaccessible to traditional parametric modeling. They position this work as a necessary step toward realizing the scientific potential of upcoming large-scale surveys (such as LSST and Euclid), which are expected to discover over 100,000 strong gravitational lenses. The authors modestly note that while the method is effective, small biases remain in complex systems (e.g., multiple deflectors), and future work will address these and explore latent-space formulations to reduce computational costs.