Probabilistic Inversion of Strong Gravitational Lenses with Flow Matching Priors

The 9th Workshop on Machine Learning and the Physical Sciences (ML4PS), 2026

Tianao Li1,5, Jimena González1,5, Xinfeng Xu1,3,5, Aggelos K. Katsaggelos1,4,
Rebecca Willett2,5, Aleksandra Ćiprijanović2,4,5, Tjitske Starkenburg1,5, Emma Alexander1,5
1Northwestern University, 2University of Chicago, 3Zhejiang University,
4Fermi National Accelerator Laboratory, 5NSF-Simons AI Institute for the Sky (SkAI)
Overview of the Blind-FlowSGS sampler
A high-level overview of Blind-FlowSGS, a principled blind posterior sampler using flow matching as a plug-and-play prior. Our method iterates between a source image likelihood step, a source image prior step, a lens mass likelihood step, and a lens mass prior step. We show sampling trajectories of the source galaxy on the right.

Abstract

Strong gravitational lensing reveals important information about the universe, but reconstructing the source image and lens properties requires solving an ill-posed inverse imaging problem. Generative image models serve as powerful data-driven priors in Bayesian inverse problems. Diffusion models have been used widely, but they are computationally costly. Flow matching has emerged as a promising alternative to diffusion models, with the advantage of straighter probability paths that enable easier training, efficient sampling, and better-quality data generation. In this work, we present Blind-FlowSGS, a principled blind posterior sampling method using flow matching priors, enabling efficient posterior sampling for non-blind and blind inversion of strong gravitational lenses. On simulated Euclid data, we demonstrate our method's improvement over diffusion-based methods in both efficacy and efficiency. Our method shows promising capabilities to accurately process strong lensing data from modern sky surveys at scale.

Citation

@article{li2026probabilistic,
    author    = {Li, Tianao and Gonz\'alez, Jimena and Xu, Xinfeng and Katsaggelos, Aggelos K. and Willett, Rebecca and \'Ciprijanovi\'c, Aleksandra and Starkenburg, Tjitske and Alexander, Emma},
    title     = {Probabilistic Inversion of Strong Gravitational Lenses with Flow Matching Priors},
    booktitle = {The 9th Workshop on Machine Learning and the Physical Sciences (ML4PS)},
    year      = {2026}
}

Acknowledgements

We gratefully acknowledge the support of the NSF-Simons AI Institute for the Sky (SkAI) via grants NSF AST-2421845 and Simons Foundation MPS-AI-00010513. This material is based upon work supported by the U.S. National Science Foundation under Award No. 2542022. This research used the Delta advanced computing and data resource, which is supported by the National Science Foundation (award OAC 2005572) and the State of Illinois. Delta is a joint effort of the University of Illinois Urbana-Champaign and its National Center for Supercomputing Applications. This work was produced by Fermi Forward Discovery Group, LLC under Contract No. 89243024CSC000002 with the U.S. Department of Energy, Office of Science, Office of High Energy Physics. The authors would like to thank Yu Sun for helpful discussions.