Trillion Token Context. No, Really — Anima Anandkumar & Benedikt Jenik, Accelerated Understanding

0:01 What's up everybody? I'm R.J. I'm here with my co-host Brandon for a latent space lightning AI for science podcast. A week ago, we released a fantastic episode with Ana Anan Kumar. It got a great reception, probably helped by the fact that she brought her company Accelerated Understanding Out of Stealth days before we released. But when we recorded a month ago, she couldn't talk about her company. So now we have Anima back on the pod with her co-founder Benedict Yanik to tell us about accelerated understanding and the things that she had to keep under wraps when we were recording before. By the way, just a little thing, Anima was just named Time magazine by Time magazine as one of the 100 most influential people of the year and we want to talk about that too. But first, Benedict, welcome.

0:50 Anima, welcome back. >> please tell us about accelerated understanding. >> Yeah. Essentially what we have been thinking about is we've been looking at all the successful projects that Anima has been doing all those like the catheter the weather model all those things she talked about on that last podcast and then we figured the thing that has happened in language like if you look back in language and that's now ancient times before the recent AI revolution people were training models for specific things like we had a spellchecking model there was a translation model. There was an autocomplete model and then OpenAI came along and was like, let's just put everything into one model and that turned out to work phenomenally well and exceed pretty much everything else those specialized models. And so now we were asking ourselves, can we do the same thing for physical simulation and for physical understanding?

1:50 Yeah, you know that universality and scale that we've seen play out for language. What is that counterpart for the physical world and that's the bet that accelerated understanding is making? >> The first question that comes to my mind with this so it's a you know basically you're trying to apply a bunch of different domain data to a single model build sort of the god model version physics version or physical world version. The question in my mind that comes up is while language there's lots of different kinds of language but they it's all language where but for weather simulation and catheterss and nuclear fusion reactors maybe the physics itself is very different and therefore there's not a lot of transfer.

2:37 So what evidence or why do you think that this is a still a transferable problem? Yeah, I mean first of all, of course, there many more challenges when it comes to the physical world, right? There isn't just a readymade internet of all these kinds of physical data available. So, real data is just not there. And also, if you think about scientific discovery, it's about doing something new. So, by definition, it cannot be there in the training data.

3:06 And so, this reliance on just purely datadriven AI is not going to be enough. And that's where you know adding the laws of physics is really critical. So think of like you know we can use existing simulators and everything else to start training our model but then the cherry on top is going to be that self-improvement using the laws of physics itself. So physics will be our guiding principle throughout and that's also something that's common across different areas whether it's like fluid in a catheter or in a rocket it's the same principles of fluid dynamics they're different regimes and that's parameterized by reol's numbers and similar such things but there is something that is shared among them and that's what we've seen this play out that across different domains if you think energy semiconductor aerospace you know it looks very different but there are common underlying principles but even if you then say that there are different mathematical equations like PDEs across dist different systems why are they common you know there's still commonality like there is variation in time there could be some kind of energy conservation you know there is aspects like object permanence there's causality so there's implicitly a lot of common features even across Ross physics that are gone by different equations. So that's how you see across different domains across different mathematical models there's a lot of shared features that these neural models can pick up.

4:44 >> I see. And and so what what is some of the most interesting? So actually where do we stand with building this model? How far have you gotten so far? >> Yeah, so we have actually been training for a bit more than a year. We've done all the usual architectural experiments you need to do. The interesting thing here even is you need to reinvent a lot of the things because the shape of our data is so different. Like if you look at for example language, you have tokens and those tokens grow in one dimension.

5:17 You can fit like a million of them in leading models. for us it grows in more dimensions and that is kind of like the the special thing here is we don't want to cut any corners. So we are like let's make everything as it is in the real world. So we have space that grows in three dimensions. It's not compressed like in video models. It's actually staying in those three dimensions and then you also have the roll out over time. And there again people like to take the easy path do auto reagive meaning you predict one step at a time except then errors build up and you lose kind of that continuity and the understanding of what in the input actually led to the output which is really important for designing stuff. So even then for the time dimension we prefer to just model the full rollout which is what we're doing. But then you're in this place where it's suddenly four dimensions and all of those four dimensions grow independently and when you multiply those numbers out, you're very quickly in the billions or even trillions in context. And that is something we've actually achieved with our models. Like we're able to train up to a trillion context input. We're able to train with like even inputs outputs both trillion context length. We are able to do inference at 5 trillion context. Like those are ridiculous numbers when you come from language and we're not even doing most of those like if you look at video models they do a bunch of tricks to make it more efficient. They like average over pixels. They patch them. We're not doing those things. We're still able to do this ridiculous context. But now going back on what have we done? How are we making it happen? Fitting this stuff is non-trivial. training this stuff is non-trivial. Like you may have heard of those standard techniques like FSTP where you shard but then you reassemble a layer in a GPU to do the training.

7:16 except our layers are so big you can't reassemble it inside a GPU. Our data samples are so big they don't fit an accelerator or even a full node. So we had to reinvent this whole sharding infrastructure, the whole sharding strategy to be able to train those models at scale. And we've done so we've trained done hundreds of training runs. We've trained up to trillion parameter models. So we've really shown this stuff to take off. The other thing that we've made sure is that we're actually benefiting from multifysics like not just doing everything big that was done small before but actually putting those things together. So we've on purpose picked a number of areas of physics that we believe are very diverse in their features, in their challenges and put them in the same model, done the full 4D roll out for those and we're able to train them.

8:07 >> And in fact I was going to add that it turns out that having the model of the same size with multiple areas of physics does better than giving all of those parameters to each single physics. So if you had separate models and made them big enough as the original one, it still is worse. And so that's saying that it's benefiting from shared learning. It's not just we make the model bigger and it gets better. That's true. But adding more areas of physics helps it to learn better in all of them.

8:38 And that's the same kind of emergent learning we've seen play out in language. In your episode, we talked a lot about neural operators and we talked about in particular some of the inductive biases you put into your models like Forier neural operators where you go in for space where you operate on a sphere and then you you have an inductive bias about geometry. How does transfer learning work when you have are now trying to do many different geometries? like are you actually now limiting yourself to say like we're only doing you know rectangular geometries or is there actually some now tricks where you can actually integrate different types of different classes of PTE like maybe with different time derivatives or things like that >> yeah so we cannot tell you every exact detail of our architecture since that's proprietary but what I can tell you is neural operators do form the basis because that's how we can make this resolution invariant. So Benedict talked about a five trillion context length at inference time that we've achieved and a trillion context length during training. But does that mean every single run we're going to do that?

9:47 No. Every single application needs that much resolution no right. So you know those are for the hardest physics that really requires all the details. we can give it the big context length and those that do not need that or even during the process like if you're doing design exploration all of these tasks in the beginning you don't need every single detail you don't need to burn that much compute and so that flexibility is crucial and if you see that immediately distinguishes us from other so-called world models whether it's video models vision models they all assume during training and inference it's a fixed resolution and that's Okay, for you know for our visual features, right?

10:31 Because we don't really necessarily need to zoom in more especially for gaming and entertainment. It's okay. The physics is somewhat a bit tan wavy. But you know for the engineering design, scientific discovery, we can't get away with those tricks. And so neural operators form the key to ensuring that we can be flexible at giving different context lengths which is equivalent to different resolution. On the other hand if you think about using transformer architectures that have worked so well for language that just wouldn't be able to support a 5 trillion context length no matter all the compute in the world is thrown at it. So that kind of quadratic complexity is infeasible and also unnecessary because the physical world has more structure than completely arbitrary alltoall correlation. So we have to rethink what works better for the physical world and that's what we've done very systematically like how in the frontier labs for language there is a lot of systematic experiments and lot of evolution of architecture you see from the earlier models to mixture of experts and so on you know we kind of like gone through and we continue to go through that same evolution within our accelerated understanding to ensure that we have the best architecture as well as the best hardware utilization, you know, thinking about communication, bandwidth, all those requirements. So, that code design is really critical.

12:04 >> Oh, real quick, I just want to say people who for people have not watched the original episode. we recommend you watch it if you want to learn about neural operators, I think it starts around 20 minutes. We'll have a link directly to jump to that scene where we start talking about some other things. but yeah, wanted to make sure we got that plug in real quick. yeah. So, I guess for for either of you, and I know a lot of this is probably proprietary, but can you give us an intuition for how how do you get to a trillion or five trillion tokens and how do you, you know, sort of that's both like a algorithmic question and also an infrastructure question.

12:40 >> Yeah. So, how you add end up with those large numbers is fairly simple. Like if you say I have a thousand resolution in each spatial dimension and then I care about resolving the thousand time steps exactly then I'm at a trillion so that is a really large number and then when you think on a little bit in that math how big of data is that like for example our 5 trillion run that we did the outputs were 22 terabytes and you want that kind of stuff in accelerator memory.

13:18 That's just what it is. Otherwise, it's going to be slow. So that is kind of like how you need to think about. Obviously, we can drive it all the sizes like we have done inference on huge clusters. We have done inference on like we can the small models with smaller context or even medium big models with smaller context fit on a MacBook or Mac Studio like you have that whole span. So like having this there is like like there is a bunch of scaling tricks that need to work out but the big one is you need a working set of your data that needs to fit somewhere and that is kind of like what drives all of this and then again I can talk about what you actually do but ultimately you can think of you want to interact the pieces of data with each other in some way over the mo like the throughput through the model and that means to a good degree you will also need to rely on large clusters with very reliable very good interconnect. So like a lot of those challenges that LLMs are walking more and more into where there's like both modern tricks and the good old HPC methods. We had those from the get-go.

14:40 Like we had to figure out how do you fit a sample that does not fit into a node. We had to figure out how to do an interaction within a layer where you can keep the state around in an accelerator. So all those tricks require quite a bit of scaling and quite a bit of infrastructure work that we have been able to put together. >> We're just talking about data. That's actually I think a really interesting question that we haven't really expanded upon. Can can you can we double click a bit on that? Like what does the data look like? Are you is this primarily you're doing like large scale differential equation solving as like a to generate training data? Are you getting physical data from a variety of sources? How do you integrate both physical data and also you know simulation computation data? Mhm. And that is that is actually quite an interesting point because that used to be a bottleneck especially for a lot of those narrow surrogates and you may have seen it. Company has some really nice pocket of data. You have a flashy demo but then the area they actually care about is out of sample and nothing gets deployed. So having this kind of limitation is something that we wanted to step out of. At the same time, we all know for simulating physics, the big theme is sim to real gap that you somehow need to overcome. So, how do we do both? one interesting thing with PTE is we are actually fairly confident like the math is known that when you solve a PD correctly, you're doing the physics correctly. Like obviously you still need to make sure that you're representing the task that you're trying to solve within that but you can solve the like sensial gap by making sure your skill domain is wide enough that you're guaranteed to be in distribution. Now for training itself what does that allow us to do? We can use numerical simulators to generate as much training data as we need and even further we can do what is called curriculum engineering meaning we can start with simple equations we can start with lower resolution simpler data that's which is a stark contrast to how language is trained like language you download all of the internet it's a giant mess it's shuffled all over some of it is correct some of it isn't and this is not how for example as humans would have learned like learning in order is much better.

17:11 So we're able to do that. We're able to build our own curriculum using numerical simulators. But it doesn't stop there. Like if you were to stop there, you kind of end up as everything else in machine learning does at the average quality of your training distribution. And we wanted to exceed that. That's where also self-improvement comes in. So the interesting piece is you can use those PDEs both for numerical simulators to generate data, but if you're clever about it, you can even use them as a training signal. You can check how well is my model actually doing on the PTEs themselves and use that as an additional training signal where you're suddenly in a place where you can push beyond the quality of the data in your model quality. And the difference there is compared to language where self-improvement needs something like human feedback or other reward signals that are very sparse. They just tell you yes or no, thumbs up or down. We have dense feedback because the physics laws there's so multiple of them and you can decide again a curriculum of how to arrange them but also the feedback you're getting from them is dense and so that kind of self-improvement can be even better and that's what we see in our experiments. Maybe this is already implicit in your answer your previous statements but one thing I'm curious about in physics is this idea of multiscale where like large scales are you know you need to simultaneously represent things happening at large scales and things happening at small scales and a lot of the times the difficulty of good physical representations is this you know is that you do need multiple scales which gives you this like large like orders of magnitude and understanding. how do you represent that simultaneously like how do you how do you deal with that problem?

19:06 >> Yeah. So if you see lot of our published work with neural operators that's exactly what we address right. So in terms of like if there multiple scales you also need multiple resolutions. Of course you could do everything at the highest resolution but that's extremely extremely expensive. So you're better off first collecting more data at lower resolution, maybe even cheaper solvers that coarsen and ignore the fine scale effects. They're wrong, but there are good starting points to kind of get the overall average effects, the core scale effects, and then you can fine-tune with more specialized that take into solvers that take into account those finer scale features. And neural operators have this flexibility because they can allow you to mix across different scales within the model rather than be prescribed externally like a lot of other hybrid machine learning for physics do and that you know allows us to be a lot more data efficient and much better at learning.

20:12 >> We don't have a ton of time. I want to hear a little bit what are the the applications here? What are you what are you chasing first? you've come out of stealth. So presumably there's something that you're doing more publicly with commercialization. Can you talk a bit about that? >> The way we're thinking about it is like there's two perspectives. There is the domains of physics and then there's the areas that can be served by them. So if you look at areas that are obviously very interesting right now because everybody needs improvement there the big ones to everybody are semi and energy and those are precisely right in our wheelhouse. So we can't tell you fully who we're talking to, who our customers are, but like within Sammy, you can think of a lot of problems that physics become increasingly relevant.

21:05 And the interesting thing is if you look at how a lot of stuff worked there. especially when you look at the chip design itself, it was much more a let's start in the digital, let's freeze the digital in, let's send it through some physics for a one-time check. Like the PDK dictates I have to have the following feature otherwise TSNC doesn't make it for you. And it stops there. Like it's like do we have enough headroom to push it through? But we're not taking advantage today of that headroom to get extra performance. So there's this potential to unlock much more of the skill set when you're also able to grab into what does the physics allow me to do to push the envelope a bit more and can I rearrange everything a little bit to push the envelope even further. That's one area. Obviously the semi-production process itself also has all these interesting physical challenges that we're addressing. So lots and lots of work to do. we're excited. So do I understand correctly like semiconductor digital analog code design >> like obviously you need to start somewhere and that there is like this whole pipeline that is also now improving like we are making sure that like we're bottom up enabling that physical like eventually that's where the journey is going obviously there's different companies that do different pieces we're starting from that direction because we think that direction is opening up a lot of impact where previously there was a lot of simplifications like you did wire length you did certain optimizations you're like oh do I have a thermal spike somewhere do I have an em spike somewhere that messes with me but not oh I have all this other stuff where it's actually fairly cool where there's fairly little craziness happen can I take advantage of that >> nice and and energy presumably we've talked about nuclear before I presume assume that's that's one of the areas there is are you also doing solar and other things like that >> there is both the design kind of piece which is I want a physical object with certain capabilities that's was something we're really good at there's other areas where the question is I have observations of the world what does that tell me about the world like think geodermal where you want to know where It's warm down there with the right mix of water that you can pump through. Like all of those things were or where are the critical minerals for for all the electronic devices we need? like all of those things kind of end up in our wheelhouse and you get there with you need a good amount of physics but physics at the same time like there's not that many things that are different from each other. So you like once you're at this point where your model reaches a good level of universality, suddenly a lot of doors open.

24:10 >> Just to push on that a little bit and you're seeing some of that transfer happening in the applications that you're chasing with with real commercial. >> Exactly. Like we're seeing that physical uplift happening. We mentioned it earlier. We've trained models where we were like, let's just make sure our simulator is fine, configured fine, and we can actually train it. And then we were like, oh, and let's add it to the universal one. Like, you want to always test those kind of things. Like there's this whole bunch of pipeline work that you want to make sure you're correct.

24:43 And what we observed is we had like like in terms of the performance characteristics we're seeing in the narrow one versus what we suddenly saw in the broad one, it was like broad all the way. >> We're running low on time, but maybe is the last question. What what is next for accelerated understanding? what do you want people to know? And are you you hiring, you're opening offices, like what's going on? >> Yeah, certainly all of that. I mean we are scaling our models and we continue to scale our company. So yes, you know, we are getting a lot of inbound interest since last week. So I'm just getting through that. So I apologize if I haven't gotten back to someone yet, but it's a really exciting time. And you know, and for us, this is just the beginning in the journey of the model scaling, right? If we have to go towards the biggest challenges we are facing in scientific discovery and inventions, those are really hard physics, hard phenomena. But at the same time we have to be also true to our commercial side and ensure along the way we unlock a lot of value to our customers. So it's that combination that it we are very excited about.

25:53 >> Well and actually I want to just hear for 30 seconds tell me about the time 100 thing like what what happened there? How did it go down? I' I've never received one so I just I'm curious. You know, I was just thrilled and really surprised to be included in the time 100 AI list. you know, I was also very lucky to receive the timeund impact award last year. So, I was in Dubai for that and you know, got to meet a lot of luminaries in the field and now you know this year. So, it's I think it's good what they're doing trying to bring people with different perspectives on AI together. So yeah, so that's a good thing.

26:37 >> Yeah. Well, congratulations. Welld deserved and thank you so much for for jumping back on with us and we're really excited to see how everything unfolds with accelerated understanding. >> Thanks a lot for having us back. >> Thank you. >> Take care. Bye.