Faster Chips That Don't Melt — Anima Anandkumar & Benedikt Jenik, Accelerated Understanding

What's up, everybody. I'm RJ, 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 Anima Anandkumar. 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 Benedikt Jenik, to tell us about Accelerated Understanding and the things that she had to keep under wraps, uh, 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, Benedikt, welcome. 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 spell-checking model, there was a translation model, there was an auto-complete 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?

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, um, but for weather simulation and catheters and nuclear fusion reactors, maybe the physics itself is very different and therefore there's not a lot of transfer.

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 are many more challenges when it comes to the physical world,right? There isn't just a ready-made 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. And so this reliance on just purely data-driven 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 Reynolds 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 different systems, while they're 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 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.

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. You can fit like a million of them in leading models. For us, uh, 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 rollout over time.

And there, again, people like to take the easy path to auto-aggressive, 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, um, like even inputs, outputs, both trillion context length. We are able to do inference at five trillion contexts.

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 FSDP, where you shard, but then you reassemble a layer in a GPU to do the training.

Uh, except our layers are so big, you can't reassemble them inside a GPU. Our data samples are so big, they don't fit in 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 multi-physics, 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 rollout for those, and we're able to train them.

And in fact, uh, 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. And that's the same kind of emergent learning we've seen play out in language.

In your episode, we talked a lot about, um, neural operators, and we talked about in particular some of the inductive biases you put into your models, like, uh, Fourier neural operators, where you go in Fourier 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, um, different classes of PDEs, like maybe with different time derivatives or things like that?

Yeah. So we cannot tell you every exact detail of our architecture since, uh, 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 Benedikt 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? 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, uh, you know, for our visual features,right?

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 hand-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, uh, we can be flexible at giving different context lengths, which is equivalent to different resolution. On the other hand, uh, if you think about using transformer architectures that have worked so well for language, that just wouldn't be able to support a five-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 all-to-all 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 a lot of evolution of architecture you see from the earlier models to mixture of experts and so on.

You know, we've kind of like gone through and we continue to go through that same evolution within our, uh, Accelerated Understanding to ensure that we have the best architecture as well as the best hardware utilization, you know, thinking about a communication bandwidth, all those requirements.

So that co-design is really critical.

Oh, real quick, I just want to say people who, for, um, people who have not watched the original episode, uh, 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. We start talking about some other things. Um, but yeah, wanted to make sure we got that plug in real quick. Um.

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 an algorithmic question and also an infrastructure question.

Yep. So, um, 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 times 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, uh, how big of data is that? Like, for example, our five trillion run that we did, the outputs were, uh, 22 terabytes.

And you want that kind of stuff in accelerator memory. 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, uh, all the sizes, like we have done inference on huge clusters.

We have done inference on, like, we can do 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't 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, like the throughput through the model.

And, uh, that means to a good degree, you will also need to rely on large clusters with very reliable, very good interconnects. 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.

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 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, uh, as like a generative generate training data?

Are you, um, getting physical data from a variety of sources? How do you integrate both physical data and also, you know, simulation computation data?

Mm-hmm. 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? Uh, one interesting thing with PDEs is we are actually fairly confident, like the math is known that when you solve a PDE 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 sim-to-real 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, uh, 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, us humans would have learned. Like learning in order is much better. 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, uh, 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 PDEs 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, um, statements, but one thing I'm curious about in physics is this idea of multi-scale 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 in understanding. Um, how do you represent that simultaneously?

Like how do you, how do you deal with that problem?

Yeah. So if you see, uh, a lot of our published work with neural operators, that's exactly what we address,right? So in terms of like if there are 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, uh, collecting more data at lower resolution, maybe even cheaper solvers that coarsen and ignore the fine scale effects. They're wrong, but they're good starting points to kind of get the overall average effects, the core scale effects, and then you can fine-tune with, uh, more specialized, uh, that take into solvers that take into account those, uh, finer scale features.

And neural operators have this flexibility because they can allow you to mix across different scales within the model rather than being 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.

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? Um, you've come out of stealth, so presumably there's something that you're doing, uh, 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 interestingright now because everybody needs improvement there, the big ones to everybody are semi and energy, and those are preciselyright in our wheelhouse.

So we can't tell you fully who we're talking to, who our customers are, but like within semi, you can think of a lot of problems that physics become increasingly relevant. And the interesting thing is if you look at how a lot of stuff worked there, um, especially when you look at, uh, 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, ooh, 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 skillset 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.

Um, we're excited.

So do I understand correctly, like semiconductor digital, uh, uh, analog co-design?

Like obviously you need to start somewhere and, uh, 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.

Um, 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 lengths, 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 that's, that's one of the areas. Are there, is it, are you also doing solar and other things like that?

There is both the design kind of piece, which is, uh, I want a physical object with certain capabilities. That's what something we're really good at. There's other areas where, um, the question is, I have observations of the world.

What does that tell me about the world? Like think geothermal where you want to know where it's warm down there with theright mix of, uh, water that you can pump through. Like all of those things where, or where are the critical minerals for, for all the electronic devices we need.

Uh, 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.

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 customers?

Exactly. Like we're seeing that physical uplift happening. Uh, 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. 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 this is the last question. What, what is next for Accelerated Understanding? Uh, what do you want people to know? And, uh, 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, uh, commercial side and ensure along the way we unlock a lot of value to our customers. So it's that combination that we are very excited about.

Well, and actually I want to just hear for 30 seconds, tell me about the TIME100 thing. Like what, what, what happened there? How did it go down? I've, I've, 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 TIME100 AI list. Uh, and I was also very lucky to receive the TIME100 Impact Award last year. So I was in Dubai for that and, you know, got to meet a lot of, uh, luminaries in the field and now, uh, you know, this year.

So it's, uh, 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.

Yeah. Well, congratulations. Well 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.

Yeah, thank you.

Take care. Bye.