Scientist Stories: Anima Anandkumar, Generative AI & LLMs to Accelerate Science

All right everyone, welcome to another episode of the TWIML AI podcast. I am of course your host Sam Cherrington. Today I'm joined by Anima Anandkumar. Anima is Bren Professor of Computing and Mathematical Sciences at Caltech and Senior Director of AI Research at NVIDIA. Of course, before we get going, take a moment to hit that subscribe button wherever you're listening to today's show. You can also follow us on TikTok and Instagram @twimlai. For highlights from every episode. Anima, welcome back to the podcast. It's been a bit.

Yeah, thank you, Sam. It's such a pleasure to be back and so great to see where, how TWiML has expanded its audience. It's now even on TikTok. And when we started back then in the beginning of the AI revolution to where we are now, it's so great to see TWiML be part of that journey and helping our viewers and audience to really be up to date with the latest and greatest machine learning tools.

Thanks so much. I'm really excited about catching up with you and digging into all the work you're doing around AI for science. Since it's been such a while since we last spoke, I'd love to have you just reintroduce yourself to our audience and share how you came into the field.

Yeah, certainly. It's been such an exciting journey to be part of the AI revolution, starting all the way back when AI was considered more a theoretical concept, right? People didn't think in our lifetimes it would take off like the way it has done now. But that also meant I could hone in on the theoretical foundations, from statistics, signal processing, machine learning, probabilistic models, and ask questions like, how do we extract hidden or latent variables or phenomena from data without labels? You know, what we call unsupervised learning. And I was working on tensor methods, which conceptually looks at higher order moments of data when people were mostly limiting to second order moments or linear algebraic techniques. We talked about that back then. And a lot of my work was saying, No, no, don't just limit to linear models. You know, don't just limit to spectral methods that use matrices. Go beyond, you know, think bigger. And we have the computational abilities to do that now. We can go to tensors. And that kind of nonlinear transformation is what we also see with neural networks. So we were early in that journey to go from matrices to tensors.. And from tensors, now we have all kinds of nonlinear mappings learned through all the neural networks and related tools. And so to me, to see that journey, and now last year was the year of generative AI. And so to me, who's been working on generative modeling since my earliest days in graduate school, it's so great to see that. Come into fruition and see how it's having a lot of really important practical implications.

That's awesome. And it's been quite a year for generative AI, and I'm sure we'll be talking a little bit about that. But one of the things we wanted to dig into a bit is the work that you've been doing around AI for science. Tell me a little bit about how you got involved in that.

Yeah, certainly. You know, at Caltech here, we founded AI for Science as a campus-wide initiative in 2018, way before researchers were thinking of it as a mainstream area in AI, right? Indeed, a lot of the AI development has been driven by industry, by big tech, and focuses on natural images and text. Because that's where we can easily get web-scale data. And they also have a lot of commercial value for building the next generation internet, you could argue. But when it comes to scientific domains, there's even a hard question of where do we get started? There's so many challenges that we don't see in our standard AI domains like image and text. I know that starts with the lack of data, right? You may not have large-scale data that's really necessary for deep learning. And you not only think about a very specific task, you most of the time require zero-shot generalization. So you wanna go beyond the training domain. A lot of the science is extrapolation, right? With the data I've seen, what is the possibility beyond? What are new drugs? What are new molecules? How do I create those? You don't have data on that because all the design and discovery is about finding something new that doesn't exist so far in your training. And so that's, to me, that extrapolation and generalization beyond the training data is such an important aspect. That standard machine learning wasn't capable of doing. But that's where I think now in the era of generative AI, so much of that is becoming much more possible. You know, just like we can generate avocado chairs, can we generate better drug molecules? You know, for the viewers, what I'm referring to is like Stable Diffusion or DALL-E, these models that once they learn text-to-image mapping, can now create entirely new images that were not in training data, right? Entirely new scenarios, like a chair that looks like an avocado. You probably didn't see it in training data, but if you learn the concept of what is a chair and what is an avocado, you can try and mix that together, do this compositional generation.. And I think these are concepts now that we are coming full circle, is aspects that can have huge implications in the sciences. And yeah, so a lot of unique challenges, but also I think opportunities that lie ahead when it comes to AI for science.

I think one of the areas of AI for science that's received a lot of publicity over the past year, couple years, has been around the various protein folding research. Can you talk a little bit about that and why that's interesting and important?

Yeah, absolutely. Protein folding has already been making such a big impact in our ability to quickly generate the full three-dimensional structure of proteins, right? So how do you map from sequences to now 3-dimensional structures. And this folding is so important because that's what will determine how other molecules or other proteins interact and bind. So if you now have a drug molecule, which is typically a small molecule, I mean, not all the time, but most of the time it's simpler to design a small molecule. And then the question is, how does it go and bind in a certain pocket of the protein, you know, what are the active areas in this protein, right? So all that's determined by this three-dimensional structure. And that's why for biologists, you know, this was really critical because you cannot use all the knowledge and intuitions and say, oh, this could be a good binding candidate. But until you have the three-dimensional structure easily available, how do you determine if that's good computationally. And that's where you had to go to the lab and do those experiments, which are really expensive. And now we are going even beyond the folding, right? Even just the protein by itself, what is the structure, to now how do the protein and ligands, which are small molecules, bind together. So in many cases, it's even more complicated. As always, biology is quite complicated, right? The protein, when it's binding to this small molecule, in fact changes its structure. So it's not even a static process. And so that's one of the work we've been exploring using these generative AI diffusion models is not only to model the static behavior of how a protein looks in 3 dimensions, But how does that three-dimensional structure change when it binds to the small molecules? And so to model that process, and with that we can predict better about the contacts, how well does it bind? And that's so crucial for coming up with much more realistic predictions. And that's just one example. You know, we are talking a lot about structure and this is, Quoting Frances Arnold, who is Nobel Prize winner here at Caltech for directed evolution of protein. So she is quoted in New York Times on generative AI for proteins that appeared, I think, about 2 weeks ago. And so she's asking, that's great, we are solving all these structure problems, right? We can speed up our prediction of structure, we can greatly increase the accuracy of our structure prediction. But what about function? Because ultimately what we want to understand is what is the functionalities of different proteins? And can we create, like, better proteins which have functional implications? Because, you know, if we can create better targets and create drugs for it, we can cure a lot of diseases and issues that so far have not been able to do through traditional therapeutics. And I think that's where a lot of the work that, for instance, my group is doing as well as others is to not just limit to the structure. Because if you only look at proteins and their structure and create the foundation models, you know, it's not very easy to talk about the function. But if you go to the level of the genome, you know, the DNA, right, that determines the function. If you can learn the relationship between different genes, you know, what is their relationship and how does it relate to different functionalities of the proteins? Because genes can be mapped to the proteins they generate. And by understanding relationship between the genes, can we generate better proteins that are functionally meaningful for different tasks? And that's what we've been doing now. And it has so many applications. So in the largest biological language model that we built of about 25 billion parameters, what we showed was the ability to learn at the genome level. So long sequences of bacteria and viruses, more than 110 million such sequences. So we consumed all that into a language model. So in a way, you're understanding the language of the genome.. And then from that, we can generate now new gene sequences, right? And with that, we could in fact predict new variants of coronavirus as well as, you know, the ones we had held out, the existing variants that appeared like Delta, Omicron, right? All those we could predict even though the model had never seen that before.

What was the text that that model was trained on?

Genome sequences. Got it. So think of now the language is the genome, right? So that's your now alphabet, that's your sentences. It's longer, and that's a challenge. In fact, we used a combination of language model like GPT, but with a diffusion backbone in the latent space. So it's a hierarchical model that has both diffusion and GPT components., which is necessary for these long sequence modeling. And, you know, and that really helped us to understand at the genome level these relationships. And you could also see organization in the latent space of different related genes, right? So genes that are close in that latent space are also functionally similar or related. And I think those also help us now to design better proteins.

This reminds me a little bit of work that I spoke with Richard Socher about when he was— I think he did this at Salesforce, of all places.

Are you familiar with that at all? Yeah, certainly. That was some of the early work on understanding protein, I think, sequence models, if I recall right, or like molecule models, right? So, A lot of work in the last few years has gone first at the molecular level. You know, how do we generate better molecules? How do we give like condition on various properties and generate better molecules? And then we said, we can't just look at a molecule in isolation. We have to look at how it binds to the protein, right? That's the important interaction. And that's why we went into the protein structure prediction, which is the AlphaFold and the MetaESM and OpenFold and all these tools. But then we said we can't look at the protein in isolation, we have to look at how the protein and this molecule bind. And that's the joint prediction. And we also do that dynamically because the protein structure can change as it's binding. So that's modeling the dynamics. And now we are saying all this is great, but this still doesn't tell me the function, you know, what are these proteins doing in our body or in an organism. And that's where we need to go to the DNA level. And so we need to understand the language of the genome and not only learn those relationships between genes because they map to proteins, but also potentially generate new gene sequences and through that, new proteins and new targets. And I think that's exciting because we are like really going up the hierarchy and essentially understanding the code of life, right? Life on Earth.

So I think it's very exciting times. Yeah. If the genome encodes the protein structure, how does that get—

how does understanding that get you to function? Because it's at the genome level that you're determining a lot of functionalities of the organism, right? So it's encoding more than just protein structure? Yeah, absolutely. Because that's what evolution has endowed it with. You know, all the fitness, right? Like which ones survive, which ones are good proteins, you know, that's what is encoded in all our DNA as well as, you know, the DNA from all the way from viruses to all the higher organisms. And it's complicated, right? You can't just directly map one-to-one in higher organisms. And that's why we started from virus and bacteria where it's clearer and simpler to analyze. But I think we can, as these models get bigger, as we create foundation models on long genome sequences and really get into understanding what in the latent space we can encode, right? And what does the generation of new sequences mean? And, you know, what kind of evolutionary gaps or other things they're filling. I think we can do a lot more than what we've done today.

You mentioned diffusion a couple of times in the couple of works that you've discussed. Can you talk a little bit about the relationship between how you're using diffusion and how it's used in the context of Stable Diffusion?

Yeah, absolutely. Last year has been exciting for generative AI, and Stable Diffusion has really been a trailblazer in the sense is that we have an open-source model, right, where it has spawned a whole range of an ecosystem of applications, and it can generate seamlessly new image candidates based on the text instruction. And so what the diffusion model at its core is, is a generative model that can sample new candidates, new realizations from a distribution it has learned.. And the way it does is to go from, let's say, a Gaussian distribution, which is simple to sample from, to the distribution of images, which would be hard to directly learn. But by learning this mapping, like, think of it as a slow diffusion. You go from Gaussian noise and you progressively make it look more and more like the image. It'll be a noisy image, and ultimately at the very end, all the noise would be filtered out and it would be the true image.. So that slow process, how can we model that through learning these mappings, right? Learning this transformation from a Gaussian distribution to the image. And similarly, we can do that for scientific domains as well because it can, you know, the concept is you can learn any arbitrary distribution, probability distribution. And so we can learn now the probability distribution of How does the process, you know, how do we create new genome sequences, right? What is the probability of them occurring, right? Because there is, you know, we know a lot of like, there's high correlation between different gene sequences in a genome. You know, there's just like a code where some combinations are just not valid or they're not going to survive, they're not that, you know, relevant for the real world. And that's encoded in that dataset of all known genome sequences that people have collected. And so by encoding that, we can learn now a better way to sample. And that's true for also so many other applications. Another setting we are thinking about over the last few years has been to look at all kinds of spatio-temporal processes and other phenomena in science and engineering, right? Think about how fluids move, how does our weather forecasting change, how does our ability to predict earthquakes develop. And so, so that's another area as well where uncertainty quantification is very important.

And I think tools like this could be very effective, and that's something we are working on. Are you foreseeing kind of the diffusion, applying diffusion to all the things? Is it a tool that we kind of stumbled upon for generating cool images, but it's going to have much broader implications over the coming years?

I think it's one of the tools, right? I mean, the question of which specific architecture and which specific framework can evolve over time, it was GANs. A few years ago, and we saw like challenges in optimization. And maybe one day we'll overcome that optimization challenges, right? And I do think that's important for the scientific applications too, because in sciences, we have a lot of constraints. For instance, one setting we've been extensively working on is solving partial differential equations. And so they're incorporating the physics constraints, incorporating, the fact that you should satisfy this equation. So you penalize if you don't satisfy that equation, right, are constraints. And so these kind of constraint optimization frameworks are similar to a GAN, right, because they're primal-dual optimization. And so we still haven't solved those fundamental optimization issues. In deep learning, we've kind of swept it under the rug and said, okay, this problem is hard, so let's just develop a different methodology where we don't have to tackle that at all. And I think in the short run, that gives us some gains, but at the same time, in the long run, I think that comes with its own issues. For instance, diffusion models are slow to sample because they have to gradually go from a Gaussian distribution to the distribution of images. And one of the frameworks we've developed to overcome that, to do a decoding in parallel. So instead of sequentially sampling, going from Gaussian to the image distribution, can we directly jump? And can we do it in a generalizable manner, not just overfit it in one setting and then, right, it doesn't work in other settings? And the frameworks we are using in all these settings has been what we call neural operators, which are, I think, a really important concept when it comes to scientific domain applications. They can solve differential equations like the one used in diffusion models for sampling, but they can also solve all kinds of other differential equations like in fluid dynamics, the ability to model seismic waves underground to predict earthquakes, to model carbon capture and storage as we mitigate climate change, how do we pump carbon dioxide deep underground, and how does it interact with water, these what we call multiphase flow systems. And I think that to me conceptually is different from the standard neural networks, and that becomes a really important tool for scientific domains.

Can you dig into that a little bit more? What does a neural operator look like, and how did you arrive at that?

Yeah, absolutely. I think this has been such an exciting journey because we went about asking, you know, let's say I have the problem of saying, how does the fluid dynamics evolve, right? I start the fluid and I shake it a little. How does the flow change over time? It sounds very similar to a video prediction problem. So why not just use all the tools we've already developed to do this prediction? Prediction. So it turns out the main challenge is that here the fine scales matter a lot, right? We know the turbulent phenomena like this, you can't just like filter out or smooth out the fine scales or the high-resolution information. If you do that, you will get— come up with bad predictions. On the other hand, in the natural image domain, we want to all the time filter out, right, all the pixel is just too high-dimensional and mostly useless information. We want to go from pixels to object localization. And so that means all the processes to remove all this irrelevant information and filter out, smooth out. And that's what, like, convolutional filters do because locally you're learning this ability to smooth out and extract only the relevant features like edges.. But that's different in so many of the scientific simulations like fluid dynamics because you cannot just filter out or smooth out the higher resolution information, right? You have to keep that. And the consequence of that is you cannot just learn your input to output mapping in one resolution. So think of any kind of image generation or image prediction. You always specify, pre-specify what is the resolution of the image you wanna either generate or predict on, right? It doesn't work at a different resolution. People train and they train only at that one resolution. They use the model only at that resolution. Whereas for the scientific domains, that would break down. And that's why until now, until these neural operators, where we proposed until that point, people never thought about replacing the full traditional numerical solvers. So the whole point was let's keep the solvers, let's get measurements from them, and then try to do super resolution on top of it. Still mapping from one fixed resolution to another fixed resolution. But on the other hand, we couldn't completely get rid of the numerical methods and the traditional solvers.. And so the way we went about overcoming this fundamental drawback is to come up with now framework called Neural Operator that, so once you've trained the model using some training data at certain resolution, at testing time or at inference time, you can test it at any resolution. So you can give it an input at a different resolution, even a higher resolution than what it has seen during training., and still it can make valid predictions. And that's the concept that makes this both, first of all, important for scientific simulations because people want the flexibility to choose different meshes, different sampling techniques, right? They don't want to limit to one resolution, one grid. And the other is it also gives it the superior ability to capture the fine scales, which is so important for these simulation frameworks.

And you mentioned in explaining this convolutional operator, is the idea here that the neural operator is a kind of an abstraction of that idea and the exact relationship between the thing that you apply it to and the output is learned?

Absolutely. At a conceptual level, it's very similar to the current neural networks, right? Could be convolutional neural network, transformers. It is also trying to learn a mapping from the input to output. But the difference here is that it doesn't just accept your input at one resolution. It has the flexibility to accept input at different resolutions. And which is lacking in our current standard networks. You know, if you take convolutional neural network, it learns filter at one resolution. So you give it input at a different resolution, right, it completely fails. And so this fundamentally says that we can now learn an operator which is mapping between function spaces. So we are changing our input to a fixed size, right, could be a vector or matrix like image which is of a fixed size, to one that is a function. And that function you could sample anywhere in a domain, usually the continuous domain. So if it's a fluid flow, right, and I tell you what is the domain, what is its boundary, you know, you could make your resolution finer and finer, right, because it's a continuous domain. So you don't just limit to one resolution.. And that's the ability we provide in these neural operators. Of course, the next question is, how do I make this practical, right? I mean, sure, this is a wishlist. So far, what I said is, this is what I want out of a learned mapping that is not present in standard networks. Now, how do I make neural operator actually possible that can handle these different resolutions?. And so the way we now go about doing this is through Fourier domain operations. So what we call Fourier neural operator. And you can operate the Fourier transform at any resolution, even any grid, right?

So conceptually, that has the expressivity to handle different grids, different resolutions.

And is this where the traditional numerical methods come in? It is inspired by traditional methods, which also use Fourier transform,. But the difference is we are marrying that with nonlinear transformations like in standard neural networks. That's the power that deep learning has, right? It's not just linear. You have nonlinear activations. So we are combining the principles from signal processing, signal representation theory that using Fourier transforms can represent signals at any resolution.. But now, if I learn like in the frequency domain weights and I go back to the standard domain, that's not enough, right? That's just a linear transformation. And so I now do nonlinear activations and keep doing these series of operations again and again. This way I can capture scales at different frequencies. So I can capture a big spectrum even though like each set of operations is just linear, but I have nonlinear activations in between. And so in a way, it's a marriage of the old and the new, right? The old being the properties that Fourier transforms have that, you know, people in numerical methods have been using that as a way to express signal more efficiently. But then they assume linearity that that, you know, you should be able to capture the spectrum of your input compactly through Fourier transform. And that's usually not true, right? Many of this have high frequencies. So just doing Fourier transform once isn't very efficient. But if you combine it with nonlinear transformations of multiple such blocks, it is expressive. And in fact, we show that this becomes a universal approximator for function spaces, meaning just as theoretically a standard neural network can universally approximate any function in a fixed dimension, we can now use neural operators and approximate any operator in a function space. So it can— it has the expressivity to handle any nonlinear operators like what we see as solutions of fluid dynamics and other partial differential equations, as an example.

So you've gone from wishlist to something that's more tangible, but you haven't, you know, yet demonstrated practicality.

Can you talk a little bit about practicality? Oh yeah, yeah, absolutely. And that's where, you know, the last year is where we really took this off to very practical and real-world large-scale applications. So One of them has been in the realm of weather forecasting. Our ability to predict what's going to happen in the next week or next 2 weeks is so critical, and especially extreme weather events, which are increasing in their intensity and scale as climate change intensifies. And so what we showed is that these framework using these kind of concepts from neural operator we can predict weather as good as the current numerical weather models for as long as 2 weeks, but be about 45,000 times faster. Wow. And so that's, you know, really exciting and a hard benchmark because currently numerical weather models use numerical solvers, sometimes over thousands of variables of partial differential equations.. And that can take hours on CPU clusters, right? Whereas our model works within a second, in fact, a quarter of a second on a single GPU. And we've open-sourced this model. We've made this available to the community. And I think that's very exciting to see this kind of democratization, not just in the generative AI that's so much in the news, but in these scientific domains. Enabling everybody to run their own weather model and build all kinds of downstream applications, whether it's for agriculture. Can we come up with accurate regional predictions that ingest new sensor data, new information, and that can refine the scale locally to provide better estimates? We can also ask how this information could be used in conjunction with the climate model, because not just looking at the weather of today, but what about the weather of the future? How well does it extrapolate? And these are aspects that we are right now working closely with meteorologists and climate scientists, both at NVIDIA as well as the broader community. And that's been very exciting. And that's just one of the examples. The other one that I had mentioned before has been in carbon capture and storage. So climate change is upon us, and I think a lot of scientists believe that the only way to completely tackle this is to mitigate it through frameworks like capturing carbon, right? Whether it's directly from the atmosphere, what is known as direct air capture, people are also trying that from the ocean, or it's like as it's being produced, you isolate and capture it and pump it deep underground. And so there, the phenomena that we need to model is how does the carbon dioxide pressure build up deep underground? As we interacted with water, it's a highly nonlinear gas plume evolution. And so how do we contain the pressure over several decades? And even there, we can see benefits as much as 700,000 times faster faster than numerical solvers. Because here it's a very, again, fine-scale phenomena in these wells underground. You need to model how both carbon dioxide and water, you know, this is called multiphase because you have both liquid and gas interact. And so modeling that nonlinear phenomena, but also it's not just in a single well, it's multiple such wells where there's some porosity, right? There's some permeability, they still interact. And so having a multigrid approach and the ability to capture both the coarse and the fine scales is so important in this application. And so that's just another application of these methods. And yeah, it's been very exciting. We have now weather forecasting, carbon capture and storage, you know, modeling deformation in materials, you know, how much can I stretch this material, right? How plastic it is. Is also a very nonlinear phenomena. And we can again show hundreds of thousands of times speedup over traditional solvers. We are able to do modeling of lithography process. So how do I go from, right, a mask design to what is the resist image, like what is finally being shown on the silicon wafer? And we also show the inverse problem, which is really critical in many of these applications because You wanna design a better mask to be able to create the wafer of like desirable properties, right? So even those problems are especially hard with numerical methods because you have to keep running this forward simulation, which is very expensive. And, but again, this is not data-driven, so it's very hard to make changes or explore the design space. But what we showed with our neural operator-based methods is that we can progressively self-train and improve. We can keep creating better masks and we can train on them. And with that, it can generate better candidates. And so all this, what you see with reinforcement learning or progressive self-training, all these phenomena in other general AI applications, we can bring all those tools here as well. And marry it with neural operators because that gives the right foundation to capture all the fine scales, which are very important in many of these scientific domains.

It sounds like the various applications you've described all involved kind of research efforts to apply neural operators to a particular domain. I imagine what you eventually want is a tool that you take off the shelf and assuming you know your underlying PDE models, you can readily apply this tool. Can you talk a little bit about, A, is that a fair characterization?

And B, kind of how do you think you get there? Yeah, I mean, as all these efforts, right, in the beginning, it's been very important to work with domain scientists and that's been the goal of also the AI for Science initiative when I founded at Caltech. But now it's even broader than just the Caltech community, right? We are working with national labs, we are working with other universities, closely working with also NVIDIA engineers who are enabling and helping scale up all these frameworks and open sourcing it to the community. And I think that kind of building trust in the beginning is important. Because, you know, Sally Benson at Stanford is considered a leader in carbon capture and storage, right? And we are collaborating with ECMWF, the European weather agency that, you know, that in fact created the datasets for these historical weather. And so working with them and asking what are the metrics that matter here, right? You know, it's not just the short-term, what is the long-term behavior? What are the uncertainties? And I think that aspect of the deep domain expertise being married with AI and fully integrated is important. But you're perfectly right. As we show that these proof of concepts are really solid, they've already been making impact in these domains. The next question remains, how do we generalize? How do we create foundation models for science and engineering. And yeah, and that's a mission that I'm undertaking now. And I should say it's not straightforward. It's because one is we don't have the data like what we do for text or images, right? And even there, it required quite a bit of curation effort. And here we need to think what that is. The other aspect we've been working really well on is How to incorporate the right physics constraints. What are the meta-learning and other approaches to generalize beyond one domain? And we've been seeing really good success. So I do think it's a great time to scale up, but we need to bring all these pieces together.

And that's what I hope to do this year. Yeah. Yeah.

How different does the application to weather look relative to the application to carbon capture? Yeah, that's a great question. You know, there's some architectural details that could be different. For instance, one of the newest architectures we are experimenting in the weather case involves spherical geometry, right? Because we know we wanna predict the weather on the surface of the Earth, that's a sphere. So by incorporating that geometric information, how much better can you do, right?, right? So, and that's always this balance between, you know, knowing more about the domain, incorporating the inductive bias versus generalist model that may know less about it, but with more data could do well. And, but the underlying concepts are still same. And I think there is a way to still capture the aspect that, you know, the underlying, right? If you only look at, say, the turbulence, right, there is similarity. But then on the carbon capture and storage, that interaction with water is very important to model. Whereas in the weather, there's certainly water. We have the oceans, but that's usually considered a longer timescale evolution compared to the shorter timescale wind, surface wind, and other variables. And so people conveniently kind of ignore some of the phenomena or simplify it for different domains, which I think is very important for tractability. But yeah, I do think that, you know, there is a broader range of phenomena that could be captured through a common model, and it need not be a model that's learning everything all at once, right?

It could kind of provide the seed or initialization for other models to further hone in and fine-tune on a very specific domain. When you're referring to your results with weather, for example, and you have the significant speed-up, is that operating under the same sets of variables, or are you— is the model that you're referring to simplified relative to the model that you're comparing against? And likewise, one of the big motivations here is the ability to handle, you know, different resolutions.

Are you— when you are comparing those models, are you looking at the same resolutions, different resolutions? Yeah, absolutely. First of all, we're not even currently taking all the information that's available in the historical data, right, and all the variables. So, you know, in the beginning when we started this project, we wanted to do just a rough cut. We thought, okay, let's start with a subsample of data and a selection of variables just to see what we get. And we didn't expect we didn't expect it to be producing such good results with even just that subset. And that's where I think, you know, there is a potential and we, you know, see evidence that we can even beat the current weather predictions by training bigger models on all of the data available. And the weather is a great use case because we have wealth of historical data available, you know, that's collected since the 1970s. Collected hourly, right? Whereas these numerical weather models hardly use any of that data. So they're using it to just calibrate a few coefficients or parameters in their numerical methods. And so in the sense, they're redoing these computations again and again without having the wealth of data being understood and incorporated into the calculations.. And I think that's where the, you know, impressive speed-up comes, right? One of the important reasons is because these AI-based methods can learn from data and they can learn better ways to do these numerical iterations, whereas the standard PDE solvers are forced to go to a very fine grid. As I said, we cannot ignore the fine-scale phenomena. So they have to, for guarantees, for convergence, be operating at a very fine grid, and that makes them expensive. And so our ability to learn from data, better nonlinear transformations where computations become quicker, I think is a key to the speedup. And regarding your question of, can we predict at higher resolutions? So currently the data that's available we've trained on is is a 25-kilometer spatial resolution, right? So it's 25 kilometers spaced on the, across the globe. And so we are working with regional weather agencies to see if we can get now, you know, we are in fact getting some of the higher resolution data. And then we can ask, you know, how do the predictions look? And is there even potential to further refine those? And I think, That's why this kind of a foundation model is valuable because you can go from being able to predict all across the globe to regional models where especially countries that don't have a lot of computing capabilities or technical abilities, right, could start with a reasonably good model and hone in on their region and refine in a much more cost-effective manner rather than running these numerical methods, which are extremely expensive as you get to the fine grids. So just to give you an intuition of how expensive it gets as you refine the grid for these kind of calculations, my colleague Tapio Schneider here at Caltech, who is in climate sciences, estimates that you need about 100 billion times more computing than what we have today. If we have to go to the actual finest resolution to be able to capture all these turbulence in low-lying clouds, which for climate models is the biggest source of uncertainty. And so that requires computing at the resolution of 1 meter. I'm talking 25 kilometers now, right? We need to go all the way to 1 meter spatial resolution all across the globe. And 1 second in time temporal resolution. And so that's, you know, the Moore's law is no longer upon us. And even all this acceleration and scaling will not get us to 100 billion times, right? And even if you say, oh, that's too much, what if we get somewhere intermediate? It's still too large. And that's where I think machine learning is a necessity.

Otherwise, there's no way we can tackle these deepest questions in scientific domains. With these models, you've referred to the, um, the relative speed-up, um, in the case of the weather, for example, in terms of performance, predictive performance.

How do they— this method compare to traditional numerical models? Yeah, absolutely. Uh, in the case of the weather, we are doing as well as the current weather models for as much as 2 weeks, which is considered where, right, it's very predictive. So after that and a bit longer, you kind of get to the sub-seasonal and seasonal scales where it's no longer predictable. So it becomes chaotic. And, you know, you can kind of say, right, statistical measures on average what it would look like, but not the actual trajectory, the actual whether what it is because it's just not predictive. And that's another challenge too with many of these models, right? So for short term, we can predict, but for the longer term, we can only simulate and on average get to what's known as the invariant measure, which is the attractor or the right distribution, right? Because you cannot predict exactly where the trajectory evolves.. And we are also working on using these neural operators, and we've shown their ability not only to capture short-term phenomena, but also the long-term behavior by capturing the nature of chaotic systems.

And hence, through that, be able to simulate them again effectively. Do you characterize the failure or error of these models and like compare that to the way traditional error models fail. Just thinking about, you know, weather is something that, you know, we— there's kind of this inherent, you know, we joke about how bad weather models are, right? And how, you know, difficult it is to predict the weather. But then a lot of people count on those predictions in very important ways. I'm thinking about like if we, you know, start predicting the weather in different ways and those predictions are, you know, generally better but occasionally much worse, like how do we communicate that to the users of the predictions?

And I'm just curious, you know, how that, you know, factors into the way you think about the problem. Right. And that is a big challenge, right? Like, as you said,, you know, we need the uncertainty quantification as well. You know, are we getting to the right probabilities? So for instance, if we slightly perturb the initial condition, how does my prediction change? And so that's an important notion of stability. And also because we have some uncertainty on our input, right? Like the, what we are measuring from our satellites and how that is all assimilated. Has certainly errors, right? It has uncertainties. And so this kind of what we call ensembling, because you're not just feeding in one fixed initial condition, but you perturb it and then you ask, even with those perturbations, what is the output now? And through that, can I get a measure of the uncertainty or probabilities? And so we can also calibrate our weather models, AI-based ones, similar and even better many times than the current numerical weather models because we are much cheaper, right. So these numerical weather models, you'll be surprised, use only about 50 ensemble members. So what kind of statistical averaging do we get out of 50, right? So 50 samples. And but that's so expensive, so, you know, people are using large clusters even for one single prediction and takes a few hours. So they are bound by the cost. And that's where this is another big benefit that we are seeing from AI-based models. Because it is so cheap, you can now run thousands of ensemble members like we are doing now. And we are carefully testing how does it do in all kinds of scenarios, right? We show, for instance, it can nicely capture the uncertainty around hurricane evolution. So many of the famous hurricanes, that didn't appear in the training data, but we are testing them on, you know, how not just like the single trajectory, right? Because people care about uncertainties. And that has all kinds of downstream implications. You know, does hurricane hit Florida or go into the ocean, right? That makes all the difference. And so being able to get to the right uncertainties is such a big aspect. And that's where, AI models are already proving to be superior. And by honing in all that and giving the right ensemble-level estimates very cheaply, I think will make this so valuable.

And that's what we are seeing now and working with weather scientists. Now, all these applications that you've been talking about have traditionally been, we talked about kind of the traditional numerical approaches. Those have all driven kind of the development of high-performance computing as a field. Um, and I, I guess I'm curious what, you know, if you, you know, thought of HPC as a, you know, a pie or whatever, like how much of that pie is, you know, being eaten up by AI today and in time, you know, over time, does AI, you know, eat all of that pie? Are there elements, are there things that we're doing with HPC that, you know, we don't think, or you don't think, will be, that are incompatible with AI in some way?

Or does AI, you know, just come to be a standard tool that's applied to that type of problem? Yeah, no, that's a great question. I mean, so many now national labs and other centers are rebranding as HPC AI, right, combination, which is great. And, you know, that's been our experience too, like kind of to think of that as an integrated approach to begin with, right? And that doesn't mean like keep the numerical methods as they are, keep an AI tool as a standard hammer, you know, take something that works on natural images, bring it here, right? That to me is not a winning strategy. When I say hybrid HPC and AI, it's really thinking conceptually at an algorithmic level and say, what does an integrated algorithm look like? And that's how we came up with Fourier Neural Operator, right? So we know numerical methods compute in the spectral domain that's efficient. And that's also a nice way to create a basis that works in any dimensions. And any resolution because it is having the ability to learn in a function space itself. So those, you know, thinking and that intuition, we don't need to reinvent from scratch. We already have so much that's been dealt in numerical methods, right? On the other hand, we know what's also the downside of numerical methods. They can't be data-driven.. So you can't change the computation based on if it's an easier or a harder sample, right? And you can't use all the existing computations you've done to learn and improve. And that's what AI provides. The other is all these nonlinear transformations, which is also an effect of being data-driven, that we have the ability to learn these nonlinear transformations. So I think to me, ultimately they'll all come together. But I don't want to get into the debate of what was numerical method, what was AI. You know, in that cyborg-like baby, you shouldn't be able to tell them apart.

That's kind of the whole point, that they just seamlessly integrate together. That's my hope, and that's what we are working towards.

Are there other ways that you see the kind of advances around, you know, that are happening broadly in AI kind of driving AI for scientific applications? Yeah, yeah. As we've been talking quite a bit, last year has been the era of generative AI and foundation models, right? So these big models having the ability not just to do now one targeted task like we were doing for the past decade, but really be general purpose. You know, they're not all the way, but if zero-shot, you can't get a good answer, you could do few-shot, right? And you can give some examples. You could even fine-tune. And what we've been very excited is taking one step further and asking how we can use these foundation models for decision-making. Using with reinforcement learning, imitation learning, all the tools which are conceptually the right thing to adapt to new environments and make decisions. But are very, very expensive, right? We really haven't broken beyond standard game settings and still fairly low-dimensional action space to one now that is open-world setting. And that's where I think there's been a lot of exciting work. And one of the works that we've provided as a benchmark to the community is called MindDojo. So it's a suite of thousands of tasks in Minecraft along with internet-scale information about YouTube videos of how people are playing Minecraft, how they're building structures, Wiki, Reddit articles. So all kinds of like text, image, tabular data, all the information that you can glean. And from that, can you now solve not just one task, but thousands of tasks, right? So what we call as open-ended task solving. And so, you know, if I give the instruction in text-based prompts and say, "Go mine a diamond," or, "Go build me a castle," right? So it should figure out how to do it, but it's not doing that from scratch like what we saw in AlphaGo, right? Or AlphaZero, rather. We said either it's some limited imitation data or it is from scratch, both of which are still way too expensive in these kind of domains. And that's where I think Minecraft is different from other games because it's not one goal, right? It's all about unleashing creativity and solving all kinds of new tasks, coming up with new structures. You can create a castle, you can create a flying dragon, all kinds of new structures., but you still need to understand the environment and its complexity is huge. It's not one tool, it's not killing somebody, it's not trying to negotiate and imitate what we saw in training data, right? It is having to solve and learn something new. But all the information that's available as foundation models through videos, through text will help us towards that goal. And that's where I think that benchmark, we got the outstanding paper award at NeurIPS, which I'm very proud what the team did. Congrats on that. Thank you. But it also paves the way for the community to take this as a new challenge, right? And this is in a way much, much harder than any of the game-playing bots we've seen solve. But at the same time is, to me, very timely and relevant because all these foundation models are becoming so powerful. But it's challenging them beyond their current capabilities, right?

Because it's not just learning what they've already seen, but pushing the envelope to create new worlds, new structures using all that knowledge.

And is it a formal challenge in the sense that there's a competition and a leaderboard and all that? We have the website where all the information is available. You know, we haven't launched a leaderboard because this ultimate goal is very hard, right? So people are solving, there was a DeepMind paper that solved only mining diamonds, for instance. So, and you know, we are building to end our first work showed that you can, you know, do solve multiple tasks by using what we call a MineClip, meaning you look at Minecraft videos, you connect it with text. So you create a clip-like model and use that as a way to get dense rewards, right? And so you make progress towards the goal. And that is just the first step. And so we are working on it. We know others in the community are. So as people start solving it, it'll be easy to create a leaderboard.

Still a bit, you know, out of reach, I would say, from what the community is capable of. So that's why it's the challenge for the next decade. Do you expect that the models that perform, you know, best in here will be based on kind of these foundation models, language models, you know, in some way, or is that still, you know, way too open-ended a question, or?

I mean, the foundation models are a necessity to begin with, right? Because we are providing text instructions on what to do in Minecraft. So you're describing the task in text, and it could be entirely new things that, right, this agent hasn't seen before. And so you need to go to now the foundation models to understand the text, but also understand what this image means, right? And go to the wiki and see 'Oh, this is what a hammer is. Now how do I go and pick up a hammer? Where do I find it?' So even to just get started, all this is necessary, right? And so that— because we are not pre-programming all of this, whereas if it's one fixed environment, you just give all these rules beforehand. You set what the objective is, you design the reward function beforehand, right? So, and then you let reinforcement learning do its magic. Which can take very long and can be very challenging to do, but it's all preset. Whereas in MineDojo, none of this is given beforehand. So to even get started, you need this, but it's a great test for foundation models because what can you really learn from it? You know, I give you everybody playing Minecraft before, you know, what can you glean from it? But still we are asking for something new beyond that.

You know, how well can you do problem solving?

And I think it's exciting to see how that develops. Awesome, awesome. And I feel like there's a whole interview just on, on that topic. Absolutely. And in fact, with NeurIPS, there was a workshop on foundation models for decision making, and there's just so much excitement because I think that'll provide, in my view, the right starting point, the right initialization for some of really challenging deep reinforcement learning problems. Because I would argue that reinforcement learning has shown, you know, really fantastic gains, but in highly focused and limited domains, right? It hasn't become the generalist agent like we see with, say, generative modeling. You know, you can sample from distribution through diffusion models or GPT-like models very easily, but that's not task-oriented. That's not learning online with the reward function.

And so these two have to marry together to bring that generality while being able to learn online, improve, adapt, and make better decisions.

Yeah, yeah. Well, clearly, Anima, 3+ years is too long between our conversations, 'cause there's so much to cover. Certainly been an exciting time, and I'm privileged to have a great team that is enabling me to, you know, think about such a wide range of problems. So that's been very exciting. Awesome. Absolutely. Well, thanks so much for joining us and sharing a bit about what you've been up to.

Thanks a lot, Sam.

It's been a pleasure, and hope you have a great rest of 2023. Absolutely.

You too. Thank you.