🔬“We have foundation models for language, not for physics” — Anima Anandkumar, Bren Professor of Computing

Speed and accuracy in weather models

0:00So we, you know, set out looking for interesting examples. And one of them was like weather modeling because the weather data is open source. And so given that the data was there, we were like, OK, let's just go try it. Right. And that's the beauty of it. Whenever data is available, it's really good news. But a lot of weather scientists did caution us back then. This was back in 21. And they said, no, no, no, this is so difficult. You know, there have been decades of development in traditional weather forecasting, and that's very careful, bottom-up physics-based modeling. Right. So assuming, oh, this is the fluid dynamics. Can you go predict the weather the next day and so on?

0:39And so that's how a lot of the thinking was that AI is just not going to be able to beat the decades of work in weather modeling. But to our surprise, we just went ahead. We trained them. We used neural operators to be able to effectively capture the phenomena. And then we found that it's not only accurate, it's almost as close to what the traditional weather models can do accurately, but also tens of thousands of times faster. So what would take a big supercomputer to run can now be run. And we only needed a consumer-grade like GPU.

1:16Like, you know, it was a small model. It fit very well. It's very fast. And it's accurate. And I think that just changed everybody's thinking.

Introduction of guest Anima Anandkumar

1:27Welcome to LeanSpace. This is the AI for Science section of LeanSpace. I'm Brandon. I work on RNA therapeutics using AI at Atomic AI. I'm joined by my co-host, RJ Haneke, who develops spatial transcriptomics and is the CTO and founder of Mirroromics. Today, we're excited to be joined by Anima Anankumar, the Brin Professor of Mathematics and Computer Science at Caltech. Anima has done all sorts of really cool work combining AI with basically models of the physical world and has a really diverse background.

2:02I don't think I could even remotely cover it. But anyway, I'll let Anima introduce herself. Thank you for coming on the show. Yeah. Yeah. Thank you, Brandon and RJ. It's a pleasure to be there. And I really like the term latent space because that very much figures in a lot of my work because it's really, you know, the world is latent. But yeah, just as a brief introduction, you know, I've been working in AI for more than two decades. In a way, you know, before even deep learning, when a lot of the theoretical foundations had to be built for probabilistic models, I worked on them.

2:37And then as deep learning started taking off, I also had a foot in industry until recently. So I was at NVIDIA. I led AI research there. And before that at Amazon Web Services helped found the cloud AI team and built the first cloud AI products back almost a decade ago. So, you know, like kind of having this one foot in industry and academia, I think has given me a lot of interesting perspective of how to bring theory and practice together and think of AI at large scale, but also AI that is principled.

3:11A lot of your work has been related to the modeling of physical systems using certain types of physical systems, which you model with differential equations and you help model them with using machine learning. So maybe first, let's go ahead and talk a little bit about that as a high level, but we'll get to kind of the details about the same neural operators and some of the applications like weather later.

Torch Lean and mathematical verification

3:36But first, I'm actually really curious to hear about TorchLean and how this recent work you've been doing connects with that larger research program. To me, broadly, like, you know, my thesis is AI and science, how we bring that together, right? So, you know, when I started at Celtic almost a decade ago, that's when, you know, my passion was always science, was physics. And but, you know, I was doing AI. So how to bring that together was where, you know, the first kind of foundations got laid there.

4:10And to me, like, you know, there are several aspects to that. One is people have been thinking how to use language models for science. Yes, you can do a lot of hypothesis generation. You can have ideas, but ideas are not enough, right? So you can have a lot of ideas. The bottleneck is going, testing and verifying that they work in the real world. And so this aspect is where a lot of my recent focus has been on how do we ensure that we can build AI that has guarantees that it will work in the physical world or any aspects in scientific domains.

4:51And one way to think about it is, you know, can we model the physical world and keep the physics correct? And that's where neural operators come in. The other aspect is, can we verify symbolically certain aspects? For instance, you know, if we claim that the theorem is correct, we have to go verify that. You know, that's where lean as a formal language can be useful for verification. So how do we bring that together with language models is where a lot of mathematical reasoning has been at the forefront.

5:28And so Torch Lean kind of is in that realm where we say, you know, not only that you want to verify mathematical statements, you may want to verify what neural networks themselves claim to deliver. You know, for instance, if you know, for instance, if you are now using a neural network and you want to ask whether it's going to be robust, say you want to use a neural network in a control loop, you want to control, you know, whether it's a drone, whether it's a nuclear reactor.

6:00So all of this ultimately, when we build AI systems with deep learning into control loops, we want robustness. And so now Torch Lean can help us do those verifications seamlessly. So we can now have neural networks be part of the verification loop and have confidence that we can use them appropriately. We have already discussed on the podcast, Lean, and everyone should probably be familiar with neural networks. How, but neural networks seem very unconstrained.

6:33What kinds of proofs are you talking about? Are you talking about bounds on the outputs, inputs? What can you prove with Torch Lean? Yeah. So Torch Lean is an overall framework, right? So what it really enables is that you can now write neural networks essentially in Lean. So instead of writing in like PyTorch, it's like a PyTorch-like abstraction, but you can like kind of, you know, write it in Lean. And so it can be fully formalized in Lean. And then there are several implementations.

7:04You know, we have algorithms for certified robustness like Crown. You know, those are implemented under this framework. Sorry, was it like what? Like Crown? Crown. Crown is one of the... Can you explain a little bit? So there are different ways to bound, you know, for certified robustness. You know, how tight those bounds can be, it depends on the relaxation techniques. And so, you know, without going into those, there's many such algorithms. But, you know, we are kind of like implementing them and enabling them in Lean.

7:35So we can seamlessly run both. You know, we can both first kind of write down Torch-like framework neural networks very simply, right? And then we can also make statements about them formally and verify them. So all of that can be brought together in one framework. So what's an example of a bound that you could claim? Like, so we're operating a nuclear reactor. We don't want it to melt down. What are the sort of guarantees that you could provide to the inputs and outputs that would help that not melt down?

8:09Yeah. I mean, the natural one is the certified robustness that I mentioned. So saying that if your inputs are, you know, perturbed by a certain amount, how much is the output going to be perturbed, right? This sensitivity analysis is another term. And so having those kinds of bounds for different neural architectures, so you kind of automatically get those bounds, can then help us, you know, not only train neural networks to do well in a control loop, but also worry about safety and robustness or stability.

8:41These are all part of control systems that people worry about. So that's one example of an application. So it's really more broadly, the idea is you need verification in lots of scenarios that involve neural networks. So control loops are one. Another example is, you know, we used physics-informed neural networks to, say, solve partial differential equations or come up with systems that are guaranteed to satisfy certain physical laws.

9:12But we also want to verify, for instance, that our neural network is only trained in finite precision, right? So can we overcome those requirements and what happens when we are, what are the shortcomings because we are using this finite precision? Can we also bound those? Oh, right. So those are other kinds of bounds that work in torch lane. So all aspects of, like, you know, the effect of precision, the effect of perturbation, all of these, we can, you know, we can have algorithms that are implemented in lean that can be seamlessly now part of the verification loop.

9:50So is the descriptive power of the torch lane, is that sufficient to describe basically any neural network or is there, are there constraints on that? Yeah, so it's essentially a, you know, PyTorch-like, you know, framework, right? So you can just kind of nicely define neural net layers in the same way. But the back end having, like, lean helps us formalize and prove it. And for, like, transformer architecture, for example, is it reasonable to prove these kinds of bounds on a very large neural network?

10:23So, you know, there is the aspect of one is, like, kind of having the framework, right? The other is scalability. So lean still has a lot of shortcomings there. It's CPU-based. And, you know, it's not, like, getting that onto the GPU has a lot of nuances there. So, you know, a lot of work needs to be done. So what we've started with is a framework, you know, making that more efficient, especially at a very large scale, requires still a lot of work to be done.

10:54Yeah. But that's true broadly for lean as well. I'm just trying to understand, like, how I imagine this. So if I were to take, let's say, numerical analysis classes, you know, on graduate level numerical analysis class, you have a differential equation, you have some discretization error or something, and you bound, like, given these properties, I can bound the solution, right? So solving some of these physics-based or AI-based solutions to differential equations, I think, historically, has been kind of the wild west. So I think you mentioned physics-inspired neural networks.

11:27Really cool idea. It'd be fun to talk about that a little bit. But I know that sometimes they are particular and that people don't always work. And I think people don't always know when they will or won't work. I mean, I'm not an expert, but I'm just wondering if that's in your experience. And what I'm wondering is, like, has this helped you understand, like, the domain of applicability for PINs? And is that sort of, like, the goal is, like, you can rigorously say, like, this solution will converge?

11:57Or is there not necessarily the same concepts of convergence in a controlled way for neural networks?

Physics informed neural nets and neural operators

12:03Yeah, so, you know, like, physics-informed neural nets are about, like, saying that, you know, I write down, like, a PDE, partial differential equations, and hopefully the optimization succeeds and I get the answer, right? And, of course, if, like, optimization was not at all an issue, this would be universal. You solve everything, you know, we're all happy, but that's not the case. And so, optimization ends up being usually very difficult, especially for problems that are time-dependent, meaning it's not just stationary.

12:37You also have time, and the time component in many cases could be turbulent, like in the case of fluid dynamics. You know, you kind of, like, if you run it long enough, it can become chaotic. So, you really, you know, have, like, very small, fine-scale effects matter. And so, in those cases, just trying to solve a partial differential equation at all times is just hopeless. Like, you know, this is not an optimization landscape that, you know, I think we'll, you know, we can have any handle on.

13:13And this is where the idea that, from scratch, we would be able to solve these equations using a neural net is not possible. So, pins don't work everywhere. And our idea of neural operators came as a way to overcome this, right? So, saying, you know, we can't rely just on physics constraints alone to come up with answers. We have lots of data available. You know, I'll talk about the weather example where we even collect data, right? So, we don't just solve equations and have synthetic data, but we also have real data by observing the weather as one example.

13:51So, why not make use of all of the data available? So, we don't just rely on trying to solve partial differential equations and other physical problems from scratch, because it's really the data-driven approach that makes it possible to get quick answers. And so, with neural operators, we can bring both of them together. We can have all the data that's available. We can utilize it. We can add physical constraints. And then that overcomes the limitations that pins face.

14:21Can you give a little bit more intuition on the difference there and why that is possible? So, I heard you mention, you know, with pins, you're basically just baking the physics constraints into the neural network, but that this becomes unstable over time or other variables. Whereas, if you add a little bit of data, like, I can kind of intuitively understand why that might help, but can you give a little intuition for what's going on? What's the difference here? So, with a pin, like, you know, every instance of an equation, you solve from scratch, right?

14:55At least in the classical sense. So, you start, you take the specification of what equation you want to solve, and you hope that the optimization landscape succeeds, which in many cases it doesn't. Whereas with neural operators, what we do is we, you know, have lots of data. So, we have a training phase, we teach it how to come up with solution for different instances of equation. And so, just as in other supervised learning at test time, you can now ask, you know, can you come up with an answer?

15:26And you can still have physics constraints as a way to guide that. So, you know, it can be both data-driven and physics-informed together. But the benefit is because we have data, you know, it's like you're not stuck in an optimization landscape, right? So, you know what the answers are during training. So, you are now at a better chance to come up with the right answers even at test time. My understanding is that a neural operator is a function fit to data or a neural network, you know,

16:02learns to fit functions to data. Is that a good intuition here? Yeah. So, you know, neural operators are, in that sense, similar to, you know, it's the same as neural networks, right? You're learning on data. But the difference is neural operators are, you can think of it as a generalization of neural networks. So, with standard neural networks, the inputs and outputs are a fixed size. So, in language, we have fixed vocabulary, we fix what the input and output are. And same with images in computer vision, in videos, we assume a fixed resolution and we always, you know,

16:36our inputs and outputs are always at that fixed resolution. We can't change it post-hoc. Whereas with a lot of this physical data, the idea is our world is inherently multi-scale. So, you should not be, like, deciding beforehand what the resolution is. You know, maybe you have, like, weather data available only at course resolution. But really, the actual phenomena is happening at a finer scale, right? And maybe you want to, after that, incorporate either additional data at finer resolution or add in physical constraints at finer resolution.

17:10So, we should be having that flexibility and we should really think of the world not at these fixed resolution, but one that's happening infinitely, you know, that one real world happens at that infinite resolution. And that's what neural operators enable because they model inputs and outputs as continuous functions that can be infinitely resolved, that can have infinite discretization. And now we can have, you know, at inference time, you can give it now inputs and ask for outputs at any resolution.

17:46So, you're not just limited to the resolution of training that we see in standard neural networks. And that's what neural operators enable. So, neural operators enable us to zoom in and out as we like. So, obviously, that, as stated, that any function, that's under constraint, right? You could have many, many functions that fit the data. It would be easy to overfit. So, how do you regularize that? Yeah. Certainly, like, you know, if you're asking about making predictions at a higher resolution than what's seen,

18:20like what we call zero short super resolution, you're kind of making some guesses, right? And that's what these models are doing. They're trying to regularize and kind of smoothly extend to higher resolution. But, of course, if you now give it the model additional information in terms of, let's say, a physical loss. So, you could give it partial differential equation constraints, conservation loss. And you can now enforce them at a finer resolution than the data you have. Then there's more guidance in a way.

18:51So, that way it can now come up with the right answers even at higher resolution because you're, you know, giving it constraints at higher resolution. And so, that's how we can ensure that these physics-informed neural operators can work at higher fidelity and higher resolution than even the training data that was available. My understanding is a lot of your work uses a particular kind of neural operator, a Fourier neural operator.

Fourier neural operators and efficiency

19:21So, Fourier is a dual domain. It is extended across the entire domain of the inputs. That's a lot of jargon, maybe. Can you give some intuition for why is that important? How does that help? What I mentioned, neural operators as a class of models that allow us to have any resolution input and any resolution output, right? And learns the mapping between them. You know, that's really called an operator.

19:52So, the mapping between function spaces. So, that's the reasoning behind the name neural operator. And, you know, Fourier neural operator was one of the early setups or architectures we came up with. And the reason why that's been so successful is because it kind of strikes a nice trade-off between efficiency and expressivity, right? So, why is the Fourier space a good one? The Fourier space allows us to, you know, it's a dual space, like you mentioned, but it really allows us to capture non-local phenomena, right?

20:29So, meaning something that's like non-local in the Fourier domain could be even efficiently captured. And a lot of phenomena like we see in nature, whether it's fluid dynamics, material deformation, quantum chemistry, it's all, you know, there's a lot of them are non-local. You know, the differential equation, like the derivative is local, but the inverse of it is you're kind of doing essentially integration. It's non-local, right? So, the solutions are non-local and these models are able to capture that.

21:03But at the same time, doing Fourier transform is efficient and it kind of like nicely captures a lot of inductive bias we see in many of these natural phenomena. But this doesn't mean that we are capturing the world entirely in the Fourier basis, right? It's not a linear representation in the Fourier basis, which is what classical numerical methods do. We add non-linearity just as in transformer and other neural nets in between Fourier layers.

21:36And we also add residual collections. So, all of these architectural aspects that are inspired by other neural nets that work well in other neural nets, bringing that together really kind of helps us get best of both the world. So, you can think of like if we were to use transformers and we require a very high resolution, it would become untenable because of the quadratic complexity and all-to-all connections. On the other hand, if you did that with Fourier transforms, we have like quasi-linear complexity and still we have global connections in a way we can model these non-local phenomena.

22:17And so, that's why it's a nice middle ground. So, that allows you to learn from what is happening on the other, like it's talking whether like what's happening in Chicago may have some impact on what's happening in San Francisco. Well, maybe not, but that's the idea. Yeah. So, that's the idea. And time, like kind of, yes, at this point, maybe local, but eventually they have an impact in other locations. And yeah, so both in space and time, we want to capture that dependence. Yeah. So, what happens today in Chicago will have an impact in a month in San Francisco or something like that.

22:53Yeah. So, you know, so there is like both the short-term and the long-term effects. So, in a short-term, like we think about predictable weather, but longer-term, we're talking about climate, right? Right. So, what happens? You may not be able to say precisely, you know, what happens in Chicago, what will happen in San Francisco. That's like the butterfly effect. On the other hand, we can kind of give averages. You know, if there's heat wave in this kind of overall region, you know, we kind of have one idea that it's going to be higher than average temperature.

23:27So, those are the aspects we can capture together. From an architectural standpoint for all the AI engineers here, are we just talking about doing all the work in the Fourier domain, but it's basically the same neural network, but I'm just operating in the Fourier domain? Or is there other aspects that are required in order to do this properly? So, think of it, I guess, maybe the easiest way to think about it is, you know, if you think of a transformer architecture, instead of like the, you know, attention map, you now have the Fourier, but you still have other non-linearities.

24:03You have like, you know, the residual, you have, you know, many other parts of the architecture still there that give it like expressivity. And we have lifting to higher dimension, like, you know, in a channel space to give it more expressivity. So, all of those kind of best principles are still available, but the Fourier helps us capture that all-to-all, you know, for dependence without requiring very huge complexity that I see transformers.

24:33That makes sense. The other advantage is that it gives you the natural multiscale, what's sort of a implicit cutoff, is, you know, the sort of, if you have a signals background or physics background, you might ask, you know, in linear, if you're doing everything linearly, there's a maximum frequency and, you know, above that you can't represent anything. But how does that, how does adding these other architectural changes in a non-linear domain actually affect, you know, your choices of frequency bounds? Yeah, yeah, no, that's a great question.

25:05And that's where the expressivity comes in, right? Otherwise, if you're just taking a Fourier transform of a signal and trying to represent it, you know, that's what numerical methods have also attempted to do. And that requires very fine discretization, and that's why it's very expensive to do the simulations in a classical way. And instead, if you want to move away from that and say we want to learn the features, which is what deep learning is all about, then we cannot force it to be only in the Fourier domain. We have to give it non-linearity to figure out what the right basis for it or the best basis to represent the signals are.

25:41And so that's the kind of, like, nice kind of combination we have, that it's like all these non-linearities will help it kind of, you know, find the right latent space. No pun intended. This is the right latent space. This is the right latent space. And so, you know, and if you do Fourier in that latent space, you know, that may be a more efficient way to represent. So that's one way of thinking because, you know, first of all, we are lifting the signal to more dimensions.

26:13Even if the signal is two or three dimensions, we are now lifting it to much higher dimension. So in that space, the idea is it's easier to learn. And we're doing it as a non-linear lifting, right? So there's already a latent space there. And then we are doing further non-linear transformations in between our Fourier transforms. So that means we are saying, yes, you know, maybe with this limited number of frequency modes, it's not expressive enough. But when I add non-linearities, I can, you know, I can kind of nice, more nicely capture that.

26:47Yeah.

Balancing classical math and deep learning

26:48So you started your career back before neural networks were, I guess, taken off, right? So I think back then people really did think a lot about, you know, appropriate basis sets and, you know, function expansions and orthogonal polynomials or whatever. How does that evolution from, you know, your research standpoint, like as the community has evolved from that to, oh, just screw it, throw it all in. It seems like you still believe in at least some of those concepts as being guiding principles.

27:18Do you think that there is actually some lessons to be taken from, you know, classical mathematical, like rigorous mathematical techniques that you can use those techniques to actually help improve modeling of the real world, even if you still are just throwing the kitchen sink at things? No, I think it's a nice, I think there's a trade-off involved. I mean, it's funny, my undergraduate thesis more than two decades ago now was on fractional Fourier transform, right? And so, yes, I mean, by themselves, like, you know, that wasn't enough to do computer vision, but I was curious, okay, what are these techniques and how well do they work?

27:57And so, you know, I'm completely with you that we cannot just force ourselves to use Stone Age techniques or classical techniques, right? I mean, so we have to have feature learning, we have to have flexibility, expressivity, you know, they have to be easily optimized. So all of these aspects are very important with deep learning. But when it comes to the physical world and physical data, it's never going to be as plentiful as we see with language models, because we are, you know, our weather model like had about like 50,000 samples, right?

28:3250,000 samples of fairly high resolution, like global weather maps, but it's nothing like what we see with language. And in other domains, it's even less because it's so expensive to simulate and the real data may just not be available. And so here we have to think about the inductive biases more, we have to add in the physics constraints, cannot be just reliant on data. And that's where I think a little bit more thinking of the architectural design comes up.

29:04The other aspect is computational complexity. So think about language, it's just one dimension. And even there, the context length, you know, we are getting to millions and we are struggling, right? I mean, on the other hand, now we are thinking about not just 2D, 3D, even 4D, you know, 3D and time. And if each of the dimension is even a few hundred grid points, which is where, you know, industrial scale starts at like a thousand grid points in each dimension, we're talking like hundreds of billions to even a trillion context length, right?

29:40So forget ever having a transformer for anything of this scale, all of the world's compute will not be enough. And first of all, they all have to be co-located to be able to ever do this. So that's why we need other architectures. But I would push back a little bit, right? We have the vision and video language models, right? And they use, they basically learn a mapping. But the resolution is very low. So that's the key, like for the physical world, the resolution, what we require, I mentioned like thousand by thousand and by thousand by thousand.

30:12So there, you know, if you count that, that's like already in hundreds of billions. Yeah. So, you know, we are not doing that high resolution when we think about images and videos currently. I have a friend here. And the video is also like autoregressive. So it's essentially only like you only need to do the next step. Right. Yeah. But you're learning, I mean, like generally you're learning a code book, right? So you have, you're kind of learning the bias of the latent space or the real world to the latent space.

30:44And so that if there is like a compression that you can do from the physical world into the latent space, then, you know, these autoregressive techniques have been successful. Yeah. But the idea is, you know, a lot of these autoregressive and, you know, techniques for vision and video models are for mostly like, you know, looking good. Right. So they are like not for very precise simulations and they're, you know, having that higher resolution and details is really important.

31:18And so we need to at least take in the data of that high resolution. Right. So we need to be able to process that and reason over them. And so this is where a lot of the bottleneck is because we, you know, cannot afford to just throw away everything and say, oh, let's just like have 100 grid points in each dimension or 50 grid points because there just isn't enough detail to correctly model phenomena like fluid dynamics, plasma, how materials deform. So all of this requires high fidelity and for that we need high resolution.

31:53My understanding, you have a thesis that AI needs, you know, to incorporate the physical world into it in order to scale and be accurate going forward. Many people have this thesis. You are somewhat unique in that you have several examples of applying our operators to the physical world in this way. And then it seems like you're constructing a thesis around your experience here. So can you share with us some of the really interesting and exciting looking things that you've done using neural operators and other techniques?

Open source weather forecasting with ForecastNet

32:32Yeah. I mean, you know, for us, when we started with like neural operators for partial differential equations, but also more broadly, you don't even need to assume they're partial differential equations. Right. Could be any spatio-temporal or data at multiple scales. So we, you know, set out looking for interesting examples. And one of them was like weather modeling because the weather data is open source. It's available called the Aero5 from the ECMWF, the European Agency for Global Weather Modeling.

33:05And so given that the data was there, we were like, OK, let's just go try it. And that's the beauty of it, whenever data is available, it's really good news. But a lot of weather scientists did caution us back then. This was back in 21. And they said, no, no, no, this is so difficult. You know, there have been like decades of like development in traditional weather forecasting. And that's very careful, bottom-up physics-based modeling. Right. So assuming, oh, this is the fluid dynamics, can you go predict the weather the next day and so on.

33:36And so that's how a lot of the thinking was that AI is just not going to be able to beat the, you know, decades of work in weather modeling. But to our surprise, we just went ahead. We trained them. We used neural operators to be able to effectively capture the phenomena. And then now we, to our surprise, we found that it's not only, you know, accurate, it's almost as close to the, what the traditional weather models can do accurately, but also tens of thousands of times faster.

34:11So what would take a big supercomputer to run can now be run. And we only needed a consumer grade like GPU. Like, you know, it was a small model. It fit very well. It's very fast. And it's accurate. And I think that just changed everybody's thinking. And after that, DeepMind, Huawei, many others followed us a year later, released their own models. We were the first to actually open source our weather model for CastNet and do it permissively. So that's what allowed companies, weather agencies, everybody to build on us.

34:45And so, you know, it's been a really interesting revolution to see that the weather models are now out there, weather agencies are adopting them. And it allows us to now have small weather agencies in the global south, for instance, have the same kind of fidelity that very big agencies were in the past only able to do, right? So it's democratizing weather modeling. And so that's just one example of where there's been very quick, rapid progress and a paradigm shift in terms of saying that, oh, now we can have AI as a reliable way to do weather modeling.

35:24I saw that, you know, the models that you mentioned, you made an insight that nobody else had and that people were able to devise other mechanisms to kind of follow behind you. But there was some sort of shift in thinking that was required here. And was it simply, we believe that there's enough structure in this data to learn and that people are just doing it wrong and people found other ways to learn the structure, but that your method was very...

35:56So let me clarify, right? So there is, you know, first of all, the very first work was to just say that, you know, look, traditionally this has been done with trying to solve partial differential equations, each time doing it again and again. Whereas AI learns from data, learns patterns and can be just as accurate, but fast. Yeah. And then the next iterations was to say, you know, how do we make it even more accurate? Right. And there's the aspect that, you know, there is the short-term weather, like what is predictable for the next two weeks.

36:30And then there's the long-term, you know, going to sub-seasonal to ultimately climate modeling. And traditionally what people did was to have different models for these different scenarios. So there's a different kind of system that works for short-term and other system works for long-term. But to me, there's only one Earth. You know, if you want a foundation model, if the claim is that it should be able to do both very short-term as well as very long-term together. And that's where in ForecastNet 3, the latest iteration of the model, we're able to do both.

37:03And that's because you also, you know, incorporate the spherical geometry of the Earth. And so with a lot of the other architectures that have been getting fairly good accuracies for the short-term weather, when you run them for like several months to even year, even before that, it just very quickly blows up, right? Because it assumes the world is a rectangle, which it isn't. And so incorporating all of the geometry and that information into neural operators means that we can faithfully run the same model also longer till and make this into a climate model.

37:44This is where the Allen AI Institute has now built climate models based on our neural operator architecture. And that's the only one that works as an AI emulator, right? None of the other architectures work for climate because climate requires us to assume the world is a globe and that if you are repeatedly rolling out, you kind of keep that information. Whereas if it's like a narrow surrogate, that's what I consider a weather model, you just narrowly look at a few metrics, many different architectures should do the job, right?

38:20But if you are asking one architecture to do a range of different tasks like a foundation model, that's where incorporating the geometry of the Earth, which is that it's a sphere, and using neural operators as an efficient way to do that enables us to accomplish that. Yeah, so ForecastNet, you said, trained on 50,000 data points.

Spherical geometry and climate simulation

38:43Can we just talk about what does this look like? What does the data input look like? What are you actually trying to predict from here? And then what is sort of the large scale, you said you're going from weather to climate. What does it look like to do that generalization? Because I could imagine if you have 50,000 data points, these are some high resolution in North America, then they might even depend on the local geography. Like, if you are always modeling, you know, Kansas, is this going to transfer to, you know, let's say the Swiss Alps or something?

39:16And then does that transfer to, you know, the Himalayas? So first of all, to clarify, we are training it on the global weather model, right? So we have all the information around the Earth, and then we are, you know, asking it to predict, like, given the current weather, like, say, wind conditions, humidity, and so on, what happens, you know, in an autoregressive way, and it's every six hours. So what happens in the next six hours, and so on, and you roll out, and you train the model to predict.

39:48And so, you know, you can have potentially the same model predict forever, right? But the predictability window is like the weather. And if you want to go beyond, you have to do what we call ensembles, meaning you have, like, a probabilistic estimate of what happens in several months to years, and that's how you get a climate model. Ah, okay. So you have, you actually have an ensemble of these local predictors, and then you use the, some sort of statistics on the ensemble to get a climate prediction.

40:20Yeah. So you kind of, like, have several rollouts. Essentially, you have several trajectories of rollouts, and then you get, you get probabilistic. And so the efficiency is actually pretty key here. Yes. And that's why, that's the biggest bottleneck with traditional climate modeling, that it's so expensive to do even one single run. You have to do long trajectories, very high resolution. And, you know, that's why we don't have a lot of very high resolution ability to do climate change predictions, for instance.

40:50So how do you validate climate? So you do lots of rollouts, and you look at sort of how things evolve kind of in aggregate. I mean, you said butterfly effect. Locally, I think we can believe even a perfect climate model, or perfect weather model would only give you maybe two weeks before it sort of, it becomes non-computable. So you do lots of rollouts, there's some chaos, you average all these things. How do you actually validate that this works over long enough times, time scales?

41:21Yeah. And, you know, it's a tricky question, right? So, for instance, you have to kind of ensure that you satisfy all of the physical constraints. And if you just do a standard rollout, that is, you know, likely not going to happen. And so some of the ongoing research we're doing is how do you kind of enforce the right physical constraints as you do the rollouts? Like, you don't want to, you know, completely wash out the fine details, because then it's not accurate. But on the other hand, if you keep them, you may be physically they're invalid.

41:54So this is still an open problem. And that's what makes it difficult, that you want AI to be fast, and you want to be able to do these long climate simulations, and at the same time be able to have full confidence in them. But these are things we are working on now. I don't know if this falls into weather or climate, but we recently just had one of the most extreme heat waves in the history of modern climate data, right, in the just right here in the kind of southwest United States.

42:25I'm wondering, did you, I don't know if you were involved in this regularly, but do you know if you or anyone actually modeled that or predicted that correctly? Yeah. Yeah. So we have, you know, I don't have information on this specific one, but we've tested in our latest forecast net three model extreme weather events of all kinds, right? And that's the key, like, you know, that's where you need probabilistic answers. So having just one deterministic output and saying that this is the weather is not enough when we are looking at extreme events.

42:57So we need careful probabilistic calibration, and we show that we are able to capture those well. And I think that was the surprise, even in our very first attempt, that we visualized certain hurricanes and storms, and it was able to do well, which is very surprising because you would think that rare events are not something AI would do well, right? It would do well on typical events. But I think this is where more broadly the lesson is the physical world may be more forgiving because,

43:28you know, where there are extreme events like hurricanes that have very specific physical signature, right? So it's like extreme, but in a very specific way. So maybe you don't need as many samples because the physical world has a lot of structure. And that's something we see this again and again, that there is a lot of structure in so many other examples, you know, I talk about plasma and fusion reactor. You know, we barely have a few thousand samples, but we are able to accurately predict events like disruption very well.

44:02And we are able to do that a million times faster than what traditional simulations were able to do. To me, yes, all these sound very surprising, but it's because I think the nature helps us a lot. You know, it has a lot of latent space structure that I don't think traditional numerical methods are able to uncover because they are focusing more on correctness, that in any scenario you should be able to solve these equations. On the other hand, with AI, it's learning from data,

44:34it's uncovering the structure, it's uncovering how easy or kind of tractable these problems are. And that's what we are seeing in many cases. Yeah, I've had a similar analogy. So from, you know, my own domain is probably closer to computational biology, but, you know, alpha fold is the obvious, like really exciting development in the community. So the solving protein structure prediction, and of course all the caveats of what was actually solved, I think we discussed this in a previous episode with the Boltz team.

45:06I encourage listeners to listen to that if they want more. But one of the, I think, key points about protein structure is that it really is constrained by physics. And that's why it was in some sense one of the few big wins in the field of biology, which is otherwise very complex. And we've had trouble, generally speaking, making a lot of success. And it seems like problems which are solved by differential equations or, I'd say, modeled well by differential equations have a lot more room for also integrating these techniques as a general form.

45:39I don't think there was a question there. Yeah, absolutely. So getting back to my question, we have climate or weather and climate. We have plasma, we have biology. And I know that you prepared for us a few visualizations. Can you just share with us, what does that look like? You know, so the visualizations, and then what is the thread that runs through here? And I think maybe the listeners will already have a hint about that,

46:11but I'd be really excited to see that. Yeah, I can certainly share some of them. I mean, this one is just kind of showing that, you know, we have world at different scales, right? And these are examples of phenomena happening at different scales, from atomic to protein to even planetary scales, like the weather we talked about. And, you know, we need to capture all of that. That's what neural operators are designed to do and to kind of, you know,

46:45feed in data at these different scales. And that's really like the aspect that makes a lot of the physical world problems hard, that need for fine scale. You know, we talked about how a lot of traditional computer vision video models are just designed to make things look visually good. And that requires low enough resolution. It's tractable, autoregressive. It's enough that, you know, it works out. It's short enough videos. But that's not how a lot of the physical simulation for, you know,

47:18industrial scale, high fidelity work. You really need high resolution. The, you know, atmosphere is one example. Like you, depending on the resolution you observe, different phenomena can be captured. So you just miss that out if you don't have that high resolution. And now the question is, with AI, can we do this much faster than what we could with traditional simulation? And this one with neural operators is kind of showing that, you know, if you use a standard neural network

47:50and you had a fixed number of pixels like you're seeing on this side and you zoom in, it gets blurry, right? That's the end of it. There's nothing beyond those fixed resolution that you can capture. But the idea is with neural operators, because it's a function space representation, meaning you can keep zooming in, you can add the relevant details either by giving it data at higher resolution or physical constraints at higher resolution, you can kind of bring that multi-scale phenomena together.

48:23So where are the physical constraints? I mean, I assume physical constraints here are like local simulating, you know, fluid equations or some sort of hydrodynamic? Yeah, could be, right? So it could be of any nature. The idea is now you can add like conservation laws, for instance, in an incompressible fluid. You can add like material deformation, like how things stretch. So, or it can be a full partial differential equation. So that's also an interesting question we've been researching, how is the curriculum of different physics?

48:56Like you mentioned, you know, some physics may be very hard to impose or add as a loss function. Others may be easier. So you also need to kind of, you know, think about what to impose, you know. And the intuition is here is that when I'm adding a physical constraint, I'm adding it to the loss function. Is that more or less what's happening? Yeah, because that's what is tractable, you know, making it a hard constraint is not tractable, whereas adding it as a loss function. And of course, there's still the balancing of that loss with the data we have.

49:27So we have to, you know, do that in the appropriate way. Right. Yeah. So as I mentioned, this is the example of the weather model, where here we are showing how we are able to capture like atmospheric rivers, which is the phenomena we see here in California. You know, resulting in life storms. I think we have one expected later this week. We'll hold you to that.

49:56So the idea of like, you know, why show this is this kind of global phenomena, right? These are like thousands of miles wide. So you really need non-local models that capture these very large span phenomena and do that accurately. And that's what our neural operators are able to do. So the training data for this is, you were talking about this a bit before, but I'm still like wondering, what is this, this is weather satellites? Or is there ground-based data?

50:26Is it some hybrid of the two? It's kind of a combination of different sources. So it's what we call re-analysis data. So this is historical weather data that is in a way re-analyzed, meaning that the sad rock satellite data is combined with essentially what the physics solvers tell you together, assimilated. And so this is made available by the weather agencies and we can train on them. So you're saying that they take a low-resolution data set,

50:57which is compiling all of the world's data set, you know, all of our meteorological data we have across the world. And then they do short time simulations using physics-based, you know, classical techniques to fill in the details. Yes. And you can do that over short time spans, but as you go farther, it breaks down very quickly. So you're amortizing that across all the, everyone would have to do that. And so somebody does it and then you are able to take advantage. Yeah.

51:27I mean, this is data, right, that's already prepared. But the idea is already this data assimilation with physics kind of makes our model physics informed implicitly. So it's able to kind of, you know, keep that information and that's why maybe that's one reason maybe it does well on even extreme weather events. Yeah. So this is just showing that we, our model is available in ECMWF, which is the weather agency, like, you know, the European weather agency.

51:58And so this was launched like more than two years ago, but, you know, I think fall 2023. So, you know, I think ECMWF making these AI-based weather models available to the public to me was a very big step because that's where, you know, everybody could see what's happening. There were several hurricanes, like, for instance, there was Hurricane Lee, and that's where the public could see what are these weather models doing. For instance, our forecast net was able to correctly predict that the hurricane

52:33making the landfall several days earlier compared to the standard weather forecasting models. And so the idea that these models could be very good for extreme weather events and do early prediction, you know, both for human lives, for economic costs is a very big deal. And so that's when the public kind of got a lot more, I think, buy-in from weather scientists because of how well it was doing in these events. And this is what I was talking about in an ensemble prediction,

53:07both for extreme weather or if you're thinking about climate, it's not just about looking at one trajectory, right? Because, you know, unless you're somebody with a Sharpie somehow saying where the hurricane is going to go. No, pun intended. But, you know, what you really want is the probabilistic prediction, meaning, you know, I'm going to try different adding noise levels to my initial condition, what the, when the weather, when the hurricane is forming in the Caribbean,

53:38I'm going to add some noise because anyway, it's noisy. I don't know truly what the measurement there is. And then I'm going to look at what happens to the possible hurricane tracks. And then I can come up with the probability of landfall in different regions. And that's how I can do risk assessment. And so this is where it gets even more expensive for traditional weather models, because you have to run all of these ensembles. And now AI weather models being so fast, tens of thousands of times faster,

54:09means we can now do very large ensembles. And this is a very big improvement in terms of what we can do for risk assessment. Have you gone through and done, let's say, looked over the historical hurricane maps and then tried to do ensemble predictions and calibrated how often your predictions are? Like a cone? Yeah. So in forecast net three paper, there are, you know, we have metrics of like extreme weather events and ensemble prediction. And in fact, we've trained the model to do good ensemble prediction.

54:42And so this is where the calibration matters for these kind of events. What was the sort of key insights or developments in forecast net three and versus two versus the first version? Yeah. So the first version was kind of the, you know, the using like the Fourier neural operators, but we didn't incorporate the spherical geometry, right? In this next version, we said, I think, you know, it's important that the world is a sphere, because first of all, otherwise it's distorted. So you kind of are not predicting it in a...

55:12Sorry, dumb question, but if it wasn't spherical, what did you do? Like a Mercator projection or something? Yeah, the standard, like kind of the, you know, like as all, and all the other weather models do the same, right? So they just kind of have the standard projection and then, you know, predict the weather. And which is okay for short term prediction. But when we, you know, our goal was to have the same model also do longer term. And that's when incorporating the spherical geometry added this additional stability, we could do longer rollouts.

55:42And then in forecast net three, the idea was, it's not just about deterministic prediction. We want to get ensemble predictions right. So we have to train them based on this objective, that we get the probabilistic predictions correct as well. How long are you predicting out? So, and how many rollouts are you doing? Yeah, so the, you know, rollout is how long you predict, right? So each step is six hours and then you predict for however long you want. You know, you just have to roll out.

56:13Sorry, how many examples in the ensemble do you have? So, and again, that's our choice. We can have like ensembles of different levels. So we, I think it's like a few tens or something like is what we are currently, you know, shown, but you can do much larger too. And that's adequate to get out how far? My intuition is the longer you want to predict, the more. So not necessarily. It's really like about, again, calibrating the ensembles and ensuring that they have the right spread rather than, you know.

56:47Okay, so you have tens of these models or examples in your ensemble and that even with a very, very long rollout, that's adequate. So again, like, you know, there's, as I said, a lot of still outstanding questions to do very, very long rollouts, right? Because you do need to incorporate the physical constraints in a way to ensure that. That's something that we are actively researching now. But these models that we have are able to do the longest rollouts

57:19compared to any of the other weather models that completely ignore spherical assumption and a range of other things. So when you say incorporate physical laws for climate over long times, I mean, what does that look like? Because there's a lot of local conservation, which may be just broken if you take an ensemble average, even though any given snapshot respects that. No, the idea is to make sure you look at each ensemble member and it's respecting the physics. Okay, okay. You're not assuming the ensemble,

57:51you're not deriving a, like, coarse-grained, like, equivalent of a probability or something? No, because then you would lose that, you know, resolution and the feather. Okay. That's a little confusing to me. Wait, so it's not an average. How are you combining the ensemble? No, you are doing the average, but you're predicting each one. Oh, you're predicting each one separately. Okay, okay. So each one is independently satisfies these constraints, but the ensemble does not. Yeah. And so that's how you ensure physical biology.

58:21Yeah. So when you go on a sphere, you operate in, do you use basically spherical harmonics or some sort of? Yes. Okay, yeah. You have a spherical basis for it. Exactly. Which is actually very natural with Fourier NetFlow. It's probably much harsher if you're doing other. Exactly. So that's where the, you know, like, the Fourier saves us. Yeah. You can incorporate these geometries well. And so I think be very faithful to, you know, what the globe is. Yeah, our EBS is much more natural than, like, a Mercator projection or whatever other.

58:53Yeah, which is, like, Greenland becomes huge. So that is a different story. But, yeah, but I think that this is where I think the aspect of, you know, more broadly incorporating more of geometry and information about the domain becomes a lot more important for the physical world, right? So this is me, again, emphasizing that we need to incorporate more of the structures because one is the data is limited. And the other is a lot of what we are asking is extrapolation, you know,

59:25to go beyond than what it's trained on. You know, we're just training it to predict the next six hours and maybe do a little bit of multi-step fine-tuning for autoregressive rollouts, right? So we're not training it to do very long, like, a climate simulation because that's just too expensive. But we hope magically it works well. And it cannot if you just say, I'm just going to put a standard transformer or whatever else there, and it won't work out. So we added more of the domain constraints like spherical geometry.

59:59We add in maybe more of the physics in certain ways. And that's where it becomes more interesting algorithmically as well. You know, there's more involved design here. So the time scale that you train on, how long is that? To predict for the next six hours. Oh, so only six hours. Yeah, and a little bit of multi-step fine-tuning. I see. Yeah, I said that. Okay, understood. So which is very surprising. Yeah, that is very surprising. I would have expected it was weeks or months. Yeah. No, and then it kind of just works well,

1:00:31even for like now we are showing for several months that it's able to do that. The number of steps is on the... Hundreds or thousands. Yeah, and that your Fourier basis, it's many times the sort of bass harmonic. I mean, this is like, that's in space, right? So we're talking rollout is autoregressive in time. No, but like in time, maybe I'm misunderstanding here, because it's in the Fourier domain, right?

1:01:01No, no, in time it's not. That's what I'm saying. It's autoregressive in time. Oh, oh, I understand. Okay. So it's space it is. Yeah, yeah. Interesting. Yeah. Okay. What's the angular resolution? At least in this scenario. In other cases, we also have... Time is also represented in the Fourier domain, and that's a question as well. Can we do that? But in this example, it's autoregressive. Got it. What's the angular resolution that you use for this in the spherical version? Yeah, so all of the data that's available

1:01:33is like, I think, a quarter, like 0.25 degrees. In terms of L, maybe? I mean, in terms of like a spherical harmonic frequency? So this is like, oh, you mean like how many modes we utilize? I think we... So for that resolution, we essentially utilize, I think, most of it, only a few of them we'll leave out. I forget the details, but... I'm just curious, like, what is the actual angular resolution on the globe that you are resolving to? Or maybe just the physical resolution?

1:02:05I mean, that's what... Because that's the data we get, right? So right now, the data we get is like a quarter, like 0.25 degrees. Oh, okay, like 0.25 solid angle. I think kind of comes out to like, I don't know, 700 by a few thousand, like, resolution. So, but this is already standard, like, kind of process. Yeah, yeah. And just trying to understand, like, how large is the basis do you need to represent that? Yeah, I mean, that's really depends on the resolution.

1:02:36And the idea is, you know, right now, our weather data is just limited by this resolution. But if you could, you know, you could, like, kind of do synthetic climate simulations of even higher resolution, right? And that's kind of the next thing on how to combine these together. Do you think you can predict with super resolution, basically resolution lower than the data provided? Again, like, yes, we can always predict them with the neural operators. But, you know, you do want to incorporate more of the physical constraints

1:03:07to ensure that they are valid. Okay, so can we talk about some of the other? Yes, yes, I know. It's a lot. So this is just showing, like, how, you know, what I described that on the left, where the world is being assumed, it's a rectangle, it blows up very quickly. And on the right, because we assumed the world was a sphere, it kept rolling it out and it kept being stable, so. I still see a little bit of a singularity there, is that right?

1:03:37And it's still like, you know, so the idea is, yes, because it's a very long rollout and we have no guardrails of physics. We are not, you know, kind of projecting it to the right physics, right? This is full extrapolation. But the idea is the sphere assumption stabilizes it to a much greater extent. Compared to the left, it's much better. Yeah. But if you're in the South Pole, you're still not going to get a good forecast. Exactly, so it's that poles are the hard part, so.

Modeling plasma evolution in fusion reactors

1:04:03So this is an example of the fusion reactor. So this is a tokamak. And we are able to model the complex plasma evolution and do this a million times faster than what we could do with traditional simulations. And this was, in a way, we're creating a digital twin of the plasma, right? And then we can, you know, do further things. Like right now, we are, as a next step, looking at, like, control, but with a full valid physics, like being able to prevent disruptions,

1:04:36ideally, and make fusion sustainable. So are you simulating MHD equations here? Or, sorry, Mnito-hydrodyneinous equations? Yes. Okay, yeah. And then so for context disruption in this case is this phenomenon which plagues plasma physicists where at some point your entire plasma collects in a little tiny beam and then shoots a strong, you know, right through your containment vessel and throws it up. And it can damage the reactor. And that's a big bottleneck because then you have to kind of shut it down before that happens.

1:05:07And then plus it's no longer possible to have a sustainable fusion. So there's a lot of open challenges here, but the idea is, you know, it's very expensive to go to physical experiments. The more you can capture that in the digital twin, but ensure physical validity, the more you can even do design and other considerations in the digital realm, you know, we can hopefully make advances. And these are the first steps towards that. The goal is that if you have one of these events

1:05:39that you can somehow adjust the magnetic field so that it contains that and stabilizes. Yes. And that's the next step we're doing now. We are looking at like designing both the control and the simulation together. Are you working with a specific lab? I'm just curious. So this one was with the UK Atomic Energy Agency. Okay. And now we're also working with a few others here in the U.S. as well. So we are, you know, kind of getting the information

1:06:10from many different approaches of fusion itself. So this is the tokamak. We're also working with stellarators. We are working with different... The stellarators are tricky. Yes. But the idea is ideally, you know, like our goal is to be able to design them in the digital twin. So can we come up with good designs that would make it maybe more practical? And so that's, I think, also a good thing as an AI person and much more like, you know, agnostic and not picking a winner beforehand, right?

1:06:42Like I like to work with different approaches, you know, and see whether AI can accelerate all of them and then we can kind of, you know, not prematurely rule out one approach over the other. So that's what AI enables us to be more kind of taking risks and exploring different approaches as opposed to in the physical world, trying to build any of these. You kind of have to cut a lot of the risky ones and say, I'm only going to do this because this is the most likely to work.

From theoretical foundations to applied science

1:07:14I noticed over your career, you started out spending a lot of time on, you know, really theoretical foundations, and mathematics of machine learning. And maybe, I don't know, something like six, eight years ago, you started working really, working a lot in applications and branching out in a diverse set of problems. What sort of prompted that shift in your approach and what you're looking at? I mean, so you're still working on very hard math problems as well. Like for example, the TorchLean work,

1:07:44but the applications have really grown. And I'm wondering what prompted that? And like, what were some of the lessons you've learned since then? Sure. I mean, to me, it's like, you know, I feel like I've grown along with AI, right? So when AI was, you know, in this where neural nets were not working because there wasn't enough data and all kinds of other reasons, you know, then you kind of have to build the theoretical foundations and try to hope that that leads you to a place where, you know,

1:08:16you get algorithms to work, right? And, you know, back then, like Tensor Methods was with that idea that, you know, pre-deep learning, we still want a structure. We have probabilistic models like late and Dirichlet allocation for topic modeling. And solving those were hard, but now Tensor Methods gave us a way to be very practical. It's parallel and can be done at large scale, but still has nice theoretical basis. So that was where, you know, we're starting off. And then as deep learning started taking off

1:08:48and we could see that it works well in practice. And yes, there is a little bit of maybe theoretical understanding, but not a whole lot because of the way, how complex it is. To me, theory should not be a constraint, right? It should be an enabler. And so that's where a lot of like the exploration was to make this work well in practice and over to Amazon Web Services, then NVIDIA. So really like making things work at scale and really kind of getting hands dirty, right?

1:09:19It was kind of like where a lot of the development is. And now I see a full circle because a lot of purely data-driven approaches in a way seeing saturation, right? So now we want to ask, okay, either make them more hardware efficient, right? There's a lot of now room to kind of say, can we now, you know, make them much more energy efficient or hardware efficient? So that's one aspect. But the other is areas like this where in the physical world, we don't have enough data.

1:09:50We are asking for hard extrapolation. You know, we want to think of doing discovery. By nature, it's about extrapolation. So we will never have data about a new discovery, right? That's my definition. And so there we need to, again, go back to thinking in principled ways. Whether it's architecture design, algorithm design, the right loss functions. So we need to be much more mindful. So I see that coming a full circle because all of the things that work with deep learning,

1:10:21let's take them, but make them a bit more principled. There are several other applications which seem very natural. I'm wondering if you forked on these, or did I just miss some papers? If I did, I'm sorry. So some examples are design of like electromagnetic circuits, I think is a big one, or maybe not a big one, but I think we'll be coming up in the near future. Design of, let's say, materials, design of, let's say, dissipation and heat sinks, or any sort of like fluid flow.

1:10:52I guess I'm going through what differential equations do I know? Electromagnetism, I know diffusion equations, you know, MHD, yeah. I'm wondering some of the other domains that are...

Inverse design and future foundation models

1:11:04Yeah, I mean, to me, there is just endless possibility, right? So there, you know, as like you can just have this work on any data, and we have several other examples. So this was like, you know, being able to ask, can we sequester carbon dioxide underground and model how carbon dioxide expands or, you know, what is the pressure buildup in these reservoirs? And, you know, can we kind of model how they migrate over several decades?

1:11:36And so this one, we were able to do much faster than what traditional simulations could do. I mean, the other aspect is being able to do all kinds of geometric shapes, like, you know, being able to model aerodynamics in cars, planes, and so on. And so, again, this is a nice example of a latent space because you can transform a car or any other shape to a donut and then model on the donut and then transform the donut back to the car. You didn't turn it into a coffee cup. Is that the classic joke?

1:12:07Your donut and your coffee cup. Exactly. So the idea of like a latent space to handle all kinds of different geometries and be able to capture the physics there in the latent space well means we can now have a model that generalizes across a lot of different geometries. And my understanding that maybe the larger vision here is that you can train a foundation model in the sense of being able to model many different physical phenomena

1:12:38with the same model. And so you may fine tune or there may be some kind of prompt that you give it to have it understand the particular geometry, but that you, you know, sort of on all these different physical problems you train and then, and then you have your particular one and you're able to model that very effectively. Yeah. I mean, that's really the future, right? Because we have foundation models for language, maybe vision, but not for physics. So, you know, the idea is instead of like right now what we've seen are narrow surrogates

1:13:09and we're trying to broaden their scope more and more, but ideally we have much broader models that can work on a range of phenomena, but also multi-physics. So not just have like one single physics, but coupled physics. The real world has all of the physics kind of coming together in coupled ways. So can we bring all that together? So that's one aspect, like, you know, have foundation models that can do design, that can do simulation. But the other aspect that's really interesting is the inverse problem, right?

1:13:40So can I now not just simulate, but ask what is the best design? And then these kinds of like models can like do simulation, but you can even do that implicitly and come up with the best design rather than in the earlier era of it was humans trying to come up with design, then you go and try to simulate or go to the wind tunnel, whatever physical testing and validate that. But now you have AI come up with optimized designs, but you have the guardrails of physics.

1:14:11So you have models that are accurate in physics. You have the confidence they work well. So you're kind of able to do that as well in the same model. Is there a reason, have you seen any evidence that you talked about these like sort of multi-physics being able to transfer or that you may be able to generalize to sort of unseen physics? So, I mean, you know, like the physics by nature, if it's completely unseen,

1:14:42it's not possible to transfer, right? I mean, if you're saying that we're going beyond the standard model, there's absolutely no data. That's not possible. But if you're asking about like, you know, for instance, like, you know, there is the, like, say, I've like, you know, shown it examples of like just how the heat propagates and the other examples of how the material like stretches. And now there's coupling, like because of heat, there's also stretching or kind of the joint phenomena.

1:15:14You could like now hope to fine tune with much fewer samples because it kind of individually knows this phenomena. Then combining them together, maybe it can't do it from scratch because that's still too much to ask. It's highly nonlinear and coupled, but it can do it with fewer examples. And we've seen evidence of that in a lot of our papers that you're able to kind of essentially build up a curriculum. And that's what we see again and again in many of these examples that, you know, the real world,

1:15:44we can kind of control a lot of curriculum and say, you know, let's kind of build in like modules and put them together. And that's what it now allows us to do in a systematic way here. I guess the design aspect, I don't know if we wanted to show very quickly. So this one was like, you know, looking at like designing the mask for inverse lithography, meaning now this is an inverse design problem. And we are also able to do that for designing gates

1:16:16and quantum dots. This is like nonlinear photonics. And all of this, what is common is the idea that, you know, there's a forward model that is simulating the physics. But now what we want is the inverse design, like the problem that we can optimize the best design. And humans are usually not good at this, right? We are not good at like looking at highly nonlinear phenomena and say, oh, somehow maybe this combination of all

1:16:47these gates coming together helps pull the electrons together in a quantum gate. And so our collaborators were struggling to do that manually. And with AI, we are now able to come up with very efficient designs, but also those we know actually work because we have already the simulation as part of the loop saying that they work well. So I think these are examples where we see that it's not just about simulation, it's about really novel designs

1:17:17and novel discoveries that enable us to move the needle of innovation itself. Each one of these examples takes a lot of domain knowledge. How could somebody take your basic research if a domain expert and quickly get started applying neural operators and the other frameworks that you've developed to their problem? Yeah. You know, neural operators are an open source library. It's extensively already adopted. It's part of the PyTorch ecosystem.

1:17:49It's, you know, used by a number of not only researchers, but also in companies. We have a lot of documentation there. So I encourage people to go there. We have like, you know, many different architectures, examples, recipes. So I think that's a great place to get started. You recently joined the UN Scientific Advisory Board. I know we're running out of time, but maybe just can you quickly give a bit of the story behind this and what you hope to accomplish? Yeah. You know, I'm really honored to be part of that advisory board

1:18:20for the UN and in these tricky times with a lot of geopolitics there, you know, which again, I'm not the expert on that. But when it comes to, you know, aspects, especially related to AI, having scientists in the room is something that, you know, I think is very important. I hope I could have an unbiased view and try to provide scientific evidence for any aspect, right? We want to think about how AI impacts globally, like, you know, how do we ensure the benefits of AI

1:18:50reach everybody? How do we democratize access to AI? How do we ensure the unintended consequences and harmful impacts can be controlled? I think these are just the beginning aspects. Of course, the other side when it comes to weather models, I'm already excited, like, you know, there is a push to seeing how we can have better weather and climate modeling. So then our food, you know, like using weather for better agriculture. So all these aspects are also where UN

1:19:21has a lot of dedicated agencies and people on the ground across the world. So I'm looking forward to contributing and being part of this.

UN scientific advisory board and compute bottlenecks

1:19:31You know, looking, you know, looking at your career and how you, you know, talk and what you work on, it seems like you very much are a person who likes to solve concrete problems. You don't like to philosophize about things, which, and you're also seem, I think, maybe more optimistic than a lot of people in the AI space. You have a very hope, like hopeful view of the world. I think not, it's not honestly always true. What are the ways that you can uniquely bring that viewpoint to the board versus maybe some, you know,

1:20:01you know, thank you. I, you know, to me, I think, as I said, I tried to be unbiased and as a scientist and as a scientist, I think that there's a lot of beneficial aspects of AI that are sometimes missed when we think of only the harmful impacts, right, and especially that is with respect to AI for science because a lot of regulatory frameworks equate AI with language models and, yes, language models can, you know, manipulate people,

1:20:32can have all these kinds of harmful impacts that we should think about controlling. But AI for science is different. So I think this one-size-fits-all is where a lot of problems come up. So we have to be mindful that there is, you know, AI that can change the world with new discoveries and we should enable people around the world to not only benefit from them, but also be able to do research, you know, have access to AI that they can go

1:21:02innovate and use them in interesting ways. One question that we have been trying to ask every guest is if you could pick a bottleneck in your domain that you could magically remove, what would that be and why? More compute. You know, that's the, I know that's an easy one, maybe a lazy one, right? Because, you know, and you know, of course our compute that we have is growing so much more than even a few years ago,

1:21:33thanks to NVIDIA, thanks to others. You may be placed there. So no, no comments on that. But what I mean by that is also like for research enabling more and more compute, you know, it's very important. I know there are national labs building more supercomputers, you know, hoping that we can have more compute for research. But I think, you know, without that, we cannot experiment, we cannot innovate. I think this is a part that I push a lot and, you know, I think I cannot

1:22:03emphasize that it's so critical. If you had a call to action or something that you would like people to do or think about or learn about, what would that be? Yeah, so, you know, you can go to neural operator libraries so you can kind of hands-on, play with different architectures, recipes, you know, look at use cases, but also think about like, you know, AI for science as not just language models and agents. Yes,

1:22:34that's one aspect of it, but ultimately, you know, those are still like external wrappers in a way, right? Until we have AI that fully understands the physical world, not just as symbols, but as one that can simulate and design and control based on that, you know, there's a big piece missing. So, that's the other aspect that I think that people should really think about AI for the physical world in this way. Anima, this has been so fascinating.

1:23:05I'm excited to check out neural operators myself. I have some ideas in my head already. I really appreciate you taking the time to sit down with us. Thank you, Arjit. Thank you, Brandon. I really enjoyed it. We really dug deep into a number of things, so I appreciate you doing that. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you. Thank you.