At the Computer History Museum: building ML inference for astrophysics that respects and exploits physical models, with applications to time-domain surveys.
Day within Feb 26-27 approximate.
“What could be more important than understanding the universe? And doing this with ML actually, as we'll wind up seeing, pushes the envelope of what is needed algorithmically and potentially even from the hardware perspective.” – Joshua Bloom
“What we decided to do is replace grad students and experts, because they don't scale, with ML.” – Joshua Bloom
“How do we impose our knowledge of physics into the learning process directly into the architecture?” – Joshua Bloom
“It's important for astronomers and domain scientists in general to be connected deeply with the cutting-edge work of computer scientists and statisticians, but it is also tremendously important, I think, for computer scientists and statisticians to learn from us the kinds of questions we ask of our data, because our data looks different.” – Joshua Bloom
Well, thanks very much. I thought I'd start off with this quote from Jim Gray, who was a consummate data scientist and obviously had a very good sense of humor. He recognized something important decades ago: that it was great working with astronomers and our data, because our data is noisy, it's messy, it's streaming, it's heterogeneous, and it's dirty. And if you made a mistake as a computer scientist or a statistician as you try out new algorithms or new approaches or new hardware on our data, you're not leaking PII, there's no billion dollars of damage to your company, and nobody dies. So indeed this is a wonderful thing about working with astronomers. On the flip side, astronomers love working with statisticians and computer scientists
really for two reasons. Number one, the pace at which algorithms and new hardware are coming out that are able to directly impact what we do in our daily lives in discovery and inference — which in some sense is a big part of this talk today — is just wonderful. And astronomers are very well known for co-opting tools. So the most famous astronomer you might think of is Galileo, who heard about this thing called the telescope, which was meant to look over the horizon for enemy ships, and he said, what if I just do this? And you know the rest of the story. But also, the second reason why we like working with computer scientists and statisticians is that, with all deference to my academic friends, oftentimes in
that world we feel like they're not working on important problems. So what could be more important than understanding the universe? And doing this with ML actually, as we'll wind up seeing, pushes the envelope of what is needed algorithmically and potentially even from the hardware perspective. So this talk is really dedicated to bringing all of you up to speed on some of the really interesting domain problems that we have in astronomy that we think we can start applying ML to. And in fact, as you'll see, it's becoming part of our daily lives, and what we're doing is starting to recognize there's this very deep connection between the way in which machines learn and the physics that we actually do with these learning systems.
So for us, this is the most exciting plot I think I can show you, just to summarize why we do what we do, at least in time-domain astrophysics. In the middle what you have are basically the light curves — this is brightness versus time — of some of the most important explosive events that we know about. Type Ia supernovae, for instance, as you see in the blue curve, are used for measuring the accelerating expansion of the universe. And if you want to understand how stars die, if you want to understand where all the metals in the universe come from, you also have to understand type IIP supernovae. So there are concerted efforts for many people to be looking for these things and use them as cosmological probes and for our understanding of
the physics of how atoms are created. But then there are these other objects which are up on this curve, like the so-called neutron star neutron star mergers. And this plot was made many years ago, before we actually found one of these electromagnetic counterparts to a gravitational wave event, and it turns out that the observations that we got look very similar to the theoretical curve that you see there. But then there are these other curves that are on this plot, like the pair-instability supernova, and when a neutron star runs into a red supergiant like Betelgeuse there are some really interesting things that we might be able to see. And so these are the unknown, or the known unknowns, but obviously by definition we can't
put up the unknown unknowns. In general what we're trying to do is, in a swarm of lots and lots of data, satisfy all these different scientific constraints simultaneously. And if what you're trying to do is write down some optimization metric, you would try to say, how can I maximize the science for the discovery of these objects, but then also our efficient and effective use of follow-up machinery? Just finding one of these things in the sky doesn't really mean anything unless you can actually follow them up with other telescopes and do deeper introspection. This is obviously easier said than done in many senses. There are sociological impediments for us to be able to say that we're going to be able to do this optimization, even if we write
down this metric, because our optimization metric may be very different than somebody else's. But this is sort of the grand goal: to be able to find these things that we know about and find things that we don't know about but may actually push the envelope for science. And what looms tremendously large for us is the so-called Large Synoptic Survey Telescope, which will be acquiring more data, at least at optical wavelengths, than we've ever had before. This comes online in just a few years. It's a billion-dollar facility that all you American taxpayers have already paid for, but it's going to be producing something like 20 terabytes of raw data per night, tracking 18 billion objects on the sky simultaneously, basically observing the entire night
sky that's available from the southern hemisphere every three days, and producing something like 150 teraflops are going to be required for this first data release, and the final catalog will be of order tens of petabytes. So for us to be able to mine that data in real time and effectively use all the resources that we can get at it to find these objects of interest and then follow them up is really our grand goal. So the agenda for this talk is to introduce you to the centrality of ML in some of our everyday practices — a lot of this is supervised learning. I'll introduce you to some of the things that we've been doing for a number of years and then talk about some of
the new directions in semi- and self-supervision. In particular I'm talking about what we've started working on in the context of astronomical time series, and then also talk about an imaging problem as well that's interesting to us, and then get into how we're starting to figure out how to bake our physical understanding of these various different phenomena into the learning process itself. So one thing is using ML to be able to do discovery at scale; another is being able to imbue what we already know about the universe into the learning process to accelerate it, to do better, and lastly with a couple of statements about so-called likelihood-free inference. So as I've already said, ML in some sense powers a lot of
discovery of what's already going on in astronomy at scale. And this is one of the first, I think, prime examples of not just doing ML on offline data and saying my scaling curve's better than your scaling curve, but actually putting ML into production. This is this challenge that we have when we get lots of imaging data. We're trying to find a new object in the sky, and the way in which you do that is you take a big median stack of all the images in that part of the sky — I call that the static image — take a new image that just came in, align the two of them, subtract them, and what you get is what remains. And if there's nothing there, then
you should just see noise; if there's something new, then you should see something like the rogue's gallery at the bottom there, what we call real. But as it turns out, because the atmosphere is turbulent and we have a lot of noise in our instrumentation in general, we wind up having not perfect alignments, and so most of our subtractions lead to the so-called bogus detection. So these are not real astrophysical objects; this is just an artifact of the data processing. So what people did in the past isn't so different than what a lot of my colleagues were doing just about 10 years ago when they were looking at astronomical images, which is of course hire grad students to look at the data and decide, is this
real or bogus? It's pretty remarkable, this idea of using people, experts, to look at the data and then opine on that. It started hundreds of years ago in astronomy. We had a big data problem, as more data was coming off of telescopes even back then. And just as a side note, that's Henrietta Swan Leavitt seen in the back there. She's one of the most important figures in modern cosmology for some of her major discoveries, in this same room, and there's actually a new play about her that you should see; it's called Silent Sky. Anyway, what we decided to do is replace grad students and experts, because they don't scale, with ML. This is obvious to all of you in the room.
Obviously this allowed us to create a fast and parallelizable version of all that, and transparent and deterministic statements about all these different postage stamps that were flying off of the telescope. And this is a massive needle in the haystack problem that we had to attack. I think one of our biggest discoveries that most people outside of at least our immediate fields got to know about was the discovery of the nearest type Ia supernova — which I said before is the most important object for cosmology — in 25 years. And we did this in an ML-assisted way, where we basically presented to humans a ranked ordered list of the most interesting objects in the sky as they were flying off of the telescope. And at the
time this meant that we were able to discover this type Ia supernova about 11 to 12 hours after its putative explosion, which was days earlier than it had ever been done before. Now it turns out that this supernova got so bright that if you'd had binoculars and you looked at the right place, you would have gotten photons from the supernova hitting your eye, which is just remarkable. So it would have been discovered by amateurs days later or maybe weeks later. But what was so special about being able to recognize in real time that we had an interesting object is that we're able to throw the world's resources at it and get some novel science out of that. We learned a lot about the progenitors,
the objects that actually make type Ia supernovae, in a way we couldn't have done by any other means, even if we had hundreds or thousands of grad students doing this kind of discovery. Another place where ML is actually starting to have an impact is in the search for Planet 9. Many of you of course think that Planet 9 is Pluto; we won't get into a large debate about that. There is growing evidence that if you look at the orbits of long-period comets, there looks to be a gravitational perturbation which will be more massive than Pluto, and this could actually be the real Planet 9, and we just haven't found it yet. But what many of us think is that Planet
9 has already been imaged somewhere on the sky, because the entire sky has already been imaged, but it's too faint to see in any one detection waiting in one image. So what you have to do is stack up these images. The problem is Planet 9 is moving over time, and so for us to actually find this thing we actually have to stack a whole bunch of images where we basically shift them along the orbit of Planet 9. The problem is we don't know the orbit of Planet 9, so we have to guess it, and as you can imagine that adds to a massive complexity of this data cube that we're trying to find a new object in. So what we've
done is started a search using old data at LBL using the NERSC computers, where we had to basically look at every five to ten sigma detection in our very large data cube, and that led to hundreds of billions of possible candidates. But what we did is we trained an ML, basically zero-one, classifier on a bunch of synthetic data that we had created, that allowed us to get a very good accuracy, throwing out essentially 99 percent of the false positives so we could actually keep a good fraction of the true positives. The problem was that most of those models were too big to fit in the amount of RAM that we were allocated. We need to do this at scale, and
importantly, there were so many candidates that we didn't have enough data space at NERSC to be able to save all of the candidates to disk and then process it afterwards to figure out which ones might actually be that Planet 9. So what we realized we had to do is create a whole server farm of basically these little tiny apps that you could ask a question, send it a postage stamp and say, should I save this data or not, is it a possibly interesting candidate? And we had to get this round-trip time down to ten milliseconds. So we couldn't even save it on disk; we had to put it in a TCP/IP packet and distribute it throughout the NERSC supercomputer. And so
one of the interesting things that this brings up is the need not just for high-quality models, but, as many of the other speakers have spoken about today, they need to start thinking about the importance of fast models and ones that are deployable and ones that are versioned. So that's some of the bread and butter, in a supervised sense, of where we've been. What I want to tell you about now is some of the pain points that we've had with the traditional classification approaches. This is an example over a very large swath of the sky of 50,000 variable stars that are changing in time, and there's just one object that we pull out of that. Unless you're really good
at doing Fourier transforms in your head, you won't know that the periodicity is about a half-day; that's a classic star called an RR Lyrae. What we had to do to classify all of these stars was to be able to essentially build a whole bunch of features where we imbued our knowledge of how all these different types of variable stars change, and do random forests. This was sort of 10 years ago when we started working on this. As you know, when you're doing traditional featurization it leads to a lot of hand-coded feature engineering, and the compute generally will scale with the number of features that you have to do. And in our case we had a very small number of labels. So
one of the things, if you're ever talking to an astronomer who says we have a big data problem in the context of inference and classification, it's actually a small label problem — which, those two things are not completely inconsistent with each other. But in this case, out of that 50,000 objects, we had 25 classes of variable stars we were interested in, but we only had eight hundred and fifty or so labels for that. So we had a very small number of labels and we had to figure out how to bootstrap, and we used a bunch of active learning techniques to be able to get a bigger and bigger sample. Anyway, this is quite challenging. So what we've recently recognized is that rather than do our hand-coded
features, why don't we just throw this into an autoencoder? And for those that have seen this before, it's a classic autoencoder idea, where we take this encoding of the original light curve, we compress it down to a small bottleneck, maybe 64 numbers or eight numbers, and then we have a decoder which tries to get back the light curve that we had before, and then you build a loss function on that, you backprop, and you wind up getting auto-learned features. And what we wind up using is the features that are in that bottleneck layer, throwing that into a random forest, and we wound up getting best-in-class accuracy on a whole bunch of different data sets. So this got us
pretty excited. Now here's where the interesting thing comes in. We didn't just use an off-the-shelf LSTM; we actually had to modify the architecture, because there's no notion within an LSTM of the fact that data could be acquired at slightly different times. There's sort of a normal cadence beat that's assumed, but astronomical data is taken irregularly, so we had to figure out a way to have the network be able to handle irregularly sampled data. We also have noise on our data, so we wanted to make use of the noise properties of the data in the loss function. And importantly, as I said before, we don't have a lot of labels, so we wanted to do feature learning and we wanted
to build this network without any label. So this was a completely self-supervised way of building up labels over not just the corpus of 850 but over a much larger corpus, and as you can imagine that's one of the reasons why we did so well. And as a side note, I'll say one of the things we're exploring now is our ability to find anomalous or new types of objects in this feature space, in this bottleneck space. Another interesting problem that we're working on is denoising autoencoders. On the left-hand side is an image that you actually get out from the Hubble Space Telescope. All the scruff you see there are cosmic rays, charged particles which are hitting the detector, and what you'd
like to see is the thing on the right, which is the cosmic-ray-free image of a whole bunch of galaxies — in this case it's a galaxy cluster. So we used a modified U-Net architecture, again with a bottleneck at the very bottom there, where we take one of these postage stamps that has these cosmic rays in them and we try to predict a mask, and then we do a second task, which is, given the mask and given the original image, we want to get the beautiful inpainting final version of that. And one of the things that's interesting and exciting for us is that if you look at the activations in the top layers of this network that we were building,
it actually learned the Laplace transform, which is the current state of the art of what people do now to find cosmic rays, but of course it also learned a whole bunch of other interesting kernels directly from that data. And so we got good answers; at least visually it's very pretty, it seemed to do incredibly well, and then compared to the other state of the art — what would be a talk here without ROC curves? Here's your ROC curve, false positives and true positives — we actually bested it not just in quality of the results relative to the best in class but in speed, because as many people have said to us, well, it's great that you have a better cosmic ray detector, but unless it's faster and easier to use, I'm
not going to use it. And then on the inpainting task we did better than the traditional median mask inpainting and by harmonic interpolation. So in the last part of my talk I want to shift over a little bit to where things are going and where our big interests lie now in ML, and this is what I'll call physics-informed ML. And there's this great paper, if you haven't read it, which is called “Why Does Deep and Cheap Learning Work So Well?” by Max Tegmark and company at MIT, where he said, we have this image that's, let's say, a hundred by a hundred or a thousand by a thousand image, there's 256 possible values in a grayscale image, and so if
you're trying to determine between cats and dogs, the available state space is 256 to the million, and yet with a pretty small network, if I give you a bunch of images of cats and dogs I can actually distinguish that. And the point in this paper is to say that the reason why networks that don't have that much capacity can do so well is because they're learning what's physically plausible; it's not having to search over that massive space to get a good classification answer. So this becomes the point of departure for a question that we might ask: how do we impose our knowledge of physics into the learning process directly into the architecture? And this is starting to be done in a bunch of different places.
In the context of computer vision you have these spatially invariant transforms, so if I rotate an image I still get the same activations out. In the context of high-energy physics they're building networks that are actually QCD aware, so when you throw in data from the LHC it winds up being invariant to all the possibilities. And in quantum chemistry, like the picture you have there, if you're trying to do some notion of protein folding for instance, the same molecule you have on the left is the one on the right, and if you're trying to predict forces it shouldn't matter what that actual orientation is. So building architectures that are aware of these different kinds of symmetries is critically important. And one of the
papers that we just submitted kind of helps answer some of this question, which is whether we can find these embeddings and these network architectures that conform to our known understanding of the problem. We just submitted a paper to ICLR where we're making use of the fact that we have periodic data in the case of pulsation of RR Lyrae and the kinds of things I was telling you about with variable stars, where they repeat over time. So once you know the period, then all the data you get at another period basically doesn't add much information, and so we tried to build up an architecture that was aware of the fact that as you wind up going backwards in phase or forward in phase
eventually you wind up wrapping around yourself. So we're essentially doing convolutions in polar coordinates instead of Cartesian coordinates, and we'll see in the next couple of days if this gets accepted or not; if any of you are the reviewers, be kind. So another thing we're starting to do is ask the question, can we imbue the kinds of things that we care about, if we're trying to do simulations of variable stars, into the latent space in a variational autoencoder? And in particular what we'd like to do is start with our real light curves on the left-hand side, where we know not just the time histories, brightness as a function of time, but also the labels, and we also
know in some cases what their temperatures are, what their masses are, etc. We'd like to basically teach the network to be able to make objects that are like that, and then we'd like to sample from that latent space so that we can create lots of instantiations of it, so we can optimize our telescopes to find more objects like those. I won't go into the details of what this architecture is, only to point out that it's a traditional VAE except for the fact that we're injecting, in that yellow box there, the labels and the physical parameters both before the bottleneck and then after the bottleneck, so that at test time when we want to generate new objects we can basically just start from
a random Gaussian sample from that, add in the parameters that we care about, and actually produce realistic sources. And one of the things that we had to think hard about is how we did this, and, because we're trying to make and simulate real light curves, we had to do some architecting of our loss function and do some curriculum learning, where that last component there in the term also wound up allowing us to build a KL divergence between the uncertainties in the predicted magnitudes, or the predicted light curves, and what came out of the VAE. And as a result what we're able to do is walk around not just in latent space, like you see on the left-hand side, but also walk around
in something like temperature space, that allowed us to give these realistic light curves that came out. So this is still a bit of a work in progress, but when we get this working and we're able to deploy this, it means that we'll be able to build simulations of the variable-sky universe in a way that people have never been able to do before. Of course this sort of generative modeling isn't just in time series, and for our kind of science, many people in astronomy are starting to do this in the context of cosmology. What you see on the right-hand side are a bunch of instantiations of little slices of dark matter halos on very large swaths of the sky. The data
at the top is the real data that came from simulations on supercomputers, and the data you see on the bottom is generated data. Now, they might look sort of similar, it might sort of look like it has similar noise properties. Well, one of the clever ideas here as you're building these generators — and this was done again with a VAE — is to not just look at the images that come out but to do regularization on the aggregate properties of those images, in this case two-point correlation functions or other sorts of things like you see in the top left plot. What you want to do is not just produce pretty images, you want to actually produce realistic physical kinds
of summaries of that data, and that's what this group was doing. And then on the last note, what I'll say is we're also starting to get interested in inverse problems. So one thing is to be able to generate data; another is to say, given data, can you actually get for me the parameters that might have been used physically to generate that data? So this is an inference problem: I'm trying to take data and trying to get posteriors in a Bayesian sense on the parameters that matter to us. And one of the really exciting things that's coming out of particle physics is this deep connection between the simulations, which are extremely expensive to run, and so-called likelihood-free inference techniques, where
you're simultaneously drawing from priors on what you think your parameters are going to be, and you wind up directly building up something that can do density estimation that can give you out posteriors. So what I'll conclude with is, in some sense, the summary slide that says that where we're going now is trying to bake the physics into the entire process. In some sense, starting on the far left, you have featurization — that's what we've always been doing for the last 20, 30 years, taking our knowledge of these objects that we care about and building handcrafted features and then basically letting traditional learning algorithms learn on that. But now we're starting to build these symmetry-preserving layers — that's number two — and we're also thinking
about imposing sparsity, sort of a way of encoding Occam's razor into the size of these bottleneck layers, and then we're doing loss function curation, and we're doing distributional loss enforcement over the entire aggregate. So this is the state of where I think this field is going; it's tremendously exciting for us. So with that I'll leave up my conclusions here. What I will just say is that, again, it's important for astronomers and domain scientists in general to be connected deeply with the cutting-edge work of computer scientists and statisticians, but it is also tremendously important, I think, for computer scientists and statisticians to learn from us the kinds of questions we ask of our data, because our data looks different, and the kinds of
ways in which we want to get these answers may be very different than what you wind up seeing in industry. So with that I'll say thank you and happy to take a question or two. [Applause]
Questioner: Thank you for that awesome talk. We have a question. Thank you for your talk, it was very informative. So I do not know what your research focus is, but I know that many astrophysicists — but not only astrophysicists, also computer scientists — have worked on astronomical time series data using Bayesian nonparametrics, but it seems like they are giving up on those ideas they've been working on for 20-30 years. What is your take on Bayesian nonparametrics?
Josh Bloom: Well, so for instance in the context of time series, Bayesian nonparametric models are incredibly important. We're using Gaussian processes, for instance, to homogenize our datasets, where what I didn't tell you is that there's a variable length to a lot of these light curves that we get — sometimes we get 50 data points, sometimes you get a thousand. Doing that in a traditional batch learning sense makes it really hard to throw in different sized sources, so we're actually using these sorts of techniques to get lots of instantiations, from a data augmentation perspective, of fixed-length light curves. So that's really helpful. This certainly isn't to say that this is the only approach for doing inference; there are lots of other ways to do it. And I think one of the things that scientists tend to
be pretty good at — or they ought to be pretty good at — is using the right tool for the right problem, and we've identified a number of problems where ML and these new self-supervision techniques actually, I think, are pretty useful and pretty important.
Questioner: Let's thank Josh one more time.