Andi Gu: GIGA-Lens: A Fast Differentiable Bayesian Inference Framework for Strong Lensing
Channel: NERSC
Date: 2022-12-20

GIGA-Lens: A Fast Differentiable Bayesian Inference Framework for Strong Lensing
 Andi Gu, Harvard University

December 13, 2022
12:00-12:30 Pacific Time

Strong gravitational lensing systems constitute a powerful tool for cosmology.  They are uniquely suited to probe the low-end of the dark matter mass function and test the predictions of the cold dark matter model beyond the local universe. Multiply lensed quasars (and supernovae in the near future) are being used to provide independent constraints on the Hubble constant.  In this talk, we present GIGA-Lens: a gradient-informed, GPU-accelerated Bayesian framework for modeling strong gravitational lensing systems, implemented in TensorFlow and JAX. The three components, optimization using multi-start gradient descent, posterior covariance estimation with variational inference, and sampling via Hamiltonian Monte Carlo, all take advantage of gradient information through automatic differentiation and parallelization on GPUs. The average time to model a simulated system on four Nvidia A100 GPUs is 105 seconds. The robustness, speed, and scalability offered by this framework make it possible to model the large number of strong lenses found in current surveys and present a very promising prospect for the modeling of O(10^5) lensing systems expected to be discovered in the era of the Rubin Observatory, Euclid, and Roman Space Telescope.