FlowLensing: Simulating Gravitational Lensing with Flow Matching
Hamees Sayed, Pranath Reddy, Michael W. Toomey, Sergei Gleyzer
TL;DR
FlowLensing tackles the bottleneck of simulating high-fidelity gravitational lensing images at scale by learning a direct mapping from astrophysical parameters to lens images using flow matching. The approach uses a Diffusion Transformer backbone to model the velocity field $v_theta$ and supports both discrete dark matter classes and continuous parameter conditioning with classifier-free guidance, enabling fast, controllable sampling. Results show substantial inference speedups (over 200×) while preserving image fidelity and enabling accurate downstream classification (AUC = 1.00) and regression (R^2 up to 0.945). This enables scalable dark matter studies in upcoming surveys like Euclid by enabling rapid, physically plausible simulations of lensing with substructure.
Abstract
Gravitational lensing is one of the most powerful probes of dark matter, yet creating high-fidelity lensed images at scale remains a bottleneck. Existing tools rely on ray-tracing or forward-modeling pipelines that, while precise, are prohibitively slow. We introduce FlowLensing, a Diffusion Transformer-based compact and efficient flow-matching model for strong gravitational lensing simulation. FlowLensing operates in both discrete and continuous regimes, handling classes such as different dark matter models as well as continuous model parameters ensuring physical consistency. By enabling scalable simulations, our model can advance dark matter studies, specifically for probing dark matter substructure in cosmological surveys. We find that our model achieves a speedup of over 200$\times$ compared to classical simulators for intensive dark matter models, with high fidelity and low inference latency. FlowLensing enables rapid, scalable, and physically consistent image synthesis, offering a practical alternative to traditional forward-modeling pipelines.