National Energy Research Scientific Computing Center (NERSC) NERSC 2022 Early Career Achievement Award Seminars, 13 Dec 2022

Previous Meeting Next Meeting

⏯

youtube image

►

From YouTube: Andi Gu: GIGA-Lens: A Fast Differentiable Bayesian Inference Framework for Strong Lensing

Description

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

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

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

A

uh Welcome congratulations! Great.

B

Thank you, um hi everyone thanks so much for the um kind introduction and for this honor um really excited to be giving this talk. So today, I'm going to be talking about Giga lens, um which is our fast differentiable Bayesian lens modeling framework um that we've called Giga lines um yeah. So before I get started. I just want to um talk about.

B

You know the wonderful collaborators I've worked with from a variety of places, ranging from Lawrence lab to NSF um to um Max, Planck um yeah, so their input and their guidance has been really invaluable to me and I just want to take a second to acknowledge their guidance and I'd also like to thank, of course, the nurse team for all their support and resources like none of this would be possible without, um like the the support and the resources they provided.

B

Okay. So what exactly is uh strong gravitational lensing? um It's this phenomenon, that's predicted by Einstein's general relativity, where there's a chance alignment between um Earth, some foreground Galaxy and some background Galaxy, um so only one in ten thousand galaxies actually turn out to be gravitational lenses. So these are very rare phenomena.

B

um They're caused by this massive elliptical galaxy known as the lens, so the the Galaxy in between Earth and the source um it warps space time to bend light from this background, Galaxy, which we call a source um and the bent light, forms an image, that's highly warped and it typically makes large arcs um and shows like multiple images.

B

So um on this slide, you can show you can see two typical examples of a gravitational lens, so the one on top is a structure known as an Einstein ring which occurs when there's like a high degree of circular symmetry um and the Earth the lens and the source are all almost perfectly collinear.

B

um So you get this like very ring-like circular structure, um on the other hand, below we have something else called an Einstein cross. So this is another very generic or like common structure, you see for gravitational lenses um where you have the um lens Galaxy sitting right in the middle and then the four Source Galaxy, uh the four uh counter images of the source sitting outside okay. So aside from being interesting to look at what are these lenses really good for um in terms of like a scientific use case?

B

Well, first of all, they can be used to test predictions of the Lambda CDM model, which is a cosmological model that you know. People generally have a strong decay of faith in um particularly they can characterize the nature of potential Dark Matter halos, in which the London Galaxy is embedded or even Halos, on the line of sight between us and the source. Second, they can be used to infer the Hubble constant, which is an extremely important constant in astrophysics, indicates that a multiply land Supernova explodes behind the um the lensing Galaxy.

B

um This is because the two images will actually appear to explode at different times, and we can use this delay or this time difference to infer the Hubble constant.

B

um Finally, when multiple different sources are lens, these systems can also be further used as a pro for cosmology and dark energy, especially at hybrid ships.

B

Okay, so before I get into the lens modeling framework um and the work that we did, I think it makes sense to give a little bit of an overview of um the basic mathematics and physics behind gravitational lensing. So I'm just going to try to give you um a high level understanding of what's going on um so the the central equation here is really the known as the lensing equation, which I've written on the slide, and it expresses a fairly simple, geometric fact, which I've Illustrated in the diagram.

B

So here the s in the diagram marks the source. The disk on the circular thing in the middle marks, the lens and the O marks. The Observer, which is us or Earth and Theta, is the angle at which we see the image and Alpha in this diagram. Is this quantity known as the deflection angle, um because it basically measures how much the light from the source has been Bent by the lensing Galaxy um so for different types of galaxies?

B

The functional form of alpha May differ depending on the mass distribution, um and since it turns out that dark matter consists of a large fraction of the mass. In a galaxy, um the typical number is something like eighty percent.

B

um That's why modeling lensing can help us understand dark matter, because dark matter is believed to really dominate this lensing effect. um So with all that, aside, like the simplest, archetypical, Lens model is known as a singular isothermal sphere. uh This is basically where one assumes that all the stars in the galaxy can be modeled or described as free atoms in a gas of uniform temperature um and for systems belonging to this class of lens. The deflection angle is actually analytically, expressible, it's very simple.

B

um It's actually just a constant, um so this constant is known as the Einstein radius and it's a parameter that depends on the mass and the redshift of the lens or so the distance of the lens from us.

B

um The reason that this is called a radius is that if a source sits perfectly behind the lens, then the image that will be formed will have a radius that's exactly equal to the Einstein radius.

B

um So with the lensing equation, it's actually quite easy, I guess conceptually to simulate lensing, because if we know the parameters of the source in the lens, then we can compute the deflection angle. Hence we can compute the image, um so this is known as the forward problem, because we're kind of going into forward direction from the ground truth to observables.

B

The image on the right is an instance of a simulated lens where um I pretend I know that the ground Street parameters for the source and the lens, and then um we basically try to simulate it with realistic, observing conditions, so there's a noise and there's effects from the telescope, etc, etc.

B

um So, although it's not perfectly symmetric, um I've roughly indicated the Einstein radius on the figure um which um gives the order of magnitude for what the Einstein radius of this lens is. Okay. So now what is our problem? Well, our problem is pretty typical, which is that we don't have access to the ground truth parameters.

B

um We only have the observable, which is typically just a noisy image of the lensing system, um so to move towards our science goals that we mentioned on previous slides. So such as dark matter, testing or Hubble constant inference, we need the ability to do several things so using the observed image, we want to learn about the light profile of the source Galaxy.

B

We also want to learn about the distribution of matter in the London Galaxy and potentially even Dark Matter Halos near the lens or along the line of sight, um to maybe put it in another way in the language of Bayesian statistics. What we actually have is an inference problem right so, as you might anticipate we'll, be making quite extensive use of the Bayes equation, which just expresses the posterior distribution of the parameters as a product of the likelihood Empire.

B

um But I just want to note um you might catch me slipping into like a more physics language and just calling this problem modeling, but modeling and inference I'll will refer to the same thing.

B

Okay, so I'm going to describe a high level strategy for like how we're going to tackle this inference problem. So the first thing that we need to do is write down some sort of model for um the physics, um so this includes a functional form for the source and the lens light, as well as the lens Mass.

B

um This, this Mass model is actually always translated into a deflection angle and that's where we use it in practice for the modeling um but yeah. So the next step is then to write down a model for the probability. um What that really means is uh we need to find a prior for the parameters as well as the likelihood, um so the likelihood typically will assume that the noise at each pixel in the image is gaussian. So the log likelihood really just looks like a squared error or a Chi Squared up to some multiplicative factors.

B

um So speaking, Loosely all we're trying to do is match uh the simulated image to The observed image, and this matching is measured by a squared error and all the Bayesian formalism does is to kind of make this rigorous and calculate and to characterize the uncertainty for the parameters after we've done our fit. Basically, so once we've written down a physical and a probabilistic model, we're pretty much ready to sample from the posterior, which you know seems simple enough.

B

um The problem is this: um when we start looking at how to sample we're faced with a number of hard problems. Well, the the main one is that most regions of a high dimensional uh parameter space have a Vanishing likelihood. So if we use some out of the box, Monte Carlo sampler, like mcmc, we're likely to waste a lot of time in regions of the parameter space that we really shouldn't be looking at, because they have Vanishing likelihood.

B

So our solution, which is an approach first introduced to us by another open source package, it's called monstronomy- um is we essentially look for the global maximum of the posterior density first and if we can find that, then, and also assuming the posterior is unimodal, which is something I'll talk about later.

B

um But assuming it's unimodal in the sense that, like all the other local maximal, the posterior have very like low likelihood anyways, then we can just initialize our sampler at this point of a high posterior density and it'll do much better.

B

um So this actually leads to another problem, which is how do we identify the um the global maximum, because the parameter space is highly non-convex and therefore is it's hard to search over right? So the figure on the slide kind of roughly illustrates that problem, which is that when you have a non-convex loss function or a non-convex loss landscape, which is, in our case the posterior density, um depending on where you start you'll end up with different solutions of varying quality. It's it's hard to determine.

B

You know whether or not you found the global maximum foreign.

B

So with these um with our goals and these problems in mind, then we have three main digital. When we were trying to come up with a new lensing package, um the first that we required fast simulation um so as I kind of hinted at simulation, is a key to forward modeling. And if you want a fast inference code, then you need a fast simulation code.

B

um Luckily, most popular lens models are analytic, um so their their deflection angles can be calculated via a series of simple Matrix operations, and this is what I mean by the lens models. Are linear algebraic right. So, since this is the case, we can actually use gpus for fast simulation, because we know gpus are exceptionally good at elementary operations on large matrices.

B

So next we need some way to guide us through high dimensional parameter spaces. um I'll make the case that without some guide known as zeroth order methods, modeling can be very inefficient.

B

um So we found like zeroth order methods to be an issue, because there's virtually no performance guarantees associated with these and convergence can be very, very slow. um So what we've chosen here is we've made our guide of local one. So specifically, we've chosen the gradient.

B

um This is actually particularly fitting because, as I mentioned previously, lens models are analytic and therefore they're differentiable. So we can easily access the gradient via automatic differentiation, which is a technology that's now widely available in most lineological libraries.

B

um So, finally, we want a robust framework, which means that it produces consistent results across many different runs, so it doesn't fail often and also doesn't require us to carefully tune the initialization to get the right answer.

B

um So our approach is pretty simple: we hedge our bets, basically so as a general design principle we'll use many candidate Solutions whenever possible, so the failure probability will vanish very quickly with the number of candidates we use. So I'll try to elaborate on that further in coming slides, but to maintain many candidate Solutions, we'll use parallel computation with um gpus okay.

B

So this is why we developed our framework, which is known as Giga lens uh standing for a gradient, informed, ngpu, accelerated lens modeling framework. um We wrote this library with these three digital in mind and, as some of you may have realized, the popular tensorflow is a really good choice for most of the features that we wanted. um So our code is fully vectorized and fully differentiable using tensorflow's Auto diff, and it also runs natively on gpus.

B

So, furthermore, we actually wrote a parallel branch in Jax that has further support for distributed computing using multiple gpus, which we were conveniently able to test out on the new Pro motor system. So we should also note that we use another Library called tensorflow probability, which is very useful for some of these, like Bayesian uh modeling, Primitives, okay, great okay, so now to get down to the actual modeling procedure used in this Library.

B

So, as I mentioned before, our strategy is to First identify the global maximum of the posterior density and then to use this to inform how we initialize our sampler. So we do this with multi-starts gradient descent, which is so what that means is we'll initialize. Many different candidate solutions by just sampling from our prior um this prior should be broad enough to include the global maximum um and then next we'll use a bijector to unconstrain various parameters in our Lens model.

B

I know that might sound kind of confusing, but to make it to make it concrete. Let's, let's take the Einstein radius. um So since this is a radius, it should always be positive, um but it's undesirable to like manually constrain the parameter to always be positive in our gradient descent, um like these kind of hard walls are not very natural. So what we do is we use bijectors that map parameters into an unconstrained space and then optimizing that unconstrained space so for an Einstein radius.

B

What that means is we just take the log of Einstein radius and optimize over that, and that's that can be unconstrained.

B

um The next thing that we do is we localize each individual candidate solution using gradient descent with respect to the posterior density, we'll typically use atom within a yield learning rate, which is pretty typical setup for non-convex, optimization problems and then, finally, after a number of iterations, we'll stop and select the best performing solution and use this as the um the maximum. A posteriority estimate.

B

Okay, so how well does this actually work in practice? Let's look at an example system um on the left. I've shown the trajectory for two parameters, or specifically the error trajectory of two parameters with respect to time or the number of gradients descent iterations.

B

um The solutions that have been bolded are the ones that converge to the correct solution.

B

um So you'll note that the failure probability seems to be pretty high like it looks like you know, for each individual particle, um the failure probability is roughly 96 is what we actually found here um and the test we did for our paper. We found roughly 98 Upper bound on the failure probability, but when you, um when you aggregate this across 300 or 500 candidate Solutions, then the failure probability becomes vanishingly, though, because at least one of those particles with high probability is going to find the um the correct solution. Basically so on the right.

B

I've shown the loss trajectory for the best performing six candidate Solutions over time and you'll see that by the 300th iteration roughly, the particles seem to have come into the neighborhood of the optimum, because after this they exhibit a very smooth geometric convergence, which is characteristic of loss, trajectories in convex regions of the last function, so they've kind of reached the Basin around the global Optimum.

B

So this is actually what informs our cutoff for the number of iterations that we use in gradient descent, okay, um so before sampling, we still have one more important, auxiliary step, which is that we want to estimate the covariance um which will be important later. So we find this estimate by using a technique known as variational imprints.

B

um This procedure is essentially one where you um you fit in onsat's distribution, which I'll call our surrogate posterior to the true posterior. um So if we take a look at the diagram, let's say our true posterior. Is this funny looking object like the one in blue um that we don't really have a good handle of? um If we want to find a rough characterization of it, we might use a simpler analytic distribution indicated by Q Z, it's in green and try to fit this distribution to the ground distribution.

B

um So the fit or the degree of agreement between these two distributions is measured by the KL Divergence between the two. um So for us, since we want the covariance we're going to make our anzots or our surrogate a multivariate normal um and then basically try to fit the covariance and the mean to the underlying posterior.

B

Okay, so it turns out that the KL Divergence up to an additive constant, can be written as something else known as an Evidence lower bound.

B

um It looks kind of messy, but all it is is really just an expectation of a difference in log densities, um so the gradient is actually similarly expressible I'm not going to go into details into the details, but um yeah, there's very simple: Monte Carlo estimates for both of these quantities right so to optimize the KL. What we're going to do is um the following right, so we initialize the um the mean of our surrogate posterior at the uh map estimate from the previous gradient descent step.

B

um The next thing that we do is we minimize the elbow or the evidence lower bound um by approximating the two expressions, um the elbow and the gradient of the elbow using Monte Carlo estimates. um So what that really means is we sample many different datas from our surrogate posterior? We calculate some gradients average and then update on gradient descent and then, finally, after we've minimized our KL, the final optimized covariance Matrix is an estimate for the true cover rates.

B

um So if you want a frequentist interpretation of what we're doing here, um this can kind of be understood as taking the Hessian of the log posterior at the mode, um and this actually works pretty well um just as well as VI um in low dimensional parameter spaces, but in high Dimensions we find that the hashing becomes unstable and then VI is significantly more reliable, okay, so, finally, we're prepared to sample.

B

um Since we have gradient information natively with Google Lens, we can take advantage of this to use a sampling technique known as hamiltonian Monte Carlo. So this is a sampler that was proposed to handle the inefficiency of mcmc and high dimensions.

B

um If you want to understand um like what this inefficiency comes from, um it origin originates in this idea of a typical set. So when we're sampling, we really want to spend most of our time in this area, known as a typical set um and in high Dimensions this, uh the space can be highly non-trivial.

B

It can have a very funny um structure, so um just uh like a some sort of diagram, it could could be like this that you can think of it in this way um and while mcmc will make lots of rejected proposals that exit this typical set, um HMC is able to use gradient information to just smoothly. Follow the the curvature of the typical set basically um and you get very efficient sampling it this way.

B

um So uh you know hamiltonian Monte Carlo is not a perfect sampler you they require some hyper parameters and an initialization configuration um so, most importantly, we'll have to set a parameter known as the mass Matrix to be the inverse of the covariance Matrix from VI. So this is really to give you some intuition. The purpose of this is to inform the sampler of the various scales associated with each parameter, and these things can vary significantly for our problem.

B

Right, like some parameters, could have a variance of 10 to the minus 3 and some could have a variance of one. So these vastly different scales HMC does not handle that well, or this like anti anisotropy, well um without being given a mass Matrix.

B

um So once you give it this Mass Matrix, then HMC is basically able to adapt to that um other hyper parameters for HMC are just adaptively tuned over the burn interface.

B

Okay. So, to give you an idea of the performance improvements, we're able to extract with HMC um I'll, compare it to another popular Ensemble sampler in astrophysics, known as EMC, which performs very similar to mcmc. So you can roughly think of the two as the same. um But let's take a look at the trace plots for just one critical parameter, which is Einstein radius.

B

We see that every single iteration um HMC is drawing pretty much an independent sample, which is ideal Behavior, whereas for EMC we get a highly correlated behavior and it seems like it takes many different iterations before we get a new independent sample. So this idea can be kind of Quantified with metrics like the effective sample size which measures the number of independent samples drawn in a sampler, and you see that HMC is able to do several orders of magnitude better than EMC in terms of effective sample size, foreign.

B

Okay, so with that that kind of completes our modeling pipeline um once HMC is done, sampling we've drawn many samples from the posterior, and then we can do inference um so I've shown on this slide. Modeling results for one typical system, which is very close to a perfect Einstein ring.

B

um The corner plot on the left shows the distribution of HMC samples, as well as the vi posterior. So, interestingly, we can see that variational inference did a pretty good job of capturing. The typical scales of the posterior to observe the scale of the Einstein radius is 10 to the minus 3 and for the for gamma the parameter on the right. The scale is almost 0.1, so very, very different scales here, differing by almost two orders of magnitude and VI, is able to actually correctly identify that.

B

um So, however, it's not perfect, obviously, because we found that our posterior itself is not perfectly gaussian, despite how the marginal distribution might appear. So, for instance, to illustrate this, we looked at a random two-dimensional cross-section of our posterior and found um non-gaussianity in the HMC samples, which is shown right over here, um whereas of course the variational inferences gaussian.

B

um So on the right, I've shown some diagnostics for just basic timings um or sorry on the right. I've shown some Diagnostics, which have the observed image, the reconstructed image, as well as the residual of the two. um So you can see we get almost a perfect reconstruction and the residual is basically just a noise, um and the execution time in total is over just over 100 seconds on four, a 100 gpus with um promoter.

B

So, as you can see, variational inference is actually the most costly step, which is something you know we'll be trying to address in the future foreign so to verify the correctness of our modeling code. The last thing that we did is we simulated a hundred random systems um completely blinded and just model all of them um and then compared basically what our pipeline inferred with um the actual ground truth and we found no bias as well as well calibrated uncertainties so as in the mean Chi Squared, for each of these was close to one.

B

This was this was like a pretty important test in both the speed and the robustness of our code.

B

Okay, so I think finally, I'm excited to show modeling results on a real observed system. um So this is a system that we observed with the Hubble Space Telescope.

B

um What we have here is a pretty classical um lens uh in terms of the geometry, so we have a quad, so that's highlighted in the orange and then we have a double which is highlighted in the white, um so we have two sources being lens making this a very interesting system, and we see that Giga lens did pretty well modeling this system. So we get a Chi Squared of roughly one. Basically.

B

Okay, great so for next steps. um Our paper reporting on these results was published. The archive and also has been accepted for publication in the astrophysical journal. um What to do next? Well, the next steps that we're considering are whether or not we can actually improve speed by I mean we assume that we will be able to improve speed, but whether or not we can get our code to run with two GPU nodes.

B

So this would mean eight gpus on pro model, which is an exciting prospect for us, um and another question is: could we eliminate the vi step, which is the most costly step by using mass Matrix adaptation for HMC um and in terms of science goals? We want to start applying our modeling code to more real observed data from the Hubble Space Telescope, um so yeah that wraps up my talk. um Thank you for listening and I'd be happy to take any questions.

A

Oh great, thank you Andy great great talk. Do we have any questions.

A

Well, I guess: I'll start with one, so um you know you laid out at the beginning, so the things you'd like to be able to do with with the work in gravitational lenses, one of which was determine the um the form of the the lensing distribution of matter. Yes, and so, where are you with being able to make um make statements or draw conclusions about that sure.

B

Sure that's a great question, um so the as soon as you get into like questions of you want to determine the form you really have to introduce many more parameters to kind of model that uncertainty in the form, um which is something that kind of this modeling framework was written to addressed like it was written to work well in a high dimensional parameter spaces.

B

um So that's currently a work in progress. um So I guess I can't say anything about that. Yet so we don't have results, but it's something that we're thinking about.

A

Okay, all right thanks.

C

um Other questions.

A

So I guess I have one for you: I see um that you're now pursuing a degree in Quantum and Quantum science requirement of information, science.

B

How are you continuing.

A

To work on this project or what is the.

B

A

um What is your focus, these days, yeah.

B

Yeah so I think I'm I'm, definitely focusing more on um the quantum side of things.

B

um I think for this project, I've, taken on more of like a maintain our role and I kind of I, I guess meet with the team once every month or once every two weeks um just to discuss future directions, but in terms of like actually active technical work, I think I've taken a step back, um but yeah.

A

C

So yeah we meet within the uh once every as any said one to two months and uh as we pursue uh actual real systems to to model, and we also getting Spectra data, um it's needed uh too. So what Andy has done with this framework is uh um without the distance information, it's all geometrical and so now we're getting reshift information which would allow us to actually measure the distance to the lens so the background source, and then we can actually get the physical parameters. That's the next step and Andy's helping us with that.

A

Well, so is James Webb um offering new opportunities for you in this space.

C

um Yeah we'll probably be uh proposing uh in the next cycle, uh but meanwhile uh Desi and the other spectroscopic Ops it Desi, which is also um processing State owners, is giving us spectroscopic data so that we can model our Hubble systems.

A

A

Okay, well, thank you if there aren't any other questions, I want to thank um all of our speakers today. This was really uh very interesting and I I know I really enjoyed it. So thank you and um you know best wishes and and keep letting us know um interesting things that you're doing um in the future. So thank you thanks. Everybody. Thank you.

A

Thank you very much.