Transcription
Much of the progress in large language models has been driven by scaling. A couple of years before ChatGPT was a thing, Jared Kaplan and co-authors from OpenAI released their preprint "Scaling Laws for Neural Language Models." If I was to really simplify their findings, then I would just say that increasing the size of your models helps your test loss go down, provided that you also increase the number of training tokens. And this makes intuitive sense. The larger your model, the more data you need to take advantage of those extra parameters. Otherwise, you wouldn't really be using your compute in an optimal manner, or worse, you could run into overfitting.
So, what is optimal? Let's say that you have a model that you're happy with, but you need to scale it up. The authors found that if you want to use your compute optimally, then an 8x increase in your model size should come with a 5x increase in your data set size. At least that holds true for the specific conditions reported in the paper. And I think it's fair to say that the language modeling community has somewhat diverged away from those specific conditions. But nonetheless, that guided OpenAI and the broader LLM community on how to budget, uh, for their ridiculously long training runs. If you increase your model size, then you're increasing the GPU hours that you need. And likewise, if you increase your data set size, it's the same deal. So if you have a particular test loss that you're chasing, then you can project the amount of compute that you need. Then you know how much to beg from your VCs.
So my question to you is, between your model size, data set size, and compute, what do you think is the limiting factor of these? If you think it's data, then you would be right. We have pretty much exhausted the internet. Pablo Vobos and co-authors show in their paper that there is a much slower growth in human-made internet-based data than what language models have been using. And this point was also stressed by Ilia Sotsskova's keynote at Europe's 2024. An upper bound on the data set size means that there is an upper bound on how much useful compute we have.
So a logical question might be, is there a way to decouple compute from model size or from data set size? Well, we've already seen Mixture of Experts being used quite effectively where scaling the model size doesn't necessarily change the compute. And this works perfectly well as long as you have more data to compensate. So, is there a way to decouple data from compute? If you thought of reasoning models, then you would be on the right track. And there are a few different ways to elicit reasoning from simply prompting your model to chain of thought with long rollouts and self-checking.
Say we ask our model a simple math question and we let it stream text. Most humans would probably have a mental lookup table to figure out that 6 * 4 is 24, and then you can divide that down by 2 to get 12. This brings on the first problem with reasoning models. You're forced to extend your context. The more context, the greater the risk of forgetting critical pieces of information. And I'm sure many of us have had these multi-turn conversations with language models that span days or months or however long, only to watch them hallucinate and forget basic information and just, I don't know, generally themselves. So while this math equation doesn't necessarily illustrate it, you could imagine that generating long blocks of tricky code will bring on its own special challenges.
Secondly, as your questions get tougher, then your model might need to undergo a few rollouts before hitting the right answer. And once it does hit that right answer, then we reward it. And we might pick some of the incorrect answers and supply a penalty for those. You might not want to penalize all of your wrong answers, as that could lead to imbalance. So, what could we possibly do to increase the chances of getting that correct answer? Well, you could increase your model size and your data set size, but then you run into a bit of a catch-22, right? Your reasoning process still has an upper limit set by the capability of your base model.
If I take a data set of multiple-choice questions and pass it through my base model and my post-trained model repeatedly up until 1,024 times, the y-axis will show me the accuracy of my model. And as long as I get at least one correct answer across the K samples, then we count it as a correct answer. And we can run this across multiple models. And you'll notice that the base model really does set a ceiling on performance. One school of thought is that reinforcement learning doesn't teach your model anything new. Rather, it just amplifies pre-existing knowledge that might be buried very deep within your model and all the while trying to suppress incorrect responses. But then the problem here is that the pre-trained model sets an upper limit on reasoning performance.
The third problem is that the model operates on your vocabulary. Why might that be a problem? My parents are from Iran, and in Persian culture, there's this kind of overly polite ritual of give and take. So, just as an example, say I go to my auntie's house, and she offers me tea, she offers me desserts. I will always decline out of respect. It doesn't matter how hungry or how thirsty I am. And she's going to insist like a dozen times, and I will say no every single time. This is so ingrained in Persian culture that it has its own dedicated word, "taarof." There is no equivalent word in the English language. And, and what that tells me is that reasoning in your vocabulary space cannot possibly be the most optimal mode of doing so, right? Because different concepts are going to take a different number of tokens to completely capture. And I suppose what especially sucks there is that you're completely underleveraging all of the trillion pre-training tokens that you have available to you for reasoning. Right? In pre-training, there are only two scaling dimensions, right? There's your model size and your data set size. Reasoning is just treated as this kind of post-hoc process. It's, it's really an afterthought.
With that, we get to the point of this video. What we've discovered here is that merging reasoning with pre-training addresses all of these problems and gives us a third axis of scaling. That requires an alternative architecture, which we present in our recent paper, "Scaling Latent Reasoning via Looped Language Models."
So let's start by addressing, what is a looped language model? A standard transformer will take an input and generate an output. In a looped transformer, we start the same way, but before generating the output token, the model is going to take the latent vector and pass it through an exit gate. It's going to ask, "Is this legit, or should we try again?" If the exit gate is happy, then we terminate. We move on to the next token. If the exit gate isn't happy, then that latent vector is going to be looped back around to the input of the model. And the process is repeated until the exit gate is satisfied. We're no longer operating on the vocabulary, and we don't need to generate a chain of tokens, which ultimately compresses your KV cache. Which means we don't need a pre-trained model because this is optimized during pre-training. You're suddenly taking advantage of the trillions of tokens available to you.
We've dropped four models in total: Oro 1.4 billion and 2.6 billion, each with their own thinking variants. And when comparing them against state-of-the-art language models that are considerably larger, we're effectively performing on par. And what you're looking at are the results of our 2.6 billion parameter model up against Qwen3 and Gemma 3. And it's worth noting that Gemma 3 is nearly 5x larger and still underperforms against Oro 2.6b. Qwen3 is three times larger, but it was also trained on almost three times more tokens.
This isn't the first time we've seen loop structures or dynamic reasoning at all. ChatGPT, for example, takes a varying amount of time to execute based on the complexity of your input prompt. It's just that much of the underlying details are opaque to us. Um, unless you're at OpenAI, then you really don't know what's going on under the hood. Though, the best guess we have is that it is running Chain of Thought just on the vocabulary of the model. Additionally, the Universal Transformer from 2019 applied looping, and we've seen variants emerge at small scales ever since then, but this is the first time that we've seen it pushed to truly industrial scales with 7.7 trillion training tokens. And while that is a lot, we still have some internet left.
Let's dig into some of the lower-level details. How does the early exit mechanism work? Well, let's assume that the model is trying to generate a token. Once an output embedding is generated, it's passed to an exit gate. And this is very simply a dense layer with a sigmoid activation. And that can sort of be interpreted as the gate's instantaneous probability of exiting at that step. The output of the sigmoid function will be bounded between 0 and 1. But is that good enough? Well, the problem is, if we were to loop four times through the model, each of these is going to have some probability of exiting. And when we add these up, it's not going to equal to 1. The most tempting approach could be to just apply softmax or whatever normalization hack that you prefer. But unfortunately, we can't forecast into the future. I don't know the probability of exiting at loop one until I've completed all four loops.
So, let's break it down with an example. We run the first loop, and our sigmoid function tells us there's a 32% chance of exiting. Excellent. Now, we run the next loop, and this time the sigmoid function might give us 51%. However, the fact that we even got to the second loop was a result of making it past the first loop. So, this 0.51 is conditioned on the survival of the first loop. And the probability that we survived up until this stage is just 1 minus 0.32, or 68%. So, if we want an unconditional probability, or the probability mass function, then we take the probability that we survived loop one and we multiply it with the probability that loop 2 is terminated. That is the survival of loop one multiplied by the exit probability of loop two. We accumulate the unconditional probabilities. And what's really nice about this is that it's automatically bounded between 0 and 1. This still doesn't guarantee a probability distribution because we're not guaranteed to actually hit one. But the thing is, if we reach the maximum number of allowable loops, then we can force an exit at that final step. Mathematically, we're assigning the remaining probability mass to the final loop. The unconditional probability is converted into a cumulative density function, and that value is then thresholded. If the CDF at a given loop is greater than the threshold, then we exit. And if you hit the final step, then by default, you'll be forced to exit and then proceed to the next token.
Our first implementation of this honestly didn't work whatsoever. So the model basically learned to reward hack. What you're seeing is the final loop dominating every other loop. And at first, we thought, "Damn, everything needs so many loops. We are total geniuses. We have solved AGI." But it turns out that's not quite the case. To understand why, let's take a quick look at how the model is trained. During training, we don't stop when the exit gate tells us to. We instead run a full rollout of all possible steps. We calculate what the loss would have been if we stopped at step one, step two, step three, and so on. Then we combine those losses into a weighted average. And the weight for each step is simply P_t, the probability that the model actually decides to exit there. But here's the problem. When we start training, pure randomness means that one specific exit gate is going to start with a slightly higher probability than the others. Let's just say that it's our final loop for now. And because that step has the highest probability, then on average, it's going to contribute the most to the total loss. The model then updates its weights to minimize the error, but it places heavy emphasis at that exit step. And this creates a self-reinforcing cycle. The model gets better at exiting at that step. It becomes more confident. So the probability of exiting at that final step goes up. And therefore, that loop, that final loop is going to now dominate the loss function even more. Eventually, the model ignores all other steps and collapses into always exiting at that one point. And this is what we observed during training. Whatever the exit loop was at the start would pretty much dominate the rest of the training iterations. And this is illustrated here by showing you the probability of exiting in this reward-hacked case.
The solution to this turned out to be quite simple. We encourage the model to spread out its probability across steps. This is done by adding an entropy regularization term to the loss function, which penalizes the model and adds to the overall loss if the distribution deviates away from a uniform distribution. I'll also do a variable substitution. P_t represents the probability of exiting at a step. But just so we're clear, it's parameterized by phi because there are learnable values in the exit gate. And the output is also conditioned on whatever the input token is. The first term represents the exit distribution. And the second term is your prior distribution. That's what we want to match. And we apply the KL divergence to encourage the exit distribution to match the prior distribution. Now, of course, not everything needs a loop. So we modulate the strength of this term with beta. This idea comes from PonderNet from Google DeepMind, where they applied a geometric distribution to encourage early exiting. But when we tested this out, we found that it led to undertraining of later steps. So we imposed a uniform prior instead, and these are the results that we got. The x-axis represents the number of training steps, and the y-axis is the loss, where lower is better. We sweep across all of these distributions, and ultimately, the uniform distribution does a better job than the geometric distribution.
At this point, the looping mechanism works effectively. But keep in mind, every step of that loop is still adding more computation and adding more memory. Each loop has to store its own KV cache. Looped models are kind of weird when it comes to KV caching. On the one hand, there are a bunch of constraints on how we can use the cache. And these constraints are different between training and inference, and between prefill and decoding. But on the other hand, the additional loops per token gives us more flexibility in how we can use the KV cache. So let's break it down. During training and prefill, the model has access to the full sequence of tokens. And if speed is our priority, which it often is, the fastest thing we can do goes like this. We take all of our tokens, denoted x_1 and x_2 and so on. We run the first loop in parallel for all tokens. And this means that the second token has access to the KV cache of the first token, but only for the first loop. We then run the second loop in parallel again for all tokens. The KV cache can effectively be passed forward through your sequence, but only up until the second loop. But this goes on. Some of you might see a bit of an issue here. Let's say that for token one, the model wanted to exit at loop 3. And what would honestly make more sense is to use the KV cache of all layers within the third loop and then pass that to all loops of the following token. But that would kill the parallel nature of training and prefill. You would have to run your model sequentially, token by token, and it would be so slow that it's impossible to train on enough tokens for decent performance at most of the things that we care about.
How about decoding during inference? Well, we have a few options here. Let's start with the default option, which is what all of the results in the paper use. We can't start inference for the second token until the first token has finished processing. And once we start processing token 2, then we could use the KV cache of the corresponding loop. The main reason is to stay consistent with how the model is trained. But we also tested three other cases: using the KV cache from the exit loop only, and this makes logical sense because the exit loop is what gives us the token that we end up using; averaging the KV cache from every loop before feeding it to the next token; and finally, using the KV cache from the first loop. Our tests showed the following results. Using the KV cache from the first loop sucked. All other cases appeared to perform quite similarly, which is kind of cool, given that two of these cases diverge from how we trained the model. So that's something that might be worth probing a little bit. And that covers what is probably the most interesting low-level details of the model.
What comes next is the full training pipeline. And, and this is where it turns into a bit of a monstrous engineering effort. And I've got to say, my PhD student, Ridger, he was, he was hustling. Like, he was, he was really grinding. He was waking up from nightmares about loss spikes. He'd get out of bed at midnight and make sure that the training run was still stable. So, yeah, I don't know. Send him some condolences on LinkedIn or whatever. But this is mostly pretty standard stuff, and I will defer you to the paper for specific details. But to save a little bit on compute, the first pre-training phase involved optimizing the 1.4 billion model on 3 trillion tokens. The model was then forked into two pathways. For the 2.6 billion parameter model, we duplicated the non-embedding layers of the smaller model. And so this kind of resembles a 2x loop pass. And then we relaxed the weights so that the larger model could be trained. And the training data also increased in quality as the phases progressed.
Now it's time for the results. We've shown you how far these base models can stack up against larger state-of-the-art models. From here, I want to show you how the thinking versions of the models go on more challenging datasets along with some really interesting tests that helped us probe some theoretical questions about Oro. Some of these benchmarks are Olympiad and competition-level math and considered quite challenging. For MMLU, we show the 10-pass accuracy and one-pass for all other benchmarks. We're comparing against Qwen3 and Deepseek Distilled. And once again, Oro is performing far better than equivalently sized models such as Qwen 1.7 billion and Deepseek Distilled 1.5 billion. And it's mostly on par with the 7 to 8 billion parameter variants. In particular, these are the benchmarks where Oro wins despite being approximately 1/3 of the size.
And that's all well and good, but we wanted to dig into the model and gain a better understanding of when looping helps. A reasonable question might be, is there a certain number of optimal loops, and is extrapolation beyond your number of training loops possible? So on these benchmarks, we trained on a maximum of four loops, and it was kind of cool to see that there are a couple of benchmarks that do benefit from looping beyond four steps. Although in other cases, the performance did degrade a little after a while, but it shows that overlooping is probably safer than underlooping, at least in these specific cases. On the more challenging benchmarks, we likewise swept across four loops and then we extrapolated out to eight loops, and it turns out that we hit optimal performance at 3 to 4 loops. Going beyond that caused rapid performance degradation. So, to be honest, it's hard to make any strong claims other than looping seems to help.
But wouldn't it be nice to understand why does looping help so much? And this leads me to the final set of results and what I personally think are the coolest results in the whole paper. If we want to make strong claims about when looping helps, then we need total control over the tasks that the model is evaluated on. And two tasks that might be worth testing include: one, can a looped language model memorize information more effectively? And two, can a looped language model understand and manipulate information more effectively? This is where the work by Turuan Alanju and Shaoli Shu on the physics of language models comes in handy. I personally interpret these two questions as memorization versus understanding. That helps me make the distinction clear, but it's more precise to refer to the first task as knowledge storage and extraction, and then the second task as knowledge manipulation. They propose testing both of these things using highly controlled synthetic datasets. And my description will simplify the process a bit, but it goes something like this.
For storage and extraction, the dataset consists of biographical information: their name, their birthday, their workplace, and so on. The test is then asking the model to recall some of these facts. We tested a large variety of models across multiple loops. First, a 1 million parameter model with one loop. Then we cycled over four loops. Damn, no improvement. We could try and increase the model size for a single loop and then test across four loops, but there's still no improvement. I don't know, maybe we need to train on larger datasets. We could go up to 50,000 samples, but we're really seeing that there is negligible variance between one loop and four loops. And this holds true across all parameter scales and all dataset scales that we tested. The conclusion: looping does not seem to help knowledge capacity. And that shouldn't be surprising because looping doesn't add parameters.
How about knowledge manipulation? This is where the model must go beyond retrieval and instead operate and reason on those stored facts. A few different operations are tested, but here's one example. So let's see our performance. On the x-axis, we have the number of training steps in thousands. On the y-axis, we have the accuracy. And one of the constraints is that no chain of thought is allowed. The model needs to immediately deliver the correct answer. With one loop, the accuracy saturates pretty quickly and doesn't go beyond 14%. Two loops give us a big improvement, and then when we go to four loops, that performance jumps even further. And so, what this shows is that knowledge manipulation is where the gains of looped language models come from. This extra looping provides more opportunities for more internal computation, and I think that's a really powerful result.
So we've seen inference-time scaling enhance the capability of commercial large language models. But they all seem to rely on teaching the models to think only after the base model has been trained. And what I think we've shown here is that the base model itself can do better. And if we're injecting multi-step thinking, or or looping, into the pre-training pipeline, then not only does that lift your performance, but we've also shown the types of tasks that can benefit from that. And this obviously helps the large language models, but I think there's a lot to be said about the small language models as well. The, the stuff that can only fit on a mobile device because, I don't know, you're memory constrained or something. Um, and I think that that's what looping can really help, right? It can help us lift the parameter efficiency and how we use the weights of a model. Now, I don't want to reach too much here, but this kind of echoes how the brain works, right? When we learn something new, we don't undergo neurogenesis, right? Rather, we just learn how to use the pre-existing neurons and synapses a little more effectively.