# RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory

> RecurTrace adds loop-time attention memory and oracle-distilled halting to frozen LLMs, improving MathQA accuracy by 2.2 points at matched compute while scaling gains with model size.

- **Source:** [arXiv](https://arxiv.org/abs/2609.03379)
- **Published:** 2026-10-24
- **Permalink:** https://picx.dev/p/mZtufK
- **Whiteboard:** https://picx.dev/p/mZtufK/image

## Summary

## 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 $L$ transformer layers. A contiguous block of $k$ middle layers, $\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 $e$ be the block input (output of layer $s-1$). Each loop after the first re-injects the block input:

$$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 \leftarrow h + \text{LMA}_\ell \left(x_\ell^{(t-1)}, M_\ell^{(t)}\right) \tag{2}$$

$$x_\ell^{(t)} \leftarrow \text{layer}_\ell(h) \tag{3}$$

where $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 $W$ 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_\ell^{(t-1)}$ and a window of $m$ past states $X$:

$$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}$$

$$\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** $-\beta_h(t - t_j)$ with learnable per-head slope $\beta_h$ (initialized to ALiBi values) enables extrapolation to more loops than seen in training.
- **Scalar gate $\gamma$** and **token-wise gate $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 $p_t = \sigma(w^\top z_t + b)$ where $z_t$ is the pooled block state. It is supervised by an oracle that marks "continue" only when deeper loops help:

$$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 $T_{min}$ and budget $T_{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)

| Method | Acc (%) | Correct | Loops | Δ vs T=2 |
|---|---|---|---|---|
| Fixed T=1 | 49.96 | 1199 | 1.00 | −114 |
| Fixed T=2 | 54.71 | 1313 | 2.00 | 0 |
| Fixed T=8 | 53.62 | 1287 | 8.00 | −26 |
| Fixed T=16 | 47.54 | 1141 | 16.00 | −172 |
| ACT | 49.96 | 1199 | 1.00 | −114 |
| PonderNet | 49.96 | 1199 | 1.00 | −114 |
| ACT + floor | 54.71 | 1313 | 2.00 | 0 |
| PonderNet + floor | 54.71 | 1313 | 2.00 | 0 |
| CALM (top-prob) | 53.04 | 1273 | 4.15 | −40 |
| CALM (margin) | 54.12 | 1299 | 5.63 | −14 |
| LoopUS-Conf | 55.25 | 1326 | 3.21 | +13 |
| TaH-Mismatch | 55.67 | 1336 | 2.12 | +23 |
| **RecurTrace** | **56.92** | **1366** | **2.04** | **+53†** |

†Exact McNemar $p < 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)

| Task | Fixed (T=8) | Best Fixed | Adaptive | Loops |
|---|---|---|---|---|
| GSM8K | 5.5 | 9.2 | 10.4 | 1.72 |
| MATH | 18.3 | 19.1 | 21.5 | 1.88 |
| AQUA-RAT | 40.2 | 43.3 | 44.1 | 1.14 |
| ARC-Challenge | 63.8 | 67.9 | 69.7 | 1.06 |
| HotpotQA | 37.2 | 38.1 | 39.6 | 2.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.4 | **54.7** |
| Classic acc (1.7B) | 64.0 | **66.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)

| Configuration | NLL Δ (nats) | MathQA (%) |
|---|---|---|
| Full RecurTrace | −0.03 | 54.7 |
| – QK-normalization | +0.05 | 52.8 |
| Loop-distance bias frozen | −0.01 | 53.9 |
| Gates open at init | +0.02 | 53.1 |
| Window W=1 | −0.01 | 53.5 |
| Window W=2 | −0.02 | 54.2 |
| Window W=5 | −0.02 | 54.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.

---

_Markdown view of https://picx.dev/p/mZtufK, served by PicX — AI-generated visual whiteboard summaries of research papers._