Project · Astrophysics × Deep Learning

GravLens

Generating and classifying strong gravitational lensing imagery: a DDPM for synthesis, and a residual network for substructure classification

Harishwar Rao J.

GravLensDiffusion GravLensNet

Overview

Strong gravitational lensing is one of the few direct probes we have into dark matter substructure. The trouble is data: real lensing systems are rare, and the traditional route to more examples, physical simulation, is computationally expensive and slow.

GravLens is two halves of the same problem. GravLensDiffusion asks whether a generative model can stand in for simulation and produce usable lensing imagery cheaply. GravLensNet asks whether substructure can be read off those images automatically at all. Together they were my first end-to-end diffusion project, built from scratch on free Colab GPUs.

GravLensDiffusion: Generation

A DDPM implemented closely following the original denoising-diffusion paper[1], trained on 10,000 strong lensing images. The aim was data augmentation: generating high-fidelity lensing images to counter scarcity and class imbalance in the real dataset.

DDPM sampling progression — Reverse diffusion over 1000 timesteps: 100 random noise samples resolving into 100 generated strong-lensing images.

Setup

Model: DDPM with a U-Net denoiser, implemented from scratch.
Data: 10,000 strong lensing images.
Training: 40 epochs on Google Colab's free GPU tier.
Evaluation: Fréchet Inception Distance over 100 generated samples.

27.93 FID
(100 samples)

10k Training
images

40 Epochs
(compute-limited)

Limitations. An FID of 27.93 is above the ~20 that is usually considered acceptable, and the constraints behind that number are worth naming: FID was computed on only 100 samples, training ran for 40 epochs because of GPU quota, and inputs were downscaled from 150px to 28px spatial resolution to fit in memory. More samples, more epochs and higher resolution would each be expected to bring the score down.

Training progression — Sample quality across training.

GravLensNet: Classification

A custom 20-layer ResNet[2] that sorts strong lensing images into three substructure categories, trained on 30,000 images with data augmentation, L2 regularisation, and a learning-rate scheduler that drops the rate when validation loss plateaus.

89.6% Test accuracy
(best model)

30k Images
train/val/test

0.968+ AUC across
all three classes

Metric	Value
Test-set accuracy (best model)	89.6%
AUC, Class 0 (`no` substructure)	0.987
AUC, Class 1 (`sphere`)	0.979
AUC, Class 2 (`vort`)	0.968

ROC curves for the three substructure classes.

Sample lensing images — Representative images from the classification dataset.

References

J. Ho, A. Jain, P. Abbeel. Denoising Diffusion Probabilistic Models. NeurIPS, 2020.
K. He, X. Zhang, S. Ren, J. Sun. Deep Residual Learning for Image Recognition. CVPR, 2016.
U-Net denoiser implementation inspired by DDPM-Pytorch.
FID computed with pytorch-fid.

DDPM ResNet Astrophysics Scientific ML