Bay.Area.AI: AI for Science with Neural Operators, Jean Kossaifi
Recording: Bay.Area.AI: AI for Science with Neural Operators, Jean Kossaifi
hi everyone so yeah we had a lot of great talks today I think it was really good combination of like this covering uh various aspect of efficiency like what enabled first the recent advancements of deep learning was kind of data hardware and algorithms coming together and so yeah we heard a lot about the hardware and so I kind of want to talk about the AI part so yeah as Alexi mentioned um I'm a senior research scientist at Nvidia and I just want to present some of the works we've been doing on applying AI for actual problems in science so just to kind of get a sense of the audience did anyone hear of neural operators before cool oh nice did anyone try neural operators in any project or cool so I guess hopefully everyone can get something from the talk so yeah essentially the starting point is that a lot of phenomenons that we're interested in are modeled by Continuum equations known or not for example um Navia STS equations which describe the motion of viscous liquids um and the the solution of these equations are functions so we typically for example have the state of the physical system at a time T and we want to describe how it will look like at time t+ one and so yeah we may have discrete observations to learn from but really we want to learn what happens at the Continuum and so as a motivating example I kind of want to take the example of uh weather forecasting where we may we have the state uh how the weather looks like globally right now and we want to look at how it will be in an hour 6 hours um maybe a week from now and the key thing is that often it's not sufficient to just do a prediction we also want to do probability so if you look at weather for example as you I'm sure all well know weather forecasting is not super accurate and so what we want to look at is embling like what is the probability of a certain event happening and this is especially true if you look at extreme events like a hurricane and you want to know what will be the trajectory of this hurricane and so on and so forth and it's also a really interesting problem for us because we were just talking about distributed uh training well this is really relevant here because the scale is huge so um weather happens at multiple scales and so for example if you want to resolve things like convection you need to look at maybe a 25 kilomet scale um if you want to resolve storms accurately you need to have a much more detailed state so that would be maybe 1 kilometer and then if you want to resolve clouds and have more accurate prediction you want to go to 100 m scale ideally eventually we' want to go to one meter scale the issue is that of course this is not a sustainable growth of compute and it's completely out of reach right now and both in terms of compute and storage actually that would be petabytes of data to to handle so even right now if we want to train even at we right now we're at more or less 25 to 1 kilm scale both for numerical solvers and AI driven models and so already this requires actually Distributing models on hundreds of gpus and kind of yeah um using these latest models so this is where neural operators come in um traditionally neural operators uh neural networks regular neural networks work on inputs and outputs at fixed resolution for example you have an image to image um problem you you take your units you train at fixed resolution and you predicted this fixed resolution but once you've trained your model you cannot do inference or you cannot easily retrain at a different resolution you could maybe upsample but you canot lose all the new information and so the idea of neural operators is to actually generalize neural networks to instead of doing mapping between finite dimensional space is doing mapping between functions so you take as input now a function and you output a function and you learn an operator to map between these function spaces and so if you have this you can now query at any point in the domain um and one property we want to have is uh what we call this discretization convergence which means like for example here if you have an air foil and you have an air flow around it as you refine your mesh you want the operator you're learning to converge to the same solution you don't want to get a different solution at each resolution and so kind of to summarize the properties we want of a neural operator so the first thing as I mentioned is we want to learn a solution operator that Maps between functions we want to decouple the resolution and number of parameters which means that you want to be able to train the same model on different resolutions without increasing the number of parameters and you would want the resolution to stay coherent between these different resolutions and then once you have a model you want to be able to find unit or inference it at a different resolution and finally you want to be able to query at any points so the framework in which we do this is called the neural operators and the idea is that in deep neural networks we use um linear mappings for example convolutions of fully connected layers which are followed by you know normalization uh nonlinearities Etc and you stack these layers you build this deep neural networks and you learn a mapping from your inputs to your output the idea here is the same except we use um integral linear linear operator which are frequently used in approximation Theory and that kind of map they are linear mapping from one function to another function so I kind of want to subdivide in the rest of the talk in multiple use cases depending on the geometry of your problem so the first case is if if you have a regular grid say for example a 2d or 3D image uh or something on a 2d domain and in that case you can use um the fast for transpose which will make things very efficient so the idea is you take your kernel integral at the top and we look at a specific case of the kernel integral which is the convolution operator so instead of having a general kernel of so V here is your input function so you're mapping your input function V of Y to a new function uh that you will query at point x here so you integrate your input function with some kernel and here we simplify the kernel to just be a convolution and so you may know that you can realize this efficiently with a for transpose transform you take the for transform of your kernel the for transform of your input function and then the convolution simply becomes a point-wise multiplication in the spectral domain and the advantage of this is you can if you have a regular input you can do this very efficiently with the fast for transpose transform I don't know why I keep saying transpose um maybe I blame their previous talk um and so yeah so now basically instead of having a o of n squ complexity you can have an O of n logn which is much more efficient and so in this case the architecture will look very similar to a regular deep neural network except now um we have these neural operator blocks so you have your forier layer which taking your function you transform it uh into the the spectral domain there you truncate the high frequency to remove the noise and kind of have something that to have a fixed number of bases um you do your point-wise operation you multiply by uh your weights which here model your kernel and then you do an inverse transform to get back into the original domain and you stack a bunch of these layers with some nonlinearities in the middle to learn a mapping between your functions and so the advantage is that now you can uh here for example we I show I'm showing some example on a nav stock problem so fluid dynamics and we trained on the we trained our neural operator our for neural operator on 64x 64 resolution but we can test it directly on 256x 256 resolution without any retraining and the nice thing is that you could say okay I can train my neural network a regular neural network at a small resolution and just interpolate but if you do this your neural network has never seen the high frequency information you had and so as you can see on the right I guess you can't see maybe you can see the cursor so the if you look at the energy Spectrum um on the left these are the low frequencies so you cannot learn them well because you've seen them during training but on the right these are wave numbers that were not SE during training and so if you now use interpolation your neural network kind of behaves erratically because it doesn't know what to do with the new frequencies while the neural operator can have generalized as well to unseen frequencies um yeah so for anyone that knows me uh I really like tensor methods I worked a lot on them so I couldn't do a talk without mentioning tensor method so they're still relevant here um essentially the core idea is that a lot of the data we manipulate in modern machine learning is high dimensional and so by considering this structure in the data and so considering the data as multi-dimensional array we can extract more from it and so the idea is very simple so a matrix is just um let's say a 2d array of numbers you can think of a tensor as an N dimensional array of numbers and then instead of doing matrix multiplication you do tensor contraction so you multiply for example here Matrix with a tensor and when you have this you can generalize know Concepts like Matrix decomposition so here um you can express a matrix as a a sum of rank one matrices so a sum of outer product of vectors and you can generalize these two tensors by by expressing them as a sum of rank one tensors which is the sum of outer products of three or more vectors and so you have many kinds of tensor decompositions and they're very relevant here because the we have Global convolutions so they are they are parameterizing um the convolution over your entire input so they have kind of size Heights by width by maybe Time by input channels by output channels you already have a huge tensor and so by expressing them in factorized forms you can drastic Bally reduce the number of parameters so in the for neural operator I presented before we apply this idea and we factorize the weights in the spectral domain so we express them as a TENS of factorization and then during the forward pass you can directly multiply your activation with the factors of the decomposition and so essentially we're able to um improve the performance of the model with just a fraction of the model so the the one on the right was the original for neural operator um in the middle we actually just brought some existing known techniques from Deep learning to neural operator realm like properly doing things normalization Etc and then at the bottom is the tens rise version so we have better performance with a fraction of the parameters because these tensor factorization kind of act as a regularizer on the on the method on the model so you can learn more efficiently so yeah we're able to compress um up up to let's say 500 times without any lot of performance and if you if you decide to just trim your model instead of using factorization performance uh decreases very quickly so this is a really principled way to do this compression and it also helps in the low data regime so actually if you have less training samples tensorization helps you to prevent overfitting so anyway that was a quick parenthesis on how we can bring tensor methods to this um and I want to go to the next use case which is now if your domain is still kind of regular but instead of uh being on a regular 2D grid we now learning on a sphere and so this is the weather case I was presenting uh earlier so here we were trained on some um the essentially gold standard of weather forecasting it's called era five it's a reanalysis data um it has several decades uh of data it's a 25 kilm scale and oh yeah that's essentially the the best data we can have have for training now and so there the idea is exactly the same as the for neural operator except now our data is on a sphere so we want to use this and instead of doing a Foria transform we can now do a um spherical harmonics transform and so instead of having s and cosine as the basis functions because now we're representing a function so the basis of our space Also function we use um spherical harmonics so circular functions on the sphere and so this a very natural way to represent uh things on on the sphere and as you can see on the left we have a Transformer based model that actually does very well at the beginning but if you do auto regressive inference so it means you take your forecast let's say you're predicting the weather 6 hours from now you take your forecast and you refit it to your model to get a prediction for 12 hours and as you enroll your model it quickly you quickly get these weird patterns that essentially you can see the patching that the model was doing while um the the spherical harmonics base model because it uh and C the geometry of the problem is able to do prediction on hundreds of days without going completely crazy um so yeah and the advantage of our model is that it more or less matches the accuracy of current numerical solver of like weather forecasting models but it is a lot faster um so here like 45,000 times faster and so this is really crucial if we want to do things like um ensembling and for example this is the hurricane Harvey in 2017 um and this is we can see the observe the actual trajectory of the hurricane versus what was predicted and so the issue is if you just look at one prediction you will get something that's typically very wrong there's a chaotic part in the system so you can't really well predict it so what we would want instead is to run our models many times like thousands of times and then build an ensemble of possible trajectories and so from these we can now not only have a predicted trajectory but also the probability of it happening in different places and so concretely for example right now the models are really slow they take minutes to run while our models runs in milliseconds so you can run thousands of ensembles much faster and so this allows you to here on the left if you have only 50 Ensemble maybe you will only capture the kind of you always have a sort of Goan distribution so here it's like only a small range of wind speeds and you need many more ensembles to actually be able to model the um on the Tails of the distribution so the extreme winds and so these are actually the ones we're most interested in predictive because those could be the most destructive or those we couldn't PR from four because we didn't have this trajectories cool and so the last thing I want to talk is if you don't have any structure at all in your data so previously we were doing nice um fluid dynamics on the Taurus everything was regular even for weather it's all nice on a sphere but there are cases where we just have arbitrary geometries and we still want to be able able to solve things um so one case is um aerodynamic simulation so for example if you have a car geometry and you want to predict the pressure on the surface of the car or more generally the drag of your car ideally if you're designing a car you would like to be able to play with the geometry and then know how this affects the drag without having to stick your car in a wind tunnel or without having to wait a few hours for your simulation to run and so this this is a use case where it would be useful to have a surrogate model an AI surrogate model that can do this uh efficiently and so this is the idea of the graph neural operator and the idea here is to be able to have a model we can query at any point and that can works at different levels of discretization so you can now take your car and give it give uh this point Cloud as input to your model at different discretization and get the prediction at any point so this idea also of discretization convergence can be really useful if you don't have that much compute then you can start training at a course resolution and give few samples at high resolution as you get it so this saves you like data and computes and so the idea is once again based on this um linear kernel integral except now we have a you have your function still V of X the same as before but now you have have a query point so let's say you want to query for example the pressure at one point on the surface of the car you now build a neighborhood so a ball around this point and you integrate simply in the ball so all the points around all the neighbors of this query point and so you do your kernel integral only in this neighborhood and the size of this neighborhood Remains the Same regardless of the discretization of your car and so the idea of the whole model which here we called Geno was first you have a Geno layer that takes your um arbitrary geometry you query it on a regular grid so you have an arbitrary geometry you query it on a regular grid so now you can use your fno which is very computationally efficient because it relies on fft and at the end you have another uh gno layer to kind of decode your uh representation and you which you can query at any point so then you can again query your model on all the points on the surface of the car predict the pressure and then from this predict the uh the drag so these are just some visualization of the pressure on the surface of the car you can see the ground Truth versus the predicted pressure the error is quite low and we found that we can get results with a 20 about 25 26,000 speed up over GPU based optimized um comput like fluid Dynamic solvers and so this is just to visualize this speed up you can see in red and orange is the this model um and this is a cost trate of accuracy so you can see the time so it's basically the error in the drag that the car drag as a function of time and you can see that with our model you can get a as good an accuracy as the open foam solver which is GPU accelerated in much less time so actually this is all from me so just to summarize I think um yeah this AI for science is really an exciting application of deep learning and it really takes everything we've learned before whether it's like how to build the architecture how to scale the models um how to use very large amounts of data and kind of applies it to really real world problem which is really interesting and neural operators are really principled way to do this so yeah I hope uh you guys are interested in trying it and because uh Alexa mentioned this is an open source only conference um everything is open source and actually it all uses py TCH as well so the neural operator Library just provides uh layers and models to run these in py to and tensor provides all the tools to kind of manipulate factorized tensors as regular P pych tens tensors and torch harmonics is a library that that just provides a py torch differentiable spherical harmonic trans transform so great that's all from me thank you very much and any questions yeah thank you that was great um you showed us towards the end that if you already know the geometry you can figure out what the basis transformation needs to be so that you can do the operator what if you don't know the geometry can you optimize the basis transformation just having like a pile of data that's a really good question I I guess the most General way would be this graph neural operator that I presented at the end then you just take any geometry and you you're literally just learning the kernel so that would be the the most General otherwise you could also try to parameterize this kernel by another Network for example you could have an MLP like I'm sure you're familiar with um Nerf for example like implicit representation you could have a a model that takes as input your coordinates x y and that kind of maps them to the corresponding kernel so there all sorts of Tricks possible but the the key thing here is that the fft is really efficient you can do it in N log in and that's the main reason we use it so you totally can learn the basis but it would be harder and you wouldn't get the speed up and also the the key Insight here is that whenever you have expert knowledge I think it's always good to use it it's like additional inductive bias that's why convolutional neuron networks work well on images I think yeah it's the same on every problem okay not thank youone again thank you