Physics Informed Inversion For Deep Space Imaging
Jack Walsh
20jwalsh@greystonescollege.ie
Greystones Community College
Gravitational lensing can be viewed as a powerful natural telescope to observe distant galaxies, but inverting lensing distortions outside parametric modeling remains a complicated previously unsolved inverse problem. We present AstroRim, a Recurrent Inference Machine (RIM) jointly trained with a physics informed differentiable forward lensing operator, to reconstruct unlensed source plane images from simulated strong lens observations. Using a 120,000+ sample synthetic dataset per model with complex multi component lens models such as: Single Isothermal Ellipsoid (SIE), Navarro–Frenk–White (NFW), Single Isothermal Sphere (SIS), External Shear, and extended Sersic sources at 96×96 resolution, AstroRim achieves a Structural Similarity Index (SSIM) of 0.95 ± 0.02 and Mean Squared Error (MSE) of 4.4×10⁻⁴ on unseen test data. We demonstrate stability improvements via mixed precision, gradient clipping, and Exponential Moving Average (EMA), and discuss pathways towards real data application on HST, JWST, and Euclid imagery. Our pipeline offers a scalable, physics informed approach to delensing.
Gravitational lensing by massive foreground objects distorts and magnifies background sources, enabling high resolution studies of distant galaxies. However, inverting these distortions to recover source morphology is greatly challenging: multiple source configurations can produce identical lensed images. Traditional parametric inversion methods are computationally expensive and sensitive to model assumptions, as well as a severe lack of generalization, meaning each RIM or model developed must be finetuned on simulations with the same mass profiles as the lens, making inversions lengthy. Recent advances in deep learning, particularly Recurrent Inference Machines (RIMs) and Physics Informed Neural Networks (PINNs), offer mass simulation based data driven inversion by iteratively refining source estimates through learned gradient steps. Yet, most prior work uses a fixed forward operator or pretrained physics models that are subject to overfitting and poorly adapted to data provided outside their simulation scripts. We propose AstroRim, which jointly trains the RIM and a differentiable PINN based lensing operator end to end, leveraging physical constraints while maintaining dataset flexibility. To our knowledge, this is the first delensing pipeline to jointly train a RIM and learned Differentiable Physics Informed Neural Network (DPINN) based lensing operator end-to-end.
-
Joint RIM + Physics Operator: End to end training of both modules for consistent gradient flow.
-
Complex Simulations: Varying datasets with dual mass (SIE+NFW+SIS+SHEAR) and 1–6 Sersic sources, enabling realistic multi component lens tests.
-
Stability Enhancements: Integration of mixed precision (AMP), gradient clipping, EMA, and learning rate warmup to mitigate training instabilities.
-
Quantitative Benchmarks: SSIM and MSE metrics demonstrating high fidelity reconstructions on synthetic test data. Reconstructions of validation data are also provided throughout the training process to check model adaptability.
-
Verification Via Real Data: Using databases such as the Harvard-NASA JWST and HUBBLE libraries combined with the CASTLES gravitational lens survey, each model of AstroRim produced is tested on real data to identify where our simulations and model architectures are lacking.
- Recurrent Inference Machines (RIMs): Originally applied in MRI and deconvolution, RIMs learn iterative update rules for inverse problems. Morningstar et al. (2019) applied RIMs to lensing with fixed forward models.
- Physics Informed ML: Differentiable simulators (e.g., Caustics, Lenstronomy) have been used to embed physical laws. Joint training remains underexplored in lens inversion.
- Alternative Approaches: CNN based regression of lens parameters (Hezaveh et al. 2017) and variational inference techniques achieve fast inference but often lack source-plane detail recovery.
- Lens models: SIE + NFW + SIS + Shear via Lenstronomy.
- Sources: 1–3 Sérsic profiles with varied amplitude, size, ellipticity, and position.
- Grid: 96x96 pixels, normalized coordinates [-1,1] (to 98%).
- Dataset: 250k+ pairs (85% train, 15% val) 250-pair test set.
- Forward Operator: CNN (Conv 9×9 → ReLU → Conv 5×5 → ReLU → Conv 3×3).
- RIM Core: Gated Conv RNN (hidden_dim=96), 15 inference iterations.
- Loss: MSE primary, SSIM tracked; per-image metrics monitored.
- Optimizer: AdamW + LR warmup (10% epochs) + ReduceLROnPlateau.
- Stabilization: AMP, gradient clipping (max_norm=5), EMA (0.999).
- Hardware: 6GB GPU + 2080ti GPU, 150 epochs, evaluation every 5 epochs.
Metrics: SSIM = 0.95 ± 0.02, MSE = 4.7×10⁻⁴ on test data.
Visuals: High-fidelity reconstructions matching ground truth across varied lenses and sources.
Morningstar et al. (2019) demonstrated the application of a RIM for gravitational lens inversion using a fixed, parametric forward model (ray tracing through an SIE lens) paired with a separate CNN for mass estimation (Morningstar et al., 2019). However, their approach does not cotrain the physics operator, limiting adaptability to modeling errors, and leading to inevitable failure in the context of real data. In contrast, AstroRim jointly optimizes a differentiable CNN based lensing operator and the RIM, enabling end to end correction of forward model mismatches and yielding higher fidelity reconstructions. Another angle is that of generalization. Morningstar et al. (2019) was able to produce high quality reconstructions, however these reconstructions were all using the same lens types and not accounting for variance in object location and size, as well as providing the model with data outside of the pure lensed information during evaluation.
- 96-dim RIM = higher SSIM, near hardware limits.
- 64-dim RIM = more stable training.
- Joint training accelerates convergence but can oscillate—mitigated with AMP, EMA, gradient clipping.
- Synthetic-to-real gap remains.
- Plans: add PSF/noise, scale to higher resolutions, real HST/JWST validation.
- Scale to 128×128 and multi-band images.
- Validate on real HST/JWST datasets.
- Add uncertainty estimation and RGB channels.
- Expand RIM capacity with more resources.
AstroRim shows that jointly training a RIM with a forward operator yields high-fidelity delensing. With real data validation, it could be a valuable astrophysical imaging tool.
Thanks to ML4Astro community, Lenstronomy, and SciPy developers.
Code: AstroRim GitHub
Sample Images Figure 1: Simulation example 
Figure 2: Recontruction of a lensed image compared the ground truth 
Figure 3: Recontruction of a lensed image compared the ground truth 
Figure 4: Newer simulations for inital training of the S2R model 
- Morningstar, W., et al. (2019). Deep recurrent inference for gravitational lensing. arXiv:1901.01359.
- Hezaveh, Y., et al. (2017). Fast parameter estimation in strong lensing using CNNs. ApJ.
- Shu, Y., & Bolton, A. (2020). Lenstronomy: Multi-plane lensing and ML-ready simulations. JOSS.
- Pinciroli Vago, N.O., Fraternali, P. (2023). DeepGraviLens: Neural Comput & Applic, 35, 19253–19277.
- Birrer, S., & Amara, A. (2018). lenstronomy: JOSS.
- Ruff, A. J., et al. (2011). Adaptive Optics Imaging of Galaxy-scale Gravitational Lenses. ApJ.
- Wagner-Carena, K., et al. (2021). Variational Inference for Gravitational Lens Mass Modeling. A&A.
- Li, Z., et al. (2022). GNN-Lens: ApJ.
- Nightingale, J. W., et al. (2019). Skylens: A&C.
- Serjeant, S., et al. (2019). Euclid Strong Lens Working Group Simulations. A&C.
- Zhang, T., et al. (2021). Uncertainty Quantification in RIM-based Inversions. NeurIPS AstroML Workshop.