Unrolling looped LMs past training depth: reasoning +2.97pp, knowledge -10.69pp

What Makes Recurrence Effective in Looped Language Models?

Xinlin Zhuang, Siyuan Wang, Imran Razzak, Weiyang Liu

cs.LG, cs.CL

2026-09-29

Unrolling looped LMs past the training horizon gains +2.97pp reasoning but loses 10.69pp knowledge; history-state injection plus timestep conditioning makes extrapolation robust.

What problem this solves

Looped language models (LoopLMs) re-execute a shared stack of Transformer blocks, buying depth with parameter sharing: no new parameters, and each extra pass at inference adds compute. That makes them a parameter-frugal route to test-time scaling. Huginn and related work established feasibility, but evaluations mostly stayed at the fixed loop count used in training. Whether unrolling further still helps, which blocks should loop, and what signals the loop needs between iterations were all open questions. This paper runs controlled experiments across three inference regimes (under-unrolling, the training horizon, extrapolation) on knowledge and reasoning task groups.

Method

Notation first. A looped decoder is p+s×K+c: a Prelude of p independently parameterized blocks, a shared core of s blocks executed K times in training, and a Coda of c blocks. With p=c=0 the whole network loops (BaseLoop); adding boundary blocks gives CoreLoop. Physical depth p+s+c sets the parameter count, effective depth p+r·s+c sets per-token compute, and varying r at inference moves the budget. Every comparison holds four things equal: physical depth, training effective depth, inference effective depth, and the training pipeline. Models are randomly initialized on Llama3.1-1B and Qwen3-0.6B architecture configurations and pretrained from scratch on FineWeb-Edu at the Chinchilla ratio, with Muon on matrix parameters. Benchmarks split into a knowledge group (SciQ, ARC-Easy, PIQA) and an eight-task reasoning group (ARC-Challenge, WinoGrande, HellaSwag, ProofWriter, CLUTRR, BBH subtasks, and others); ProofWriter proof depth and CLUTRR chain length provide difficulty knobs.

Conditioning is the focus of the second half:

Results

The headline: unrolling past the training horizon keeps improving reasoning while degrading knowledge.

SettingMetricResult
BaseLoop 2x10 (Llama), L(r) 20 to 40Reasoning28.52 to 31.49 (+2.97pp, test-time only)
SameKnowledge62.80 to 52.11 (-10.69pp)
NonLoop 20x1 (equal training FLOPs)Reasoning30.25 (compute-matched reference)
BaseLoop 5x4 at L(r)=30Reasoning31.88, above NonLoop
Qwen 4x7, history + Loop Gating at D=84Overall37.30 vs 36.91 for the next-best pairing; +1.16pp over its stronger component

Two stratified results cut against intuition. ProofWriter gains at every proof depth when extrapolated to L(r)=24, about +7.5pp at depths 3-5, while CLUTRR gains concentrate on short chains (31.6 at length 2, 23.4 at length 3) and shrink or turn negative on longer chains. Harder instances do not automatically benefit more. Configuration matters as much as depth: at the same training depth of 20, 2x10 peaks early and decays, 10x2 stays flat, and only the balanced 4x5 and 5x4 keep extrapolating. Effective depth alone does not predict behavior.

Conditioning results are equally sharp. Dense initial-state injection collapses knowledge accuracy under extrapolation; lightweight variants avoid the collapse but add nothing over BaseLoop. History-state injection with dense maps and a window of one step rescues the collapsed curve back near the NonLoop reference, while windows of two or four erode the gain. Timestep gains depend on the configuration: at L(r)=56 on 7x4, Branch Gating reaches 38.46 overall (BaseLoop 36.59, AdaLN 38.01), yet the same schemes give small or negative gains on 4x7 and 2x14. The final combination, validated on Qwen 4x7: channel-wise history injection (w=2) plus Loop Gating is the strongest pair, and adding initial-state injection brings nothing further. Geometry analysis suggests why: this pair shows 18.6% and 10.4% higher cross-loop interaction than the two initial-state pairings, coupling that a static anchor cannot provide.

Why it matters

For anyone working on efficient inference or test-time scaling, this reads as a design manual. Three rules land directly: balance physical depth against loop count, since lopsided configurations stop extrapolating; keep a Coda to protect knowledge under under-unrolling and weight the Prelude for reasoning extrapolation; condition with channel-wise history injection plus scalar-level timestep gating, which costs a few vectors. The paper also draws a boundary for the parameters-for-compute trade: under equal training FLOPs, recurrent variants exceed the non-recurrent reasoning bound purely with extra test-time compute, a real lever for small models on hard tasks. To be clear about magnitude, this is incremental empirical work, with individual gains of one to three points at scales up to 1B.

Limitations

The authors list three: compute caps experiments at about 1B parameters under a fixed token budget; extrapolation multiples are set manually, with no adaptive per-input loop count; the geometry metrics, including computational interaction, remain correlational observations. Two more concerns from a close read: the knowledge group has only three tasks, with PIQA deliberately assigned to the knowledge side, so group means are sensitive to that split; and conditioning benefits swing with configuration (4x7 and 7x4 prefer different gating granularities), with no account that predicts in advance which pairing works where.

Terms

Source

Related papers

All paper explainers