[Paper Review] Strong-lensing source reconstruction with variationally optimised Gaussian processes

Konstantin Karchev, Adam Coogan+1 more|arXiv (Cornell University)|May 20, 2021

Gaussian Processes and Bayesian Inference3 references4 citations

TL;DR

This paper presents a differentiable, GPU-accelerated Bayesian strong-lensing pipeline using variational inference with Gaussian processes to reconstruct high-fidelity source images and infer lens parameters. It achieves unbiased, noise-limited reconstructions and precise uncertainty quantification in under an hour, enabling scalable training data generation for substructure inference in dark matter studies.

ABSTRACT

Strong-lensing images provide a wealth of information both about the magnified source and about the dark matter distribution in the lens. Precision analyses of these images can be used to constrain the nature of dark matter. However, this requires high-fidelity image reconstructions and careful treatment of the uncertainties of both lens mass distribution and source light, which are typically difficult to quantify. In anticipation of future high-resolution datasets, in this work we leverage a range of recent developments in machine learning to develop a new Bayesian strong-lensing image analysis pipeline. Its highlights are: (A) a fast, GPU-enabled, end-to-end differentiable strong-lensing image simulator; (B) a new, statistically principled source model based on a computationally highly efficient approximation to Gaussian processes that also takes into account pixellation; and (C) a scalable variational inference framework that enables simultaneously deriving posteriors for tens of thousands of lens and source parameters and optimising hyperparameters via stochastic gradient descent. Besides efficient and accurate parameter estimation and lens model uncertainty quantification, the main aim of the pipeline is the generation of training data for targeted simulation-based inference of dark matter substructure, which we will exploit in a companion paper.

Motivation & Objective

• Develop a fast, differentiable strong-lensing image simulator compatible with modern deep learning frameworks.
• Create a statistically principled, computationally efficient source model using layered Gaussian processes with adaptive scale selection.
• Enable scalable Bayesian inference over tens of thousands of lens and source parameters using stochastic gradient variational inference.
• Generate high-fidelity, realistic training data for simulation-based inference of dark matter substructure in future work.
• Quantify uncertainties in both lens mass distributions and source light with minimal bias and high resolution.

Proposed method

• Implement a fully differentiable, GPU-accelerated strong-lensing image simulator using PyTorch and KeOps for efficient linear algebra operations.
• Introduce a multi-layered, variationally optimized Gaussian process kernel that models multiple spatial scales in the source plane simultaneously.
• Use fixed and adaptive inducing points: fixed grids for large-scale structure and movable inducing points for fine details, reducing parameter correlation.
• Apply stochastic gradient variational inference (SVI) with a proposal posterior that captures lens-source correlations, enabling scalable posterior approximation.
• Optimize hyperparameters via stochastic gradient descent, with the evidence lower bound (ELBO) used as the objective for joint parameter and hyperparameter learning.
• Incorporate pixellation effects directly into the GP model to ensure physical consistency with observed image resolution and noise.

Experimental results

Research questions

• RQ1 Can a differentiable, end-to-end strong-lensing simulator enable fast and accurate Bayesian inference for complex lens systems?
• RQ2 How well can a layered Gaussian process model reconstruct high-resolution source images while quantifying uncertainty in the presence of noise?
• RQ3 To what extent does the choice of GP layer configuration and inducing point placement affect reconstruction fidelity and parameter inference?
• RQ4 Can variational inference with a structured proposal posterior yield conservative and unbiased posterior estimates for lens and source parameters?
• RQ5 How effectively can this pipeline generate realistic, high-fidelity training data for downstream simulation-based inference of dark matter substructure?

Key findings

• The pipeline achieves noise-limited source reconstructions on high-resolution mock observations, with residuals at the noise level.
• Lens parameter estimates are unbiased, with precision limited only by the resolution of the observation.
• The method is robust to GP layer configuration: adding many small-scale layers has only a marginal impact on residuals.
• Hyperparameter optimization effectively suppresses noise from high-frequency GP components, preventing overfitting.
• Variational posterior uncertainties for source parameters are conservative, as validated analytically and numerically in one-dimensional tests.
• The runtime scales linearly with model complexity, enabling full convergence in a few hours and reasonable posteriors in ~10 minutes on modern GPU hardware.