Scale By The Bay 2021 : Anima Anandkumar, Convergence of AI, Simulations and HPC
Recording: Scale By The Bay 2021 : Anima Anandkumar, Convergence of AI, Simulations and HPC
great so i think you know we are in the golden age of ai right then what we are now seeing is increasing impact in so many domains and to me what's really exciting is the impact it's having on the sciences uh and in the area of scientific simulation which has traditionally required uh large-scale hpc and um you know what can ai provide there right what can we do better with ai that's what i hope to give an overview today and uh yeah and that's i think the question right ai has been one of the most powerful technological forces we've seen and it's mind-boggling how the complexity has grown over the last few years if you count the how larger the model has gotten it has grown 30 000 times in the last five years right so we now have a new form of moore's law which is about ai complexity doubling every two months so with this exciting skill what is possible right i guess the largest developments have been in the language domain but can we also have those developments in the sciences indeed the alpha fold is one of them where scale has given us right the best model compared to any of the previous ones the question is can we do that also in other problems in sciences and one important challenge with that is extrapolation right so going beyond what we've already seen or what the model is trained on how well can it accomplish and that's really important for the sciences and indeed i don't need to remind you that for this we need both data and large-scale computing right the design of the algorithms is driven by also the availability of the data and the computing scale and the question is in the sciences you know can we get large enough amounts of data and this is a challenge especially if the simulations are so expensive can we you know do it even larger than um what uh the training data is right you may not be able to for instance simulate for a molecule of say size hundred or thousand so can you do it on smaller molecules and extrapolate and that's what makes it different from standard supervised learning that we see uh in most of the standard ai problems and that's where the other dimension of the scale comes in uh on the hpc front we have seen a lot of exciting gains in terms of our supercomputers getting more powerful right with each generation but the difference is unlike traditional hpc which is involving 64-bit floating points which is very high precision ai requires much lower precision and for that much lower precision uh if we can do better right then we have a lot more flops available and that's where now we are measuring the flops of our super computers in terms of the ai computations are no longer the traditional high precision computing and so that's where we can now use the hardware much more efficiently than with traditional computing and the question is on top of it can ai further enhance and give us orders of magnitude speedups while retaining the accuracy of traditional simulations and that's what i'm here to address what would be the algorithmic developments needed to enable this and just as a pointer to where we are in terms of the power needed in our super computers right we are now in the realm of exo flops right so the recently announced world's fastest supercomputer would have 20 extra flops of ai and so we are getting now this large scale and the question is what is the right class of ai algorithms that could make an impact in challenging problems such as climate weather material sciences and so on and so this is where um for ai for science right how do we replace these traditional algorithms with uh ai hopefully which can be much faster while retaining accuracy the biggest challenge is that most of these phenomena are highly complex multi-scale and possibly chaotic right so meaning you have to really go to the fine scales to be able to learn this phenomena or understand this with traditional methods that's why these discretizations have to be at the fine scale the question is with enough training data can ai learn this multi-scale phenomena without having to go to find discretizations for solving each problem and another aspect is also extrapolation in terms of the nature of training data right so if we can train only on small molecules and directly extrapolate to large molecules then again we can make this technique practical because the traditional methods may be too slow even on our world's fastest supercomputers to be run on larger molecules to give you an example the fundamental equation that uh governs um right the properties of molecules is the schrodinger's equation so it's looking at the quantum level interactions between atoms in the molecule and uh by current estimates in order to compute this exactly on a molecule of size say hundred atoms it will take longer than the age of the universe on the current fastest supercomputers so that's what we are talking about this sheer importance of or just infeasibility of being able to do brute force computations for getting some of these complex phenomena and the question is can ai make the right kind of approximations implicitly and be able to give reasonable answers to these challenging questions so the first part of what i'll talk about is what we call neural operator learning which is uh the result of having an interdisciplinary team um come together right from applied math material sciences and ai of course and the question we asked was you know we have partial differential equations can we use ai methods to solve them accurately and indeed to frame this problem it appears to be a standard supervised learning problem right so the input is different examples of in initial and boundary conditions and the output are solutions to those given conditions right so if this is indeed the case then uh you know we can use any standard neural network train it and that's you know that should work fine so why is it more nuanced than that so the difference is unlike standard setups where we have a fixed sized input right so such as an image of fixed size or a vocabulary of fixed size with the pdes what we have instead are functions right so it's a function to function mapping so we should be able to solve the equation at any different discretization so grids are different sizes different kinds of uh right shapes of grids here so all this we should still be able to solve and the traditional solvers can do that today right they can um you know handle these uh different resolutions and so now the question is can we design ai methods that can do the same and so that's why having a standard neural network is not good because it is for instance the convolutional models are learning filters at a fixed resolution and don't easily generalize to another resolution and the other difference we have to think about is right you're not just solving one instance of the problem you're ultimately mapping functions to functions and this is known as operator meaning these inputs no longer are of a finite dimension but potentially infinite dimension right because you can have like now the input as continuous uh within a fixed you know defined domain so given this how do we now design methods that can handle changes in resolution and grid and that was our motivation to go beyond um but again the main intuition of where we start is like kind of asking in a standard neural network what we have is typically linear operations followed by non-linearity right and these linear operations are of a fixed sized [Music] say matrix operation or convolution operation so the question now is how do we extend this to infinite dimensions and the main principle here is what is the solution of a linear partial differential equation look like and what it looks like is a global convolution where you're integrating with this kernel which is known as the greens function to give you an intuition if it's a simple heat diffusion equation the function would look something like that right in terms of how the heat propagates so for simple linear partial differential equations we can write this down as just a global convolution but instead for more general nonlinear pdes such as turbulent fluid flows it won't be this simple and so instead what we do is define this neural operator as non-linear compositions of such integral operators right so we compose this global convolution layers with non-linearity and by doing that we can now approximate any general operator and so then to make this tractable what we do is to use the fourier transform to solve this convolution efficiently so meaning instead of directly trying to learn that convolution which is hard we transform it into the fourier domain and then we learn the weights in the fourier domain so think of this as now we are filtering out frequencies in the fourier domain and uh changing their weighting and we also only filter out the low frequency modes in each step but because we compose with non-linearity we can obtain the full non-linear spectrum as well or rather high frequency spectrum as well and so a very simple approach that draws a lot of inspiration from classical methods that use the fourier transform but at the same time having the flexibility of representing any uh general operator and having the ability to learn this efficiently and so the main intuition of why fourier layers are much more effective than convolutional layers with learned filters is that they can capture global correlations much more effectively right whereas the filters are small in standard convolutional neural networks so that's one advantage and the other is this ability to compute this global convolution efficiently and having the property of handling different resolutions and grid sizes and so that's the overall framework and what's the outcome what we can show is the ability to train on just low resolution data right so just say the 64 by 64 sized images and directly evaluate it on much higher resolutions and so this is zero shot right so training time we've never seen this higher resolution evolutions but we are now able to also accurately resolve them whereas standard uh neural networks would not have this ability and so that's the benefit of this principled approach that we are resolution invariant and by doing that we can capture the underlying physics much more effectively and this we can actually validate by looking at the frequency spectrum right the idea is if you only have low resolution data we only see this part of the spectrum right this other part is never seen before but still this our model the fourier neural operator is able to much more closely approximated compared to say unit and so that's the benefit of having uh you know this framework that is able to handle different resolutions and we see the benefit of this in many other applications um just to name a few um in the previous one uh i forgot to mention the speed up we got was about thousand times and that's because the solver the traditional solver is also gpu based and uses fourier transform right so fast solvers exist for this case and hence uh you know we even then we get such a big speed up and the main intuition of why is for you know these solvers they still need to do fine discretization as they move forward in time whereas we can do much bigger gems using the machine learning method right we don't need that fine discretization and we can be much faster and in other methods such as like modeling plasticity this is a complex multi-scale phenomena and we have solvers at different scales because it's too complicated to do all at once and so the fine scale we replace the uh traditional solver with the neural operator and what we see is a big speed up here and and so that also shows like in some of the more complex problems the benefits we get from these methods is even larger another such example is carbon dioxide means capture and mitigation so you know one of the hypothesis is that climate change uh will be so untenable that one of the solutions would be to actually capture carbon dioxide and store it deep underground right and so to see if this is feasible we need to accurately model how will the gas saturate and the pressure builds up in this underground uh you know and this is complex because this is a multi-phase flow right there are different phases here and what we see again is a big speed up and here what we did was we added like a residual layer which did have like a unit along with the fourier neural operator which could capture both the fine and core scales effectively and we see a big speed up over traditional methods another example is understanding the molecular dynamics of um you know viruses and this one we did it on the covet virus and what we see is the ability to understand right like the different regions like the idea is like the conformal changes that happen in different parts of um the covet like you know is important for it to be able to replicate effectively so understanding this is really important again the details i want to you know have you here but the what it shows is the versatility of this method in being able to handle so many different uh right uh time varying right high resolution spatial variations and being able to get big speed ups over traditional methods we also see the impact on transformer architectures where we show that you can replace self-attention with a version of this fourier neural operator so effectively what it's doing is to enable token mixing much more effectively by making it as a continuous um right of object rather than discrete so it's making a continuous approximation and through that we see that we get the same right quality as the full attention but be much better in terms of the number of flops and so that's also something we see with the clustering of the representation in these layers and so this shows that there can be many diverse use cases for this methodology and talking about scale we are integrating all this into the nvidia simnet which will help us scale up right on djx and bigger systems and be able to solve highly complex real world problems and incorporate right both forward inverse data assimilation all these aspects together so as another example of what you know algorithmic development is needed for ai for science i'll talk about right computing the quantum mechanical energy of the molecules and this is really important for drug discovery right indeed that involves so many complicated steps here uh but the first part is of course the search you know how do you search and come up with candidates of good molecules and that space is enormous it's 10 to the 60 right so the question is can we um you know come up with good search strategies but also be able to accurately compute what is the right property prediction of these molecules and to do that what we did was to combine domain knowledge uh so these molecular orbitals are calculations that are fixed right that say how much each atom has influence over others in the molecule and so on top of this we build the graph neural network uh to learn and predict uh the energy right in this molecule and so the main reason why we use molecular orbitals is now this becomes transferable so even if you're trained only on small molecules like we did with just under 40 atoms we can directly make this method work on much larger molecules right 10 times larger than what we've trained and that's because these orbitals capture the dependencies within a molecule irrespective of how big it is so we can avoid the distributional shift between training and testing by incorporating this domain knowledge and so we see that in the outcome here that our method uh right has accuracy similar to these gold standard traditional methods but be three orders of magnitude faster so that's the benefit we get that with ai we can retain the accuracy while being thousand times faster and this is possible because we combine domain knowledge with the flexibility of graph neural networks and also compared to other machine learning methods we are very efficient in terms of the data needed right because remember we need training data from these traditional methods that's also very expensive and so the lower amount of data we can use the better right so and that's where from having the orbital information as domain knowledge helps us work with lower amounts of data and so with this we can now explore a broad range of properties right so i've just mentioned a few here that's important to chemists geometry optimization being an important one and compared to like right traditional methods ours is much faster while also being more accurate and you can see that here with the match um and recently we also use this method to understand how the covet virus right interacts in aerosols right i mean this is so important because this is what makes it so infectious and again not a lot of details here but the main aspect is that the orbnet calculations can make it much more accurate in terms of the energy predictions compared to the standard methods because the methods like dft which are the gold standard are way too expensive to be run on molecules of this scale yeah so to conclude here i truly believe that ai for science is the future of science and also the future of ai right but we need to be thoughtful because so many times this needs extrapolation and zero shot generalization right so it may not be your standard supervised learning method uh we need to ask uh you know is this testing going to be very different from the training data and in case of partial differential equations the reason was because we need the flexibility to operate on different resolutions and different grid sizes and the other aspect we found with orbnet is using the right amount of domain knowledge right that is cheap and quick to compute and building graph neural network on top of it can greatly enhance its ability to transfer from small molecules directly to much larger molecules thank you