Differentiable Strong Lensing for Complex Lens
Modelling

Luca Biggio∗
ETH Zürich

Aymeric Galan
EPFL

Austin Peel
EPFL

Georgios Vernardos
EPFL

Frédéric Courbin
EPFL

Abstract
Strong lensing is a stunning physics phenomenon through which the light emitted
from a distant cosmological source is distorted by the gravitational field of aforeground object distributed along the line of sight. Strong lensing observations
are important, since, from their analysis, it is possible to infer properties of boththe light-emitting source and the lens. In particular, precise lens modelling allows
for the extraction of precious information on the distribution of dark matter ingalaxies and clusters, which can provide tight constraints on several cosmologicalparameters. In this work, we consider the case where a comprehensive closedform parametric model of the lens potential is only partially available, and we
propose to model missing mass along the line-of-sight with a deep neural network.
We incorporate the network within a fully differentiable, physically sound strong
lensing simulator, and we train it via maximum likelihood estimation in an end-toend fashion. Our experiments show that the model is able to effectively interact withthe other components of the simulator and can successfully retrieve the underlying
potential without any assumption on its form.

1

Introduction

In strong lensing, the light emitted from a distant cosmological source (e.g. a quasar or galaxy) is
deflected by the gravitational potential exerted by a foreground object (e.g. an additional galaxy orcluster of galaxies) along the line of sight. This deflection leads us to observe an image that is adistorted version of what we would observe if the lensing object were not present [1]. Observationsof strong lensing systems represent precious probes of the structure of the lens object, namely theirdark matter content and its spatial distribution (see [2] and references therein). They can also be used
to infer the properties of very distant cosmological objects owing to the fact that the lens magnifies
the source light, allowing us to see objects that would otherwise be too faint [3].
Several efforts [4, 5, 6] have been made to use machine learning tools to extract information fromlensing images, including both the parameters of the lens and of the source. Typically, convolutional
neural networks (CNNs) [7] are used for this purpose, given their ability to process grid-like data
structures. This class of techniques has resulted in dramatic gains in inference efficiency compared toclassic Markov-chain Monte Carlo methods. Their main downside is that they need significantly large
datasets for their training procedure. This often causes a performance degradation when the same
network is tested on real data, which present inevitable discrepancies compared to simulated ones.
A different line of research leveraging the power of modern automatic differentiation (AD) tools [8]has recently emerged [9, 10]. In these works, the use of AD enables the optimization of parametersof the lensing system via gradient descent, and paves the way to gradient-based posterior samplingtechniques such as Hamiltonian Monte Carlo [11].
In this work, we build on the work of [9, 10] and [12, 13], and we propose to infer the form of
unknown perturbations to the lensing potential with a neural network in model-agnostic fashion. We∗

Corresponding author: luca.biggio@inf.ethz.ch

Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021).

do so by building a parameteric model combining physically principled and black box components ina unique differentiable strong lensing simulator. We introduce the basic principles of strong lensing
along with our neural network-based parametrization of potential perturbations in Section 2. In
Section 3, we present our preliminary results on two relatively simple lensing systems and in Section
4 we summarize our conclusions and we discuss future research directions.

2

Differentiable Strong Lensing

In the formalization of strong lensing, one usually starts from the so-called thin lens approximation,
which assumes that all the mass is contained into a single plane. This is justified, since the dimensionsof the lensing object are much smaller than the distances, DLS and DL (see Fig. 1), separating the lensfrom the source and the lens from the Earth, respectively. Under this approximation, the following
equation relating the source-plane coordinates, ~x, and the image-plane coordinates, ξ~ = (ξx , ξy ),holds2~~x = ξ~ − α~ (ξ),(1)

where α~ is the deflection angle determined by the projected gravitational potential Ψ of the lensing
object. This lensing potential can be written as an integral along the line of sight of the 3D potential
ψ as
~ ≡
Ψ(ξ)

Z ∞
ψ (ξx , ξy , z) dz,

(2)

−∞

where z represents the coordinate orthogonal to the planes. From general relativity considerations, it
can be shown that the deflection angle is related to the projected potential by the following equation

~ =
α
~ (ξ)

~
~ ~ Ψ(ξ)
DLS 2 ∇
ξ.
DS c2 DL(3)

It is common to split the potential into two terms: one describing a macroscopic smooth componentΨM , which, given the source light, reproduces the observed multiple images, and a “substructure”
component Ψsθ which describes local effects due to any additional mass along the line of sight,
either internal or external to the main lensing potential. In this work, we use a neural network withparameters θ to approximate Ψsθ . We define the substructure component on a fixed coordinate grid
and we train the network to map input locations in the grid to the corresponding pixel values. The
final potential is then given byΨ = ΨM + Ψsθ .

(4)

The last step needed to generate simulated data is to apply the distortion introduced by the potential
to the lens light intensity Isrc . This can be done by exploiting the fact that lensing conserves surface
brightness, and thus the image of the system in the lens plane, Ilens , is simply obtained by evaluating
the source light distribution, Isrc , on the lens plane. Hence, we have
~ = Isrc (ξ~ − α~Ilens (ξ)~ (ξ)).

(5)~ will depend on the parameters θ of the network through the deflection angle α
Note that Ilens (ξ)~,
which, according to Eq. 3, depends on the lens potential defined by Eq. 4. Therefore, in order to findthe parameters θ, we exploit the fact that our entire pipeline is fully differentiable and allows us tocalculate gradients of a likelihood function measuring the distance between observed data and thegenerated image Ilens . Our simulator is written in JAX [14] and our network is implemented in Flax[15], allowing us to exploit the AD engines at the core of these tools.
2

The notation used in [10] is followed hereafter.

2

Sourc
e Pla

Imag
e Pla

ne

ne

Earth

Figure 1: Illustration of the geometry of a strong lensing system under the thin lens approximation.

3

Experiments

We test our method on two different types of perturbations: a singular isothermal ellipsoid (SIE)and an external shear. The source light is modelled as a Sérsic ellipse, whereas the macroscopic
potential is described by an SIE. More details on the form of these models and on the ground truth
parameters characterizing all the components of the lensing system, including the ground truth
perturbations, are reported in the Appendix. Generated images further include instrumental andatmospheric effects to mimic real data. More specifically, after generating the image of the lensed
source, we convolve it with a gaussian point spread function and add random gaussian noise toaccount for measurement errors. For both experiments, we use the same model consisting of a simple
fully-connected neural network with three hidden layers, each with 200 neurons. We model the fullpotential on a pixellated grid and, similarly to [16], the network takes as input a pixel entry in theimage plane and outputs the corresponding intensity of the perturbation. Contrarily to [16], though,we use the Swish activation function [17], since we noted that other choices (e.g. ReLU and sine)
led to instabilities in the calculation of the gradients of the final potential with respect to the input
locations. In our experiments we found it beneficial to interrupt the training of the network after
200 epochs, since longer training resulted in overfitting the observational noise (see Appendix). The
Adam optimizer [18] with learning rate 0.001 was used to minimize the mean squared error between
simulated image and data.

Hybrid model
3

0.8
0.7

2

0.6

1

0.5

0

Ground Truth image
3

0.8
0.7

2

0.6

1

Masked noisy data
3
1

0.4

0

0.4

0

1

0.3

1

0.3

1

2

0.2

2

0.2

2

3

0.1

3

0.1

3

2

0

2

0.0

2

0

2

0.6

2

0.5

0.0

0.8

0.4
0.2
0.0
2

0

2

Figure 2: Results on the SIE perturbation. (Left) Output image from our hybrid model. (Middle)
Ground truth denoised lensed image. (Right) Noisy image data.

Hybrid model

Ground Truth image

0.8

Masked noisy data

0.7

3

0.7

3

2

0.6

2

1

0.5

1

0

0.4

0

1

0.3

1

2

0.2

2

0.2

2

3

0.1

3

0.1

3

2

0

2

0.0

0.6
0.5
0.4
0.3

2

0

2

0.0

3

0.6

2
1

0.4

0
0.2

1

0.0
2

0

2

Figure 3: Results on the shear perturbation. (Left) Output image from our hybrid model. (Middle)
Ground truth denoised lensed image. (Right) Noisy image data.
3

As shown in Figs. 2 and 3, the final lensed images generated by our hybrid simulators match theground truth denoised images very closely. This performance results from the ability of the neuralnetwork to infer the shape of the perturbation given noisy data (see Figs. 4 and 5 for the SIE and theshear perturbation respectively).
True

Reconstructed

Absolute residual

0.12

0.8
80

0.7
0.6

60

0.5
0.4

40

0.3
0.2

20

80

0.8

80

60

0.6

60

40

0.4

40

20

0.2

20

0

0.0

0

0.10
0.08
0.06
0.04
0.02

0.1
0

0

20

40

60

80

0.0

0

20

40

60

80

0

20

40

60

80

0.00

Figure 4: Results on the SIE perturbation. (Left) Ground truth perturbation. (Middle) Predicted
perturbation. (Right) Absolute residual between prediction and ground truth.

True

Reconstructed

Absolute residual
0.6

0.6
80

0.4
0.2

60

80

0.4
0.2

60

0.2
0.4

20
0

0.6
0

20

40

60

80

0.08

60

0.06

40

0.04

20

0.02

0.0

0.0
40

80

40

0.2
0.4

20

0.6
0

0

20

40

60

80

0

0

20

40

60

80

0.00

Figure 5: Results on the shear perturbation. (Left) Ground truth perturbation. (Middle) Predicted
perturbation. (Right) Absolute residual between prediction and ground truth.
As a sanity check, we also assessed the performance of our network on fitting the perturbation withnoise-free data. As shown in the Appendix, the model retrieves the ground truth potential with a very
high level of accuracy.

4

Conclusions

In this work, we have proposed a new fully differentiable strong lensing simulator employing neuralnetworks to learn unknown physics components in the lensing potential. We have shown thatthe method is effective in inferring the shape of both an SIE and shear perturbations, resulting in
final simulated images in good agreement with observed data. Our method presents a number ofinteresting properties: first, it uses neural networks as black-box models only to learn unknown partsin the modelled systems; all the remaining components are either fixed or constrained by physicallymotivated parametric assumptions. Second, the network acts within a physically sound simulator,
satisfying the well-known mechanisms underlying the strong lensing effect. Third, once trained, thenetwork can interpolate between the pixels in the image plane and predict the potential values on
points it did not use as training inputs. On the other hand, the current implementation of our method
lends itself to a number of extensions and presents some limitations. Most importantly, due to its
high expressivity, the network might overfit the training data (see Appendix for an example of thisbehaviour) by learning the noise, resulting in unrealistic potentials. Future work should focus on
introducing regularization terms on the loss function penalizing physically unreasonable solutions. It
will also be important to test the method on more complex and unstructured perturbations that might
result in more subtle distortions in the final image. Additionally, we plan to extend this approachin such a way that a similar neural network to that used in our experiments can be included at any
level of the strong lensing systems (e.g. to enhance the source model). We conclude by highlightingthat the choice of a neural network as a black box function approximator is not the only one andthat the power of our method mainly resides in its fully differentiable nature, making gradient-based
optimization on the model parameters possible.
4References[1] Tommaso Treu. Strong lensing by galaxies. Annual Review of Astronomy and Astrophysics,
48(1):87–125, 2010.
[2] R. B. Metcalf, M. Meneghetti, C. Avestruz, F. Bellagamba, C. R. Bom, E. Bertin, R. Cabanac,
F. Courbin, A. Davies, E. Decencière, and et al. The strong gravitational lens finding challenge.
Astronomy Astrophysics, 625:A119, May 2019.
[3] Philip J. Marshall, Tommaso Treu, Jason Melbourne, Raphaël Gavazzi, Kevin Bundy, S. MarkAmmons, Adam S. Bolton, Scott Burles, James E. Larkin, David Le Mignant, David C. Koo,
Léon V. E. Koopmans, Claire E. Max, Leonidas A. Moustakas, Eric Steinbring, and Shelley A.
Wright. Superresolving Distant Galaxies with Gravitational Telescopes: Keck Laser Guide Star
Adaptive Optics and Hubble Space Telescope Imaging of the Lens System SDSS J0737+3216. ,671(2):1196–1211, December 2007.
[4] Yashar D. Hezaveh, Laurence Perreault Levasseur, and Philip J. Marshall. Fast automated analysis of strong gravitational lenses with convolutional neural networks. Nature,548(7669):555–557, Aug 2017.
[5] Laurence Perreault Levasseur, Yashar D Hezaveh, and Risa H Wechsler. Uncertainties inparameters estimated with neural networks: Application to strong gravitational lensing. The
Astrophysical Journal Letters, 850(1):L7, 2017.
[6] Warren R. Morningstar, Yashar D. Hezaveh, Laurence Perreault Levasseur, Roger D. Blandford,
Philip J. Marshall, Patrick Putzky, and Risa H. Wechsler. Analyzing interferometric observationsof strong gravitational lenses with recurrent and convolutional neural networks, 2018.
[7] Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D.
Jackel. Backpropagation applied to handwritten zip code recognition. Neural Computation,
1(4):541–551, 1989.
[8] Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark
Siskind. Automatic differentiation in machine learning: a survey, 2018.
[9] Marco Chianese, Adam Coogan, Paul Hofma, Sydney Otten, and Christoph Weniger. Differentiable strong lensing: uniting gravity and neural nets through differentiable probabilistic
programming. Monthly Notices of the Royal Astronomical Society, 496(1):381–393, May
2020.
[10] Konstantin Karchev, Adam Coogan, and Christoph Weniger. Strong-lensing source reconstruction with variationally optimised gaussian processes, 2021.
[11] Handbook of markov chain monte carlo. May 2011.
[12] Simon Birrer, Adam Amara, and Alexandre Refregier. Lensing substructure quantification in
RXJ1131-1231: a 2 keV lower bound on dark matter thermal relic mass. Journal of Cosmology
and Astroparticle Physics, 2017(05):037–037, may 2017.
[13] D. Bayer, S. Chatterjee, L. V. E. Koopmans, S. Vegetti, J. P. McKean, T. Treu, and C. D.
Fassnacht. Observational constraints on the sub-galactic matter-power spectrum from galaxygalaxy strong gravitational lensing, 2018.
[14] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal
Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao
Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[15] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas
Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2020.
[16] Vincent Sitzmann, Julien N.P. Martel, Alexander W. Bergman, David B. Lindell, and Gordon
Wetzstein. Implicit neural representations with periodic activation functions. In Proc. NeurIPS,2020.
5

[17] Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions, 2017.
[18] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
[19] J. L. Sérsic. Influence of the atmospheric and instrumental dispersion on the brightness
distribution in a galaxy. Boletin de la Asociacion Argentina de Astronomia La Plata Argentina,
6:41–43, February 1963.
[20] R. Kormann, P. Schneider, and Matthias Bartelmann. Isothermal elliptical gravitational lensmodels. Astronomy and Astrophysics, 284:285–299, 03 1994.
[21] Y. Yao, Lorenzo Rosasco, and Andrea Caponnetto. On early stopping in gradient descentlearning. Constructive Approximation, 26:289–315, 2007.

6

A

Checklist
1. For all authors...
– Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [YES]– Did you describe the limitations of your work? [YES]– Did you discuss any potential negative societal impacts of your work? [N/A]– Have you read the ethics review guidelines and ensured that your paper conforms tothem? [YES]
2. If you are including theoretical results...
– Did you state the full set of assumptions of all theoretical results? [N/A]– Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
– Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [NO]. This
paper presents only preliminary results and it is part of a larger project close to
its completion. Code and documentation for our library will be published soon.
– Did you specify all the training details (e.g., data splits, hyperparameters, how they werechosen)? [YES]. Details on network architecture, data structures and optimizers
are included both in the main paper and in the Appendix.
– Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [N/A]– Did you include the total amount of compute and the type of resources used (e.g., 518
type of GPUs, internal cluster, or cloud provider)? [YES]. Details are reported in the
Appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
– If your work uses existing assets, did you cite the creators? [N/A]– Did you include any new assets either in the supplemental material or as a URL? [NO]– Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [N/A]– Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
– Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A]– Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A]– Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]

7

B

Appendix

B.1

Experimental details

Our unperturbed lensing system is described by an elliptical Sérsic source light profile [19] and an
SIE mass profile [20]. The first can be fully described by 7 parameters: the surface brightness I0 , thesemi-major axis half light radius Rsersic , the Sérsic index nsersic , the two eccentricity components es1
and es2 and the coordinates of the centre xs0 and y0s . An SIE mass profile is described by 5 parameters,
namely the Einstein radius θE , the eccentricity components el1 and el2 and the coordinates of thecentre xl0 and y0l . Our choices for these 12 parameters are listed in Table 1.
Finally, a shear potential is used to model additional angular structure in the final lensed image andresults from unstructured additional matter distribution either within the cluster or distributed along
the line of sight. It is parameterized by the two shear components, γ1 and γ2 . Our choices for theparameters of the perturbations are listed in Table 2.
For all our experiments, we select a pixel size value equal to 0.08 arcsec (similar to data obtainedwith the Hubble Space Telescope) and each image is of size 100 × 100 pixels. All our experiments
have been performed on the CPU of an Intel(R) Xeon(R) Gold 5120 processor.
Parameter Name ValueI0
10.0
Rsersic
1.2 arcsec
nsersic
1.5
es1
−0.09
es2
0.03
xs0
0.4 arcsec
y0s
0.15 arcsec
θE
1.6 arcsec
el1
0.15
el2
−0.04
xl0
0.0
y0l
0.0
Table 1: Parameter choices of source and lens models. The source light profile amplitude is given inarbitrary units, as it represents a simple scaling of pixel values.

Parameter Name Valuep
θE
0.15 arcsec
lp
e1
0.0
elp
0.3
2
xlp
1.9
0
lp
y0
−1.4
γ1
−0.1
γ2
−0.1
Table 2: Parameter choices of the lens perturbations.

B.2

Results without noise

In this section, we repeat the experiment with the SIE perturbation with the difference that we remove
the noise from the data. As shown in the figure below, the network infers very well the ground truth
potential. The final image is indistinguishable from the true one, and this can be understood bylooking at the gradient components of the perturbation, which match the truth almost perfectly.
8Hybrid model

Ground Truth image
0.8
0.7
0.6
0.5
0.4
0.3
0.2
0.1
0.0

3
2
1
0
1
2
3
2

0

2

0.8

3

0.7

2

0.6

1

0.5

0

0.4

1

0.3

2

0.2

3

0.1
2

0

0.0

2

Figure 6: Results on the SIE perturbation with noiseless observations. (Left) Output image from our
hybrid model. (Right) Noise-free image data.

True

Reconstructed

80

1.0

80

60

0.8

60

0.8

20
0

20

40

60

0.020
40

20
0

0.0

80

0.025

0.015

0.4

0.2
0

0.030

60

40

0.4

20

40

60

0.005

0.0

80

0.010

20

0.2
0

0.040
0.035

80

1.0

0.6

0.6

40

Absolute residual

1.2

1.2

0

0

20

40

60

80

0.000

Figure 7: Results on the SIE perturbation with noiseless observations. (Left) Ground truth perturbation. (Middle) Predicted perturbation. (Right) Absolute residual between prediction and groundtruth.

Partial Derivative w.r.t x (Prediction)
3

Partial Derivative w.r.t y (Prediction)
0.2

3

0.2

2

2
0.1

1
0

1

0.1

2

2

0

2

2

1

3

0

0

0.0

1

0.1

2
3

0.2
2

0.1

1

0.1

2

0.2

2

0.0

0

2

3

0.1

1

0

Partial Derivative w.r.t y (Ground-Truth)
0.2

2

0.2

3

Partial Derivative w.r.t x (Ground-Truth)
3

0.1

2

0.2

3

0.0

0

0.0

1

0.1

1

2

0.2
2

0

2

Figure 8: Results on the SIE perturbation with noiseless observations. The top (bottom) row shows
the partial derivatives of the predicted (ground-truth) perturbation.
B.3

The risks of Overfitting

In this section, we investigate the negative effects of overfitting on the final reconstructed images.
Overfitting results from the network learning the noise in the image. Looking at the figures below,
9this effect manifests itself in the appearance of multiple artifacts in the final lensed image as wellas in the potential and in its gradients. As shown in the paper, we empirically observed that thiseffect can be mitigated by applying an early stopping criterion [21] during training. As discussed in
the conclusions, future work will focus on introducing meaningful regularization terms in the loss
function in order to constrain the network to output more physically sound predictions.

Hybrid model

Ground Truth image

3

0.7

3

2

0.6

2

1

0.5

1

0

0.4

1

0.3

2
3
2

0

0.7
0.6

Masked noisy data
3

0.8

2

0.5

1

0.6

0

0.4

0

0.4

1

0.3

1

0.2

2

0.2

2

0.1

3

0.1

3

0.0

2

0.8

2

0

0.2
0.0

0.0

2

2

0

2

Figure 9: Results on the SIE perturbation with the network in the overfitting regime. (Left) Output
image from our hybrid model. Note the numerous artifacts generated by the network. (Middle)
Ground truth noiseless image. (Right) Noisy image data.

True

Reconstructed

Absolute residual
1.0

0.8
80

0.7
0.6

60

0.5
0.4

40

0.3
0.2

20

80

0.8

60

0.6

40

0.4

20

0.1
0

0

20

40

60

0.0

80

0

0

20

40

60

80

0.175
80

0.150
0.125

60

0.100
40

0.2

20

0.0

0

0.075
0.050
0.025
0

20

40

60

80

0.000

Figure 10: Results on the SIE perturbation with the network in the overfitting regime. (Left) Groundtruth perturbation. (Middle) Predicted perturbation. Note the numerous artifacts generated by the
network. (Right) Absolute residual between prediction and ground truth

Partial Derivative w.r.t x (Prediction)

Partial Derivative w.r.t y (Prediction)

3

0.3

3

0.4

2

0.2

2

0.3

1

0.1

1

0.2

0

0.1

1

0.0

2

0.1

0.0

0

0.1

1

0.2

2

0.3

3

3

0.4
2

0

2

Partial Derivative w.r.t x (Ground-Truth)
3

0.2
2

0

2

Partial Derivative w.r.t y (Ground-Truth)
3

0.10

2

0.10

2
0.05

1

1

0.05

1

0.00

0

0

0.05

2

0.00

1

0.05

2
0.10

3
2

0

2

0.10

3
2

0

2

Figure 11: Results on the SIE perturbation with the network in the overfitting regime. The top(bottom) row shows the partial derivatives of the predicted (ground-truth) perturbation.
10