galaxy-prior

Galaxy Diffusion Prior 🌌

This repository contains a score-based diffusion prior on galaxy images. It is used to perform Bayesian image reconstruction (denoising, deblurring, etc) of simulated telescope images in an educational workshop.

• Model type: score-based diffusion model (VE-SDE)
• Domain: galaxy surface-brightness images
• Resolution: 64×64 pixels (single-band)
• Author / trainer: Gabriel Missael Barco

---

1. Training data

The model was trained on mock galaxy images produced by applying dust–radiativetransfer post-processing to galaxies in the TNG cosmological magneto-hydrodynamical simulations (https://www.tng-project.org/):

Connor Bottrell, Hassen M. Yesuf, Gergö Popping, Kiyoaki Christopher Omori, Shenli Tang, Xuheng Ding, Annalisa Pillepich, Dylan Nelson, Lukas Eisert, Hua Gao, Andy D. Goulding, Boris S. Kalita, Wentao Luo, Jenny E. Greene, Jingjing Shi, and John D. Silverman. IllustrisTNG in the HSC-SSP: image data release and the major role of mini mergers as drivers of asymmetry and star formation. MNRAS, 527(3):6506-6539, January 2024. doi:10.1093/mnras/stad2971

We downsample to $64\times 64$, convert to $\mu \text{Jy} / sr$, and normalize each image to $[0,1]$. If you use this model, please cite data source appropriately.

---

2. Code base and package

The model was created and trained using the score_models package by Alexandre Adam:

• GitHub (dev branch used here): https://github.com/AlexandreAdam/score_models/tree/dev

• Method reference:

Alexandre Adam, Adam Coogan, Nikolay Malkin, Ronan Legin, Laurence Perreault-Levasseur, Yashar Hezaveh, and Yoshua Bengio. Posterior samples of source galaxies in strong gravitational lenses with score-based priors. In Machine Learning and the Physical Sciences Workshop, NeurIPS 2022, January 2022. URL https://arxiv.org/pdf/2211.03812.

Barco, G. M., Legin, R., Stone, C., Hezaveh, Y., & Perreault-Levasseur, L. (2025). Blind Strong Gravitational Lensing Inversion: Joint Inference of Source and Lens Mass with Score-Based Models. Machine Learning and the Physical Sciences Workshop @ NeurIPS 2025. Retrieved from https://arxiv.org/abs/2511.04792.

Please cite if you use this model or the associated code.

---

3. Files in this repository

.
├── checkpoint_step0434907_043.pt
# Trained diffusion model weights - Last checkpoint
├── hyperparameters.json  # Training and architecture hyperparameters
└── README.md             # This file