: Fast Bayesian Inference for Strong Gravitational Lens Modeling

0000-0003-2748-7333]A. Gu
Department of Physics, University of California, Berkeley, Berkeley, CA 94720
Department of Electrical Engineering & Computer Sciences, University of California, Berkeley, Berkeley, CA 94720

0000-0001-8156-0330]X. Huang
Department of Physics & Astronomy, University of San Francisco, San Francisco, CA 94117
Physics Division, Lawrence Berkeley National Laboratory, 1 Cyclotron Road, Berkeley, CA 94720

W. Sheu
Department of Physics, University of California, Berkeley, Berkeley, CA 94720
Department of Electrical Engineering & Computer Sciences, University of California, Berkeley, Berkeley, CA 94720

G. Aldering
Physics Division, Lawrence Berkeley National Laboratory, 1 Cyclotron Road, Berkeley, CA 94720

0000-0001-8156-0330]A. S. Bolton
NSF's National Optical-Infrared Astronomy Research Laboratory, 950 N. Cherry Ave., Tucson, AZ 85719

0000-0002-5828-6211]K. Boone
DiRAC Institute, Department of Astronomy, University of Washington, 3910 15th Ave. NE, Seattle, WA 98195

0000-0002-4928-4003]A. Dey
NSF's National Optical-Infrared Astronomy Research Laboratory, 950 N. Cherry Ave., Tucson, AZ 85719

A. Filipp
Technische Universität München, Physik-Department, James-Franck-Straße 1, 85748 Garching, Germany
Max-Planck-Institut für Astrophysik, Karl-Schwarzschild-Str. 1, 85748 Garching, Germany

0000-0002-9253-053X]E. Jullo
Aix-Marseille Univ, CNRS, CNES, LAM, Marseille, France

0000-0002-4436-4661]S. Perlmutter
Department of Physics, University of California, Berkeley, Berkeley, CA 94720
Physics Division, Lawrence Berkeley National Laboratory, 1 Cyclotron Road, Berkeley, CA 94720

0000-0001-5402-4647]D. Rubin
Department of Physics & Astronomy, University of Hawaii, Honolulu, HI 96822

0000-0002-3569-7421]E. F. Schlafly
Lawrence Livermore National Laboratory, 7000 East Avenue, Livermore, CA 94550

0000-0002-5042-5088]D. J. Schlegel
Physics Division, Lawrence Berkeley National Laboratory, 1 Cyclotron Road, Berkeley, CA 94720

0000-0002-9063-698X]Y. Shu
Max-Planck-Institut für Astrophysik, Karl-Schwarzschild-Str. 1, 85748 Garching, Germany
Ruhr University Bochum, Faculty of Physics and Astronomy, Astronomical Institute (AIRUB), German Centre for Cosmological
Lensing, 44780 Bochum, Germany

0000-0001-5568-6052]S. H. Suyu
Technische Universität München, Physik-Department, James-Franck-Straße 1, 85748 Garching, Germany
Max-Planck-Institut für Astrophysik, Karl-Schwarzschild-Str. 1, 85748 Garching, Germany
Academia Sinica Institute of Astronomy and Astrophysics (ASIAA), 11F of ASMAB, No.1, Section 4, Roosevelt Road, Taipei
10617, Taiwan

Xiaosheng Huang, Andi Gu
<EMAIL_ADDRESS><EMAIL_ADDRESS>
We present : 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 massive parallelization on graphics processing units (GPUs).
We test our pipeline on a large set of simulated systems and demonstrate in detail its high level of performance.
The average time to model a single system on four Nvidia A100 GPUs is 105 seconds.

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

$\mathcal{O} (10^5)$ lensing systems expected to be discovered in the era of the Vera C. Rubin Observatory, Euclid, and the Nancy Grace Roman Space Telescope.

§ INTRODUCTION

Strong gravitational lensing systems are a powerful tool for cosmology. They have been used to study how dark matter is distributed in galaxies and clusters <cit.>,
and are uniquely suited to probe the low-end of the dark matter mass function and test the prediction of the cold dark matter (CDM) model beyond the local universe <cit.>.
Multiply lensed supernovae (SNe) are ideal for measuring time delays and because of their well-characterized light curves, and in the case of Type Ia, with the added benefit of standardizable luminosity [Refsdal, 1964, Treu, 2010, Oguri & Marshall, 2010], provided microlensing can be accurately characterised [Yahalomi et al., 2017, Foxley-Marrable et al., 2018]. Furthermore, SNe have the benefit of fading, so for these systems, lens models can be validated using images that are uncontaminated by bright point sources [Ding et al., 2021].
In recent years, strongly lensed supernovae,  both core-collapse [Kelly et al., 2015, Rodney et al., 2016] and Type Ia [Quimby et al., 2014, Goobar et al., 2017, Rodney et al., 2021], have been discovered.
Time-delay measurements from multiply imaged supernovae <cit.>,
combined with measurements from distance ladders <cit.> and lensed quasars <cit.>, can be an important test of the tension  between measured locally and the value inferred from the Cosmic Microwave Background <cit.>.

The introduction of neural networks to identify gravitational lens candidates in imaging surveys has

been transformational <cit.>.
In our recent work,
we discovered over 1500

new strong lenses [Huang et al., 2020, Huang et al., 2021]

in the Dark Energy Spectroscopic Instrument (DESI) Legacy Imaging Surveys [Dey et al., 2019] by using residual neural networks.
This trend will accelerate in the era of the Vera C. Rubin Observatory Legacy Survey of Space and Time (LSST), Euclid, and the Nancy Grace Roman Space Telescope,
when $\mathcal{O}(10^5)$ strong lensing systems are projected to be found in the next decade [Collett, 2015].
As the lens search efficiency has dramatically improved,
there has been significant development on strong lens modeling as well.

For example, the widely-used
[Birrer & Amara, 2018] is a fully-fledged (including, for example, a suite of plotting routines) and highly versatile lens modeling package.
However, the computational cost for lens modeling remains high.
For instance, Rojas et al., 2021 modeled

41 systems with a single deflector in $gri$ bands from the Dark Energy Survey using a pipeline based on .
They used Sérsic profiles for lens and source light, and for the lensing potential, a singular isothermal ellipsoid model  and external shear.

They reported that modeling a single system took 4.3 hours on average.

In addition, uses particle swarm optimization (PSO). PSO is an easily parallelizable heuristic algorithm for non-convex optimization. However, it does not have any guarantee of convergence, especially in high dimensions (see <ref>).
[Foreman-Mackey et al., 2013], a popular MCMC algorithm in astrophysics, including being used in ,

relies on sampling techniques that also show undesirable behavior in high dimensional parameter spaces (see <ref>).
For lens modeling using high resolution images, often more complex lens and source models

will be used, and this can significantly increase the dimensionality of the parameter space.

For strong gravitational lenses to realize their full potential as an effective probe for cosmology,

it is crucial that we address these issues and

make the process of modeling lensing systems

robust, considerably faster, and scalable to high-dimensional parameter spaces.

In this paper, we present a Bayesian lens modeling framework that fulfills these requirements.

We describe our gradient-informed, GPU-accelerated () framework in <ref>.

In <ref>, we
demonstrate the performance of our pipeline on a large set of simulated systems.
We discuss our results and conclude in <ref>.

§ LENS MODELING

In this section, we introduce the [Our framework will be made public upon acceptance of this paper.] modeling framework.

In the initial stages of the development of our pipeline, served as a helpful guide,
specifically, in its approach of using optimization to find a region of high posterior density from which the MC sampler can be initialized. We believe our framework represents a significant improvement upon the  modeling pipeline in terms of speed, optimization, and sampling.
Our entire framework is

implementated in both TensorFlow [Abadi et al., 2015] and JAX [Bradbury et al., 2018]. Complete integration with either of these libraries confers significant advantages.
It enables seamless execution of our code on graphics processing units (GPUs), which

can achieve much faster gravitational lens simulation and modeling.
Even modest, freely available GPUs

are capable of performing basic linear algebra operations
(which are at the core of lens modeling codes) one to two orders of magnitude faster than a typical CPU.

In addition, our tight integration with TensorFlow allows us to use the TensorFlow Probability [Dillon et al., 2017] library for probabilistic modeling (which also provides support for JAX). Conveniently, this library has already implemented
advanced statistical methods such as variational inference and adaptation algorithms (for step size and trajectory length; Hoffman et al., 2014, Hoffman et al., 2021) in Hamiltonian Monte Carlo.
These features play a central role in our pipeline.

In <ref>, we describe our lens model, both its physical and probabilistic aspects.
Obtaining the gradient for the posterior with automatic differentiation is presented in <ref>.
Next, we detail the three main steps in our modeling pipeline.
In <ref>, we find the global optimal values for the lensing parameters using multi-start gradient descent.
In <ref>, we estimate the covariance matrix around the global optimum using variational inference.
This allows us to sample the posterior distribution efficiently with Hamiltonian Monte Carlo in <ref>.

§.§ Model Specification

§.§.§ Physical Model

Given a lens model fully described by a set of parameters $\Theta$, the predicted counts per second at arbitrary positions on the image plane, $\mathcal{I}_{model}(x,y; \Theta)$, can be determined by evaluating the deflection angle at $(x,y)$, using this deflection to ray trace <cit.> onto the source plane, setting the surface brightness of the lensed source at $(x,y)$ to the corresponding surface brightness on the source plane (via surface brightness conservation), then adding the lens light, and finally convolving with the PSF[The sky brightness, $I_{sky}$, can also be included in the model parameters $\Theta$. In this work, we assume the sky brightness has been subtracted.]. To demonstrate how our modeling pipeline works, we simulate a reference system (<ref>) using . The lens model used for this simulation consists of an elliptical power law[We note the EPL model is equivalent to the power law elliptical mass density profile <cit.>.] (EPL) mass model for the lens [Tessore & Metcalf, 2015] and external shear. The EPL model is characterized by the surface mass density in units of the critical density, or convergence,

\begin{equation}
\kappa(x_{\agu{lens}},y_{\agu{lens}}) = \frac{3-\gamma_{epl}}{2} \qty(\frac{\theta_E}{\sqrt{qx_{\agu{lens}}^2+y_{\agu{lens}}^2/q}})^{\gamma_{\agu{epl}}-1},
\end{equation}

where $\theta_E$ is the Einstein radius, $\gamma_{epl}$ is the mass profile slope,
the coordinates $x_{\agu{lens}},y_{\agu{lens}}$ are aligned with the major and minor axes of the lens, and $q$ is the axial ratio. The transformation between the image coordinates $(x,y)$ and the lens-centric coordinates $(x_{lens},y_{lens})$ is
\begin{equation}
\mqty[x_{lens} \\ y_{lens}] = \mqty[\cos \phi & \sin \phi \\ -\sin \phi & \cos \phi] \mqty[x-x_{epl} \\ y-y_{epl}] \label{eqn:coord-trans}
\end{equation}

where $\phi$ is the position angle and $x_{epl}$ and $y_{epl}$ are the lens center.
The deflection angle for this model can be written in terms of the Gaussian hypergeometric function, which can be calculated iteratively, converging to a high degree of accuracy in a small number of iterations. The truncation error reaches $\lesssim 10^{-16}$ within 35 iterations for $q \gtrsim 0.5$ [Tessore & Metcalf, 2015].

To include the effects of the local environment, we include an external shear with a deflection angle

\begin{equation}
\alpha_{ext}(x,y) = \qty(\gamma_{ext,1} x + \gamma_{ext,2} y, \gamma_{ext,2} x - \gamma_{ext,1} y).
\end{equation}

Finally, for our simulations, we model our lens and source light with an elliptical Sérsic profile [Sérsic, 1963]:

\begin{equation}
I(x_{\agu{light}},y_{\agu{light}}) = I_0 \exp(-b_n \qty(\qty(\frac{\sqrt{qx_{\agu{light}}^2 + y_{\agu{light}}^2/q}}{R_{eff}})^{1/n} - 1)),
\end{equation}

where $b_n=1.9992n - 0.3271$, $R_{eff}$ is the effective radius (half-light radius), $n$ is the Sérsic index, and $x_{\agu{light}},y_{\agu{light}}$ are

aligned with the major and minor axes of the light profile. The transformation between image coordinates $(x,y)$ and light coordinates $(x_{light}, y_{light})$ is identical in form to <ref>. We will use $(x_l, y_l)$ to denote the lens light center (in this work, the lens light center is not fixed to the lens mass center) and $(x_s,y_s)$ to denote the source light center. Finally, in lens modeling, eccentricities are often used through the reparameterization,

\begin{equation}
(\epsilon_1,\epsilon_2) = \frac{1-q}{1+q} \qty(\cos(2\phi), \sin(2\phi)),
\end{equation}

where $\phi$ is the position angle. We will use $(\epsilon_{epl,1},\epsilon_{epl,2})$ to denote the lens mass eccentricities, $(\epsilon_{l,1},\epsilon_{l,2})$ for the lens light eccentricities, and $(\epsilon_{s,1},\epsilon_{s,2})$ for the source light eccentricities. Similarly, $R_l$ and $R_s$ will denote the lens and source light effective radius.

In practice, $\mathcal{I}_{model}(x,y; \Theta)$ is vectorized to be evaluated on a grid of pixels simultaneously. This grid is supersampled by some integer factor $k_{super}$ (in this work, $k_{super}=2$) from the image coordinates, so there is an additional step in evaluating $\mathcal{I}_{model}$ that consists of downsampling (by averaging) the predicted image by $k_{super}$ back to the image coordinate grid.

§.§.§ Probabilistic Model

Our probabilistic model

comprises a likelihood function $\mathcal{L}(\Theta; \mathcal{I}_{obs}) \equiv p(\mathcal{I}_{obs} \mid \Theta)$ and a prior $p(\Theta)$, where $\Theta \in \mathbb{R}^d$ are the lensing system parameters, and $\mathcal{I}_{obs}$ is the observed image (in units of counts/sec).
The likelihood function requires a preprocessed

observed image $\mathcal{I}_{obs}$ as well as specification of the background (sky) Gaussian noise $\sigma_{bkg}$ and exposure time $t_{exp}$. For this work, we define:

\begin{equation}\label{eqn:model-def}
\begin{gathered}
\log \mathcal{L}(\Theta; \mathcal{I}\agu{_{obs}}) = -\frac{1}{2} \sum_{x,y} \qty[ \frac{\qty(\mathcal{I}_{obs}(x,y)-\mathcal{I}_{model}(x,y; \Theta))^2}{\sigma_{tot}^2(x,y; \Theta)} + \log\qty(2\pi\sigma_{tot}^2(x,y; \Theta))] \\
\sigma_{tot}^2(x,y; \Theta) = \sigma_{bkg}^2 + \frac{\mathcal{I}_{model}(x,y; \Theta)}{\mathcal{G} \cdot t_{exp}},
\end{gathered}
\end{equation}

where $\mathcal{I}_{model}(x,y; \Theta)$ is our forward model for the observed image, $\mathcal{G}$ is the gain (in $\mathrm{e^-}$/count), and the sum is over all pixels $(x,y)$ in the observed data.
The second term in the variance map arises from Poisson shot noise: since the expected value of the electron counts is $t_{exp} \mathcal{I}_{model} \mathcal{G}$, the Poisson variance is therefore also $t_{exp} \mathcal{I}_{model} \mathcal{G}$, and so the Poisson variance of the model image is $\frac{t_{exp} \mathcal{I}_{model} \mathcal{G}}{\mathcal{G}^2 \cdot t_{exp}^2} = \frac{\mathcal{I}_{model}}{\mathcal{G} \cdot t_{exp}}$. The form of this likelihood assumes independent per-pixel noise. However, the variance map $\sigma_{tot}^2$ can be specified to account for correlated pixels and to incorporate prior information about the noise at each pixel. For example, an alternative form for $\sigma_{tot}^2$, defined in terms of the observed image, is sometimes used: $\sigma_{tot}^2(x,y) = \sigma^2_{bkg} + \frac{\mathcal{I}_{obs}(x,y)}{ \mathcal{G} \cdot t_{exp}}$ , so that the total variance is independent of $\Theta$.
This is done for computational efficiency.
For example, it allows for linear light profile parameters to be solved for in closed form. We note, however, this definition for the variance may induce biases at low signal-to-noise ratios [Horne, 1986].[Another issue with using $\mathcal{I}_{obs}$ to define the noise map $\sigma^2_{tot}$ is that it may have negative pixel values due to Gaussian noise. In , the negative pixels are set to zero, which is not the most rigorous way to estimate the Poisson noise. Despite the slight difference in these two approaches,
for our simulated systems (<ref>), the results using both are virtually identical.] Therefore, we opt to use the more rigorous definition.

In our pipeline, this incurs little additional computational cost.
We also point out that the first term in the log-likelihood corresponds to
\begin{equation}
\chi^2 = \sum_{x,y} \frac{\qty(\mathcal{I}_{obs}(x,y)-\mathcal{I}_{model}(x,y; \Theta))^2}{\sigma_{tot}^2(x,y; \Theta)},
\end{equation}

and the second term is the normalization factor.

A lensing system simulated using  is shown on the left.
The parameters for this system are taken directly from the  starting guide Jupyter Notebook [Birrer, 2021].
Although  uses a Gaussian point spread function (PSF) for this system, we use a more realistic PSF calculated for the Hubble Space Telesopce () WFC3 F140W band using  [Krist et al., 2011] with an original pixel scale of 0.13. We then subsample this PSF to a scale of 0.065, as shown on the right. This system will be used to show the step-by-step workflow of our modeling pipeline. It will henceforth be called the reference system.

The prior for the parameters is typically defined as a product of independent distributions, each of which can be tuned either to reflect
prior knowledge about lensing parameters, or to match a physical understanding of the particular system that is being modelled (e.g., the Einstein radius can typically be estimated to within 20% from visual inspection).
Given our physical understanding of the model, our prior ought to vanish for certain regions of parameter space, such as $\theta_E<0$.
In our modeling, we make use of mappings that naturally enforce these constraints. We describe these mappings in <ref>.

In this work, we use a “simulation distribution" to generate a full set of 22 parameters for 100 lensing systems, and the prior distribution used to model these lenses (see <ref>) is a broadened version of this simulation distribution (note some rows contain two parameters):

\begin{equation}
\begin{split}\label{eqn:sim-dist}
\text{Lens mass}: &
\left\lbrace
\array{@{}r@{\quad}l@{}r@{}}%
\theta_{E} &\sim \exp(\mathcal{N}(\ln 1.25, \scale{0.25}{0.4})) & \hspace{78 pt} \text{[Einstein radius(\twopr)]} \\
\gamma_{epl} &\sim \mathcal{TN}(2, \scale{0.25}{0.5}; 1, 3) & \text{[Mass slope]} \\
\epsilon_{epl,1}, \epsilon_{epl,2} &\sim \mathcal{N}(0, \scale{0.1}{0.2}) & \text{\qquad [Lens mass eccentricities]} \\
x_{epl}, y_{epl} &\sim \mathcal{N}(0, \scale{0.03}{0.06}) & \text{[Lens mass center(\twopr)]} \\
\gamma_{ext,1}, \gamma_{ext,2} &\sim \mathcal{N}(0, \scale{0.05}{0.1}) & \text{[External shear components]}
\endarray \right. \\
\text{Lens light}: &\left\lbrace
\array{@{}r@{\quad}l@{}r@{}}%
R_{l} &\sim \exp(\mathcal{N}(\ln 1.6, \scale{0.15}{0.25})) & \hspace{65 pt}  \text{[Lens S\'ersic radius(\twopr)]} \\
n_{l} &\sim \mathcal{U}(\scale{2}{0.5}, \scale{6}{8}) & \text{[Lens S\'ersic index]} \\
\thinspace \epsilon_{l,1}, \epsilon_{l,2} &\sim \mathcal{TN}(0, \scale{0.05}{0.1}; -0.15, 0.15) & \text{[Lens light eccentricities]}\\
x_{l}, y_{l} &\sim \mathcal{N}(0, \scale{0.01}{0.02}) & \text{[Lens light center(\twopr)]} \\
I_l &\sim \exp(\mathcal{N}(\ln 300, \scale{0.3}{0.5})) & \text{[Lens half light intensity]}
\endarray \right. \\
\text{Source light}: &\left\lbrace
\array{@{}r@{\quad}l@{}r@{}}%
R_s &\sim \exp(\mathcal{N}(\ln 0.25, \scale{0.15}{0.25})) & \hspace{65 pt} \text{[Source S\'ersic radius(\twopr)]} \\
n_s &\sim \mathcal{U}(0.5, \scale{4}{8}) & \text{[Source S\'ersic index]} \\
\thinspace \epsilon_{s,1}, \epsilon_{s,2} &\sim \mathcal{TN}(0, \scale{0.15}{0.3}; -0.5, 0.5) & \text{[Source light eccentricities]}\\
x_s, y_s &\sim \mathcal{N}(0, \scale{0.25}{0.5}) & \text{[Source light center(\twopr)]} \\
I_s &\sim \exp(\mathcal{N}(\ln 150, \scale{0.5}{0.9})) & \text{[Source half light intensity]}
\endarray \right.
\end{split}
\end{equation}

where $\mathcal{U}(a,b)$ is a uniform distribution with support $\qty[a,b]$, $\mathcal{N}(\mu, \sigma)$ is Gaussian with mean $\mu$ and standard deviation $\sigma$, and $\mathcal{TN}(\mu, \sigma; x_{low}, x_{high})$ is a truncated Gaussian with support $\qty[x_{low}, x_{high}]$. For the distribution parameters, we use the notation $\scale{a}{b}$ to indicate that the simulation distribution uses the parameter $\textcolor{red}{a}$ while the prior uses $\textcolor{blue}{b}$. For instance, when generating our dataset, we sample $x_{epl}$ from $\mathcal{N}(0, 0.03)$, while during modeling, our prior for $x_{epl}$ is $\mathcal{N}(0, 0.06)$.

Note that $\exp(\mathcal{N}(\ldots))$ is a log-normal distribution, used for 5 parameters. Since these distributions are not as intuitive, we show them in <ref>.
The simulation distribution for the light amplitudes $I_l$ and $I_s$ have been chosen in such a way that the typical signal-to-noise ratio of the arcs is 100 (with a range between 30 and 200), and is comparable to the amplitudes used in the starting guide Jupyter Notebook [Birrer, 2021].

Probability density for five components of the simulation and prior distribution that are log-normal.

Although the distributions in <ref> can be specified in any format, in this work, we make a number of deliberate choices. Most importantly, for parameters such as ellipticities (e.g., $\epsilon_{epl,1},\epsilon_{epl,2}$) and centers (e.g., $x_{epl},y_{epl}$), we use normal distributions to reflect their rotational symmetry. For instance, in <ref>, we show how a normal distribution for the ellipticities results in a uniform distribution for the position angle $\phi$, whereas independent uniform distributions for the ellipticities break rotational symmetry, resulting in a non-uniform distribution for $\phi$ and a highly non-Gaussian distributions for $q$.

The corresponding distributions for $\phi, q$ when the ellipticities are normally distributed (blue) and uniformly distributed (orange).

§.§ Obtaining Gradient Information

For the models discussed in <ref>, $\mathcal{I}_{model}(x, y; \Theta)$ is a differentiable mapping with respect to the model parameters, $\Theta$. Therefore, so is the posterior density (<ref>) within the support of the prior[Although the gradient is continuous everywhere, it may vanish in certain regions of the parameter space due to extreme misalignment of the lens and source. We avoid this by setting our prior in a way such that, $\sqrt{(x_{s}-x_{epl})^2+(y_s-y_{epl})^2} \sim \theta_E$.]. This is essential. As we will show in <ref> and <ref>, gradient information about the posterior is

necessary to successfully achieve fast convergence for complex lens models. Although the posterior is theoretically differentiable, to obtain gradient with numerical differentiation (i.e., finite difference) would be prohibitively expensive.

Therefore, we exploit the fact that the mapping $\mathcal{I}_{model}$ is composed of a sequence of differentiable operations (e.g., convolution) and use automatic differentiation <cit.>.

Crucially with AD, which is implemented in both TensorFlow and JAX, the additional computational cost of evaluating the gradient is independent of the number of parameters [Baydin et al., 2018]. The full gradient $\pdv{\mathcal{I}_{model}}{\Theta_i}$ that we are interested in can be calculated in approximately the same amount of time it takes to calculate $\mathcal{I}_{model}$ itself.

Although we can calculate the gradient efficiently with AD, there are regions of parameter space where this gradient is undefined as a result of having hard boundaries for the parameters. This can have a number of undesirable effects for modeling.

For example, any algorithm with an iterative update that has a finite step size for the parameters may result in certain solutions being updated to positions outside the support of the prior.
Therefore we instead use smooth, invertible functions (termed “bijectors”), $g$, to map from an unconstrained space to the support of the prior, $\text{supp}(p)$.
For instance, the Einstein radius has a semi-infinite support $\mathbb{R}_{>0}$, so we use the exponential map

\begin{equation}
\begin{split}\label{eqn:exp}
g_{\theta_E}: %\mathbb{R} &\rightarrow \mathbb{R}_{>0} \\
z &\mapsto e^z
\end{split}
\end{equation}

For other parameters, such as the mass profile slope $\gamma$, that may have a finite support $\qty(a,b)=(1,3)$,
a convenient bijector to use is the sigmoid mapping:

\begin{equation}\label{eqn:sigmoid}
\begin{split}
g_\gamma: %\mathbb{R} &\rightarrow \qty(a,b) \\
z &\mapsto a + \frac{b-a}{1+e^{-z}}
\end{split}
\end{equation}

The joint bijector $g: \mathbb{R}^d \rightarrow \text{supp}(p)$ is typically constructed component-wise from the bijectors for each parameter.[Although $g$ can be specified in any desired form, we find that simply constructing it component-wise from the standard bijectors (<ref>) is sufficient. This component-wise construction is also inexpensive to evaluate, since each of the standard bijectors are elementary.] That is, $\Theta=g(\tilde{\Theta})$,
using the notation $\tilde{\cdot}$ to henceforth denote quantities in the unconstrained space.
This allows us to optimize or sample over the unconstrained space $\mathbb{R}^d$, rather than manually enforcing constraints on the parameters.

Since the volume element in the unconstrained space is different from the constrained parameter space, the prior needs to be modified,

\begin{equation}
\log \tilde{p}(\tilde{\Theta}) = \log p(g(\tilde{\Theta})) + \log \abs{J},
\end{equation}

where $J$ is the Jacobian of $g$ evaluated at $\tilde{\Theta}$.
The likelihood remains unchanged,

\begin{equation}
\log \tilde{p}(\mathcal{I}_{obs} \mid \tilde{\Theta}) = \log p(\mathcal{I}_{obs} \mid g(\tilde{\Theta})).
\end{equation}

Thus for the posterior,

\begin{equation}
\log \tilde{p}(\tilde{\Theta} \mid \mathcal{I}_{obs}) = \log p(\Theta \mid \mathcal{I}_{obs}) + \log \abs{J}. \label{eqn:reparam-post}
\end{equation}

The constrained parameters that correspond to physical quantities in the lens model (e.g., $\theta_E$) are hereafter called physical parameters.
With the posterior distribution thus reparameterized, it is straightforward to compute its derivatives using AD.

§.§ Maximum a Posteriori Estimate

When sampling from a high dimensional posterior using Monte Carlo (MC) algorithms, since all but a small fraction of the parameter space has a vanishing posterior density, arbitrarily chosen initializations for these samplers will take a long time to converge to the posterior distribution.
Therefore, it is necessary to identify a region of high posterior density from which MC samplers can be initialized. Typically, the maximum a posteriori (MAP) estimate

serves this purpose well:

\begin{equation}
\tilde{\Theta}^*_{MAP} = \underset{\tilde{\Theta} \in \mathbb{R}^d}{\text{argmax}} \ \log \tilde{p}(\tilde{\Theta} \mid \mathcal{I}_{obs}) = \underset{\tilde{\Theta} \in \mathbb{R}^d}{\text{argmin}} \ \qty(-\log \tilde{p}(\mathcal{I}_{obs} \mid \tilde{\Theta}) - \log \tilde{p}(\tilde{\Theta})).
\end{equation}

Here and below, the notation $\cdot^*$ will denote a local optimum,
and the notation $\cdot^*_{MAP}$ will indicate that this local optimum is also globally optimal. Finding $\tilde{\Theta}_{MAP}^*$ is an optimization task with the objective function $f$ being the negative (unnormalized) log posterior density:

\begin{equation}
f(\tilde{\Theta})= -\qty(\log \tilde{p}(\mathcal{I}_{obs} \mid \tilde{\Theta}) + \log \tilde{p}(\tilde{\Theta})). \label{eqn:obj-fun}
\end{equation}

There are different approaches to global optimization for non-convex, multimodal functions such as <ref>.
For example,  uses particle swarm optimization (PSO). Although this method is attractive because it requires only the evaluation of the objective function $f$ (and not its gradient),

a major weakness is that there are “little to no guarantees” [Sengupta et al., 2018] for finding the global, or even local, minimum.

We take the approach of gradient descent,
which has a number of advantages over

heuristic optimization techniques such as PSO.
Most importantly, gradient descent can at least guarantee convergence towards a local minimum.
Furthermore, when close enough to a local minimum, $\tilde{\Theta}^*$,
gradient descent approximately achieves a geometric convergence rate.
That is, $f(\tilde{\Theta}^{(k)}) - f(\tilde{\Theta}^*)$ is upper bounded by $\order{\beta^k}$, where $\beta < 1$ and $\tilde{\Theta}^{(k)}$ is the candidate solution on the $k$th iteration, called the $k$th iterate [Nesterov, 2014].

To achieve the global optimum, we disperse a large number of samples $n_{MAP}$ throughout a wide region of the parameter space and carry out gradient descent on each of these samples [Martí, 2003, György & Kocsis, 2011].
Since the iterates of these samples will quickly and reliably converge towards local minima, we only need to ensure that $n_{MAP}$ is large enough such that at least one of the samples
reaches sufficient proximity of the global minimum. The loss function $f(\tilde{\Theta})$ is multimodal in the sense that there exist multiple local minima[However, the posterior is not multimodal in the sense that each of these local minima has vanishing posterior density relative to the global mode. Although each of these local minima have a large effect on the performance of MAP, for sampling, they are irrelevant.], but we find that a moderate number of samples, $n_{MAP} = 300$ and $K_{MAP}=300$ iterations of gradient descent is sufficient for consistent identification of the global optimum. After $K_{MAP}$ iterations, we take the best of the $n_{MAP}$ samples to be the MAP estimate $\tilde{\Theta}^*_{MAP}$.
As we show in <ref>, many samples, even some that start far away from the global minimum, converge to the neighbourhood of the correct solution. If the gradient descent were run for more iterations, each of them would reach the same, globally optimal solution (see <ref>). Typically $\sim 5\%$ (and at least $1\%$, for systems with a low signal-to-noise ratio) of the samples converge to the global optimum, with the other samples eventually converging to local optima. This demonstrates the robustness of the multi-starts gradient descent method.

Each of the $n_{MAP}$ samples are initialized by sampling from the prior: $\tilde{\Theta}_i^{(1)} \sim \tilde{p}(\tilde{\Theta})$ for $i=~1,\ldots, n_{MAP}$.
Finally, similar to the training of neural networks, we use the Adam optimizer [Kingma & Ba, 2017] with an initial large learning rate[The learning rate is used as a multiplier for the gradient when updating parameters in gradient descent, with the simplest implementation being $\tilde{\Theta}^{(k+1)}_i=\tilde{\Theta}^{(k)}_i-\alpha \grad f(\tilde{\Theta}^{(k)}_i)$. The Adam optimizer is slightly more complex, and rescales the components of the gradient before applying the update. We refer readers to Kingma & Ba, 2017 for a more detailed description of the Adam update rule.] $\alpha=10^{-2}$ to accelerate

“learning” and escape spurious local minima, and decay it to $\alpha=10^{-3}$ over 300 iterations to help the optimization converge and avoid instabilities [You et al., 2019]. This is reflected in the loss trajectories shown in <ref>, which show a period of rapid improvement in the first $\sim 200$ iterations toward the neighborhood that surrounds the global mode, followed by approximate geometric convergence.

As noted in <ref>, with AD,
the gradient can be calculated at virtually no additional computational cost.
Furthermore, after initializing $n_{MAP}$ samples, the optimization can be done simultaneously by virtue of our code's parallelization.
Combining these two performance enhancements, this step of our modeling pipeline is fast:
to find the global optimum, it takes just 17 seconds (see <ref>, <ref>) to run 300 iterations of gradient descent with $n_{MAP}=300$ (<ref>).

The error trajectories of the lensing parameters (i.e., the difference between the $k$th iterate and the ground truth) for the reference system (see <ref>) over the course of gradient descent. Trajectories that end in $\chi_\nu^2 \leq 1.01$ (where $\chi_\nu^2=\chi^2/\text{DOF}$) are shown with thicker lines, all of which converge to the global minimum (see <ref>). We note that a wide range of samples converge to the correct solution (see text).

Six of the best performing loss trajectories, ranked, over 800 iterations of gradient descent. We have run 800 iterations (after which the total $\chi^2$ difference is $\lesssim 0.5$) solely for demonstration purposes. In our pipeline, however, we terminate the gradient descent at the 300th iteration, because at that stage, the best performing trajectory always ends up at the global minimum.
Also, with the $y$-axis in $\log$ scale, note the approximate geometric convergence of each solution starting around the 200th iteration,
after the initial steep descent.

§.§ Variational Inference

After finding the MAP estimate, $\tilde{\Theta}^*_{MAP}$, it is necessary to do an intermediate step of analysis before sampling via Monte Carlo.
Specifically, we estimate the posterior covariance matrix of the lens parameters $\Sigma$. This covariance estimate plays an important auxiliary role for Monte Carlo sampling: it sets a scale for each of the parameters that is used to define a proposal distribution in the sampling step, as we will show in <ref>.
We find this covariance estimate by using variational inference <cit.> to fit a multivariate normal $\mathcal{N}(\tilde{\mu},\tilde{\Sigma})$ (called the “surrogate” posterior) with a probability density $\tilde{q}(\tilde{\Theta}; \tilde{\mu}, \tilde{\Sigma})$ to the true posterior.[This roughly corresponds to calculating the posterior mode and evaluating the Hessian of the log posterior density around the mode.] This means minimizing the Kullback-Leibler () divergence between the two distributions:

\begin{equation}
\tilde{\mu}^*_{VI}, \tilde{\Sigma}^*_{VI} = \underset{\tilde{\mu}, \tilde{\Sigma}}{\text{argmin}} \ \kl (\tilde{q}(\tilde{\Theta}; \tilde{\mu}, \tilde{\Sigma}) \mid \mid \tilde{p}(\tilde{\Theta} \mid \mathcal{I}_{obs})) \label{eqn:vi}
\end{equation}

Since the posterior density $\tilde{p}(\tilde{\Theta} \mid \mathcal{I}_{obs})$ is intractable, we decompose the :

\begin{equation}
\begin{split}
\kl (\tilde{q}(\tilde{\Theta}; \tilde{\mu}, \tilde{\Sigma}) \mid \mid \tilde{p}(\tilde{\Theta} \mid \mathcal{I}_{obs})) &= \mathbb{E}_{\tilde{\Theta} \sim \mathcal{N}(\tilde{\mu}, \tilde{\Sigma})}\qty[\log \tilde{q}(\tilde{\Theta}; \tilde{\mu}, \tilde{\Sigma})-\log \frac{\tilde{p}(\mathcal{I}_{obs}, \Theta)}{p(\mathcal{I}_{obs})}] \\ &= \underbrace{\mathbb{E}_{\tilde{\Theta}}[\log \tilde{q}(\tilde{\Theta}; \tilde{\mu}, \tilde{\Sigma})-\log \tilde{p}(\mathcal{I}_{obs},\tilde{\Theta})]}_{\text{ELBO loss}} + \underbrace{\mathbb{E}_{\tilde{\Theta}}[\log p(\mathcal{I}_{obs})]}_{\text{Independent of $\tilde{\mu},\tilde{\Sigma}$}} \label{eqn:elbo}
\end{split}
\end{equation}

where $\mathbb{E}_{\tilde{\Theta}}$ denotes the expectation value with respect to the surrogate posterior $\tilde{q}$.
Therefore, minimizing the divergence is equivalent to minimizing the evidence lower bound (ELBO).
This is tractable since $\tilde{p}(\mathcal{I}_{obs},\tilde{\Theta})$ can be expressed as the product of the prior $\tilde{p}(\tilde{\Theta})$ and likelihood $\tilde{p}(\mathcal{I}_{obs} \mid \tilde{\Theta})$, each of which are readily available. The gradient of this loss is also expressible as an expectation [Ranganath et al., 2014]:

\begin{equation}
\grad_{\tilde{\mu}, \tilde{\Sigma}} \kl = \grad_{\tilde{\mu}, \tilde{\Sigma}} \text{ELBO} = \mathbb{E}_{\tilde{\Theta}}\qty[\qty(\log \tilde{q}(\tilde{\Theta}; \tilde{\mu}, \tilde{\Sigma})-\log \tilde{p}(\mathcal{I}_{obs}, \tilde{\Theta})) \grad_{\tilde{\mu},\tilde{\Sigma}} \log \tilde{q}(\tilde{\Theta}; \tilde{\mu}, \tilde{\Sigma})] \label{eqn:elbo-grad}
\end{equation}

In practice, at each iteration, we use AD to calculate $\grad_{\tilde{\mu},\tilde{\Sigma}} \log \tilde{q}(\tilde{\Theta}; \tilde{\mu}, \tilde{\Sigma})$ and approximate the expectation in <ref> with a finite number of samples[This way of doing variational inference is sometimes termed stochastic variational inference [Hoffman et al., 2013] because of the stochasticity induced by the finite number of samples $n_{VI}$ used at each iteration.] $n_{VI}$ drawn from $\mathcal{N}(\tilde{\mu}, \tilde{\Sigma})$. This forms an estimator for the true gradient $\grad_{\tilde{\mu},\tilde{\Sigma}} \kl$, which we use to do gradient descent with the Adam optimizer to minimize the ELBO in <ref>.
We note that the covariance matrix $\tilde{\Sigma}$ is constrained to be positive semi-definite. In keeping with our method of using unconstraining bijectors in <ref>, we use a Cholesky bijection mapping unconstrained real vectors to positive semi-definite matrices, and optimize over this unconstrained space.[The bijector uses the fact that any covariance matrix can be written in terms of its Cholesky decomposition $\Sigma=L L^T$, where the Cholesky factor $L$ is a lower triangular matrix with a non-negative diagonal. The unconstrained space of real vectors $\mathbb{R}^{d(d+1)/2}$ is then mapped to a covariance matrix by first reshaping the vector into a lower triangular matrix, exponentiating the diagonal entries (which defines a valid Cholesky factor $L$), then multiplying $LL^T$ to find the corresponding covariance matrix.] We initialize our VI with the MAP estimate $\tilde{\mu}^{(1)}=\tilde{\Theta}^*_{MAP}$ and a diagonal covariance matrix $\tilde{\Sigma}^{(1)}=10^{-6} \mathbb{I}$, based on the intuition that it will be easier for VI to approximate the true covariance starting from an underestimate.

Based on a coarse optimization for the reference system, we run VI for $K_{VI}=1000$ iterations using $n_{VI}=500$, with the learning rate being increased quadratically from $\alpha=0$ to $\alpha=10^{-3}$ over 500 iterations.[We start from $\alpha=0$ to allow Adam to adjust its first and second order moment estimates <cit.>.] This slow increase in learning rate is to avoid initial instabilities in the optimization that may result from the crude initial guess for the covariance $\tilde{\Sigma}^{(1)}$. In our TensorFlow implementation of the modeling pipeline, we use the TensorFlow Probability methods for variational inference, whereas in our JAX implementation, we calculate <ref> directly.

The resulting best-fit distribution is only an approximation, since the true posterior is not necessarily Gaussian. <ref> show an instance of a posterior with non-Gaussian marginals, indicating the true posterior is not Gaussian. Even when the marginals of the posterior appear to be Gaussian, this does not necessarily imply the full posterior is jointly Gaussian (see <ref> and Dutta & Genton, 2014). This is consistent with the fact that the VI posterior does not always exactly agree with the true posterior (i.e., HMC samples). Nonetheless, we find that the VI covariance matrix is almost always a sufficiently good estimate of the true covariance.

§.§ Hamiltonian Monte Carlo

For the last step, sampling, we will use Hamiltonian Monte Carlo <cit.>.
HMC relies on gradient information about the posterior distribution, and is known to have several advantages over gradient-free MCMC samplers (e.g., ).
Conveniently, this gradient information is readily available via AD.
Using HMC, we achieve highly efficient sampling.
These results are shown <ref>.
Furthermore, we compare the performance of HMC and the widely used sampler in <ref> and show that in high-dimensional spaces HMC is strongly preferred.

§.§.§ Sampler Configuration

We sample from the posterior $\tilde{p}(\tilde{\Theta} \mid \mathcal{I}_{obs})$ in the unconstrained parameter space with HMC and convert back to samples of the physical parameters $p(\Theta \mid \mathcal{I}_{obs})$ using the bijector $g$.
We initialize $n_{HMC}$ chains[We use multiple chains, $n_{HMC}=50$, first because our framework naturally lends itself to parallelization, so it is more efficient to sample for $750$ iterations using $50$ chains rather than sampling for $750 \cdot 50$ iterations using one chain. Using multiple chains is also necessary to evaluate sampling diagnostics such as $\hat{R}$ (see end of this section).] by sampling from the surrogate posterior calculated by VI (namely, $\mathcal{N}(\tilde{\mu}^*_{VI},\tilde{\Sigma}^*_{VI})$),
and as is typical for Monte Carlo samplers, run $n_{burn}$ burn-in steps before sampling $n_{sample}$ times.

Correctly configuring the HMC sampler is important to realizing its advantages.
There are three hyperparameters in HMC: the step size $\epsilon$, number of leapfrog steps $L$, and the mass matrix $\tilde{M}$ that defines the momentum distribution.
Much work has been done to adaptively set the first two hyperparameters $\epsilon$ and $L$. We use the methods introduced by Hoffman et al., 2014, Hoffman et al., 2021, as implemented by TensorFlow Probability, to automatically tune (“autotune”) the hyperparameters during the first 80%[As recommended by Tensorflow Probability (see <https://www.tensorflow.org/probability/api_docs/python/tfp/mcmc/DualAveragingStepSizeAdaptation>).] of burn-in phase. They are then fixed in the remaining burn-in steps, since their adaptation generally prevents chains from reaching the stationary distribution (hence the use of only the first 80% of burn-in steps for adaptation). Part of this tuning is adjusting the step size $\epsilon$ to achieve a target acceptance probability for each proposal: a step size that is too small will result in slower sampling, but a step size that is too large will result in too many rejected proposals. Optimal values for the target acceptance probability range between 0.6 to 0.8 [Betancourt, 2018]; we use 0.75. The remaining hyperparameter, the mass matrix $\tilde{M}$, defines the momentum distribution, and this provides a way to inform the HMC algorithm about the scales and correlations of the parameters. We can significantly improve the sampling efficiency by setting $\tilde{M}$ to be the inverse covariance matrix of the posterior <cit.> — this is called “preconditioned” HMC.
This was the main purpose of the previous section, <ref>:
we set $\tilde{M}=(\tilde{\Sigma}^*_{VI})^{-1}$, the inverse of the inferred covariance matrix from VI

as defined by <ref>. There have been proposals <cit.> to adapt the mass matrix $\tilde{M}$ on the fly during the burn-in steps, which may render the VI step irrelevant. For now, these methods are not yet well-tested (however, they may be incorporated in future work), hence VI remains a necessary step of our modeling pipeline. Note that a secondary use of the VI step is that each of the $n_{HMC}$ chains are initialized (see <ref>) by sampling from the VI posterior, $\mathcal{N}(\tilde{\mu}_{VI}^*, \tilde{\Sigma}_{VI}^*)$.

We show in <ref> the posterior samples for our reference system generated using $n_{HMC}=50$ chains and $n_{burn},n_{sample}=250,750$. As with the VI step, these hyperparameters were roughly tuned on the reference system.
We report two metrics that are widely used in the statistics literature to measure the degree to which our sampler has converged. These are known as the effective sample size () and potential scale reduction factor (PSRF), $\hat{R}$ [Gelman & Rubin, 1992]. The former measures the effective number of independent samples we have drawn from the posterior by accounting for autocorrelation within each chain, and the latter is the ratio of the average within-chain variance to the variance of the pooled samples across all chains. A large and a $\hat{R}$ that is close to 1 indicates convergence has been achieved (in Gelman & Rubin, 1992, it is suggested that an appropriate condition is $\hat{R} < 1.1$).

figureA corner plot of the posterior samples for the reference system (see <ref> for the definition of each parameter). We show only the lensing parameters, marginalizing out all light parameters (marginalized posteriors for all 22 parameters are shown in <ref>).
The posterior samples are first obtained in unconstrained space from both VI and HMC, and then converted to physical parameters by applying the bijector $g$ (see text).
The samples and $0.5$, $1$, $1.5$, $2$ $\sigma$ contours (corresponding to roughly $12\%$, $39\%$, $68\%$, $86\%$ of the probability mass), for both VI and HMC are shown in blue and grey, respectively, and the ground truth is in red.
In the top right inset, we show the model reconstructed image, residuals (normalized by the square root of the noise variance map), and the reconstructed source (together with the caustic shown in green and critical curve in red) using the Bayesian mean estimate. Despite the approximately Gaussian marginal distributions on the corner plot for the true posterior, we show a random cross section of the posterior that is banana-shaped (top inset), demonstrating that the full posterior is not perfectly Gaussian (see text). Since the VI ansatz is a multivariate Gaussian, this is consistent with the fact that the marginals for the VI posterior do not entirely coincide with those of the true posterior.

figureMarginalized posterior samples for all 22 parameters of the reference system. As in <ref>, the ground truth (input values) are marked in red. We report the effective sample size () and potential scale reduction factor $\hat{R}$ for our posterior samples.
Note that for each parameter, $\ess > 26000$ and $\hat{R} < 1.01$.
This is achieved with just under 36 seconds of HMC sampling (<ref>).

§.§.§ Comparison of HMC and

Foreman-Mackey et al., 2013 implemented an affine-invariant ensemble sampler, .
It is a popular MCMC algorithm in astrophysics.
This is the sampler that  uses.
Here we compare the performance of HMC with ,
which is gradient-free, by applying both to the reference system (<ref>). To make the comparison as fair as possible, for sampling, we initialize the sampler with the recommended configuration, as detailed in Birrer, 2021, and for our HMC sampling, we initialize the sampler as detailed in <ref>. Furthermore, we run our pipeline on a single A100 GPU and on a single CPU. uses uniform priors for each parameter whereas we use the prior described in <ref>. However, we have found that the difference in priors has virtually no effect on the sampling results. Finally, for both modeling pipelines, we use the supersampling factor $k_{super}=2$ and the PSF shown in <ref>.

We take two axes of comparison between HMC and .
First, we observe that our sampling process is significantly more efficient than , as evidenced by the rate at which HMC generates independent samples, $\sim 40 \ \ess/\text{iter}$ ($\sim 300 \ \ess/\text{sec}$, on a single A100 GPU)
whereas for  it is $\sim 0.2 \ \ess/\text{iter}$ ($\sim 0.04 \ \ess/\text{sec}$) (see <ref>). Second, we compare the convergence of the two samplers. We find that although both sampling methods agree in terms of their central values, they exhibit dramatically different convergence behavior.
In <ref>, we show that individual  chains tend to devolve to random walks. This random walk behavior manifests itself in three ways. First,  makes slow progress exploring the posterior, whereas HMC draws virtually independent samples each iteration, traversing the posterior very efficiently. Second, compared with HMC, we observe high inter-chain variance in , evidenced qualitatively by the differing marginal distributions for each of the individual chains, and quantitatively by the substantially higher $\hat{R}$ for  (<ref>).
Third, in <ref> we find that the autocorrelation time for HMC is

much lower than that of : within just 10 iterations, the autocorrelation shrinks to negligible levels, compared to , which has a characteristic autocorrelation lag of $\sim 300$. The empirical autocorrelation at lag $\tau$ for a single MC chain $\qty{f_n \mid n=1,\ldots,N}$ is defined by [Sokal, 1996]:

\begin{equation}
\begin{gathered}
\hat{\rho}(\tau)=\hat{c}(\tau)/\hat{c}(0) \qc{\text{where}} \\
\hat{c}(\tau)=\frac{1}{N-\tau} \sum_{n=1}^{N-\tau} (f_n-\mu_f)(f_{n+\tau}-\mu_f) \qq{and} \mu_f=\frac{1}{N} \sum_{n=1}^N f_n.
\end{gathered}
\end{equation}

We emphasize that this autocorrelation is independent of the iteration number.
That is, burn-in does not remove autocorrelation, nor does running a chain for a very long time.

figureA comparison of HMC with  performance for the Einstein radius.
We illustrate 5 randomly chosen chains (from 50 total chains) for both HMC and  over 750 sampling iterations. The times
shown are for HMC run on a single A100 GPU and run on a modern CPU. Note the poor interchain mixing in , which leads to a high $\hat{R}$.  In particular, for , $\hat{R}$ is much higher than the recommended value of 1.1.
This is representative behavior for all other physical parameters.

figureAutocorrelation of the Einstein radius, $\hat{\rho}(\tau)$, using HMC and  (see text).
This is representative behavior for all other physical parameters.
Note the vertical scale for HMC is one tenth of that of .

Our investigation of revealed undesirable characteristics even for low to moderate dimensional spaces, as shown above in the case of the reference system with $22$ parameters. While in this regime, it is possible that through tuning and longer sampling time, higher quality convergence can still be achieved using , Betancourt, 2018 pointed out that in higher dimensions, any gradient-free sampler is likely to be much less efficient compared with HMC. Furthermore, [Huijser et al., 2017] found that in moderate ($\sim 50$) to high ($> 100$) dimensions, affine-invariant ensemble samplers (such as ) can have more severe problems. They showed that for high dimensional posteriors, in addition to slow convergence, an affine invariant sampler can misleadingly appear to converge even when it has not.
In strong lens modeling, it is critical to avoid this pernicious behavior, since models for high resolution observed data that use complex light profiles such as shapelets [Birrer et al., 2015], wavelets [Galan et al., 2021], or pixelization [Nightingale et al., 2021] can easily have $\gtrsim 50$ parameters.
The modeling of perturbations to the smooth lensing potential,
whether due to dark matter subhalos or line-of-sight halos will require even more.
If we wish to fit these sophisticated models to observed data, we must be able to do robust inference in spaces of moderate to high dimensions.

§.§ Pipeline Summary and Hyperparameter Settings

Our pipeline is a sequence of three steps

with the ultimate goal of producing a collection of samples from the posterior distribution from which robust statistical inferences can be made.
We summarize these three steps in <ref>, and report the hyperparameter and initialization settings that we used for the reference system.

From our experience of using , the PSO initialization usually needs to be at least somewhat close to the optimum. With multi-start gradient descent, we find this to be unnecessary. While samples that start near the optimum are virtually assured to reach it, as expected, those that start far away can often succeed as well (see <ref>). This suggests that multi-start gradient descent has a much weaker dependence on initialization than PSO. In the next section, we will show the application of our pipeline to 100 simulated systems. We find that the MAP initialization in <ref> does not need to be adjusted to successfully model these systems, providing further evidence that multi-start gradient descent is not sensitive to initialization, so long as the prior is broad and $n_{MAP}$ is sufficiently large. The remaining initializations (for VI and HMC) do not need to be changed either. Furthermore, the hyperparameters in <ref> that were chosen for modeling the reference system have also been found to suffice for the 100 simulated systems in <ref>.

The total execution time for our reference system is 6 minutes on a cutting-edge A100 GPU (available through NERSC Perlmutter early access[<https://www.nersc.gov/systems/perlmutter/>]). On a GPU node on Perlmutter, which has 4 A100 GPUs[Currently, only the JAX implementation of our pipeline supports distributed computing over multiple GPUs, due to the lack of support for distributed computing on TensorFlow (outside of neural networks).], it takes 105 sec.  (<ref>).

§ RESULTS

To demonstrate the performance of our lens modeling pipeline,

we simulate a sample of 100 systems using  (see <ref>).
The parameters for these systems are sampled from the simulation distribution defined in <ref>. Our prior, also defined in <ref>, has the same center as the simulation distribution, but has been broadened considerably so that it is less informative.

A sample of 100 lenses simulated using . We include the effects of Gaussian noise with standard deviation $\sigma_{bkg}=0.2$, Poisson shot noise with an exposure time $t_{exp}=100 \ \text{sec}$ with $\mathcal{G}=1$ (assuming observations), and the PSF (see <ref>). The pixel scale is $0.065\twopr$ and the cutout size is $5.2\twopr \times 5.2 \twopr$ (80 pixels by 80 pixels).

We apply our modeling pipeline as described in <ref> to each of these systems and show the excellent agreement with the ground truth (input values) in <ref>. The hyperparameters listed in <ref> were roughly tuned (to the appropriate order of magnitude) on the reference system, and left unchanged when modeling the sample of 100 simulated systems. The average time to model one simulated system is comparable to the reference system  (see <ref>).

Moreover, we find that our pipeline consistently exhibits favorable MC convergence (<ref>): even the largest $\hat{R}$ for any parameter over all 100 simulated systems was $1.017$, and the smallest ESS was $26822$, an order of magnitude higher than the typical value with .

figureDifference between the recovered parameters and ground truth (input values) for the 100 simulated systems in <ref>. The points are the mean of the posterior, and the uncertainties correspond to the 68% highest posterior density interval. Note that as $n_l$ and $n_s$ increase, their uncertainties increase as well. This is because the light becomes more compact at higher Sérsic indices, resulting in higher degeneracy between the Sérsic indices and the half-light radii.

Summary statistics of lensing parameters
Parameter     Mean error     $\mu_z$     $\langle \hat{R} \rangle$     $\max \hat{R}$     $\langle ESS \rangle$     $\min ESS$
$\bm{\theta_E}$        -0.00026     -0.04 ±0.09                          1.001                1.013                      35465            30846

$\bm{\gamma_{epl}}$         0.01608      0.12 ±0.08                          1.001                1.017                      35407            28045

$\bm{\epsilon_{epl,1}}$         0.00235      0.08 ±0.09                          1.001                1.011                      35590            29438

$\bm{\epsilon_{epl,2}}$        -0.00159      0.01 ±0.10                          1.001                1.006                      35505            31213

$\bm{x_{epl}}$         0.00031     -0.10 ±0.10                          1.001                1.004                      35617            33434

$\bm{y_{epl}}$         0.00082      0.06 ±0.09                          1.001                1.008                      35569            32745

$\bm{\gamma_{ext,1}}$         0.00088      0.07 ±0.09                          1.001                1.012                      35542            27926

$\bm{\gamma_{ext,2}}$        -0.00060     -0.06 ±0.09                          1.001                1.010                      35382            30456

$\bm{R_l}$        -0.00203     -0.09 ±0.09                          1.000                1.007                      35768            33695

$\bm{n_l}$        -0.00321     -0.05 ±0.08                          1.000                1.008                      35777            33752

$\bm{\epsilon_{l,1}}$        -0.00038     -0.21 ±0.10                          1.000                1.003                      35444            33593

$\bm{\epsilon_{l,2}}$        -0.00006     -0.01 ±0.10                          1.001                1.003                      35495            33247

$\bm{x_l}$         0.00003      0.03 ±0.09                          1.001                1.003                      35687            33086

$\bm{y_l}$         0.00006      0.13 ±0.08                          1.000                1.003                      35641            32861

$\bm{I_l}$         1.20330      0.07 ±0.08                          1.000                1.007                      35758            33603

$\bm{R_s}$         0.00599      0.13 ±0.09                          1.000                1.005                      35497            32109

$\bm{n_s}$         0.01192      0.04 ±0.09                          1.000                1.005                      35613            32820

$\bm{\epsilon_{s,1}}$        -0.00255     -0.09 ±0.10                          1.000                1.004                      35556            31342

$\bm{\epsilon_{s,2}}$         0.00282      0.05 ±0.10                          1.001                1.005                      35711            31810

$\bm{x_s}$        -0.00087     -0.06 ±0.09                          1.001                1.017                      35443            26822

$\bm{y_s}$        -0.00092      0.03 ±0.08                          1.001                1.017                      35470            29876

$\bm{I_s}$        -0.09109     -0.05 ±0.09                          1.000                1.003                      35580            33416
We show the errors for the 22 lensing parameters in <ref>. Mean error denotes the average difference between the recovered parameters and the ground truth for the 100 simulated systems.
The notation $z$ denotes errors that have been scaled by the posterior standard deviation. For example, for a given system, if the posterior mean and variance of the Einstein radius are $\mathbb{E}[\theta_E],\mathbb{V}[\theta_E]$, and the ground truth Einstein radius is $\bar{\theta}_E$, then $z[\theta_E]=(\mathbb{E}[\theta_E]-\bar{\theta}_E)/\sqrt{\mathbb{V}[\theta_E]}$. We report the average (over all 100 systems) scaled error $\mu_z$ for each parameter, and find that they are all consistent with zero bias. We also report statistics for the MC convergence diagnostics, including the mean and extremal values. Specifically, for any given parameter, the $\max \hat{R}$ and $\min ESS$ values are the largest $\hat{R}$ and smallest $ESS$ for that parameter across all 100 simulated systems.

§ DISCUSSION AND CONCLUSION

In this work we present a new framework for modeling strong gravitational lenses that is robust, efficient, and scalable to high-dimensional parameter spaces. We achieve this via algorithmic improvements and extensive use of two technologies. For the former, we use multi-start gradient descent in place of PSO, and HMC augmented with VI in place of .
For the latter, first, massive parallel processing on GPUs allows us to simulate thousands of systems at once, orders of magnitude faster than existing lensing codes that use CPUs.
This fast simulation capability is key

for efficient forward modeling.
Second, automatic differentiation provides access to gradient information that is a highly valuable guide for each step in our pipeline,
at virtually no additional computational cost.

We have demonstrated our pipeline's performance on a large set of simulated systems. We make a reasonably general choice for our lens model (EPL + external shear, with lens and source light modeled with Sérsic profiles) in this work.
But we emphasize that our modeling methodology is an overarching framework.
The capabilities described above are applicable to any parameterized lens model.
For instance, if we opt instead to use shapelets [Birrer et al., 2015] as a source light model, only $\sim 10\%$ more computation time is needed.
More importantly, as we showed in <ref>, a gradient-informed modeling pipeline is necessary to do rigorous statistical inference on models with many parameters.
Fifty-one of the lensing systems that we discovered in Huang et al., 2020, Huang et al., 2021 have been observed with the Hubble Space Telescope (ID: 15867; PI: Huang). We will apply the framework to model a subset of these systems and report the results in an upcoming publication (Gu et al. in prep.).

In this work, we have developed the core components for a gradient-based lens modeling framework.
There is much room for expansion within this framework. For instance, although we did not find significant multimodality in the posterior for the model we consider in this work (i.e., all local modes have vanishing posterior density compared to the global mode), it is unclear whether this will still be the case for more complex lens and source models. We believe this can be addressed using more advanced samplers (using HMC as a substrate) such as adiabatic Monte Carlo [Betancourt, 2015], parallel tempering [Earl & Deem, 2005], or annealed importance sampling [Neal, 2001]. The latter is also capable of estimating normalizing constants, which enables the computation of Bayes factors. This is necessary for model comparison, which is particularly useful for tasks such as the modeling of subhalos and line-of-sight low-mass halos.

Finally, the execution time can very likely be significantly shortened from the 105 seconds reported in this work via a combination of technological and algorithmic improvements.
For the former, we plan to use 8 A100 GPUs, and expect that this will bring the execution time to roughly cut in half, bringing the total time to below 1 minute. In addition, further improvement on GPU speed is almost a certainty. For the latter, on one hand, advances in mass matrix adaptation for HMC [Stan Development Team, 2021] may allow the VI step to be eliminated, potentially offering up to a factor of two speed gain.
On the other hand, the VI step can be improved to fit the posterior exactly [Kingma et al., 2017, Papamakarios et al., 2018], allowing HMC to be eliminated from the pipeline.

This framework and its further improvements make it possible, for the first time, that the $\mathcal{O}(10^5)$ strong lenses expected to be discovered in the next generation surveys can be modeled on a reasonable time scale.

§ ACKNOWLEDGEMENT

This work was supported in part by the Director, Office of Science, Office of High Energy Physics of the US Department of Energy under contract No. DE-AC025CH11231. This research used resources of the National Energy Research Scientific Computing Center (NERSC), a U.S. Department of Energy Office of Science User Facility operated under the same contract as above and the Computational HEP program in The Department of Energy’s Science Office of High Energy Physics provided resources through the “Cosmology Data Repository” project (Grant #KA2401022).

X.H. acknowledges the University of San Francisco Faculty Development Fund.
Support for program 15867 was provided by NASA through a grant from the Space Telescope Science Institute, which is operated by the Association of Universities for Research in Astronomy, Inc., under NASA contract NAS 5-26555.
S.H.S. thanks the Max Planck Society for support through the Max Planck Research Group. E.J. acknowledges funding from Excellence Initiative of Aix-Marseille University - A*MIDEX, a French “Investissements d'Avenir” program (AMX-19-IET-008 - IPhU). Y.S. acknowledges support from the Max Planck Society and the Alexander von Humboldt Foundation in the framework of the Max Planck-Humboldt Research Award endowed by the Federal Ministry of Education and Research. This research used resources of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02-05CH11231 using NERSC award HEP-ERCAP0021270.

TensorFlow [Abadi et al., 2015],
TensorFlow Probability [Dillon et al., 2017],
JAX [Bradbury et al., 2018],
Optax [Hessel et al., 2020],
lenstronomy [Birrer & Amara, 2018],
emcee [Foreman-Mackey et al., 2013],
Matplotlib [Hunter, 2007],
seaborn [Waskom, 2021],
corner.py [Foreman-Mackey, 2016],
TinyTim [Krist et al., 2011],
NumPy [Harris et al., 2020]

[Abadi et al., 2015]
Abadi, M., Agarwal, A., Barham, P., et al. 2015, TensorFlow: Large-Scale
Machine Learning on Heterogeneous Systems.

[Barkana, 1998]
Barkana, R. 1998, , 502, 531, 10.1086/305950

[Baydin et al., 2018]
Baydin, A. G., Pearlmutter, B. A., Radul, A. A., & Siskind, J. M. 2018,
Journal of Machine Learning Research, 18, 1.

[Betancourt, 2018]
Betancourt, M. 2018, A Conceptual Introduction to Hamiltonian Monte Carlo.

[Betancourt, 2015]
Betancourt, M. J. 2015, Adiabatic Monte Carlo.

[Birrer, 2021]
Birrer, S. 2021,  Starting Guide.

[Birrer & Amara, 2018]
Birrer, S., & Amara, A. 2018, Physics of the Dark Universe, 22, 189,

[Birrer et al., 2015]
Birrer, S., Amara, A., & Refregier, A. 2015, , 813, 102,

[Birrer et al., 2020]
Birrer, S., Shajib, A. J., Galan, A., et al. 2020, , 643, A165,

[Blei et al., 2017]
Blei, D. M., Kucukelbir, A., & McAuliffe, J. D. 2017, Journal of the American
Statistical Association, 112, 859–877,

[Bolton et al., 2006]
Bolton, A. S., Burles, S., Koopmans, L. V. E., Treu, T., &
Moustakas, L. A. 2006, , 638, 703, 10.1086/498884

[Bonvin et al., 2017]
Bonvin, V., Courbin, F., Suyu, S. H., et al. 2017, , 465, 4914,

[Bradač et al., 2008]
Bradač, M., Allen, S. W., Treu, T., et al. 2008, , 687, 959–967,

[Bradbury et al., 2018]
Bradbury, J., Frostig, R., Hawkins, P., et al. 2018, JAX: composable
transformations of Python+NumPy programs, 0.2.5.

[Broadhurst et al., 2000]
Broadhurst, T., Huang, X., Frye, B., & Ellis, R. 2000, , 534,
L15, 10.1086/312651

[Brooks et al., 2011]
Brooks, S., Gelman, A., Jones, G., & Meng, X.-L., eds. 2011, Handbook of
Markov Chain Monte Carlo (Chapman and Hall/CRC),

[Çaǧan Şengül et al., 2021]
Çaǧan Şengül, A., Dvorkin, C., Ostdiek, B., & Tsang, A.
2021, Substructure Detection Reanalyzed: Dark Perturber shown to be a
Line-of-Sight Halo.

[Çaǧan Şengül et al., 2020]
Çaǧan Şengül, A., Tsang, A., Diaz Rivero, A.,
et al. 2020, , 102, 063502, 10.1103/PhysRevD.102.063502

[Cañameras et al., 2020]
Cañameras, R., Schuldt, S., Suyu, S. H., et al. 2020, , 644, A163,

[Collett, 2015]
Collett, T. E. 2015, , 811, 20, 10.1088/0004-637X/811/1/20

[Dey et al., 2019]
Dey, A., Schlegel, D. J., Lang, D., et al. 2019, The Astronomical Journal,
157, 168, 10.3847/1538-3881/ab089d

[Diaz Rivero & Dvorkin, 2020]
Diaz Rivero, A., & Dvorkin, C. 2020, , 101, 023515,

[Dillon et al., 2017]
Dillon, J. V., Langmore, I., Tran, D., et al. 2017, TensorFlow Distributions.

[Ding et al., 2021]
Ding, X., Liao, K., Birrer, S., et al. 2021, , 504, 5621,

[Duane et al., 1987]
Duane, S., Kennedy, A., Pendleton, B. J., & Roweth, D. 1987, Physics Letters
B, 195, 216, 10.1016/0370-2693(87)91197-X

[Dutta & Genton, 2014]
Dutta, S., & Genton, M. G. 2014, Journal of Multivariate Analysis, 132, 82,

[Earl & Deem, 2005]
Earl, D. J., & Deem, M. W. 2005, Physical Chemistry Chemical Physics, 7, 3910,

[Foreman-Mackey, 2016]
Foreman-Mackey, D. 2016, The Journal of Open Source Software, 1, 24,

[Foreman-Mackey et al., 2013]
Foreman-Mackey, D., Hogg, D. W., Lang, D., & Goodman, J. 2013, Publications of
the Astronomical Society of the Pacific, 125, 306–312,

[Foxley-Marrable et al., 2018]
Foxley-Marrable, M., Collett, T. E., Vernardos, G., Goldstein, D. A., & Bacon,
D. 2018, , 478, 5081–5090, 10.1093/mnras/sty1346

[Freedman et al., 2019]
Freedman, W. L., Madore, B. F., Hatt, D., et al. 2019, , 882, 34,

[Freedman et al., 2020]
Freedman, W. L., Madore, B. F., Hoyt, T., et al. 2020, , 891, 57,

[Galan et al., 2021]
Galan, A., Peel, A., Joseph, R., Courbin, F., & Starck, J.-L. 2021, , 647,
A176, 10.1051/0004-6361/202039363

[Gelman & Rubin, 1992]
Gelman, A., & Rubin, D. B. 1992, Statistical Science, 7, 457 ,

[Gilman et al., 2021]
Gilman, D., Bovy, J., Treu, T., et al. 2021, Monthly Notices of the Royal
Astronomical Society, 507, 2432–2447, 10.1093/mnras/stab2335

[Goldstein & Nugent, 2017]
Goldstein, D. A., & Nugent, P. E. 2017, , 834, L5,

[Goldstein et al., 2019]
Goldstein, D. A., Nugent, P. E., & Goobar, A. 2019, , 243, 6,

[Goldstein et al., 2018]
Goldstein, D. A., Nugent, P. E., Kasen, D. N., & Collett, T. E. 2018,
, 855, 22, 10.3847/1538-4357/aaa975

[Goobar et al., 2017]
Goobar, A., Amanullah, R., Kulkarni, S. R., et al. 2017, Science, 356,
291, 10.1126/science.aal2729

[Grillo et al., 2015]
Grillo, C., Suyu, S. H., Rosati, P., et al. 2015, , 800, 38,

[György & Kocsis, 2011]
György, A., & Kocsis, L. 2011, Journal of Artificial Intelligence
Research, 41, 407, 10.1613/jair.3313

[Harris et al., 2020]
Harris, C. R., Millman, K. J., van der Walt, S. J., et al. 2020, , 585,
357, 10.1038/s41586-020-2649-2

[Hessel et al., 2020]
Hessel, M., Budden, D., Viola, F., et al. 2020, Optax: composable gradient
transformation and optimisation, in JAX!, 0.0.1.

[Hezaveh et al., 2016]
Hezaveh, Y. D., Dalal, N., Marrone, D. P., et al. 2016, , 823, 37,

[Hoffman et al., 2021]
Hoffman, M., Radul, A., & Sountsov, P. 2021, in International Conference on
Artificial Intelligence and Statistics, PMLR, 3907–3915

[Hoffman et al., 2013]
Hoffman, M. D., Blei, D. M., Wang, C., & Paisley, J. 2013, Journal of Machine
Learning Research, 14

[Hoffman et al., 2014]
Hoffman, M. D., Gelman, A., et al. 2014, Journal of Machine Learning
Research, 15

[Hogg & Blandford, 1994]
Hogg, D. W., & Blandford, R. D. 1994, Monthly Notices of the Royal
Astronomical Society, 268, 889, 10.1093/mnras/268.4.889

[Horne, 1986]
Horne, K. 1986, Publications of the Astronomical Society of the Pacific, 98,
609, 10.1086/131801

[Huang et al., 2009]
Huang, X., Morokuma, T., Fakhouri, H. K., et al. 2009, , 707, L12,

[Huang et al., 2020]
Huang, X., Storfer, C., Ravi, V., et al. 2020, , 894, 78,

[Huang et al., 2021]
Huang, X., Storfer, C., Gu, A., et al. 2021, , 909, 27,

[Huber et al., 2021]
Huber, S., Suyu, S. H., Ghoshdastidar, D., et al. 2021, arXiv e-prints,

[Huijser et al., 2017]
Huijser, D., Goodman, J., & Brewer, B. J. 2017, Properties of the Affine
Invariant Ensemble Sampler in high dimensions.

[Hunter, 2007]
Hunter, J. D. 2007, Computing in Science & Engineering, 9, 90,

[Jacobs et al., 2017]
Jacobs, C., Glazebrook, K., Collett, T., More, A., & McCarthy, C.
2017, , 471, 167, 10.1093/mnras/stx1492

[Jacobs et al., 2019]
Jacobs, C., Collett, T., Glazebrook, K., et al. 2019a,
, 484, 5330, 10.1093/mnras/stz272

[Jacobs et al., 2019]
—. 2019b, , 243, 17, 10.3847/1538-4365/ab26b6

[Jullo et al., 2007]
Jullo, E., Kneib, J.-P., Limousin, M., et al. 2007, New Journal of Physics,
9, 447–447, 10.1088/1367-2630/9/12/447

[Jullo et al., 2010]
Jullo, E., Natarajan, P., Kneib, J. P., et al. 2010, Science, 329, 924,

[Kelly et al., 2015]
Kelly, P. L., Filippenko, A. V., Burke, D. L., et al. 2015, Science,
347, 1459, 10.1126/science.1261475

[Kingma & Ba, 2017]
Kingma, D. P., & Ba, J. 2017, Adam: A Method for Stochastic Optimization.

[Kingma et al., 2017]
Kingma, D. P., Salimans, T., Jozefowicz, R., et al. 2017, Improving
Variational Inference with Inverse Autoregressive Flow.

[Kochanek, 1991]
Kochanek, C. S. 1991, , 373, 354, 10.1086/170057

[Koopmans & Treu, 2002]
Koopmans, L. V. E., & Treu, T. 2002, , 568, L5,

[Koopmans et al., 2006]
Koopmans, L. V. E., Treu, T., Bolton, A. S., Burles, S., &
Moustakas, L. A. 2006, , 649, 599, 10.1086/505696

[Krist et al., 2011]
Krist, J. E., Hook, R. N., & Stoehr, F. 2011, in Optical Modeling and
Performance Predictions V, Vol. 8127, International Society for Optics and
Photonics, 81270J

[Martí, 2003]
Martí, R. 2003, in Handbook of Metaheuristics, ed. F. Glover & G. A.
Kochenberger, International Series in Operations Research &
Management Science (Boston, MA: Springer US), 355–368,

[Meneghetti et al., 2020]
Meneghetti, M., Davoli, G., Bergamini, P., et al. 2020, Science, 369,
1347, 10.1126/science.aax5164

[Metcalf et al., 2018]
Metcalf, R. B., Meneghetti, M., Avestruz, C., et al. 2018, arXiv
e-prints, arXiv:1802.03609.

[Millon et al., 2020]
Millon, M., Galan, A., Courbin, F., et al. 2020, , 639, A101,

[Narayan & Bartelmann, 1997]
Narayan, R., & Bartelmann, M. 1997, Lectures on Gravitational Lensing.

[Neal, 2011]
Neal, R. 2011, MCMC Using Hamiltonian Dynamics (CRC Press), 113–162,

[Neal, 2001]
Neal, R. M. 2001, Statistics and Computing, 11, 125,

[Nesterov, 2014]
Nesterov, Y. 2014, 1.2. Local Methods in Unconstrained Minimization (Springer),

[Nightingale et al., 2021]
Nightingale, J. W., Hayes, R. G., Kelly, A., et al. 2021, J. Open Source
Softw., 6, 2825, 10.21105/joss.02825

[Oguri & Marshall, 2010]
Oguri, M., & Marshall, P. J. 2010, , 405, 2579,

[O'Riordan et al., 2019]
O'Riordan, C. M., Warren, S. J., & Mortlock, D. J. 2019, , 487,
5143–5154, 10.1093/mnras/stz1603

[Papamakarios et al., 2018]
Papamakarios, G., Pavlakou, T., & Murray, I. 2018, Masked Autoregressive Flow
for Density Estimation.

[Pierel & Rodney, 2019]
Pierel, J. D. R., & Rodney, S. 2019, , 876, 107,

[Planck Collaboration et al., 2020]
Planck Collaboration, Aghanim, N., Akrami, Y., et al. 2020, , 641, A6,

[Quimby et al., 2014]
Quimby, R. M., Oguri, M., More, A., et al. 2014, Science, 344, 396,

[Ranganath et al., 2014]
Ranganath, R., Gerrish, S., & Blei, D. 2014, in Proceedings of Machine
Learning Research, Vol. 33, Proceedings of the Seventeenth International
Conference on Artificial Intelligence and Statistics, ed. S. Kaski &
J. Corander (PMLR), 814–822.

[Refsdal, 1964]
Refsdal, S. 1964, , 128, 307, 10.1093/mnras/128.4.307

[Riess et al., 2019]
Riess, A. G., Casertano, S., Yuan, W., Macri, L. M., & Scolnic, D. 2019, ,
876, 85, 10.3847/1538-4357/ab1422

[Riess et al., 2021]
Riess, A. G., Yuan, W., Macri, L. M., et al. 2021, A Comprehensive
Measurement of the Local Value of the Hubble Constant with 1 km/s/Mpc
Uncertainty from the Hubble Space Telescope and the SH0ES Team.

[Ritondale et al., 2019]
Ritondale, E., Vegetti, S., Despali, G., et al. 2019, , 485,
2179, 10.1093/mnras/stz464

[Rodney et al., 2021]
Rodney, S. A., Brammer, G. B., Pierel, J. D. R., et al. 2021, Nature
Astronomy, 10.1038/s41550-021-01450-9

[Rodney et al., 2016]
Rodney, S. A., Strolger, L. G., Kelly, P. L., et al. 2016, , 820,
50, 10.3847/0004-637X/820/1/50

[Rojas et al., 2021]
Rojas, K., Savary, E., Clément, B., et al. 2021, Strong lens systems search
in the Dark Energy Survey using Convolutional Neural Networks.

[Sengupta et al., 2018]
Sengupta, S., Basak, S., & Peters, R. 2018, Machine Learning and Knowledge
Extraction, 1, 157–191, 10.3390/make1010010

[Sérsic, 1963]
Sérsic, J. 1963, Boletín de la Asociación Argentina de
Astronomía, 6.

[Shu et al., 2018]
Shu, Y., Bolton, A. S., Mao, S., et al. 2018, , 864, 91,

[Shu et al., 2016]
Shu, Y., Bolton, A. S., Moustakas, L. A., et al. 2016, , 820, 43,

[Shu et al., 2015]
Shu, Y., Bolton, A. S., Brownstein, J. R., et al. 2015, , 803, 71,

[Shu et al., 2017]
Shu, Y., Brownstein, J. R., Bolton, A. S., et al. 2017, , 851, 48,

[Sokal, 1996]
Sokal, A. D. 1996, in Monte Carlo Methods in Statistical Mechanics: Foundations
and New Algorithms

[Sountsov & Hoffman, 2021]
Sountsov, P., & Hoffman, M. D. 2021, Focusing on Difficult Directions for
Learning HMC Trajectory Lengths.

[Stan Development Team, 2021]
Stan Development Team. 2021, Stan Modeling Language Users Guide and
Reference Manual, 2.28.

[Suyu et al., 2010]
Suyu, S. H., Marshall, P. J., Auger, M. W., et al. 2010, , 711,
201, 10.1088/0004-637X/711/1/201

[Suyu et al., 2013]
Suyu, S. H., Auger, M. W., Hilbert, S., et al. 2013, , 766, 70,

[Suyu et al., 2020]
Suyu, S. H., Huber, S., Cañameras, R., et al. 2020, Astronomy &
Astrophysics, 644, A162, 10.1051/0004-6361/202037757

[Tessore & Metcalf, 2015]
Tessore, N., & Metcalf, R. 2015, , 580, A79,

[Treu, 2010]
Treu, T. 2010, , 48, 87, 10.1146/annurev-astro-081309-130924

[Treu & Marshall, 2016]
Treu, T., & Marshall, P. J. 2016, , 24, 11,

[Vegetti & Koopmans, 2009]
Vegetti, S., & Koopmans, L. V. E. 2009, , 400, 1583,

[Vegetti et al., 2010]
Vegetti, S., Koopmans, L. V. E., Bolton, A., Treu, T., & Gavazzi, R.
2010, , 408, 1969, 10.1111/j.1365-2966.2010.16865.x

[Vegetti et al., 2012]
Vegetti, S., Lagattuta, D. J., McKean, J. P., et al. 2012, , 481,
341, 10.1038/nature10669

[Waskom, 2021]
Waskom, M. L. 2021, Journal of Open Source Software, 6, 3021,

[Wengert, 1964]
Wengert, R. E. 1964, Communications of the ACM, 7, 463,

[Wong et al., 2020]
Wong, K. C., Suyu, S. H., Chen, G. C.-F., et al. 2020, , 498, 1420,

[Yahalomi et al., 2017]
Yahalomi, D. A., Schechter, P. L., & Wambsganss, J. 2017, A Quadruply Lensed
SN Ia: Gaining a Time-Delay …Losing a Standard Candle.

[You et al., 2019]
You, K., Long, M., Wang, J., & Jordan, M. I. 2019, How Does Learning Rate
Decay Help Modern Neural Networks?

Below, we show modeling results for four types of typical systems: folds, cusps, crosses, and doubles.

Four archetypal lensing systems selected from <ref>: a folds, cusps, crosses, and doubles. The numbers in the top left corner refer to the ordering in <ref>. All four systems have a comparable SNR. The source location is marked with a star, the critical curves are in red, and the caustics are in green. Observe the presence of an inner critical curve and caustic in the cross system, due to the fact that $\gamma_{epl} < 2$ <cit.>.

For all four systems, the posterior mean agrees (within uncertainty) with the ground truth. In the folds (a), we point out the clear banana-shaped posterior (for similar examples with cluster lensing, see Jullo et al., 2007), as well as the weaker constraint on $\gamma$ (the standard deviation is here $\pm 0.1$, compared to the more typical $\pm 0.05$ for the other systems). This is worth keeping in mind when doing density profile slope studies. Furthermore, note the tendency for the VI posterior in (a) to underestimate the posterior scale. In contrast, for the cusp (b), the VI posterior overestimates the posterior scale for $\theta_E$. Notably, for the cross system (c), the results are qualitatively similar to the results for the reference system (which is also an approximate cross). That is, the degree of agreement between the ground truth and posterior mean is comparable to that of the reference system, and in both cases, the marginals of the VI posterior are similar to those of the true posterior. Finally, for the double (d), the marginals of the VI nearly perfectly agree with those of the true posterior.

figures/arch-0.pdf0.49(a) Folds
figures/arch-1.pdf0.49(b) Cusps
figures/arch-2.pdf0.49(c) Crosses
figures/arch-3.pdf0.49(d) Doubles
Modeling results for each of the four archetypal systems. The samples and $0.5$, $1$, $1.5$, $2$ $\sigma$ contours (corresponding to roughly $12\%$, $39\%$, $68\%$, $86\%$ of the probability mass), for both VI and HMC are shown in blue and grey, respectively.