One-sentence Summary

Tianao Li and Emma Alexander present a physics-informed deep learning framework that unrolls the ADMM algorithm with a plug-and-play neural denoiser and learned hyperparameters for PSF deconvolution of galaxy images, achieving improvements in reduced shear ellipticity error of 38.6% (SNR=20) and 45.0% (SNR=200) over classic methods and 7.4% (SNR=20) and 33.2% (SNR=200) over modern methods for weak gravitational lensing shape measurements.

Key Contributions

• An unrolled Plug-and-Play ADMM architecture for galaxy image deconvolution combines a learned ResUNet denoiser with trainable step-size hyperparameters, trained end-to-end across a fixed number of iteration layers to integrate physics-based deconvolution steps and data-driven priors.
• On simulated LSST-like galaxy images, the method reduces reduced shear ellipticity error by 38.6% (SNR=20) and 45.0% (SNR=200) compared to classic techniques, and by 7.4% (SNR=20) and 33.2% (SNR=200) compared to modern methods, while remaining robust to systematic PSF errors.
• An open-source framework supplies source code, pretrained weights, and simulation tools that allow users to generate custom datasets and retrain the model under their own survey settings.

Introduction

Accurate galaxy shape measurement for weak gravitational lensing is essential for probing dark matter, dark energy, and cosmological models, yet atmospheric and optical blur (the point spread function, or PSF) combined with sensor noise degrades the raw images. Classical deconvolution methods such as Richardson-Lucy amplify noise, and recent physics-informed deep learning approaches like ADMMNet introduce learned denoisers within a Plug-and-Play optimization but rely on variable iteration counts and separate training stages, limiting end-to-end efficiency. The authors propose an unrolled Plug-and-Play ADMM network that fixes the number of iterations and jointly trains a denoising ResUNet together with step-size hyperparameters, yielding a fast, interpretable architecture. They benchmark this unrolled scheme on realistic LSST-like galaxy simulations and demonstrate superior shear estimation accuracy, robustness to PSF mismodeling, and reduced compute time compared to prior deconvolution methods.

Dataset

The authors construct a synthetic dataset of ground-based galaxy observations, pairing clean galaxy images with realistic point-spread functions (PSFs) and noisy blurred counterparts.

• Sources and composition
  The dataset combines the COSMOS Real Galaxy Dataset with the Galsim simulation toolkit. Simulations follow LSST specifications, using real galaxy cutouts as ground truth. Each sample consists of the clean galaxy, the overall PSF, and the convolved, noise-added image.

• Processing pipeline

  • Shearing (uniform [0.01, 0.05]) and rotation (uniform [0, 2π]) are applied to raw galaxies to mimic weak lensing.
  • An atmospheric PSF is generated with a Kolmogorov model; its seeing is drawn from a realistic distribution, and additional shear (uniform [0.01, 0.03]) and rotation inject shape variation.
  • An optical PSF is built with LSST diffraction diameter and built-in ranges for obscuration, defocus, coma, and astigmatism. The two PSFs are convolved to form the final PSF, which is saved alongside the data.
  • The galaxy is convolved with the final PSF, and Gaussian noise (σ = 19.4 ADUs) is added based on LSST read noise and a sky level of 350 ADUs/pixel².
  • All images are downsampled to the LSST pixel scale (0.2 arcsec) and mean-subtracted, but brightness is not normalised so the network can adapt to SNR directly from input intensity.

• Usage in the model
  The simulated dataset is used for both training the proposed unrolled ADMM network and benchmarking classical deconvolution methods. The paper does not detail explicit train/test split sizes or mixture ratios; the data generation code and a link to the pre-generated dataset are provided in the GitHub repository, along with trained weights for LSST settings.

Method

The authors leverage a physics-informed deep learning approach to address the Point Spread Function (PSF) deconvolution problem in galaxy surveys. To solve this ill-posed linear inverse problem, they apply algorithm unrolling and the Plug-and-Play technique to the Alternating Direction Method of Multipliers (ADMM). In this framework, a neural network is tasked with learning appropriate hyperparameters and denoising priors directly from simulated galaxy images.

As shown in the figure below:

The overall pipeline consists of two main components: image simulations and the unrolled ADMM network. For the image simulations, the authors utilize the COSMOS dataset to generate ground truth galaxy images. They randomize shear, rotation, and seeing to create corresponding PSF images. The ground truth is then convolved with the PSF and noise is added to produce the observation, which serves as the input for the network.

The core of the method is the Unrolled ADMM Network, which comprises $N$N layers. Each layer executes a sequence of operations: deblurring, denoising via a neural network, and Gaussian Maximum Likelihood Estimation (MLE). A hyperparameter subnetwork is integrated into the architecture to dynamically learn and provide optimal parameters for each iteration. The shared weights across the iteration layers in this unrolled network allow the model to efficiently learn the denoising priors and hyperparameters, facilitating robust reconstruction of the galaxy images from the blurred observations.

Experiment

The evaluation setup uses simulated galaxy images and compares the proposed unrolled ADMM method with classical and learned deconvolution techniques by measuring reduced shear ellipticity error across multiple SNR levels. Experiments show that the fixed-iteration unrolled ADMM consistently achieves the most accurate shape measurements, with ablation studies confirming the importance of jointly training the denoiser and the hyperparameter subnetwork. The method offers a favorable time-performance trade-off compared to iteration-to-convergence approaches and is robust under typical PSF errors, though its sensitivity to large systematic PSF mismatch can be reduced by using fewer unrolled iterations.