Source-linked AI summary
Teaching Pretrained Language Models to Think Deeper with Retrofitted Recurrence
Sean McLeish, Ang Li, John Kirchenbauer, Dayal Singh Kalra, Brian R. Bartoldson, Bhavya Kailkhura, Avi Schwarzschild, Jonas Geiping, Tom Goldstein, Micah Goldblum
TL;DR
The paper asks how pretrained non-recurrent language models can be converted into depth-recurrent models efficiently. It retrofits recurrence through continued pretraining, using pretrained initialization, recurrence scheduling, and data curricula. The resulting models improve mathematical performance at a given compute budget while retaining strong language-modeling performance, though broader scale and domain generalization remain open.
Problem
The paper studies how to convert existing pretrained non-recurrent language models into depth-recurrent models while addressing the training cost of recurrent computation.
Method
The paper retrofits pretrained models by looping a recurrent block and uses pretrained initialization, a curriculum over recurrent depths, and healing followed by task-specific data.
Results
Retrofitted models improve GSM8K and MATH performance over base models while maintaining high accuracy on broad language-modeling benchmarks.
Takeaways & Limitations
Retrofitting recurrence is an efficient way to train performant depth-recurrent models while decoupling parameter count from test-time computation.
Takeaways & Limitations
The experiments focus on approximately 1B-parameter and 50B-token scales and primarily strengthen mathematical reasoning, so broader scale and domain generalization remain unverified.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper retrofits depth recurrence into pretrained language models to improve test-time compute efficiency without increasing parameter count. Continued pretraining with recurrence curricula preserves language-modeling performance while improving mathematical reasoning under compute constraints.
- 1 INTRODUCTION: Depth recurrence enables additional inference computation without increasing memory consumption or context size, and can support latent reasoning without bespoke reasoning traces.The paper positions recurrence as an alternative to generating many verbalized reasoning tokens or candidate solutions.
- 1 INTRODUCTION: Pretrained fixed-depth models can be converted into depth-recurrent transformers through continued pretraining, with early layers forming a prelude, later layers forming a recurrent block and coda, and intermediate layers removed.The recurrent block is looped by feeding its output back into itself, with a linear adapter combining recurrent and prelude representations.
- 1 INTRODUCTION: Initializing recurrent models from pretrained weights is more efficient than random initialization, improving loss and benchmark accuracy per training FLOP.This knowledge-transfer advantage is demonstrated for Llama-3.2-1B initialization.
- 1 INTRODUCTION: With a well-formed data curriculum, recurrence improves math-task accuracy while maintaining high accuracy across a broad suite of language-modeling benchmarks.The paper reports this pattern for retrofitted TinyLlama, OLMo, and Llama models using Common Crawl math data.
- 1 INTRODUCTION: A recurrence curriculum slowly increases the average number of recurrent iterations during training, maintaining performance while improving training speed.The curriculum is presented as a central training-efficiency mechanism for retrofitted models.
2 RELATED WORK
The related work spans recurrent transformers, model surgery, and latent reasoning. The paper distinguishes its approach from prior conversions by targeting deeper test-time recurrence without distillation or auxiliary adapters.
- 2 RELATED WORK: Prior work studies recurrent transformers with shared weights, fixed repetition counts, and adaptive-depth mechanisms for increasing computation efficiency.Universal transformers and several recurrent-transformer variants motivate depth recurrence as an established architectural direction.
- 2 RELATED WORK: Model-surgery methods make post-hoc architectural changes, including converting pretrained transformers into recurrent models.Relevant prior work includes looped-model finetuning and conversions of pretrained GPT-2, OPT, and other transformer checkpoints.
- 2 RELATED WORK: Bae et al. convert pretrained transformers using only two or three recursions, retain the base model shape, and use low-rank adapters to recover performance.Their approach also reports performance decreases with more recurrences, limiting test-time compute scaling.
- 2 RELATED WORK: The paper positions its method as avoiding distillation and auxiliary adapters while benefiting from additional test-time computation.This contrast is stated directly against Bae et al.’s conversion approach.
- 2 RELATED WORK: Latent-reasoning work includes architectures whose recurrence supports reasoning but may lack the ability to extrapolate recurrence at test time.The cited HRM discussion describes ablations reducing the architecture to a simpler recurrent model without recurrence extrapolation.
3 EXPERIMENTAL SETUP
The experimental setup defines a prelude, recurrent block, and coda, then adapts pretrained layers into this recurrent structure. Training cost is measured by FLOPs because recurrence reuses fixed parameters while increasing computation with iteration count.
- 3 EXPERIMENTAL SETUP: The recurrent architecture defines P as the prelude, R as the recurrent block, and C as the coda, with embeddings in P and unembeddings in C.The recurrent block begins with a 2h-to-h linear adapter that processes the concatenation of the prelude output and recurrent state.
- 3 EXPERIMENTAL SETUP: The model computes e = P(x), initializes s_0 with Gaussian noise, iterates s_i = R(e, s_{i−1}), and produces p = C(s_r).The equations describe the recurrent forward pass for r iterations.
- 3 EXPERIMENTAL SETUP: A configuration such as (2, 4, 2) specifies the numbers of transformer layers in the prelude, recurrent block, and coda.The paper uses tuple notation to describe recurrent-model depth allocation.
- 3 EXPERIMENTAL SETUP: Model surgery selects early pretrained layers for the prelude and later layers for the recurrent block and coda, removing intervening layers.For a 22-layer model with configuration (4, 8, 4), the selected layers are [0–3], [10–17], and [18–21].
- 3 EXPERIMENTAL SETUP: Increasing recurrent iterations raises computation while leaving the number of distinct trainable parameters fixed.The paper therefore reports recurrent-model size at different depths using FLOPs rather than effective parameter counts.
- 3 EXPERIMENTAL SETUP: Training FLOPs account for recurrent computation separately because gradients are backpropagated through at most the last eight recurrent iterations.The standard fixed-depth approximation is FLOPs = 6ND, while recurrent models split parameters according to gradient recording.
4 TRAINING RECURRENT LANGUAGE MODELS
The study develops efficient procedures for converting pretrained fixed-depth transformers into depth-recurrent models. Pretrained initialization, recurrence curricula, optimizer choice, and post-training experiments together show improved training efficiency and stronger mathematical reasoning under comparable compute.
- 4.1 EFFICIENTLY INITIALIZING RECURRENT TRANSFORMERS: Pretrained Llama initialization yields lower loss and earlier Hellaswag accuracy gains than random initialization for recurrent models.By training step 1000, the pretrained model is already leveraging recurrence, whereas randomly initialized recurrences remain near random accuracy.
- 4.2 SCHEDULING RECURRENCES: Linearly increasing the mean recurrent depth to 32 improves loss efficiency in both training steps and FLOPs, with the largest gain measured against FLOPs.Scheduling the maximum depth can be better for FLOPs alone but worse for steps, while a 1-sqrt curriculum performs comparably to linear scheduling for TinyLlama.
- 4.3 HOW TO RETROFIT RECURRENCE: Under the same training FLOP budget, depth-recurrent models with fewer parameters achieve higher accuracy on math problems than their non-recurrent parent models.This pattern is reported across TinyLlama, OLMo, and Llama experiments.
- 4.3 HOW TO RETROFIT RECURRENCE: Muon is more stable than AdamW for recurrent training, achieving lower loss while avoiding the loss spikes and NaNs seen with AdamW.The optimizer difference is much smaller for non-recurrent TinyLlama models.
- 4.3.2 RECURRENT MODELS ARE EFFICIENT TO TRAIN: Increasing inference recurrences lets recurrent models improve accuracy over fixed-depth baselines while remaining competitive per inference FLOP with fewer trainable parameters.The reported scaling behavior appears in both GSM8K and MATH experiments.
5 DISCUSSION
Depth-recurrent models offer parameter-efficient separation of train-time and test-time compute, but training remains more expensive per parameter. The paper mitigates this through pretrained initialization, recurrence scheduling, and data curricula, while leaving scaling and adaptive computation open.
- Depth-recurrent models decouple train-time and test-time compute but require more training FLOPs per parameter.
- Pretrained initialization, recurrence scheduling, and data curricula reduce the training-efficiency gap for recurrent models.
- Adaptive recurrence that assigns computation to problem difficulty remains an unresolved goal.
- The method’s evidence is concentrated at 1B parameters and 50B training tokens, mainly for mathematical reasoning.
- Future work includes better layer selection, larger-scale validation, and applying recurrence beyond mathematics.
C.1 MODEL SURGERY ABLATIONS
Model-surgery ablations examine initialization, embedding scale, layer selection, and architecture. Pretrained initialization and carefully chosen recurrent layers are favored, while extending non-recurrent regions yields little additional benefit.
- At least approximately 950 billion tokens are estimated for randomly initialized and pretrained loss curves to intersect, likely an underestimate.
- Pretrained embedding-scale choices have minimal impact in the smaller randomly initialized model.
- The selected layer configuration performs better than ShortGPT’s least-impactful-layer removal for training depth-recurrent models.
- Early layers work best in the prelude, while later layers are preferred for the recurrent block and coda.
- Increasing recurrent-block depth consistently improves accuracy over the fixed-depth model, whereas extending the prelude and coda gives minimal gains.
C.2 SCHEDULING RECURRENCES ABLATIONS
The scheduling ablations compare curriculum duration, target depth, maximum backpropagation depth, and schedule shape. Scheduling the mean recurrence is compute- and data-efficient, while 1-sqrt and linear schedules perform similarly per FLOP.
- The curriculum increases recurrence in a staircase-like progression from one recurrence to the target depth.
- Scheduling the mean of the depth distribution is efficient in both training steps and compute.
- Scheduling maximum backpropagation depth saves FLOPs but produces worse models when measured by training steps.
- The 1-sqrt and linear curricula are similar on a per-FLOP basis, motivating the choice of a 75% 1-sqrt schedule.
C.3.1 TINYLLAMA
TinyLlama recurrent models are evaluated across training and inference compute on GSM8K and MATH. They can outperform non-recurrent baselines during training and remain competitive at inference, with additional recurrences improving performance.
- The evaluation suite includes final-step accuracy and broader benchmark results for the models shown in the main experiments.
- Recurrent TinyLlama models efficiently outperform the non-recurrent baseline on MATH as training FLOPs increase.
- At inference, recurrent models are competitive with the fixed-depth baseline and can surpass it when given more recurrences.
- Recurrent models remain competitive in inference FLOPs for both GSM8K and MATH when plotted against effective parameters.
C.3.2 OLMO
For OLMo, depth recurrence improves reasoning efficiency on GSM8K and MATH: recurrent models outperform non-recurrent baselines during training and remain competitive at inference.
- The evaluation covers GSM8K and MATH, with inference-efficiency views replotted using effective parameters proportional to inference FLOPs.
- Training-step curves report GSM8K and MATH accuracy for train recurrences of 4, 8, 16, and 32.
- Table 4 reports final-step accuracy across a broad range of evaluations for the models shown in the main OLMo results.
- Recurrent OLMo models efficiently outperform the non-recurrent baseline when accuracy is plotted against training FLOPs.
- At inference, recurrent OLMo models remain competitive with fixed-depth baselines and can surpass them by using more recurrences.
C.3.3 LLAMA
The Llama experiments retrofit recurrence into a pretrained model by forming (4, 6, 4) configurations, then evaluate reasoning efficiency on GSM8K and MATH.
- The retrofit removes layers 4 and 5 from Llama-3.2-1B, leaving approximately 850 million parameters, or 87.5% of the pretrained model.
- Inference results are also replotted with effective parameters, treated as proportional to inference FLOPs, for both GSM8K and MATH.
- Training-step curves report GSM8K and MATH accuracy for train recurrences of 4, 8, 16, and 32.
- Llama recurrent models efficiently outperform the non-recurrent baseline when accuracy is measured against training FLOPs on GSM8K and MATH.
- The recurrent Llama models are competitive with fixed-depth baselines at inference and can outperform them as recurrence increases.
C.4 DATA MIXTURES
The data-mixture analysis extends recurrent-model evaluations across recurrences and compares training curricula, data quality, and broad benchmark performance.
- Table 6 extends earlier evaluations by including more test recurrences and Huginn-0125 results for comparison.
- Depth-recurrent models increase in accuracy over recurrence and achieve better accuracy with two-phase training, whereas the non-recurrent baseline slightly favors single-phase training.
- The experiments use specified optimizer settings, gradient clipping, batch sizes, and hardware configurations that vary across experimental sections.
- The appendix documents exact parameter counts for non-recurrent and recurrent models and identifies the layers taken from pretrained models.