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

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

To ensure that the model can function when we scale up recurrent iterations at test-time, we randomly sample iteration counts during training, assigning a random number of iterations r r to every input sequence ( Schwarzschild et al., 2021b ) . We optimize the expectation of the loss function L L over random samples x x from distribution X X and random iteration counts r r from distribution Λ \Lambda . ℒ ⁡ ( θ ) = 𝔼 𝐱 ∈ X ​ 𝔼 r ∼ Λ ​ L ​ ( m θ ​ ( 𝐱 , r ) , 𝐱 ′ ) . \mathcal{L}(\theta)=\mathbb{E}_{\mathbf{x}\in X}\mathbb{E}_{r\sim\Lambda}L\left(m_{\theta}(\mathbf{x},r),\mathbf{x}^{\prime}\right). Here, m m represents the model output, and 𝐱 ′ \mathbf{x}^{\prime} is the sequence 𝐱 \mathbf{x} shifted left, i.e., the next tokens in the sequence 𝐱 \mathbf{x} . We choose Λ \Lambda to be a log-normal Poisson distribution . Given a targeted mean recurrence r ¯ + 1 \bar{r}+1 and a variance that we set to σ = 1 2 \sigma=\frac{1}{2} , we can sample from this distribution via τ \displaystyle\tau ∼ 𝒩 ⁡ ( log ⁡ ( r ¯ ) − 1 2 ​ σ 2 , σ ) \displaystyle\sim\mathcal{N}(\log(\bar{r})-\frac{1}{2}\sigma^{2},\sigma) (1) r \displaystyle r ∼ 𝒫 ⁡ ( e τ ) + 1 , \displaystyle\sim\mathcal{P}(e^{\tau})+1, (2) given the normal distribution 𝒩 \mathcal{N} and Poisson distribution 𝒫 \mathcal{P} , see Figure 3 . The distribution most often samples values less than r ¯ \bar{r} , but it contains a heavy tail of occasional events in which significantly more iterations are taken.

#### Truncated Backpropagation.

To keep computation and memory low at train time, we backpropagate through only the last k k iterations of the recurrent unit. This enables us to train with the heavy-tailed Poisson distribution Λ \Lambda , as maximum activation memory and backward compute is now independent of r r . We fix k = 8 k=8 in our main experiments. At small scale, this works as well as sampling k k uniformly, but with set fixed, the overall memory usage in each step of training is equal. Note that the prelude block still receives gradient updates in every step, as its output 𝐞 \mathbf{e} is injected in every step. This setup resembles truncated backpropagation through time, as commonly done with RNNs, although our setup is recurrent in depth rather than time ( Williams and Peng, 1990 ; Mikolov et al., 2011 ) .

## 4 Training a large-scale recurrent-depth Language Model

After verifying that we can reliably train small test models up to 10B tokens, we move on to larger-scale runs. Given our limited compute budget, we could either train multiple tiny models too small to show emergent effects or scaling, or train a single medium-scale model. Based on this, we prepared for a single run, which we detail below.

### 4.1 Training Setup

We describe the training setup, separated into architecture, optimization setup and pretraining data. We publicly release all training data, pretraining code, and a selection of intermediate model checkpoints.

#### Pretraining Data.

Given access to only enough compute for a single large scale model run, we opted for a dataset mixture that maximized the potential for emergent reasoning behaviors, not necessarily for optimal benchmark performance. Our final mixture is heavily skewed towards code and mathematical reasoning data with (hopefully) just enough general webtext to allow the model to acquire standard language modeling abilities. All sources are publicly available. We provide an overview in Figure 4 . Following Allen-Zhu and Li (2024) , we directly mix relevant instruction data into the pretraining data. However, due to compute and time constraints, we were not able to ablate this mixture. We expect that a more careful data preparation could further improve the model’s performance. We list all data sources in Appendix C .

#### Tokenization and Packing Details.

We construct a vocabulary of 65536 65536 tokens via BPE ( Sennrich et al., 2016 ) , using the implementation of Dagan (2024) . In comparison to conventional tokenizer training, we construct our tokenizer directly on the instruction data split of our pretraining corpus, to maximize tokenization efficiency on the target domain. We also substantially modify the pre-tokenization regex (e.g. of Dagan et al. (2024) ) to better support code, contractions and LaTeX. We include a <|begin_text|> token at the start of every document. After tokenizing our pretraining corpus, we pack our tokenized documents into sequences of length 4096. When packing, we discard document ends that would otherwise lack previous context, to fix an issue described as the “grounding problem” in Ding et al. (2024) , aside from several long-document sources of mathematical content, which we preserve in their entirety.

#### Architecture and Initialization.

We scale the architecture described in Section 3 , setting the layers to ( 2 , 4 , 2 ) (2,4,2) , and train with a mean recurrence value of r ¯ = 32 \bar{r}=32 . We mainly scale by increasing the hidden size to h = 5280 h=5280 , which yields 55 55 heads of size of 96 96 . The MLP inner dimension is 17920 17920 and the RMSNorm ε \varepsilon is 10 − 6 10^{-6} . Overall this model shape has about 1.5 1.5 B parameters in non-recurrent prelude and head, 1.5 1.5 B parameters in the core recurrent block, and 0.5 0.5 B in the tied input embedding.

At small scales, most sensible initialization schemes work. However, at larger scales, we use the initialization of Takase et al. (2024) which prescribes a variance of σ h 2 = 2 5 ​ h \sigma_{h}^{2}=\frac{2}{5h} . We initialize all parameters from a truncated normal distribution (truncated at 3 ​ σ 3\sigma ) with this variance, except all out-projection layers, where the variance is set to σ out 2 = 1 5 ​ h ​ l \sigma_{\textnormal{out}}^{2}=\frac{1}{5hl} , for l = l P + r ¯ ​ l R + l C l=l_{P}+\bar{r}l_{R}+l_{C} the number of effective layers, which is 132 for this model. As a result, the out-projection layers are initialized with fairly small values ( Goyal et al., 2018 ) . The output of the embedding layer is scaled by h \sqrt{h} . To match this initialization, the state s 0 s_{0} is also sampled from a truncated normal distribution, here with variance σ s 2 = 2 5 \sigma_{s}^{2}=\frac{2}{5} .

#### Locked-Step Sampling.

To enable synchronization between parallel workers, we sample a single depth r r for each micro-batch of training, which we synchronize across workers (otherwise workers would idle while waiting for the model with the largest r r to complete its backward pass). We verify at small scale that this modification improves compute utilization without impacting convergence speed, but note that at large batch sizes, training could be further improved by optimally sampling and scheduling independent steps r r on each worker, to more faithfully model the expectation over steps in Equation 1 .

#### Optimizer and Learning Rate Schedule.

We train using the Adam optimizer with decoupled weight regularization ( β 1 = 0.9 \beta_{1}=0.9 , β 2 = 0.95 \beta_{2}=0.95 , η = 5 × 10 − 4 \eta=$5\text{\times}{10}^{-4}$ ) ( Kingma and Ba, 2015 ; Loshchilov and Hutter, 2017 ) , modified to include update clipping ( Wortsman et al., 2023b ) and removal of the ε \varepsilon constant as in Everett et al. (2024) . We clip gradients above 1 1 . We train with warm-up and a constant learning rate ( Zhai et al., 2022 ; Geiping and Goldstein, 2023 ) , warming up to our maximal learning rate within the first 4096 4096 steps.

### 4.2 Compute Setup and Hardware

We train this model using compute time allocated on the Oak Ridge National Laboratory’s Frontier supercomputer. This HPE Cray system contains 9408 compute nodes with AMD MI250X GPUs, connected via 4xHPE Slingshot-11 NICs. The scheduling system is orchestrated through SLURM. We train in bfloat16 mixed precision using a PyTorch-based implementation ( Zamirai et al., 2021 ) .

#### Device Speed and Parallelization Strategy.

Nominally, each MI250X chip 3 3 3 Technically, each node contains 4 dual-chip MI250X cards, but its main software stack (ROCm runtime) treats these chips as fully independent. achieves 192 TFLOP per GPU ( AMD, 2021 ) . For a single matrix multiplication, we measure a maximum achievable speed on these GPUs of 125 TFLOP/s on our software stack (ROCM 6.2.0, PyTorch 2.6 pre-release 11/02) ( Bekman, 2023 ) . Our implementation, using extensive PyTorch compilation and optimization of the hidden dimension to h = 5280 h=5280 achieves a single-node training speed of 108.75 TFLOP/s, i.e. 87% AFU (“Achievable Flop Utilization”). Due to the weight sharing inherent in our recurrent design, even our largest model is still small enough to be trained using only data (not tensor) parallelism, with only optimizer sharding ( Rajbhandari et al., 2020 ) and gradient checkpointing on a per-iteration granularity. With a batch size of 1 per GPU we end up with a global batch size of 16M tokens per step, minimizing inter-GPU communication bandwidth.

When we run at scale on 4096 GPUs, we achieve 52-64 TFLOP/s per GPU, i.e. 41%-51% AFU, or 1-1.2M tokens per second. To achieve this, we wrote a hand-crafted distributed data parallel implementation to circumvent a critical AMD interconnect issue, which we describe in more detail in Section A.2 . Overall, we believe this may be the largest language model training run to completion in terms of number of devices used in parallel on an AMD cluster, as of time of writing.

#### Training Timeline.

Training proceeded through 21 segments of up to 12 hours, which scheduled on Frontier mostly in early December 2024. We also ran a baseline comparison, where we train the same architecture but in a feedforward manner with only 1 pass through the core/recurrent block. This trained with the same setup for 180B tokens on 256 nodes with a batch size of 2 per GPU. Ultimately, we were able to schedule 795B tokens of pretraining of the main model. Due to our constant learning rate schedule, we were able to add additional segments “on-demand”, when an allocation happened to be available.

### 4.3 Importance of Norms and Initializations at Scale

At small scales all normalization strategies worked, and we observed only tiny differences between initializations. The same was not true at scale. The first training run we started was set up with the same block sandwich structure as described above, but parameter-free RMSNorm layers, no embedding scale γ \gamma , a parameter-free adapter A ⁡ ( 𝐬 , 𝐞 ) = 𝐬 + 𝐞 A(\mathbf{s},\mathbf{e})=\mathbf{s}+\mathbf{e} , and a peak learning rate of 4 × 10 − 4 4\text{\times}{10}^{-4} . As shown in Figure 5 , this run (“Bad Run 1”, orange), quickly stalled.

While the run obviously stopped improving in training loss (left plot), we find that this stall is due to the model’s representation collapsing ( Noci et al., 2022 ) . The correlation of hidden states in the token dimension quickly goes to 1.0 (middle plot), meaning the model predicts the same hidden state for every token in the sequence. We find that this is an initialization issue that arises due to the recurrence operation. Every iteration of the recurrence block increases token correlation, mixing the sequence until collapse.

We attempt to fix this by introducing the embedding scale factor, switching back to a conventional pre-normalization block, and switching to the learned adapter. Initially, these changes appear to remedy the issue. Even though token correlation shoots close to 1.0 at the start (“Bad Run 2”, green), the model recovers after the first 150 steps. However, we quickly find that this training run is not able to leverage test-time compute effectively (right plot), as validation perplexity is the same whether 1 or 32 recurrences are used. This initialization and norm setup has led to a local minimum as the model has learned early to ignore the incoming state 𝐬 \mathbf{s} , preventing further improvements.

In a third, and final run (“Main”, blue), we fix this issue by reverting back to the sandwich block format, and further dropping the peak learning rate to 4 × 10 − 5 4\text{\times}{10}^{-5} . This run starts smoothly, never reaches a token correlation close to 1.0, and quickly overtakes the previous run by utilizing the recurrence and improving with more iterations.

With our successful configuration, training continues smoothly for the next 750B tokens without notable interruptions or loss spikes. We plot training loss and perplexity at different recurrence steps in Figure 6 . In our material, we refer to the final checkpoint of this run as our “main model”, which we denote as Huginn-0125 4 4 4 \textipa /hu: gIn/, transl. “thought”, is a raven depicted in Norse mythology. Corvids are surprisingly intelligent for their size, and and of course, as birds, able to unfold their wings at test-time. .

## 5 Benchmark Results

We train our final model for 800B tokens, and a non-recurrent baseline for 180B tokens. We evaluate these checkpoints against other open-source models trained on fully public datasets (like ours) of a similar size. We compare against Amber ( Liu et al., 2023c ) , Pythia ( Biderman et al., 2023 ) and a number of OLMo 1&2 variants ( Groeneveld et al., 2024 ; AI2, 2024 ; Team OLMo et al., 2025 ) . We execute all standard benchmarks through the lm-eval harness ( Biderman et al., 2024 ) and code benchmarks via bigcode-bench ( Zhuo et al., 2024 ) .

### 5.1 Standard Benchmarks

Overall, it is not straightforward to place our model in direct comparison to other large language models, all of which are small variations of the fixed-depth transformer architecture. While our model has only 3.5B parameters and hence requires only modest interconnect bandwidth during pretraining, it chews through raw FLOPs close to what a 32B parameter transformer would consume during pretraining, and can continuously improve in performance with test-time scaling up to FLOP budgets equivalent to a standard 50B parameter fixed-depth transformer. It is also important to note a few caveats of the main training run when interpreting the results. First, our main checkpoint is trained for only 47000 steps on a broadly untested mixture, and the learning rate is never cooled down from its peak. As an academic project, the model is trained only on publicly available data and the 800B token count, while large in comparison to older fully open-source models such as the Pythia series, is small in comparison to modern open-source efforts such as OLMo, and tiny in comparison to the datasets used to train industrial open-weight models.

Disclaimers aside, we collect results for established benchmark tasks ( Team OLMo et al., 2025 ) in Table 1 and show all models side-by-side. In direct comparison we see that our model outperforms the older Pythia series and is roughly comparable to the first OLMo generation, OLMo-7B in most metrics, but lags behind the later OLMo models trained larger, more carefully curated datasets. For the first recurrent-depth model for language to be trained at this scale, and considering the limitations of the training run, we find these results promising and certainly suggestive that further research into latent recurrence as an approach to test-time scaling is warranted.

### 5.2 Math and Coding Benchmarks