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

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

Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach

Jonas Geiping 1 Sean McLeish 2 Neel Jain 2 John Kirchenbauer 2 Siddharth Singh 2 Brian R. Bartoldson 3 Bhavya Kailkhura 3 Abhinav Bhatele 2 Tom Goldstein 2

###### Abstract

We study a novel language model architecture that is capable of scaling test-time computation by implicitly reasoning in latent space. Our model works by iterating a recurrent block, thereby unrolling to arbitrary depth at test-time. This stands in contrast to mainstream reasoning models that scale up compute by producing more tokens. Unlike approaches based on chain-of-thought, our approach does not require any specialized training data, can work with small context windows, and can capture types of reasoning that are not easily represented in words. We scale a proof-of-concept model to 3.5 billion parameters and 800 billion tokens. We show that the resulting model can improve its performance on reasoning benchmarks, sometimes dramatically, up to a computation load equivalent to 50 billion parameters.

Model : huggingface.co/tomg-group-umd/huginn-0125 Code and Data : github.com/seal-rg/recurrent-pretraining

## 1 Scaling by Thinking in Continuous Space

Humans naturally expend more mental effort solving some problems than others. While humans are capable of thinking over long time spans by verbalizing intermediate results and writing them down, a substantial amount of thought happens through complex, recurrent firing patterns in the brain, before the first word of an answer is uttered.

Early attempts at increasing the power of language models focused on scaling model size, a practice that requires extreme amounts of data and computation. More recently, researchers have explored ways to enhance the reasoning capability of models by scaling test time computation. The mainstream approach involves post-training on long chain-of-thought examples to develop the model’s ability to verbalize intermediate calculations in its context window and thereby externalize thoughts.

However, the constraint that expensive internal reasoning must always be projected down to a single verbalized next token appears wasteful; it is plausible that models could be more competent if they were able to natively “think” in their continuous latent space. One way to unlock this untapped dimension of additional compute involves adding a recurrent unit to a model. This unit runs in a loop, iteratively processing and updating its hidden state and enabling computations to be carried on indefinitely. While this is not currently the dominant paradigm, this idea is foundational to machine learning and has been (re-)discovered in every decade, for example as recurrent neural networks, diffusion models, and as universal or looped transformers.

In this work, we show that depth-recurrent language models can learn effectively, be trained in an efficient manner, and demonstrate significant performance improvements under the scaling of test-time compute. Our proposed transformer architecture is built upon a latent depth-recurrent block that is run for a randomly sampled number of iterations during training. We show that this paradigm can scale to several billion parameters and over half a trillion tokens of pretraining data. At test-time, the model can improve its performance through recurrent reasoning in latent space, enabling it to compete with other open-source models that benefit from more parameters and training data. Additionally, we show that recurrent depth models naturally support a number of features at inference time that require substantial tuning and research effort in non-recurrent models, such as per-token adaptive compute, (self)-speculative decoding, and KV-cache sharing. We finish out our study by tracking token trajectories in latent space, showing that a number of interesting computation behaviors simply emerge with scale, such as the model rotating shapes in latent space for numerical computations.

## 2 Why Train Models with Recurrent Depth?

Recurrent layers enable a transformer model to perform arbitrarily many computations before emitting a token. In principle, recurrent mechanisms provide a simple solution for test-time compute scaling. Compared to a more standard approach of long context reasoning OpenAI (2024) ; DeepSeek-AI et al. (2025) , latent recurrent thinking has several advantages. • Latent reasoning does not require construction of bespoke training data. Chain-of-thought reasoning requires the model to be trained on long demonstrations that are constructed in the domain of interest. In contrast, our proposed latent reasoning models can train with a variable compute budget, using standard training data with no specialized demonstrations, and enhance their abilities at test-time if given additional compute.

• Latent reasoning models require less memory for training and inference than chain-of-thought reasoning models. Because the latter require extremely long context windows, specialized training methods such as token-parallelization Liu et al. (2023a) may be needed.

• Recurrent-depth networks perform more FLOPs per parameter than standard transformers, significantly reducing communication costs between accelerators at scale. This especially enables higher device utilization when training with slower interconnects.

• By constructing an architecture that is compute-heavy and small in parameter count, we hope to set a strong prior towards models that solve problems by “thinking”, i.e. by learning meta-strategies, logic and abstraction, instead of memorizing. The strength of recurrent priors for learning complex algorithms has already been demonstrated in the “deep thinking” literature Schwarzschild et al. (2021b) ; Bansal et al. (2022) ; Schwarzschild et al. (2023) .

On a more philosophical note, we hope that latent reasoning captures facets of human reasoning that defy verbalization, such as spatial thinking, physical intuition or (motor) planning. Over many iterations of the recurrent process, reasoning in a high-dimensional vector space would enable the deep exploration of multiple directions simultaneously, instead of linear thinking, leading to a system capable of exhibiting novel and complex reasoning behavior.

Scaling compute in this manner is not at odds with scaling through extended (verbalized) inference scaling ( Shao et al., 2024 ) , or scaling parameter counts in pretraining ( Kaplan et al., 2020 ) , we argue it may build a third axis on which to scale model performance.

#### ———————— Table of Contents ————————

• Section 3 introduces our latent recurrent-depth model architecture and training objective.

• Section 4 describes the data selection and engineering of our large-scale training run on Frontier, an AMD cluster.

• Section 5 reports benchmark results, showing how the model improves when scaling inference compute.

• Section 6 includes several application examples showing how recurrent models naturally simplify LLM usecases.

• Section 7 visualizes what computation patterns emerge at scale with this architecture and training objective, showing that context-dependent behaviors emerge in latent space, such as “orbiting” when responding to prompts requiring numerical reasoning.

## 3 A scalable recurrent architecture

In this section we will describe our proposed architecture for a transformer with latent recurrent depth, discussing design choices and small-scale ablations. A diagram of the architecture can be found in Figure 2 . We always refer to the sequence dimension as n n , the hidden dimension of the model as h h , and its vocabulary as the set V V .

### 3.1 Macroscopic Design

The model is primarily structured around decoder-only transformer blocks ( Vaswani et al., 2017 ; Radford et al., 2019 ) . However these blocks are structured into three functional groups, the prelude P P , which embeds the input data into a latent space using multiple transformer layers, then the core recurrent block R R , which is the central unit of recurrent computation modifying states 𝐬 ∈ ℝ n × h \mathbf{s}\in\mathbb{R}^{n\times h} , and finally the coda C C , which un-embeds from latent space using several layers and also contains the prediction head of the model. The core block is set between the prelude and coda blocks, and by looping the core we can put an indefinite amount of verses in our song.

Given a number of recurrent iterations r r , and a sequence of input tokens 𝐱 ∈ V n \mathbf{x}\in V^{n} these groups are used in the following way to produce output probabilities 𝐩 ∈ ℝ n × | V | \mathbf{p}\in\mathbb{R}^{n\times|V|} 𝐞 \displaystyle\mathbf{e} = P ⁡ ( 𝐱 ) \displaystyle=P(\mathbf{x}) 𝐬 0 \displaystyle\mathbf{s}_{0} ∼ 𝒩 ⁡ ( 𝟎 , σ 2 ​ I n ⋅ h ) \displaystyle\sim\mathcal{N}(\mathbf{0},\sigma^{2}I_{n\cdot h}) 𝐬 i \displaystyle\mathbf{s}_{i} = R ( 𝐞 , 𝐬 i − 1 ) for i ∈ { 1 , … , r } \displaystyle=R(\mathbf{e},\mathbf{s}_{i-1})\qquad\textnormal{for}\quad i\in\{1,\dots,r\} 𝐩 \displaystyle\mathbf{p} = C ⁡ ( 𝐬 r ) , \displaystyle=C(\mathbf{s}_{r}), where σ \sigma is some standard deviation for initializing the random state. This process is shown in Figure 2 . Given an init random state 𝐬 0 \mathbf{s}_{0} , the model repeatedly applies the core block R R , which accepts the latent state 𝐬 i − 1 \mathbf{s}_{i-1} and the embedded input 𝐞 \mathbf{e} and outputs a new latent state 𝐬 i \mathbf{s}_{i} . After finishing all iterations, the coda block processes the last state and produces the probabilities of the next token.

This architecture is based on deep thinking literature, where it is shown that injecting the latent inputs 𝐞 \mathbf{e} in every step ( Bansal et al., 2022 ) and initializing the latent vector with a random state stabilizes the recurrence and promotes convergence to a steady state independent of initialization, i.e. path independence ( Anil et al., 2022 ) .

#### Motivation for this Design.

This recurrent design is the minimal setup required to learn stable iterative operators. A good example is gradient descent of a function E ⁡ ( 𝐱 , 𝐲 ) E(\mathbf{x},\mathbf{y}) , where 𝐱 \mathbf{x} may be the variable of interest and 𝐲 \mathbf{y} the data. Gradient descent on this function starts from an initial random state, here 𝐱 0 \mathbf{x}_{0} , and repeatedly applies a simple operation (the gradient of the function it optimizes), that depends on the previous state 𝐱 k \mathbf{x}_{k} and data 𝐲 \mathbf{y} . Note that we need to use 𝐲 \mathbf{y} in every step to actually optimize our function. Similarly we repeatedly inject the data 𝐞 \mathbf{e} in our set-up in every step of the recurrence. If 𝐞 \mathbf{e} was provided only at the start, e.g. via 𝐬 0 = 𝐞 \mathbf{s}_{0}=\mathbf{e} , then the iterative process would not be stable 1 1 1 Stable in the sense that R R cannot be a monotone operator if it does not depend on 𝐞 \mathbf{e} , and so cannot represent gradient descent on strictly convex, data-dependent functions, ( Bauschke et al., 2011 ) , as its solution would depend only on its boundary conditions.

The structure of using several layers to embed input tokens into a hidden latent space is based on empirical results analyzing standard fixed-depth transformers ( Skean et al., 2024 ; Sun et al., 2024 ; Kaplan et al., 2024 ) . This body of research shows that the initial and the end layers of LLMs are noticeably different, whereas middle layers are interchangeable and permutable. For example, Kaplan et al. (2024) show that within a few layers standard models already embed sub-word tokens into single concepts in latent space, on which the model then operates.

###### Remark 3.1 (Is this a Diffusion Model?) .

This iterative architecture will look familiar to the other modern iterative modeling paradigm, diffusion models ( Song and Ermon, 2019 ) , especially latent diffusion models ( Rombach et al., 2022 ) . We ran several ablations with iterative schemes even more similar to diffusion models, such as 𝐬 i = R ⁡ ( 𝐞 , 𝐬 i − 1 ) + 𝐧 \mathbf{s}_{i}=R(\mathbf{e},\mathbf{s}_{i-1})+\mathbf{n} where 𝐧 ∼ 𝒩 ⁡ ( 𝟎 , σ i ​ I n ⋅ h ) \mathbf{n}\sim\mathcal{N}(\mathbf{0},\sigma_{i}I_{n\cdot h}) , but find the injection of noise not to help in our preliminary experiments, which is possibly connected to our training objective. We also evaluated and 𝐬 i = R i ​ ( 𝐞 , 𝐬 i − 1 ) \mathbf{s}_{i}=R_{i}(\mathbf{e},\mathbf{s}_{i-1}) , i.e. a core block that takes the current step as input ( Peebles and Xie, 2023 ) , but find that this interacts badly with path independence, leading to models that cannot extrapolate.

### 3.2 Microscopic Design

Within each group, we broadly follow standard transformer layer design. Each block contains multiple layers, and each layer contains a standard, causal self-attention block using RoPE ( Su et al., 2021 ) with a base of 50000 50000 , and a gated SiLU MLP ( Shazeer, 2020 ) . We use RMSNorm ( Zhang and Sennrich, 2019 ) as our normalization function. The model has learnable biases on queries and keys, and nowhere else. To stabilize the recurrence, we order all layers in the following “sandwich” format, using norm layers n i n_{i} , which is related, but not identical to similar strategies in ( Ding et al., 2021 ; Team Gemma et al., 2024 ) : 𝐱 l ^ = \displaystyle\hat{\mathbf{x}_{l}}= n 2 ​ ( 𝐱 l − 1 + Attn ​ ( n 1 ​ ( 𝐱 l − 1 ) ) ) \displaystyle n_{2}\left(\mathbf{x}_{l-1}+\textnormal{Attn}(n_{1}(\mathbf{x}_{l-1}))\right) 𝐱 l = \displaystyle\mathbf{x}_{l}= n 4 ​ ( 𝐱 l ^ + MLP ​ ( n 3 ​ ( 𝐱 l ^ ) ) ) \displaystyle n_{4}\left(\hat{\mathbf{x}_{l}}+\textnormal{MLP}(n_{3}(\hat{\mathbf{x}_{l}}))\right) While at small scales, most normalization strategies, e.g. pre-norm, post-norm and others, work almost equally well, we ablate these options and find that this normalization is required to train the recurrence at scale 2 2 2 Note also that technically n 3 n_{3} is superfluous, but we report here the exact norm setup with which we trained the final model. .

Given an embedding matrix E E and embedding scale γ \gamma , the prelude block first embeds input tokens 𝐱 \mathbf{x} as γ ​ E ​ ( 𝐱 ) \gamma E(\mathbf{x}) , and then to applies l P l_{P} many prelude layers with the layout described above.

Our core recurrent block R R starts with an adapter matrix A : ℝ 2 ​ h → ℝ h A:\mathbb{R}^{2h}\to\mathbb{R}^{h} mapping the concatenation of 𝐬 i \mathbf{s}_{i} and 𝐞 \mathbf{e} into the hidden dimension h h ( Bansal et al., 2022 ) . While re-incorporation of initial embedding features via addition rather than concatenation works equally well for smaller models, we find that concatenation works best at scale. This is then fed into l R l_{R} transformer layers. At the end of the core block the output is again rescaled with an RMSNorm n c n_{c} .

The coda contains l C l_{C} layers, normalization by n c n_{c} , and projection into the vocabulary using tied embeddings E T E^{T} .

In summary, we can summarize the architecture by the triplet ( l P , l R , l C ) (l_{P},l_{R},l_{C}) , describing the number of layers in each stage, and by the number of recurrences r r , which may vary in each forward pass. We train a number of small-scale models with shape ( 1 , 4 , 1 ) (1,4,1) and hidden size h = 1024 h=1024 , in addition to a large model with shape ( 2 , 4 , 2 ) (2,4,2) and h = 5280 h=5280 . This model has only 8 8 “real” layers, but when the recurrent block is iterated, e.g. 32 times, it unfolds to an effective depth of 2 + 4 ​ r + 2 = 132 2+4r+2=132 layers, constructing computation chains that can be deeper than even the largest fixed-depth transformers ( Levine et al., 2021 ; Merrill et al., 2022 ) .

### 3.3 Training Objective

#### Training Recurrent Models through Unrolling.