Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach — part 3/5

Cleaned-text code points [30981, 46703). Bibliography omitted here; source.md includes it.

We also evaluate the model on math and coding. For math, we evaluate GSM8k ( Cobbe et al., 2021 ) (as zero-shot and in the 8-way CoT setup), MATH ( ( Hendrycks et al., 2021b ) with the Minerva evaluation rules ( Lewkowycz et al., 2022 ) ) and MathQA ( Amini et al., 2019 ) . For coding, we check MBPP ( Austin et al., 2021 ) and HumanEval ( Chen et al., 2021 ) . Here we find that our model significantly surpasses all models except the latest OLMo-2 model in mathematical reasoning, as measured on GSM8k and MATH. On coding benchmarks the model beats all other general-purpose open-source models, although it does not outperform dedicated code models, such as StarCoder2 ( Lozhkov et al., 2024 ) , trained for several trillion tokens. We also note that while further improvements in language modeling are slowing down, as expected at this training scale, both code and mathematical reasoning continue to improve steadily throughout training, see Figure 8 .

### 5.3 Where does recurrence help most?

How much of this performance can we attribute to recurrence, and how much to other factors, such as dataset, tokenization and architectural choices? In Table 4 , we compare our recurrent model against its non-recurrent twin, which we trained to 180B tokens in the exact same setting. In direct comparison of both models at 180B tokens, we see that the recurrent model outperforms its baseline with an especially pronounced advantage on harder tasks, such as the ARC challenge set. On other tasks, such as SciQ, which requires straightforward recall of scientific facts, performance of the models is more similar. We observe that gains through reasoning are especially prominent on GSM8k, where the 180B recurrent model is already 5 times better than the baseline at this early snapshot in the pretraining process. We also note that the recurrent model, when evaluated with only a single recurrence, effectively stops improving between the early 180B checkpoint and the 800B checkpoint, showing that further improvements are not built into the prelude or coda non-recurrent layers but encoded entirely into the iterations of the recurrent block.

Further, we chart the improvement as a function of test-time compute on several of these tasks for the main model in Figure 7 . We find that saturation is highly task-dependent, on easier tasks the model saturates quicker, whereas it benefits from more compute on others.

#### Recurrence and Context

We evaluate ARC-C performance as a function of recurrence and number of few-shot examples in the context in Figure 9 . Interestingly, without few-shot examples to consider, the model saturates in compute around 8-12 iterations. However, when more context is given, the model can reason about more information in context, which it does, saturating around 20 iterations if 1 example is provided, and 32 iterations, if 25-50 examples are provided, mirroring generalization improvements shown for recurrence ( Yang et al., 2024a ; Fan et al., 2025 ) . Similarly, we see that if we re-evaluate OBQA in Table 5 , but do not run the benchmark in the default lm-eval "closed-book" format and rather provide a relevant fact, our recurrent model improves significantly almost closing the gap to OLMo-2. Intuitively this makes sense, as the recurrent models has less capacity to memorize facts but more capacity to reason about its context.

### 5.4 Improvements through Weight Averaging

Due to our constant learning rate, we can materialize further improvements through weight averaging ( Izmailov et al., 2018 ) to simulate the result of a cooldown ( Hägele et al., 2024 ; DeepSeek-AI et al., 2024 ) . We use an exponential moving average starting from our last checkpoint with β = 0.9 \beta=0.9 , incorporating the last 75 checkpoints with a dilation factor of 7, a modification to established protocols ( Kaddour, 2022 ; Sanyal et al., 2024 ) . We provide this EMA model as well, which further improves GMS8k performance to 47.23 % 47.23\% flexible ( 38.59 % 38.59\% strict), when tested at r = 64 r=64 .

## 6 Recurrent Depth simplifies LLMs

Aside from encouraging performance in mathematical and code reasoning, recurrent-depth models turn out to be surprisingly natural tools to support a number of methods that require substantial effort with standard transformers. In the next section, we provide a non-exhaustive overview.

### 6.1 Zero-Shot Adaptive Compute at Test-Time

We have shown that the model is capable of varying compute on a per-query level, running the model in different recurrence modes. This is after all also how the model is trained, as in Equation 1 . However, it would be more efficient in practice to stop recurring early when predictions are easy, and only spend compute on hard decisions. Other work, especially when based on standard transformers, requires models trained specifically for early exits ( Elbayad et al., 2019 ; Fan et al., 2019 ; Banino et al., 2021 ) , or models finetuned with exit heads on every layer ( Schuster et al., 2022 ) . To test our model’s zero-shot exit abilities, we choose a simple exit criterion to evaluate convergence, the KL-divergence between two successive steps. If this divergence falls below 5 × 10 − 4 5\text{\times}{10}^{-4} , we stop iterating, sample the output token, and move to generate the next token.

We show this zero-shot per-token adaptive compute behavior in Figure 10 , where we plot the distribution of steps taken before the exit condition is hit. We do this for the first 50 questions from different MMLU categories, asked in free-form chat. Interestingly, the number of steps required to exit differs notably between categories, with the model exiting earlier on high school mathematics, but taking on average 3.5 steps more on moral scenarios. As a preliminary demonstration, we verify on MTBench that this adaptivity does not significantly impact performance in a conversational benchmark setting (standard: 5.63 5.63 , early exits: 5.56 5.56 see Appendix Table 6 ).

###### Remark 6.1 (What about missing KV-cache entries?) .

Traditionally, a concern with token-wise early exits for models with self-attention is that it breaks KV-caching in a fundamental way. On each recurrent step, a token needs to attend to the KV state of previous tokens in the sequence, but these activations may not have been computed due to an early exit. A naïve fix would be to pause generating and recompute all missing hidden states, but this would remove some of the benefit of early stopping. Instead, as in Elbayad et al. (2019) , we attend to the last, deepest available KV states in the cache. Because all recurrent KV cache entries are generated by the same K,V projection matrices from successive hidden states, they “match”, and therefore the model is able to attend to the latest cache entry from every previous token, even if computed at different recurrent depths.

### 6.2 Zero-Shot KV-cache Sharing

A different avenue to increase efficiency is to reduce the memory footprint of the KV-cache by sharing the cache between layers ( character.ai, 2024 ; Brandon et al., 2024 ) . Typically, transformers must be trained from scratch with this capability. However, as discussed in the previous section, we find that we can simply share KV-caches in our model with minimal impact to performance. We set a fixed KV-cache budget for the recurrence at every token k k , and at iteration i i , read and write the cache entry i mod k i\bmod k . For example, we set a maximum KV-cache budget of 16 steps, overwriting the KV-cache of the 1st step when executing the 17th step, and so forth. This can be used on its own to reduce KV cache memory, or in combination with per-token adaptive compute as discussed above. On MTBench, this does not reduce performance (cache budget of 4 4 : 5.86 5.86 , see Appendix Table 6 ).

### 6.3 Zero-Shot Continuous Chain-of-Thought

By attending to the output of later steps of previous tokens in the early steps of current tokens, as described in the KV-cache sharing section, we actually construct a computation that is deeper than the current number of recurrence steps. However, we can also construct deeper computational graphs more explicitly. Instead of sampling a random initial state 𝐬 0 \mathbf{s}_{0} at every generation step, we can warm-start with the last state 𝐬 r \mathbf{s}_{r} from the previous token. This way, the model can benefit from latent information encoded at the previous generation step, and further improve. As shown in Figure 10 , this reduces the average number of steps required to converge by 1-2. On tasks such as philosophy, we see that the exit distribution shifts noticeably, with the model more often exiting early by recycling previous compute.

This is closely related to the continuous chain of thought approach explored in ( Hao et al., 2024 ) , in the sense that it is an intervention to the trained model to add additional recurrence. To achieve a similar behavior in fixed-depth transformers, Hao et al. (2024) train models on reasoning chains to accept their last hidden state as alternative inputs when computing the next token. Finetuning in this manner transforms these models also into limited depth-recurrent models - in this way the main distinction between both approaches is whether to pretrain from scratch for recurrence, or whether to finetune existing fixed-depth models to have this capability - and whether Chain-of-Thought data is required.

### 6.4 Zero-Shot Self-Speculative Decoding

Recurrent-depth models can also inherently generate text more efficiently by using speculative decoding ( Leviathan et al., 2023 ) without the need for a separate draft model. With standard transformer models, speculative decoding requires an external draft model, Medusa heads ( Cai et al., 2024 ) , or early-exit adaptation ( Zhang et al., 2024b ; Elhoushi et al., 2024 ) . Zhang et al. (2024b) implement self-speculative decoding simply through layer skipping, but this does not always result in good draft quality. In comparison, our model can naturally be run with fewer iterations to draft the next N N tokens in the generated sequence, which can then be verified with any desired number of iterations M > N M>N later. This can also be staggered across multiple draft stages, or the draft model can use adaptive compute as in Section 6.1 . Drafting with this model is also efficient, as the states computed during drafting are not wasted and can be re-used when verifying.

## 7 What Mechanisms Emerge at Scale in Recurrent-Depth Models

Finally, what is the model doing while recurring in latent space? To understand this question better, we analyze the trajectories { 𝐬 i } i = 1 r \{\mathbf{s}_{i}\}_{i=1}^{r} of the model on a few qualitative examples. We are especially interested in understanding what patterns emerge, simply by training this model at scale. In comparison to previous work, such as Bai et al. (2019) , where the training objective directly encodes a prior that pushes trajectories to a fixed point, we only train with our truncated unrolling objective.

Figure 11 shows the norm distance ‖ 𝐬 i − 𝐬 ∗ ‖ ||\mathbf{s}_{i}-\mathbf{s}^{*}|| between each 𝐬 i \mathbf{s}_{i} in a trajectory and an approximate limit point 𝐬 ∗ \mathbf{s}^{*} computed with 128 iterations. We show the sentence top to bottom and iterations from left to right. We clearly see that convergence behavior depends on context. We see that key parts of the question, and the start of the model response, are “deliberated” much more in latent space. The context dependence can also be seen in the different behavior among the three identical tokens representing each of the three dots. Also note that the distance to 𝐬 ∗ \mathbf{s}^{*} does not always decrease monotonically (e.g. for school ); the model may also trace out complicated orbits in its latent trajectory while processing information, even though this is not represented explicitly in our training objective.

We look at trajectories for select tokens in more detail in Figure 12 . We compute a PCA decomposition of latent trajectories over all tokens in a sequence, and then show several individual trajectories projected onto the first six PCA directions. See the appendix for more examples. Many tokens simply converge to a fixed point, such as the token in the top row. Yet, for harder questions, such as in the 2nd row 5 5 5 This is the token ”3” in a GSM8k test question that opens with Claire makes a 3 egg omelette . , the state of the token quickly falls into an orbit pattern in all three pairs of PCA directions. The use of multi-dimensional orbits like these could serve a similar purpose to periodic patterns sometimes observed in fixed-depth transformers trained for arithmetic tasks ( Nanda et al., 2022 ) , but we find these patterns extend far beyond arithmetic for our model. We often also observe the use of orbits on tokens such as “makes” (see Figure 16 ) or “thinks” that determine the structure of the response.

Aside from orbits, we also observe the model encoding particular key tokens as “sliders”, as seen in the middle of the bottom row in Figure 12 (which is the token “wrong”, from the same message as already shown in Figure 11 ). In these motions the trajectory noticeably drifts in a single direction, which the model could use to implement a mechanism to count how many iterations have occurred.

The emergence of structured trajectories in latent space gives us a glimpse into how the model performs its computations. Unlike the discrete sequential chain of reasoning seen in verbalized chain-of-thought approaches, we observe rich geometric patterns including orbits, convergent paths, and drifts - means to organize its computational process spatially. This suggests the model is independently learning to leverage the high-dimensional nature of its latent space to implement reasoning in new ways.

#### Path Independence.

We verify that our models maintain path independence, in the sense of Anil et al. (2022) , despite their complex, learned dynamics, which we discussed prior (see also the additional examples in Appendix Figure 22 ). When re-initializing from multiple starting states 𝐬 0 \mathbf{s}_{0} , the model moves in similar trajectories, exhibiting consistent behavior. The same orbital patterns, fixed points, or directional drifts emerge regardless of initialization.

## 8 Related Work Overview

The extent to which recurrence is a foundational concept of machine learning is hard to overstate ( Amari, 1972 ; Hopfield, 1982 ; Braitenberg, 1986 ; Gers and Schmidhuber, 2000 ; Sutskever et al., 2008 ) . Aside from using recurrence to move along sequences, as in recurrent neural networks, it was understood early to also be the key to adaptive computation ( Schmidhuber, 2012 ; Graves, 2017 ) . For transformers, recurrence was applied in Dehghani et al. (2019) , who highlight the aim of recurrent depth to model universal , i.e. Turing-complete, machines ( Graves et al., 2014 ) . It was used at scale (but with fixed recurrence) in Lan et al. (2019) and an interesting recent improvement in this line of work are described in Tan et al. (2023) ; Abnar et al. (2023) , Mathur et al. (2024) and Csordás et al. (2024) . Schwarzschild et al. (2021b) ; Bansal et al. (2022) ; Bear et al. (2024) and McLeish et al. (2024) show that depth recurrence is advantageous when learning generalizable algorithms when training with randomized unrolling and input injections. Recent work has described depth-recurrent, looped , transformers and studied their potential benefits with careful theoretical and small-scale analysis ( Giannou et al., 2023 ; Gatmiry et al., 2024 ; Yang et al., 2024a ; Fan et al., 2025 ) .