Teaching Pretrained Language Models to Think Deeper with Retrofitted Recurrence

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.Source: §1 · p. 1
• 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.Source: Figure 1 · p. 1, §1 · p. 2
• 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.Source: §1 · p. 2
• 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.Source: §1 · p. 2
• 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.Source: §1 · p. 2

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.Source: §2 · p. 2
• 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.Source: §2 · p. 3
• 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.Source: §2 · p. 3
• 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.Source: §2 · p. 3
• 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.Source: §2 · p. 3

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.Source: §3 · p. 3
• 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.Source: §3 · p. 3
• 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.Source: §3 · p. 3
• 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].Source: §3 · pp. 3–4
• 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.Source: §3 · p. 4
• 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.Source: §3 · p. 4

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.Source: §4.1 · p. 5, Figure 2 · p. 5

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.Source: §4.2 · pp. 5–6, Figure 3 · p. 6

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.Source: §4.3 · p. 6, §4.3.2 · p. 8
• 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.Source: §4.3.1 · p. 7

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.Source: §4.3.2 · p. 8

• Depth-recurrent models decouple train-time and test-time compute but require more training FLOPs per parameter.Source: §5 · p. 10
• Pretrained initialization, recurrence scheduling, and data curricula reduce the training-efficiency gap for recurrent models.Source: §5 · p. 10
• Adaptive recurrence that assigns computation to problem difficulty remains an unresolved goal.Source: §5 · p. 10
• The method’s evidence is concentrated at 1B parameters and 50B training tokens, mainly for mathematical reasoning.Source: §5 · p. 11
• Future work includes better layer selection, larger-scale validation, and applying recurrence beyond mathematics.Source: §5 · p. 11

• At least approximately 950 billion tokens are estimated for randomly initialized and pretrained loss curves to intersect, likely an underestimate.Source: §C.1 · p. 19, Figure 9 · p. 19
• Pretrained embedding-scale choices have minimal impact in the smaller randomly initialized model.Source: Figure 11 · p. 20
• The selected layer configuration performs better than ShortGPT’s least-impactful-layer removal for training depth-recurrent models.Source: §C.1.1 · p. 20, Figure 13 · p. 21
• Early layers work best in the prelude, while later layers are preferred for the recurrent block and coda.Source: Figure 12 · p. 21
• Increasing recurrent-block depth consistently improves accuracy over the fixed-depth model, whereas extending the prelude and coda gives minimal gains.Source: Figure 14 · p. 22, Figure 15 · p. 22

• The curriculum increases recurrence in a staircase-like progression from one recurrence to the target depth.Source: §C.2 · p. 23, Figure 16 · p. 23
• Scheduling the mean of the depth distribution is efficient in both training steps and compute.Source: Figure 17 · p. 23
• Scheduling maximum backpropagation depth saves FLOPs but produces worse models when measured by training steps.Source: Figure 18 · p. 24
• The 1-sqrt and linear curricula are similar on a per-FLOP basis, motivating the choice of a 75% 1-sqrt schedule.Source: Figure 19 · p. 24

• The evaluation suite includes final-step accuracy and broader benchmark results for the models shown in the main experiments.Source: §C.3.1 · p. 24, Table 3 · p. 25
• Recurrent TinyLlama models efficiently outperform the non-recurrent baseline on MATH as training FLOPs increase.Source: §C.3.1 · pp. 24, 26, Figure 21 · p. 26
• At inference, recurrent models are competitive with the fixed-depth baseline and can surpass it when given more recurrences.Source: Figure 21 · p. 26, §C.3.1 · p. 26
• Recurrent models remain competitive in inference FLOPs for both GSM8K and MATH when plotted against effective parameters.Source: Figure 22 · p. 27, Figure 25 · p. 28

• The evaluation covers GSM8K and MATH, with inference-efficiency views replotted using effective parameters proportional to inference FLOPs.Source: §C.3.2 · p. 29, Figure 30 · p. 31, Figure 33 · p. 32
• Training-step curves report GSM8K and MATH accuracy for train recurrences of 4, 8, 16, and 32.Source: §C.3.2 · p. 29, Figure 31 · p. 31, Figure 32 · p. 32
• Table 4 reports final-step accuracy across a broad range of evaluations for the models shown in the main OLMo results.Source: §C.3.2 · p. 29, Table 4 · p. 30
• Recurrent OLMo models efficiently outperform the non-recurrent baseline when accuracy is plotted against training FLOPs.Source: Figure 28 · p. 29, §C.3.2 · p. 31
• At inference, recurrent OLMo models remain competitive with fixed-depth baselines and can surpass them by using more recurrences.Source: Figure 28 · p. 29, §C.3.2 · p. 31

• The retrofit removes layers 4 and 5 from Llama-3.2-1B, leaving approximately 850 million parameters, or 87.5% of the pretrained model.Source: §C.3.3 · p. 34
• Inference results are also replotted with effective parameters, treated as proportional to inference FLOPs, for both GSM8K and MATH.Source: §C.3.3 · p. 34, Figure 38 · p. 36, Figure 41 · p. 37
• Training-step curves report GSM8K and MATH accuracy for train recurrences of 4, 8, 16, and 32.Source: §C.3.3 · p. 34, Figure 39 · p. 36, Figure 40 · p. 37
• Llama recurrent models efficiently outperform the non-recurrent baseline when accuracy is measured against training FLOPs on GSM8K and MATH.Source: Figure 36 · p. 34, Figure 37 · p. 36
• The recurrent Llama models are competitive with fixed-depth baselines at inference and can outperform them as recurrence increases.Source: Figure 36 · p. 34, Figure 37 · p. 36

• Table 6 extends earlier evaluations by including more test recurrences and Huginn-0125 results for comparison.Source: §C.4 · p. 39
• 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.Source: Table 6 · p. 39
• The experiments use specified optimizer settings, gradient clipping, batch sizes, and hardware configurations that vary across experimental sections.Source: D HYPERPARAMETERS · p. 39
• The appendix documents exact parameter counts for non-recurrent and recurrent models and identifies the layers taken from pretrained models.Source: E PARAMETER COUNTS · p. 40, Table 7 · p. 40, Table 8 · p. 40