Anima Anandkumar: How do we build AI to push the frontiers of scientific discovery?

[00:00:00]Thank you. Uh, thank you for having me here. Thank you, Chuck. I mean, we first met when you came to the Caltech AI and Science Conference, and I'm really glad that you've grown this theme into so many locations across California, and I'm looking forward to seeing more of that. So, >> that was a fantastic conference.

[00:00:20]>> Thank you. So more to AI for science or what I like to call AI plus science because it's really this deep integration that you know I personally have been working on for a number of years and now it's really gone mainstream. So I'm really thrilled to see that uh you know when I started at Caltech almost a decade ago there was this idea that AI could be helping in sciences although back then it wasn't very clear how that would be because deep learning would require so much data. So you know what is the real impact and now here we are a decade later asking questions about AI pushing the frontiers of scientific discovery and indeed with the language modeling era a lot of that has been framed in the lens of AI scientists right like AI coming up with new ideas maybe AI replacing scientists but I think I'm you know less interested or even skeptical that they would really be replacing the scientists. But AI not only is about collaborating with scientists or replacing their ideas, but really asking whether AI could solve the critical bottlenecks we have in scientific discovery. Because if you think about lots of areas of sciences, you know, we were talking about physics, we were talking about material discovery. What is the critical bottleneck?

[00:01:46]Anyone?

[00:01:49]What do you think takes the longest time, takes the most resources?

[00:01:56]>> Data.

[00:01:57]>> Okay. But what does even data mean? I mean, data collection is to train AI, right? But is that really the goal if you want to discover a new material?

[00:02:06]So either is H.

[00:02:09]Okay. [laughter] Let's say I have infinite money.

[00:02:15]So what do you think is the Yeah.

[00:02:28]truly vision.

[00:02:29]>> Yeah. So, right. So, that is like adding better insights to AI. But what really is the critical bottleneck?

[00:02:38]>> Yes. Thank you. Because you can come up with lots and lots of ideas, but how many of them do really pan out? And we've seen this with AI, you know, that's why a lot of especially experimentalists are very skeptical, right? Like for instance, there was a deep mind paper talking about 100,000 new materials. And then after that, I think scientists from Berkeley lab pointed out, you know, many of them were, you know, you could kind of an expert could quickly rule out and say, oh, no, no, they can't even be synthesized in the lab. But there were many others that was more nuanced and so you couldn't possibly go test all of them out. And so that's one of the dangers of AI being used for scientific discovery because if you now multiply lots of bad ideas that's like instead of finding a needle in the hay stack, you get a bigger hay, right? So that's almost counterproductive. And so this is where we should be careful in kind of you know what I believe is the overhyped field of hypothesis generation where just lots of hypothesis being generated for their own sake is not that useful.

[00:03:50]But on the other hand, what's useful is if those are verified, you know, if instead of experiments, I could like guarantee to some extent, of course, there's always could be approximations, but at least some level of confidence that they would work out in the lab. So those are the ideas that would be more highly priced. And in areas like in math, that's where we'll I'll talk about lean. You know, in fact, we did the first open-source framework combining lean and language models. So you're not just taking language models at their face value but lean will be verifying those mathematical claims but in other areas like materials or in physics you know it's not just about symbolic verification right we may need to do experiments and instead can we do simulations and if even though simulations are very expensive can we now have AI that understands and simulates the physical world at different levels and If it can be done much much faster than traditional numerical simulations, then we have the ability to not only speed up the scientific discovery process because AI is differentiable, we can also now do inverse problems like inverse design.

[00:05:07]And so we want to think about AI that can understand different scales. could be from the atomic scale, you know, understanding m bigger structures like proteins for biological modeling all the way to even planetary scales if you think about modeling the weather and climate of this planet. And so depending on the problem of interest, you could be spanning multiple scales. For instance, we did modeling of the aerosol in corona virus where at the molecular level we were thinking of how calcium ions bind to the proteins in the corona virus but then the macroscale behavior of how the transmission happens across individuals.

[00:05:48]This was in the early part of the pandemic. And so that's a multiscale modeling problem. And same with like weather modeling. you want to like be able to capture the details of fluid dynamics and that requires usually high resolution simulations especially with traditional numerical methods. So could AI help us speed up these kinds of simulations and what we see is we are able to do that many many orders of magnitude faster. In the case of weather modeling, we have models that's tens of thousands of times faster. In case of plasma simulation in nuclear fusion, more than a million times faster. So, you're getting this immense speed up because AI learns not only the phenomena at a certain one particular scale or resolution, we have these frameworks called neural operators that learn across multiple resolutions.

[00:06:48]And so that's what I will be talking about today. And of course, one challenge is we may not have all the data to begin with, but we have mathematical models. We have equations.

[00:06:59]And now can we make our AI not only depend on data but also incorporate these knowledge of these equations, right? Such as partial differential equations. So we'll call that physics informed AI. So AI doesn't just rely on the existing data. It can also incorporate either constraints from physics, symmetries and other kinds of domain properties that are very different from the bitter lesson that you see with language models and standard image models because here if you think about scientific discovery, it is an extrapolation problem, right? So you really want to go beyond what is known. So you likely don't have the training data for that because then it's not a discovery problem. So if you have to go beyond then you cannot just train a model narrowly on the data that's already available. And so where will that additional knowledge come from? It has to come from other sources like these mathematical equations like PTEES.

[00:08:02]And so as I've mentioned the neural operators can incorporate both. They can incorporate data at multiple resolutions as well as physics constraints at multiple resolutions. So to kind of conceptualize the idea is here I'm showing like the uh simulation of the clouds if in a standard neural network so neural network where the inputs and outputs are at a fixed resolution or a fixed size you zoom in it because it's a fixed set of pixels it gets blurry right there's no further detail than the resolution on which you trained your model on the other hand neural operators don't have this limitation The inputs and outputs are both modeled as functions. So you can keep discretizing this function to any resolution. So you can keep zooming into finer details either by providing more data like for instance climate simulations we have data at different resolutions available or in like DFT calculations for material modeling we have different levels of theory. So more and more finer refined uh simulations you get you can train the same model with it and you can also incorporate physics constraints. So what just the data doesn't tell you for instance satellite data is in a course resolution you can add in physics constraints of the relevant kind could be energy conservation other kinds of constraints but make them at a finer resolution than the data you have. So in that case you even exceed the accuracy that just the data by itself gives you because physics is giving you more information than purely the training data available. And so that's the concept that we can do all of this and capture multiscale phenomena that wouldn't be possible with neural networks that were just fixed to one course resolution.

[00:09:59]So let me show you some success stories over the last many years that we've been applying neural operators. Perhaps the most uh famous one is weather modeling.

[00:10:09]We created forecast net the first AI based high resolution weather model almost five years ago and it's tens of thousands of times faster. We were the first to also permissively open source the model. This was when I led the team at uh Nvidia back then and it's under the Nvidia arts 2 uh platform and it's been used by not only lots of weather agencies but also startups, big companies and so really like that changed the notion that you needed traditional weather forecasts run on big CPUbased supercomputers, right? because instead of that you could now miniatureize and run this on your local gaming PC with a consumer grade GPU and run like two week forecasts in under a minute. So this is a massive speed up of tens of thousands of times and of course the you know idea was you know will weather scientists really trust this? So how do you kind of gain their trust? And so we worked with ECMWF which is the premier global forecasting agency from Europe. Uh so they hosted our model. Uh and they were able to like make it really widely accessible to weather scientists. And that's where you could see that daytoday, not only these forecasts were lining up with what the traditional weather models did, but to everyone's surprise, when it came to extreme weather events like hurricanes, it even exceeded what traditional weather models could predict. For instance, in this case, it's a bit dark, but you can see that the one in orange, the forecast net, hour, our model was able to predict landfall several days earlier correctly compared to the traditional forecasts. So soon it became clear that AI wasn't just a cheap substitute for traditional forecast, but can exceed now both in accuracy as well as being much much faster. Right? So this is not a trade-off. This is a clear win. And so why does AI get all these benefits here? uh the speed is clear because we are not only moving from CPU to GPU but more importantly we are able to subsample in a way when we do those uh calculations in um traditional PTEES you need very very fine discretizations whereas in a neural net it's just a few layers in our neural operator which is a kind of neural net architecture we have these few layers that can give to very fast inference And but the other aspect is okay maybe you got lucky on this one hurricane you know how do you kind of systematically evaluate success on rare events like hurricanes. So the correct way to do this is when it comes to rare events is to really ask do I have a calibrated model right do can I correctly predict the tail risk because that's what we mostly care about you know we are in California it's sunny every day it's pretty easy for any model to predict that. So that's not the interesting piece. The interesting piece is do we get extreme weather events correctly and when it comes to this we want to think about creating larger and larger ensembles. So in this example of uh hurricane you can see where it originates in um the Caribbean you want to perturb it with appropriate noise distribution. So you have a spread and then you want to see where the different tracks go and then you can kind of count within a region how many of those tracks hit that and then you can assess what is the risk of making landfall in this specific region you know this specific local government can then go make policy decisions that are really critical for saving human lives. And so how do we ensure that we can correctly calibrate especially to get the tail distribution correct and this is difficult right I mean most of the time the typical events trying to get the probability distribution correct is what people hope for but can we have AI models not only be faster and accurate on typical events but also be getting the right uh tail distribution which really matters for weather forecasting and this is also a great example of working closely and talking to the domain experts because what do they really care about? The metrics they care about may not be the typical machine learning metrics and in this case the focus is been on these tail events. And what we did was to take forecastet and then do the calibration with a much larger ensemble in this case thousands of ensemble members and then be able to show that yes it does get the tail events correctly and why can't we just do the same with our current traditional weather forecast I was very surprised that operational weather forecast so the one that you see every day as official forecast uses merely about 50 ensemble members.

[00:15:26]So if you think about statistics of just 50 uh you know quantities then we know that you cannot get the tail events that accurately right. So here the speed of AI is not just that I save more energy I'm cheaper it actually can now lead to better tail risk characterization because I can do much much larger ensembles than what was possible before.

[00:15:52]So this is AI pushing the frontier in weather forecasting because of the ability to do much larger ensembles. So in this uh figure here you see that this was the Berkeley lab researchers who took our forecast net and showed that it can correctly capture like the 95th percentile heatwave and of course we was just came from Europe where there is a really massive heat wave. So these kind of like extreme events being able to capture that accurately is is very important.

[00:16:25]And then now in the latest uh version of forecast net this is the version three we said let's now train with this new loss function. So what is called CRPS is a way to capture the full distribution and of course we want to make sure it's also on the sphere because the earth is a sphere. So we are now calibrating and training the model to match the probability. So as I said this is what the weather scientists care about and we are matching uh the training to satisfy those needs. And here you see now being doing the full ensemble roll out. So you can now not only get one answer, you can get lots of different answers and you can get probabilities as a result of that and doing that very very fast. So doing this in under a minute and there is a systematic evaluation of all kinds of metrics and this is the other message I was saying from the beginning when it comes to AI for science it's not about one metric right there isn't just like oh perplexity I'm done and of course with language models too there is entire suites of uh evaluations and similarly when it comes to weather forecasting there is lots of different metrics and you want to kind of be good on all of them. And that's what we show that we're kind of on the parto front and state-of-the-art in all these different measures. And so that's what is like of what a foundation model should be doing like being good at multiple things. And then you can make it a more trustworthy weather model because when it runs in practice, when it comes to you know determining impact on human lives, you really want to be thinking about all these different metrics.

[00:18:09]So what else helps? So in as we said we incorporate the fact that the earth is a sphere. Seems very simple. We're not flat earthers. But every other weather model decided not to. They said it does not matter. I'm going to pretend the world is a rectangle. Use a standard transformer or a graph model whatever that may be. So why does it matter that we should incorporate those symmetries?

[00:18:36]Uh honestly it doesn't matter if you just run the weather forecast for a few days, right? Because it just like kind of any effects of winds is only moving a little bit further. So you might find some difference but it's not so clear.

[00:18:49]On the other hand, if you roll out for long like I did here and this model here uh you know assumed the world was a rectangle. This one didn't and this is Antarctica. So what happened you saw earlier was that the distortions start building up on the poles and so uh as you run more and more steps auto reggressively the errors build up and once you assume it's a rectangle it just goes haywire right because the world is not a rectangle and so adding symmetries matters more when it comes to extrapolation because the bitter lesson says oh ignore all that scale solves all the problem But scale cannot solve an extrapolation problem. The minute you go beyond what the training data has told you right to a bigger extent to you go significantly towards extrapolation scale is not enough. And that's where aspects like domain symmetry in this case spherical geometry was really helpful. And here you see being able to do like a year-ong roll out while having only trained to predict like the next six hours and a little bit of fine-tuning over multiple steps. So it was never trained to do a whole year roll out but it's able to keep that stability and so this kind of extrapolation is where uh symmetries helps a lot and when it comes to predicting climate which is what the Allen AI institute built on our forecast net architecture you want to do long rollouts and you need architecture like one that uses sperical constraints as well as adding other energy conservation laws and also So training on other conditions like different carbon dioxide forcings to really be able to talk about not just the present weather but possibly uh future changes in climate due to different carbon dioxide.

[00:20:46]And so far you know it showed you an example with AI replacing traditional uh forward modeling you know what we call like the modeling of what happens to weather in next time step and so on. And so instead of solving partial differential equations from ground up like what traditional forecasting did, we used AI instead to learn to forecast uh that. So that's called forward modeling. But in so many areas, what you have is an inverse problem, right? And in the case of design, the inverse problem is what is the best design that satisfies certain constraints, does the job in the best way. And so there's some re really great examples of this uh this is uh work um I did while at Nvidia where I showed that you could uh improve the mass design significantly for lithography. As we know, lithography is one of the critical bottlenecks in chip manufacturing. And uh in order to kind of uh in account for the distortions that happen, can we now optimize what the mask is uh that kind of is aware of what the lithography is going to do for it? And so you can because AI models are differentiable, the forward model of lithography can now be used to do the optimized design of masks. In the most recent uh iteration of this what we did was to extend this kind of notion for designing gates in quantum dots. I don't know where Hartmouth is but uh oh yeah there you go. [laughter] So you know this is an example of where for quantum dots it's much more sensitive even compared to classical chips to be able to do this kind of design right like to pull electrons what is the best uh uh configuration of gates more nonlinear more challenging and we now have a chip that is uh fabricated we are in the process of testing we're very excited to see like what how much u step up that AI can do compared to how humans can design chips when it's so highly nonlinear and non-intuitive.

[00:22:57]Another such example is in nonlinear photonics. Again, humans are, I think, fairly good at designing cases where the phenomena are linear, but when there are high levels of nonlinearity, you know, we aren't that good in coming up with the right intuions. Another example is the medical catheter where uh we wanted to reduce bacteria contamination. So we designed these triangular grooves and you see under the microscope of what our AI design catheter is doing. It is pushing out all these bacterial particles and none of them are really able to swim upstream into the human body. Right? So this is under the microscope. The idea is there's barely like one particle here trying to swim upstream. But we measured a 100fold reduction in bacterial contamination as a result of AI helping optimize the design. And so here again fluid flow is not an intuitive thing for humans to design but AI can really do a great job at this. So all these are examples of AI both accelerating forward modeling as I showed in the weather forecasting but also being able to come up with optimized inverse designs.

[00:24:13]And in terms of inverse problems, another great example is medical imaging. Being able to really uh you know tell us what's inside the human body and there's really no ground truth, right? Uh so you have to be kind of coming up with the right uncertainty measures as well. And in many cases we not only do better speedups but we can also learn the prior from data itself which is another benefit that standard numerical solvers are not able to do. So you can really be learning prior from data and do base and inverse and coming up with good posteriors for uh you know also modeling the uncertainty effectively.

[00:24:54]Uh so to give you an idea of how AI is not just about speeding up medical imaging but could be even coming up with better diagnostics. Uh in the case of lung ultrasound uh the problem is quite challenging because in the lung you're trying to measure the aeration level. So you're not looking for solid lines.

[00:25:13]You're looking for what are these called as bel lines where it's very blurry and you're trying to count these blurry lines to say how much of the lung is irritated. Right? So when the lung is more diseased, it gets more challenging to come up with the right level of aation. And instead we said, let's just reconstruct the lung tissue and predict which levels where it's airrated and where it's not. And so because we're doing the forward modeling as a wave propagation and we first train on synthetic data that is rightly calibrated to lung tissue that is dissected. So you can really train on large amounts of data further validate it on dissected lung tissue and then we do clinical study on humans. So you can see that full pipeline of going from sim to real and also being able to validate uh this kind of aeration maps being more accurate very clearly interpretable compared to the earlier practice which was clearly errorprone subjective and not really giving the right diagnostics.

[00:26:19]So as I mentioned, you know, I've talked about neural operators as the one capable of capturing many of these multiscale phenomena and being able to solve problems for the physical world.

[00:26:31]To very quickly uh remind you what that is, it's uh being able to model both the inputs and outputs as functions and the operator is the mapping between function spaces. So that's why we call it neural operator. And so you can now discretize the inputs and outputs at any resolution. Right? So you can query anywhere. You can have your training data or inputs be at any resolution.

[00:26:57]And the guarantee we have is a notion of consistency. So as you refine your mesh, so as you refine your input to a higher level of discretization, you get outputs that are consistent. Whereas this is not true if you took a standard convolutional neural net or a transformer, right? You don't have this nice consistency. As you add more and more details, it kind of converges to a consistent solution. But you need that for the physical world. You know, weather model, you have more satellites.

[00:27:27]You should have like solution that ultimately converges to the weather that kind of captures the continuous uh information around the planet. And so that kind of operator uh convergence is what we have with neural operators. Of course, we have extensive documentation and a library where I encourage you to go try out the code, contribute. It's all permissively uh licensed. So, you know, everybody is welcome to do that and already has extensive usage. So, in the last few minutes, um I know there was a question asked uh you know whether we should just do coarse graining with AI. let your numerical solver be can we just use AI to refine it and uh there's the other option where you throw away your numerical solver and let AI solve everything which one do you think requires less training data so the option where I have like a numerical solver I still keep it I only train my closure model with AI or course training with AI or I throw throw away my numerical solver and train end to end the model like a neural operator. Which one do you think requires less training data?

[00:28:50]How many of you think the closure model setup keeping the course grid numerical solver requires less training data?

[00:29:02]Okay. How many of you think the end toend learning requires less training data?

[00:29:09]Okay. Yeah, it's pretty pretty good mix.

[00:29:13]So, of course, you know, I kind of cheated because I didn't tell you what exactly is the kind of training data I'm going to give you, right? So, what is the kind of data this closure model would require? The closure model would require the fine grid solver data, right? In the case of fluid dynamics, it would be like this uh direct numerical solver like a fully resolved simulation.

[00:29:36]So, and it would require lots of it because the coarse grid cannot help you with this. You need the fine grid data.

[00:29:42]Whereas with a neural operator, I can make use of both. I can make use of data from my coarse grid solver and then I can fine-tune it on the high resolution data for my fine grid solver. So I can mix data from my different solvers at different resolutions.

[00:29:58]And if you only count the very expensive fine grid solver, right, the most of the training cost will be from generating data from the expensive fine grid solver. The answer is this end-to-end training with the neural operator does require less training data.

[00:30:16]And so so that's the key that you know if we are smart about how we use the training data effectively as a curriculum you don't need a lot of data of the very expensive kind. So you can go very far with the coarse grid data and especially with the neural operator learning how to mix across scales. It even learns what is missing in the coarse grid solver. Right? What are the mistakes made in the coarse grid solver and only learn the corrections to that implicitly within its layers. And we've shown also with uh our experiments that you require much less data and also be much much better in your error. And so this was an you know for many people not so intuitive that uh you know because there's a lot of notion and lot of already work using AI just to kind of augment existing solvers and I believe that's not the right approach in the long run. You may get some quick gains, but in the long run, it's better to replace AI with all the existing solvers because there'll be ultimately no real separation, right? Like they're all algorithmic ways to solve partial differential equations and you could use some amount of learning. You can use some amount of domain knowledge and fixed ways to solve them. And I think we'll see that algorithmic integration and marriage of algorithms rather than putting AI in one bucket that is separate from numerical. solvers. So that's where I think the future holds.

[00:31:44]So I will stop here. I wanted maybe a very quick, you know, teaser on verification. I had uh lots of slides on this, but I won't get into them other than to say the other piece that I didn't tell you was on being able to verify um aspects, right? In this case, verify math. And we had like lean dojo as the first open-source framework that combined lean and language models a few years ago and now the rest is history.

[00:32:12]You know, thanks to Terry, we lean is really a first citizen in math and mathematical reasoning. So we really are pushing forth a lot on lean. uh the framework that I want to advertise and I'll leave you with this slide is torch lean which is the latest iteration of using lean for verification but in this case it's for verifying neural networks themselves so think about now you've written a program for a neural net how can I quickly prove properties about it you know was it trained correctly did it have the right level of quantization what is the effect of quantization what is the robustness guarantee that this possesses and so you know especially these properties are important for instance if I use AI to replace my numerical solver right I may need a lot of good guarantees that the solver had and that's just one application many others like that so you can get certificates you can get formal checks so you can get all of this and this is available now you know in the lean dojo website so it's available it's uh you can contribute to this and the idea is we you think of like lean as not just one to verify but also to execute. So you write the program in lean where you define neural net layers. You can run it just as you do pytorch. You know we called it torch lean. It's not invoking pytorch but it has a similar kind of you know language to kind of easily write neural network layers and execute. But in the background the same intermediate representation can now be used for verification. So you get the full uh kind of like you know framework to do all the lean kind of tasks on the same neural net. And so this combination is what I believe will really bridge what we've done with mathematical reasoning with lean as well as neural network analysis uh using neural nets as uh tools for uh solving problems like partial differential equations. Thank you.