Summary (Overview)

  • RecurTrace is a method for converting pretrained LLMs into memory-augmented looped reasoners with adaptive halting, addressing two key limitations of latent recurrence: memory loss across loop iterations and fixed loop budgets.
  • Loop Memory Attention (LMA) enables each looped layer to attend to its own states from previous iterations along a "loop-time" axis, allowing the model to revisit earlier computations without adding external memory stores or significant parameters (~2.2% at most).
  • Oracle-distilled halting uses a sequence-level halting head supervised by whether a deeper loop reduces loss, avoiding the collapse of penalty-based methods (ACT, PonderNet) and the overspending of confidence-based rules (CALM).
  • On MathQA (1.7B), RecurTrace achieves 56.9% accuracy at an average of 2.0 loops, beating the best fixed depth (54.7% at T=2) by 2.2 points at matched compute and outperforming all adaptive baselines tested.
  • Gains scale with model size: accuracy improvements over same-budget fine-tuned baselines grow from +0.6 points at 0.6B to +3.4 points at 8B, with teacher-forced NLL reductions at all four scales.

Introduction and Theoretical Foundation

Background and Motivation

Standard transformers allocate the same number of layers to every input, regardless of difficulty. However, reasoning quality scales with effective depth—whether achieved through more stacked layers, chain-of-thought token generation (Wei et al., 2022), or latent-space layer iteration (Saunshi et al., 2025; Geiping et al., 2025). Latent looping is particularly attractive because it increases effective depth without emitting additional tokens, adding parameters, or modifying the decoding loop.

Two Key Problems

  1. The loop forgets: A looped block re-runs the same layers on their own output, passing forward only a single hidden state. Earlier computations must either fit into that state or be overwritten—the loop cannot "consult its own past."
  2. The loop cannot budget: Fixed iteration counts waste depth on easy inputs while providing insufficient computation for hard ones. On MathQA, a 1.7B looped model peaks at 54.7% accuracy at 2 loops, stays near 53.6% at 8 loops, and falls to 47.5% at 16 loops—excess depth actively hurts.

Theoretical Foundation

The paper builds on the encode-think-decode paradigm (Koishekenov, Lipani, and Cancedda, 2026), where a contiguous middle block of layers is repeated. The key insight is that the loop trajectory itself provides both resources needed:

  • The states computed in earlier passes are a ready-made memory.
  • The current state carries evidence about whether another pass will change the answer.

This eliminates the need for external machinery such as gated memory banks (Frey et al., 2026) or multi-slot buffers (MeSH; Yu et al., 2025).


Methodology

Architecture Overview

RecurTrace starts from a pretrained LLM with LL transformer layers. A contiguous block of kk middle layers, Bˉ={s,…,s+k−1}\bar{B} = \{s, \ldots, s+k-1\}, is designated as the looped "thinking block" (layers 12–14 for 28-layer 0.6B/1.7B; layers 15–17 for 36-layer 4B/8B Qwen3 models). The block is chosen using an angular-distance probe that identifies where representations change least between adjacent layers (i.e., approximate fixed points).

Looped Computation

Let ee be the block input (output of layer s−1s-1). Each loop after the first re-injects the block input:

h←h+αRMSNorm(e)(1)h \leftarrow h + \alpha \text{RMSNorm}(e) \tag{1}

where α\alpha is a scalar initialized to zero (ReZero style), making injection start as a no-op. Inside the loop, before each block layer:

h←h+LMAℓ(xℓ(t−1),Mℓ(t))(2)h \leftarrow h + \text{LMA}_\ell \left(x_\ell^{(t-1)}, M_\ell^{(t)}\right) \tag{2} xℓ(t)←layerℓ(h)(3)x_\ell^{(t)} \leftarrow \text{layer}_\ell(h) \tag{3}

where Mℓ(t)={xℓ(t′)∣max⁡(1,t−W)≤t′≤t−1}M_\ell^{(t)} = \{x_\ell^{(t')} \mid \max(1, t-W) \leq t' \leq t-1\} is a sliding window of the layer's own outputs from up to WW earlier loops.

Loop Memory Attention (LMA)

LMA is an attention mechanism operating along the loop-time axis only (never mixing token positions, thus preventing future-token leakage). For a query state q=xℓ(t−1)q = x_\ell^{(t-1)} and a window of mm past states XX:

Ai,j=softmaxj(Qi⊤Ki,jdh−βh(t−tj))(4)A_{i,j} = \text{softmax}_j \left(\frac{Q_i^\top K_{i,j}}{\sqrt{d_h}} - \beta_h(t - t_j)\right) \tag{4} LMAℓ(q,X)=γ⋅g⊙([∑jAi,jVi,j]iWo)(5)\text{LMA}_\ell(q, X) = \gamma \cdot g \odot \left(\left[\sum_j A_{i,j} V_{i,j}\right]_i W_o\right) \tag{5}

Key design elements:

  • Relative loop-distance bias −βh(t−tj)-\beta_h(t - t_j) with learnable per-head slope βh\beta_h (initialized to ALiBi values) enables extrapolation to more loops than seen in training.
  • Scalar gate γ\gamma and token-wise gate g=σ(MLP([q,Xˉ]))g = \sigma(\text{MLP}([q, \bar{X}])) let the model close memory on inputs that don't need it.
  • Gates start near zero, so the memory path opens gradually during training.

Variable-Depth Training

The loop count is sampled per forward pass from a clamped log-normal Poisson distribution (mean depth ≈ 3.9, ~20% of passes run 6+ loops). This lets one trained model serve any test-time depth and extrapolate beyond training depths.

Oracle-Distilled Halting

A lightweight halting head predicts continue probability pt=σ(w⊤zt+b)p_t = \sigma(w^\top z_t + b) where ztz_t is the pooled block state. It is supervised by an oracle that marks "continue" only when deeper loops help:

yt=1[min⁡t′>tLt′<Lt−δ](6)y_t = \mathbb{1}\left[\min_{t' > t} \mathcal{L}_{t'} < \mathcal{L}_t - \delta\right] \tag{6}

for margin δ\delta. At inference (Algorithm 1), the model loops until the head says stop, with floor TminT_{min} and budget TmaxT_{max}. The threshold τ\tau is chosen on held-out data (smallest mean depth without losing likelihood).


Empirical Validation / Results

Controlled Comparison on MathQA (1.7B)

MethodAcc (%)CorrectLoopsΔ vs T=2
Fixed T=149.9611991.00−114
Fixed T=254.7113132.000
Fixed T=853.6212878.00−26
Fixed T=1647.54114116.00−172
ACT49.9611991.00−114
PonderNet49.9611991.00−114
ACT + floor54.7113132.000
PonderNet + floor54.7113132.000
CALM (top-prob)53.0412734.15−40
CALM (margin)54.1212995.63−14
LoopUS-Conf55.2513263.21+13
TaH-Mismatch55.6713362.12+23
RecurTrace56.9213662.04+53†

†Exact McNemar p<0.001p < 0.001 vs. T=2 on every training seed.

Key findings:

  • ACT and PonderNet collapse to one loop due to penalty/prior pressure dominating the oracle signal.
  • CALM overspends (5.6 loops) without accuracy gain.
  • The two-loop floor prevents collapse (raises baseline from 1199 to 1313 correct); the learned head adds 53 more.
  • A correctness oracle ceiling reaches 61.0%, so RecurTrace closes ~1/3 of the 6.3-point gap.

Depth Tracks Difficulty

Hard items (wrong after one loop) average 3.10 loops vs. 2.50 for easy items (+0.6 loop gap), with hard items showing a heavier tail distribution.

Cross-Task Transfer (1.7B, per-task heads without floor)

TaskFixed (T=8)Best FixedAdaptiveLoops
GSM8K5.59.210.41.72
MATH18.319.121.51.88
AQUA-RAT40.243.344.11.14
ARC-Challenge63.867.969.71.06
HotpotQA37.238.139.62.76

Adaptive halting matches or exceeds Best fixed depth on every task while cutting mean loops by 65–87%.

Scaling Results (Fixed 2-loop depth, no halting)

  • Accuracy gains over same-budget baselines: +0.6 pp (0.6B), +1.2 pp (1.7B), +2.2 pp (4B), +3.4 pp (8B).
  • NLL reductions at all scales: −0.15, −0.08, −0.02, −0.01 nats (0.6B → 8B).
  • Both train-covered (14 tasks) and classic (8 tasks) suites gain on both metrics at 1.7B.

Ablations: Memory is Essential

Plain looping+ Memory
NLL Δ (0.6B)−0.06−0.15
NLL Δ (1.7B)+0.07−0.03
NLL Δ (4B)+0.05−0.03
NLL Δ (8B)+0.06−0.02
MathQA acc (1.7B)52.454.7
Classic acc (1.7B)64.066.4

Beyond the 8-loop training limit:

  • Pointer-chasing (no state needed): plain looping wins 89.3 vs. 84.6.
  • Symbolic state tracking (state must persist): memory wins 56.2 vs. 38.5.
  • Arithmetic (state must persist): memory wins 49.8 vs. 33.1.

Component Ablations (1.7B)

ConfigurationNLL Δ (nats)MathQA (%)
Full RecurTrace−0.0354.7
– QK-normalization+0.0552.8
Loop-distance bias frozen−0.0153.9
Gates open at init+0.0253.1
Window W=1−0.0153.5
Window W=2−0.0254.2
Window W=5−0.0254.1

W=3 is optimal, QK-normalization is the most critical component, and the learnable loop-distance bias matters.


Theoretical and Practical Implications

Theoretical Implications

  1. The loop trajectory is self-sufficient: Both memory and computation budgeting can be derived from the states the loop already produces, without external machinery. This challenges the necessity of gated memory banks or multi-slot buffers.

  2. Loss-improvement oracle > token-mismatch oracle: The distinction between "hard items" and "items that more depth can help" is crucial. TaH-Mismatch (token-mismatch oracle) achieves only 55.7% because it cannot make this distinction; RecurTrace's loss-improvement oracle captures it explicitly.

  3. Weight-tied recurrence with attention over depth: Attention along loop time is distinct from attention over token positions or over depth in sparse-expert architectures (e.g., Dreamer). It preserves the weight-tied structure while enabling memory access.

Practical Implications

  1. Safe adoption path: All base weights stay frozen; added modules start nearly inactive, so a single loop exactly reproduces the pretrained model—a safe floor that more loops build on.

  2. Wall-clock speedups: At ~2 mean loops, runtime is ~1.1× the base model (vs. 1.7× at 8 loops), making adaptive halting a real computational saving, not just a reduced loop count.

  3. Multi-scale effectiveness: The recipe holds from 0.6B to 8B with consistent accuracy and likelihood gains, with the accuracy advantage growing with model size.

  4. Generalization beyond training depth: The ALiBi-initialized relative loop-distance bias enables extrapolation to 16+ loops despite training being clamped at 8.


Conclusion

RecurTrace demonstrates that repeating a compact middle block of layers—equipped with Loop Memory Attention and oracle-distilled halting—effectively converts a pretrained LLM into a memory-augmented, difficulty-adaptive reasoner. Key takeaways:

  1. Loop Memory Attention is the mechanism that makes added depth useful, not a minor enhancement. Without it, extra loops cause representational drift (plain looping raises NLL at scales ≥1.7B).

  2. Penalty-based and confidence-based halting objectives fail on frozen backbones (collapse to one loop or overspend). Supervision from a loss-improvement oracle is the key to accurate depth allocation.

  3. One model serves every depth: Variable-depth training with the relative loop-distance bias means a single checkpoint supports any test-time depth and extrapolates beyond training limits.

Limitations and Future Directions

  • Sequence-level halting assigns one depth per example; token-level allocation is future work.
  • Memory is limited to same-position states in a fixed 3-loop window; cross-position and older states are excluded.
  • Only frozen Qwen3 decoder-only models with direct-answer targets were tested.
  • Latent looping and chain-of-thought are complementary, not substitutes—combining them is a natural extension.

Related papers