Scaling Up Test-Time Compute with Latent Reasoning with Jonas Geiping - #723

0:00I think it's clear that we can sometimes think longer about a problem without actually verbalizing the answer in our head. We don't always think in steps in our head. But even as humans, we have these two axes in which we can scale, compute or thinking in this way. And I think that's kind of interesting, right? It's not necessarily you can do this or the other. But maybe as humans, we naturally do both of those. And these models also could do both.

0:45All right, everyone, welcome to another episode of the TwiML AI podcast. I'm your host, Sam Charrington. Today, I'm joined by Jonas Skyping. Jonas is a research group leader at Ellis Institute and the Max Planck Institute for Intelligent Systems in Tübingen. Before we get going, be sure to take a moment to hit that subscribe button wherever you're listening to today's show. Jonas, welcome back to the podcast. Yeah, happy to be back. A lot has happened in the last year. A lot has happened in the past year. I'm really excited about the conversation we're about to have. You recently published a paper called Scaling Up Test Time Compute with Latent Reasoning, a Recurrent Depth Approach that fits in at a unique time when, you know, reasoning models are getting popular, you know, just after the quote-unquote deep seek moment, if we want to call it that.

1:36And this paper proposes a different approach to reasoning. So really looking forward to digging into that work. I'd love to have you spend a few minutes talking a little bit about your background and refreshing listeners who may not have heard you introduced last year. Yeah, I'm Jonas. I have an old background in mathematics, but I've been doing computer science maybe since 2016. and I've been a bit around in Germany and the US. Now I'm back in Germany and I'm leading a little research group in Germany in Tübingen where we work on safety and efficiency aligned learning. And these are really things that interest me a lot.

2:19And this project in particular is one that's really been a favorite of mine and really have been, this has been a long time that we've worked on this because we really wanted to figure out how can we train these models very differently? Or we had this idea that we could train these models very differently and it would still be interesting and still be a good model. And so because this whole backstory to this whole recurrent approach is that way back when some good colleagues of mine trained these models in 2020, 2021, we trained these recurrent models to solve mazes. I'm not sure if you have seen that.

3:02It's really a toy problem, right? But it's kind of a very intriguing algorithmic task because you train these small recurrent models on mazes of sizes like 13 by 13. So it's really, it's a tiny pixelized maze, right? And it's like, you look at it, it's so trivial what the solution is. So you train these models on these 13 by 13 mazes. And then at test time, you just run them for more recurrences. so you do you put in more compute and in a test time you can solve arbitrarily large mazes like up the way like you have they have these funny pictures in the paper where it's like uh you know the whole screen is filled like every pixel is a different pixel in this maze it's like a i don't know 8000 by 8000 pixel maze and the model just puts in a lot of compute and just solves it right and it's an interesting example because there this model really has learned the algorithm to solve mazes or to solve this particular kind of maze.

3:57Right. And I think that was very motivating for us. So you see that like you could really learn an algorithm like that. But it was always a bit restricted to like, you know, you could only learn either you learn only mazes or you learn only addition or only chess problems. These were all a bit like maybe toy tasks or very like very limited domains where could see that these approaches would scale. And what I should also say is that, so what's interesting here about recurrence here or how that comes in is I think people are very familiar, or maybe this is maybe the ML person speaking, maybe are very familiar with recurrent neural networks.

4:41Right? And this is a bit different. In a recurrent neural network, what you have is that you have a sequence and the recurrence is basically that like every time you go forward in the sequence, you do another step of computation. And this is an efficient way to compute things about sequences. But what we're doing here is actually a bit different. What we're here doing is we're doing a recurrence in really in the depth of the model. And a good analogy here is that the model doesn't have a fixed amount of layers, but it has a small number of layers that it repeats. And this relates to lots of approaches in that field, like our weight sharing, sometimes called our weight tying.

5:21in this way the model really has like an arbitrary number of layers which really means that in comparison to a recurrent neural network where you always produce more input or more output to do more compute same way in a transformer every time you want to do more compute you have to produce more outputs with this sort of like level of this recurrent depth you can produce an arbitrary amount of compute before you compute the next output so it's sort of like the the compute axis is disentangled from the output axis or from the context of the model and that's what's so interesting i mean algorithmically about this kind of approach got it got it so the immediate thought in hearing the title um and the idea of recurrence is to think oh maybe you're doing something with rnns or like splicing that kind of recurrence into a transformer it's not really like that it's um manipulating the number of layers at test time to kind of flex the amount of compute that's being uh applied to a given generation yeah exactly exactly it's a bit um it relates a bit to these analogies for maybe for an rnn or for transformer where you um

6:44For example, I think the analogy for an RNN would be that you actually don't move to the next token in the sequence, but you sort of like produce like filler outputs and then you keep on iterating on these filler outputs. That would be like where you really stop on the sequence and then you, on your point where you currently stopped, you keep on like recurring a bit and then you move forward. You mentioned how you had this paper that explored this recurrent approach as kind of background. But then I'm wondering, talk a little bit about kind of the impetus to apply that. It comes at an interesting time when, you know, we've got these reasoning models like the O1 series from OpenAI.

7:28You know, DeepSeek showed how you could take a very different approach to training those. You know, this paper came out shortly after that, but I'm suspecting you've been working on it for quite a while. So kind of situate it for us in time and, you know, what motivated you to go after this particular approach of reasoning? Yeah. So basically, this is all the way back in 2023, like Tom and me, Tom is the last author, we had just been working on a paper on training language models more efficiently. And then we've been talking a lot about these recurrent models and how they would be cool to do in language.

8:08And because we really think it would be an interesting language model that could also solve things like math and code better, if it could reason more. Because back then, people were also doing a lot of MOE models. And MOE is interesting because they're the other extreme. It's a model that has more parameters, but uses less of them. So that seems like optimal to maybe store information. And this is, in some conceptual way, it's the opposite. It's a model that has fewer parameters. So fewer parameters, but uses them over and over and over. Yes, exactly. Right. So we really wanted to figure out how we could make this paradigm of recurrent depth work with language and how we could get, how we can modify it, what a good like training objective would be to really scale this up.

8:54Because it really wasn't really shown before in Transformers at all. Like, or there's some, at least like from our, like from this like maze direction, these were always convolutional models. And it wasn't so clear how to put this in Transformer and be nice. There's some interesting parallel work on looped transformers, which goes in a related direction. But it didn't quite work as well for us, just scaling those up larger. And that relates a bit to this part that this actually also is, so like one part that's different here is that it's latent, where basically we actually have a few non-recurrent layers in the model.

9:34And then we have this recurrent part in the middle. And then we have a few more non-recurrent layers at the end. To really have this computation be in this centralized latent space. And then still have some layers that actually decode back into language. And have some layers at the start that decode from language. And then with the right objective, this really helped us a lot. And then we had some, let's say like 100 million parameter prototypes in 2023. and we're like okay what next we're academics right so we we teamed up with lots of people that know a lot more about um high performance computing and we wrote this big allocation grant for um insight is like some big allocation grant basically this compute that we then got at oak rich at oak rich national labs and but then like it's still this um it was quite an interesting journey because it's it's one thing to like train this model in pytorch and it's 100 million parameters and it's kind of nice right to actually seeing okay now we now there are 4 000 accelerators and it still has to work and they have to work together and this like funny little probabilistic objective that you wrote still has to scale and it's just to work here and it was a pretty interesting journey that i personally enjoyed a lot although it was a bit harrowing at times as well to actually make this work and run this on these cards right because we also um so this is um this was compute on frontier and frontier is the largest amd supercomputer in the world and just and then we were just given that amount of compute because the um the folks that inside were saying yeah here we have this you could have this and so um but uh 2024 was also the year where a lot of things were moving at amd and were just beginning to be supported for example uh flash attention like flash attention or something it's like a very basic algorithm for us nowadays but it took like a lot of time maybe like until like half the year maybe until like last july where i felt this was really was stable enough that we could reliably deploy it and run it and there was a lot back and forth and amd did a lot and we had some, maybe like a lot of issues, but we got around them and it did train in the end.

11:54But it was a whole journey maybe of like, maybe a year where the underlying little model didn't really change, right? Or like the, I think in the paper, that's like, yeah, right. Like in the paper, maybe that's like section three and section three is like, okay, 2023. And then all of 2024 is just section four, which is, oh, we trained it on 4 ,000 GPUs. Oh, wow. And that was kind of an interesting journey. And yeah, then it also got pretty exciting because, as you mentioned, we had these test-sam-computed reasoning models coming out in fall, and we still hadn't finished training the model. And we were saying, oh, oh, oh.

12:35And suddenly, this was in the news everywhere, and suddenly this was happening, and we still hadn't finished training the model. And then now we did finish training the model in early December last year, and that's also why we pretty much wrote up our results in January and published them. But it was a bit of a scary moment as the world caught up to us a bit in that, although that's also why the approach is a bit different. Yeah, to what degree did the way you position the results around reasoning, you know, was that impacted by the timing of when you finished the work? like was that specifically the way you were thinking about it going in or was it more like you've got this architecture let's try to scale it and see what it can do and then lo and behold like it produced results that were relevant to you know the zeitgeist at the time it was always this idea of reasoning and learning algorithms but maybe like that that framing is maybe like the more classical framing would be that it's learning an algorithm just be like this older line of papers would maybe call this learning an algorithm or a prior towards learning algorithmic reasoning to call this all just blanket reasoning maybe that's a more modern way of looking at it and the other part it's also very modern is that like in the intro now we clearly call this out that there's a um i think we made this term up but i think it fits quite well to compare sort of like verbalized reasoning which we which is how we how we um maybe the box So we put in these other approaches to what we're doing.

14:09And that's a modern distinction that we put in to claim, okay, what are we actually doing differently? And I really hadn't thought about it so much that it wasn't so clear to me that just long chain of thoughts could scale so well. But maybe this wasn't clear to many of us a year ago. But I think it's interesting to look at them in this dichotomy that one is really trying to learn these algorithms in a high dimensional space. so like before it really decodes into language and the other approach is really trying to scaffold reasoning on to language let's maybe maybe even back up a step and really dig into that because i think that that is an idea that um you know people are really excited about and that is really exciting this idea that uh with traditional language models you know the, you know, the high compute and reasoning that we're seeing, it's all like having the models think out loud.

15:09And in a sense, what this paper is proposing is that we can also have the models kind of think to themselves and then, you know, not necessarily generate, you know, all of that thinking as language. And so with that in mind, kind of talk about how you think about, you know the relationship between those two and you know the the relative merits of one versus the other and and how they might compare like i'm curious like how you think about all that yeah so maybe the first thing i think is interestingly i think this is quite natural for us as well maybe really as humans this is a bit of anthropomorphizing i hope that's fine but basically um i think it's clear that we can sometimes think longer about a problem without actually verbalizing the answer in our head.

15:58Like we don't always like think in steps in our head, especially if it's a problem that maybe relates to maybe spatial reasoning or maybe something that you have some intuitive, maybe motor planning or like planning reasoning or how you go from one place to another. These are often things that we think about and we also think about longer, but we don't really verbalize this kind of thinking. And then interestingly, we also like these models, we can still write things down, right? We can sort of extend our thinking process by verbalizing it, by writing things down in notebooks, in equations. In this way, we can externalize thinking and think even longer.

16:35But even as humans, we have these two axes in which we can scale, compute or thinking in this way. And I think that's kind of interesting, right? It's not necessarily you can do this or the other. But maybe as humans, we naturally do both of those. and these models also could do both. But what's interesting about thinking without verbalizing it is that maybe like from a very basic perspective, this just seems much more powerful because basically the model, now if we come a bit to like the theory of like what algorithms can transformers represent, basically the transformer does some calculation because it is like n layers deep.

17:25Let's say that's like 96 layers deep. So the transformers has some input. The model does some calculation and the next word comes out. And then in the next layer, it can maybe access these previous calculations, but again, only up to the 96th layer. And if the model were to do a calculation that's deeper than 96, it's kind of hard because it has to go into these previous attention parts. so like the the model is a bit limited in its depth you're kind of calling out this idea that like i mean in some ways it's parallel to the argument for test time compute in general it's like that argument is often stated as well you know look at us as humans like there are things that we you know think about and the answer comes to us very quickly and there are other things that we spend a lot of time thinking about.

18:19And so why should models have like just a fixed, you know, compute, fixed budget for, you know, response, which kind of led to the idea of test time compute to some degree and allowing models to spend more time on that compute. And what you're kind of suggesting is that there's like another dimension of that. And there's, there's another the dimension of it from a transformer architecture perspective like you're both time and compute limited and also architecturally limited and you the recurrence kind of gives you some flex in that other dimension i don't know like what the implications of that are but um it strikes me that there are some parallels there yeah yeah i think an interesting analogy to make this clearer is that in some way if we like extract a bit away from it what the transformer is doing if it's doing a very long chain of thought is that it has some hidden representation that are computed on what the next token should be.

19:25Then that whole hidden representation, which is a high dimensional vector, is thrown away and then one token is sampled out of it. And the next step, that token is embedded again to compute a whole other hidden representation, which is thrown away. And then a new token is embedded. Because we keep embedding and unembedding these representations that the model is providing. And in some way, maybe the only thing that this recurrent model is doing is that it's just not doing that. It's basically keeping, it keeps on recurring and keeps on refining this hidden representation instead of always like encoding and decoding it in every step.

20:05And so how does it, how does it know how much to do this, when to stop doing it, that kind of thing? Yeah. So interestingly, that's actually not something we have, we have done a lot in this work. So what we've done a lot is basically is show how to scale these models and how to train them to have this capability. But we actually found that the best way we could scale these models was just to train them with a randomized depth or randomized number of occurrences. Really? Okay. So basically there's some older work also called Universal Transformers, right? Because this whole thing relates to universal Turing completeness.

20:45and these models had halting sub-modules that were doing training you had to decide when to stop but those are pretty hard to scale those are pretty hard to train because you're kind of training a model to get better and at the same time you're training it to exit early and it's hard to bootstrap a model that has to get better and harder problems and has to know when to exit early and so one way we got around this was basically that like during training we have a log normal Poisson distribution. It's really a heavy tail distribution over possible steps. And we sample one of them and we just compute loss on this step and then we back propagate.

21:26And this is how we train the model. But then interestingly, at inference time, what we found is that we could just zero shot adaptive exits. For example, just by computing how much the state changes from step to step. We could just make a simple rule based on that and just exit based on the simple rule that we just put in zero shot afterwards. It was kind of exciting that this actually did work at all. It wasn't a guarantee that this. What's the specific state that you're applying that algorithm to? So we've tried this with several ones. We've tried this with the actual latent recurrence state and for example the L2 norm of that state to the previous state.

22:08But we've also done this in the output distribution. so the right we have like a well maybe we're at state 16 and so we compute the next token distribution and we can then we go to state step let's say 24 we can compute the next token distribution and then we can compare the colbyg-like divergence of these two probability distributions and if that is small then we just exit the idea being that the model is converging on some thought or idea exactly and interestingly like the second one seems a bit more practically better because what we observe is sometimes the model keeps on thinking in this in this recurrent space and we have some interesting graphics in the paper where like what is actually like how we can project that down in lower dimensions how it's thinking there but this thinking doesn't always reflect in the probability of the next token right that's sort of like a more stable state is actually what do you predict next right and it's like it's not it wasn't so clear to us beforehand, or maybe also just wasn't clear in general, whether the model would learn this specialization.

23:12Basically, this was a research question, was like, whether just with scale, would the model learn to specialize this recurrence, right? Because we have this picture in the paper how the model, for example, converges quickly when the next token prediction is easy, but the convergence is much slower if the next token prediction is hard, or if it's a token that's somewhat important for the question to understand the question. right and we just we totally didn't train for this we trained with um we take fixed sequence we take um training sequences and we assign them the same step budget for the entire sequence just because that makes sense for pre-training scale right and this like per token specialization just seemed to emerge with scale that the model actually converges quicker on easy tokens than on hard ones, which was pretty interesting.

24:04And we really didn't know it would work like this before we saw the final model. So elaborate on that. When you refer to per-token specialization, what do you mean by that? So basically, the way the model was pre-trained was not only do we pick a random step for every recurrence. But we also pick the same step count for all the tokens in an example sequence. Because basically for this to be parallelized, the model has to do the same amount of compute at every token. This is trained on sequences, 4 ,000 tokens long. So meaning that random number is at like a batch level or something. Exactly, right?

24:52We're actually saying, okay, this batch gets three steps, this batch gets eight steps, this batch gets 128 steps. But then we look at the test time and we actually see that the model now has this convergence graph that's very different on a per token level. Where on some tokens, it goes like green very quickly, which means like it converges very quickly. And in other ones, it just keeps on doing things for a longer time. And this specialization really just emerged during training. Okay, so let me make sure I'm understanding this. What you're saying is that you've trained it on just these random numbers on a per sequence level, like random number of steps.

25:32But then at test time, the model actually does apply more compute to more difficult tokens in this case. I was a bit more like dancing around saying like, does it use more compute for that? meaning because that's more nuanced it's more nuanced than that oh it's more like in the chart like in the chart in Hippie what we show actually is that we in that chart we actually run the model for the same number of compute right for all steps that's why this is like a rectangular chart and then on that chart we just see how the convergence like how easily like how quickly it's done in every token but to make that chart we actually run all the steps to completion and from that we imply that you could have exited early but the chart really is quadratic and it has a full compute so basically what I'm really saying is it can do that for most of the paper we haven't done this a lot yet maybe this is the most exact way to say it because we have this section about how you could do adaptive exits and you could exit early and that's section 6.3 in the paper and there we try this But that's also where we discovered that this is possible.

26:45And so a lot of the rest of the paper just hasn't had this yet. A lot of the rest of the paper is like, here's a query. I'm going to run the query with 32 steps. Here's a different query. I'm going to run it with 64 steps. So a lot of the paper hasn't caught up to this more adaptive paradigm yet. Where it's also something that we were surprised it would work at all. And which is something where I'm personally very excited. Thinking back to these halting modules of now that this model is trained, we can probably just fine tune these stopping modules into it again, because now the bootstrap problem is solved and now we can just fine tune these and have these exit conditions again that I learned.

27:25But we just haven't gotten to that yet or no one has, I think. I mean, even before you said this, like it was surprising to me that the random stopping worked at all. Like, and why is that better than just like, you know just maximizing um during training i'm maximizing the number of steps right uh yeah that's a good question right because they have been like they have been previous architectures for example do you not sure if you remember something like albert which was a weight shared bird variant so these and there's also some some utmoe models that have this um where basically they have recurrence but the recurrence is always fixed.

28:11It's always like you always do 8 times or you always do 16 times or something. And those also work but interestingly those models don't really generalize to fewer or more steps of compute. They really train like a fixed depth transformer where they lose performance once they go for more steps and they also don't work if they have fewer steps. and with this objective really try to generalize to maybe really an arbitrary number of steps so the model is stable and can really learn more mathematically you can really learn some fixed point or some steady state behavior and that's what this objective is trying to get it to do because otherwise you would maybe you try you put in more computer test time and then it's it would get worse again it would like veer off track and go into some direction it wasn't trained on and so like right because like actually it really is this very non-linear dynamical system that we're training here and we're trying to get it to do what we want it to do and to predict the next token and that's a bit it's a bit tricky to train and if it's for example if it's fixed then it behaves like it's a fixed transformer.

29:32Is there an interpretation of the paper that is searching in latent space for the next token? It's interesting. It's an interesting way to look at it. In some ways, yes. Especially because if you think about it, what actually is happening is that you have some initial state in the recurrence. And the initial state is always random. And then We can also visualize those and often these random states, they lie on like some, for example, if you plot PCA directions, then you will see these random states just somewhere in the middle because they don't have a significant component in any direction. So they sort of like lie in the middle.

30:14And then you see as the model recurs, it keeps on like moving these points away onto some target position that where it wants them at and where then it can predict the next token from them. and we have these kind of funny little spine charts in the end of the paper where you sort of like see this like middle column rising and then you have these trajectories going out at the end. I think this is somewhere in the appendix where you really see that like from like these like from random it then converges to this or finds this point from which to predict the next token. And what's also super exciting for us there is that this is not always just a point.

30:55It doesn't always just converge to a fixed point. Because sometimes you actually also observe the model actually rotating. It has these little orbits that it makes. And it's kind of interesting because then we have these, right? It's still very hard for me to actually visualize this 5 ,280 dimensional space. So I only can show like, right? I can only show these like first 40, first 80 PCA directions or something, right? But in those, then we can see like, okay, what is it doing? It's like often like there's like some trajectories and they all come together and then they do these little circles.

31:27And that's quite interesting, right? Because then like in later tokens, it can pick up on these circles or in these like these orbits that it had done in earlier steps. And it can do some calculation with that. Because we do know that models often implement additional algorithms with these sort of schemes where you have orbits that reinforce each other. And it was quite interesting for us, right? Because it really was an emergent behavior that it would implement some of these problems like this. It wasn't really trained to have these little orbits, but it just emerged from training that this was a good solution to predict the next tokens.

32:07there's a note in the paper that kind of compares this to diffusion models elaborate on that a bit from some perspective this really is a diffusion model I mean maybe this is a bit generic maybe because all recurrent models are or maybe both of them are recurrent models diffusion models are also recurrent models in a different way it's important to separate that from there have also been some very exciting language diffusion models recently but this one's a bit if you think about this as a diffusion model it would be that you pick one token the next token, right? And then you diffuse all the way until you found the next token and then you go one step forward and you diffuse again until you found the next token and you go one step forward you diffuse again until you found the next token this kind of analogy would work here and during development we actually tried to make this a bit more explicit because we also tried to not only start from a noisy state, but also inject noise proportional to how many steps you've done so far, like you would in a diffusion model.

33:18And this didn't really help doing pre-training, so we removed that part again. But it really shows this analogy that this recurrence is not so different from a diffusion model that would diffuse the next token. What is quite different is that the objective, the training objectives are very different. In a diffusion model also this is why discrete diffusion models or language diffusion was all a bit hard is ideally in a diffusion model you would want this process where you have a you have a continuous target and you add some noise to it and then you learn this denoising operation. That's really how large scale image diffusion models are trained.

33:59And in a language that's not so clear because you can try to denoise embeddings but it's like not so nice to do that there. And so the training objective here is still quite different from a diffusion model, but they are both related models in that. Yeah, they both diffuse the next token. Maybe that's the way to say it. Okay. So you mentioned that you started with a hundred million, I think, parameter model as kind of your test model. And then you ultimately scaled that up to a three and a half billion parameter model. did that all happen in one shot or was that an incremental process and maybe talk a little bit about the process?