Teaching Pretrained Language Models to Think Deeper with Retrofitted Recurrence

AI summary

Overview

Research area: Natural Language Processing — specifically efficient language model training, depth-recurrent transformer architectures, and test-time compute scaling.

Technical level: Advanced. The paper assumes familiarity with transformer internals (residual streams, grouped-query attention, layer norms), pretraining pipelines, and FLOP accounting.

Scope in one sentence: The paper studies how to cheaply convert already-pretrained, fixed-depth ~1B-parameter language models (TinyLlama-1.1B-intermediate-step-1431k-3T, OLMo-2-0425-1B, and Llama-3.2-1B) into depth-recurrent models by cutting out middle layers and re-training a looped core block, and measures whether the resulting models beat simply continued-training the original model at the same compute budget.

What This Paper Is About

Test-time compute scaling normally means generating more tokens — longer chain-of-thought traces or many sampled candidate answers. A less common alternative is depth recurrence: looping the same block of layers more times at inference, which spends more compute without growing memory cost, context length, or parameter count. The problem is that depth-recurrent models have so far been pretrained from scratch with many recurrent passes, which slows training dramatically (the authors cite Geiping et al. (2025), who pretrained Huginn-0125 from scratch on 800 billion tokens). This paper's goal is to retrofit recurrence into models that already exist, so the expensive pretraining is not repeated.

Key Contributions

1. Pretrained initialization beats random initialization for recurrent models. The authors show that initializing a depth-recurrent model's parameters from a pretrained fixed-depth model is significantly more efficient (in loss and benchmark accuracy per training FLOP) than the scalable random initialization of Takase et al. (2023).
2. A curriculum over recurrent depth. They propose slowly increasing the mean number of recurrent iterations during training, which preserves performance while improving training speed.
3. Converting TinyLlama, OLMo, and Llama into recurrent models with math data. Using Common Crawl math data, the resulting recurrent models achieve better GSM8K and MATH performance than the base models.
4. A "healing" phase plus a data curriculum. Because layers are removed during conversion, the authors introduce an initial period of minimal distribution shift to recover basic language modeling ability before switching to task-specific data.

Main Findings

• Pretrained init is dramatically more sample-efficient than random init. Two models were trained for roughly 120 billion tokens on FineWeb-Edu with a mean of 32 recurrences: a (2,4,2) configuration initialized from Llama-3.2-1B layers versus a random (Takase et al., 2023) initialization. The Llama-initialized model starts at a higher loss but consistently reaches lower losses; by training step 1000 it is already using recurrence to raise accuracy, whereas the randomly initialized model is at roughly random accuracy on Hellaswag. Extrapolating the loss curves in log-linear space suggests it would take at least approximately 950 billion tokens for the curves to intersect.
• Scheduling recurrences saves compute. Linearly scheduling the mean of the recurrence distribution up to a final mean of 32 over a long training period lowers validation loss as a function of both training steps and FLOPs. Measured over 48 hours on 4 MI300A GPUs (approximately 1e20 FLOPs), the FLOP-efficiency gain from scheduling is described as large, while the step-efficiency gain is small and positive. A 1-sqrt schedule was found to be as good as a linear one for TinyLlama.
• Muon beats AdamW for recurrent post-training. Using a learning rate of 5e-5 for AdamW and 0.001 for Muon, Muon was more stable and achieved lower loss; the AdamW run ended early because the loss spiked and became NaN. For the non-recurrent TinyLlama model the difference between optimizers was minimal. The authors also compare against a variant they call AdamW*, used by Geiping et al. (2025), which adds update clipping, removes the epsilon constant, and uses a different decoupling method than PyTorch's AdamW.
• Model surgery removes layers and keeps most parameters. The (4,8,4) TinyLlama configuration drops 6 layers (layers 4, 5, 6, 7, 8, and 9 with 0-indexing), leaving approximately 700 million parameters — 72.7% of the parent non-recurrent TinyLlama model. For OLMo-2-0425-1B and Llama-3.2-1B, (4,6,4) configurations remove 2 layers (layers 4 and 5), leaving approximately 900 million parameters, or 87.5% of the pretrained model's parameters. Notably, the (7,8,7) comparison showed that removing layers was efficient even from the prelude and coda.
• Recurrent models win on math at equal training FLOPs. Training (4,8,4) TinyLlama and non-recurrent models on approximately 50 billion tokens of Nemotron-CC-Math-v1-4plus, with a 1-sqrt curriculum for the first 75% of training and constant mean recurrences after, the recurrent models matched the non-recurrent baseline at small budgets and kept improving as more compute was invested, while the non-recurrent model's accuracy plateaued. This held on GSM8K for TinyLlama and on MATH for OLMo.
• More inference recurrences convert into more accuracy. When accuracy is plotted against inference recurrences, the recurrent models improve significantly over the fixed-depth baseline by using more test-time compute despite having fewer trainable parameters.
• Stronger pretrained weights transfer better. The accuracy achieved by Llama and OLMo based models was higher than that of TinyLlama based models. For GSM8K, larger training recurrences yielded larger per-FLOP gains for OLMo and Llama, but this trend did not hold for TinyLlama.
• A healing phase recovers language modeling ability. Training (4,8,4) TinyLlama models for 26 billion tokens on an even mix of FineWeb-Edu, Nemotron-Pretraining-SFT-v1-General, and Nemotron-Pretraining-SFT-v1-Math ("single phase") left the final recurrent model worse than the non-recurrent model. With "two phase" training — 26 billion tokens of FineWeb-Edu followed by the same mix, totaling 52 billion tokens — the recurrent model improved by 5% on Arc-Challenge (reaching 37.7% at test recurrence 32), while the non-recurrent model's accuracy barely differed between phases. The non-recurrent baseline showed only a small increase from two-phase training, which the authors attribute to there being no surgery to heal.
• Broad benchmark performance is largely preserved. In Table 1, the two-phase (4,8,4) recurrent model at test recurrence 32 scores 65.2 Arc-E, 37.7 Arc-C, 60.4 Hellaswag, 60.5 Winogrande, 44.8 MMLU, 73.6 PIQA, 40.0 OpenbookQA, 51.2 GSM8K, and 14.2 MATH. The text states the depth-recurrent model is only outperformed by the non-recurrent model on MMLU, and by less than one standard error. Compared with Huginn-0125, the authors report an MMLU score over 12% higher and GSM8K performance 10% higher than Geiping et al.'s published evaluation results.
• Recurrence is defined with input injection. The prelude P produces embeddings e from the input; the recurrent block R begins with a linear adapter that maps the concatenation of e and the previous state s(i-1) from 2h back down to width h; the initial state s0 is sampled from a Gaussian; and the coda C produces the output distribution after r recurrences.

Methodology in Plain English

The authors take a pretrained transformer and physically split it into three parts: a prelude (early layers, including embeddings), a recurrent block (middle layers, which get looped), and a coda (later layers, including the unembedding). Layers in between that are not kept are discarded, which is why the resulting model has fewer distinct parameters than its parent. Each pass through the recurrent block takes the concatenation of the prelude's output and the previous pass's output, projects it back to the model width with a linear adapter, and runs the block again. At the first step, random noise stands in for the previous output.

Training proceeds in a loop-aware mode. The number of recurrences used per training step is sampled from a Poisson-Lognormal distribution, whose mean is called the "training recurrence"; gradients are backpropagated through at most the last 8 passes to keep memory bounded. Training FLOPs are computed with a modified rule — (6N₁ + 2N₂)D, where N₁ counts parameters whose gradients are recorded and N₂ counts parameters used only in the forward pass — rather than the standard 6ND.

To speed things up, the authors schedule that mean over training (starting small, ramping to a final value like 32), analogously to gradual stacking for non-recurrent models. They then continue pretraining on math data, using a 1-sqrt schedule for the first 75% of training followed by a constant mean. Evaluations use a single in-context example; GSM8K is scored with flexible extract accuracy and MATH with the Minerva criteria via Math-Verify. For the general-capability experiments they mix FineWeb-Edu with Nemotron-Pretraining-SFT-v1-General and Nemotron-Pretraining-SFT-v1-Math, filtering out rows generated by reinforcement-learning reasoning models such as DeepSeek-R1 and the user/assistant tags, and they compare single-phase against healing-then-mixed two-phase training.

Why This Matters

Impact on research. The paper argues that converting fixed-depth models to recurrent ones can be treated as a small addition to a training pipeline, analogous to extending context length late in pretraining. It reports that recurrence improves math reasoning at a given compute budget compared with post-training the original non-recurrent model, and it contrasts its approach with prior work (Bae et al., 2024) that required low-rank adapters and distillation and saw performance decrease when recurring more, and with concurrent work (Koishekenov et al., 2025) that keeps all parameters, omits input injection, uses a fixed recurrence count, and does not report compute.

Real-world applications (as motivated by the paper's framing):

• Deploying reasoning models with a per-query compute dial: the same weights can be run with few recurrences for easy queries and many for hard ones, without increasing memory or context size.
• Upgrading already-deployed 1B-class checkpoints instead of paying for a from-scratch recurrent pretraining run.
• Systems where inference memory, not just latency, is the binding constraint, since deeper recurrence does not grow memory consumption or context length.
• Domains where reasoning traces are expensive to collect, since the model is trained on standard data rather than bespoke chain-of-thought traces.

Industry relevance. The recipe is applied to widely used open-weight families (TinyLlama, OLMo, Llama-3.2-1B), code and converted models are released at github.com/mcleish7/retrofitting-recurrence and on Hugging Face, and the compute-efficiency framing (loss and accuracy per FLOP, with training done on AMD MI300A accelerators in bfloat16 with Flash Attention) speaks directly to practitioners deciding whether depth recurrence is worth the extra FLOPs per parameter.

Future Directions

• Recurring deeper at test time than in training. The paper identifies as unsolved how to build depth-recurrent models that can recur beyond training-time depths to solve harder problems than those seen during training.
• Native adaptivity / learned stopping criteria. A related goal is giving recurrent models a built-in mechanism that automatically assigns the right amount of recurrence to a problem based on difficulty, letting models think deeply on hard problems while solving easy ones quickly.
• Better layer selection during model surgery. Appendix Figures 12 and 13 present the search over which layers to keep and which to discard, and the authors state that future work could identify a more optimal method than their selection, which was compared against the ShortGPT pruning method.
• Scaling up. The experiments are at the 1B-parameter and 50B-token scales, and the authors state that more experimentation is required to verify the method generalizes to much larger model and data scales.
• Beyond mathematics. The work focuses primarily on strengthening mathematical capabilities, and the authors call for extending depth recurrence to other reasoning-intensive domains.

Target Audience

Researchers and engineers working on efficient transformer training, test-time compute scaling, and latent reasoning; practitioners who want to squeeze more reasoning performance out of existing ~1B-parameter open-weight checkpoints without increasing inference memory; and readers following the line of work on depth-recurrent language models, particularly Geiping et al. (2025) and concurrent conversion methods such as Koishekenov et al. (2025). Beginners will find the architecture section heavy going, since the model definition is given in formal notation and the results depend on FLOP accounting conventions.

Authors’ abstract

Recent advances in depth-recurrent language models show that recurrence can decouple train-time compute and parameter count from test-time compute. In this work, we study how to convert existing pretrained non-recurrent language models into depth-recurrent models. We find that using a curriculum of recurrences to increase the effective depth of the model over the course of training preserves performance while reducing total computational cost. In our experiments, on mathematics, we observe that converting pretrained models to recurrent ones results in better performance at a given compute budget than simply post-training the original non-recurrent language model.