# Closing the Loop: Practical Training Recipes for Looped Language Models

> Looped language models beat parameter-matched dense transformers on all 12 benchmarks using 25x less training data, via a two-stage recipe and minimal dense-to-loop conversion.

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

## Summary

# Closing the Loop: Practical Training Recipes for Looped Language Models

## Summary (Overview)

- **Dramatically reduced training cost**: The authors develop a from-scratch training recipe for looped language models (LoopLMs) that reduces the token budget from 7.7T tokens (Ouro) to 310B tokens—a 25× reduction—while retaining strong reasoning performance, achieved through a two-stage pipeline (pretraining + mid-training) instead of Ouro's five-stage schedule.

- **Controlled comparisons isolate recurrence gains**: A 1.4B LoopLM outperforms a parameter-matched dense model trained on the same data and token budget on all 12 evaluated benchmarks, with gains of +14 points on GSM8K, +10 on MATH, and +22 on DROP. At matched inference compute, it approaches a 3.9B dense model on math and reading comprehension while using only 36% of its parameters.

- **Minimal dense-to-loop conversion**: The authors introduce a recipe for converting pretrained dense models (Qwen3-1.7B-Base) into looped models using just one learned input-mixing scalar and a smoothed exit loss—no step-specific parameters required. Looped Qwen improves over an identically continued dense baseline on every benchmark across two data regimes.

- **Stability insights**: A learning-rate warmup prevents early exit-gate collapse, and a larger KL coefficient (β = 0.15) enables gate recovery if collapse occurs. The paper also identifies that post-norm and sandwich-norm normalization schemes cause catastrophic gradient vanishing in recurrent settings.

- **Reinforcement learning findings**: RL only helps looped models when the rollout task has a mixed solve rate (neither too easy nor too hard), and the learning rate must be large enough to move the policy.

---

## Introduction and Theoretical Foundation

### Background

Transformers spend the same computation on every token regardless of difficulty. **Looped language models** address this by applying a shared transformer block repeatedly, which increases effective depth without adding parameters. With adaptive exits, they can spend extra computation only where needed (Dehghani et al., 2019; Zhu et al., 2025).

### Motivation

The prior state-of-the-art, **Ouro** (Zhu et al., 2025), demonstrated that looped models can be pretrained at scale and compete with much larger dense models. However, Ouro's recipe has three critical limitations:

1. **Excessive cost**: Five training stages requiring 7.7T tokens
2. **No use of pretrained checkpoints**: Training starts from scratch while pretrained dense checkpoints go unused
3. **Uncontrolled comparisons**: Gains are measured against dense models trained on different data, so the share attributable to recurrence itself is unknown

### Key Concepts

The theoretical foundation draws on several strands:

- **Universal Transformers** (Dehghani et al., 2019): Introduced the shared-block construction with adaptive halting
- **Entropy-regularized ELBO objective**: The exit gate is optimized jointly with the language model using a KL coefficient β
- **Recurrent depth**: Each application of the shared block constitutes a "recurrent step," with exit probabilities determining how many steps are used per token

---

## Methodology

### 3.1 From-Scratch Recipe (LoopLM)

The authors collapse Ouro's five-stage schedule into two stages:

| | Pretraining | Mid-training |
|---|---|---|
| **Initialization** | From scratch | Pretraining checkpoint |
| **Peak learning rate** | 3 × 10⁻⁴ | 1 × 10⁻⁴ |
| **Final learning rate** | 3 × 10⁻⁵ | 1 × 10⁻⁵ |
| **LR scheduler** | Warmup + cosine decay | Warmup + cosine decay |
| **LR warmup steps** | 1,000 | 1,000 |
| **Weight decay** | 0.1 | 0.1 |
| **Gradient norm clip** | 1.0 | 1.0 |
| **Optimizer** | AdamW (β₁=0.9, β₂=0.95) | AdamW |
| **Sequence length** | 4K | 4K |
| **Training steps** | 100,000 | 47,684 |
| **Training tokens** | 210B | 100B |
| **Loss** | Next token | Response tokens only |
| **Recurrent steps (T)** | 4 | 4 |
| **KL coefficient β** | 0.15 | 0.15 |

**Key decisions:**

- **Depth T = 4**: A T=8 run showed that loops 5–8 change the loss by ≤0.01 nats, while the exit gate spreads mass evenly over loops 4–8 (wasting compute). T=4 is fixed from the start.
- **Learning-rate warmup**: Prevents early exit-gate collapse. Without warmup, the gate learns to exit after the first loop, later iterations stop receiving gradient, and the model trains as a plain transformer. The LM loss does not reveal this collapse—the exit distribution must be monitored directly.
- **KL coefficient β = 0.15**: Larger than Ouro's post-stage value of 0.05. A sweep (Table 7) showed that β = 0.05 stays collapsed permanently without warmup, while β ≥ 0.10 allows recovery.

### 3.2 Recurrence and Exit Objective (From-Scratch)

The recurrence is defined as:

$$h_0 = e, \quad h_t = B(h_{t-1}), \quad t = 1, \ldots, T$$

where $e$ is the token embedding and $B$ is the shared decoder block.

A shared linear gate produces halt probabilities:

$$\lambda_t = \sigma(w_g^\top h_t)$$

The exit distribution uses a survival construction:

$$p_t = \lambda_t \prod_{s<t}(1 - \lambda_s) \quad (t < T), \quad p_T = \prod_{s<T}(1 - \lambda_s)$$

The training objective is:

$$\mathcal{L} = \sum_{t=1}^{T} p_t \mathcal{L}_t - \beta H(p)$$

with β = 0.15.

### 3.3 Dense-to-Loop Conversion (Looped Qwen)

For converting Qwen3-1.7B-Base, the authors found that naive recurrence degrades later steps. Their minimal recipe uses:

**1. Scalar input injection**: Let $e$ be the token embedding and $h_{t-1}$ the previous recurrent step output. For $t > 0$:

$$z_t = we + (1 - w)h_{t-1}, \quad w = \sigma(a)$$

where $a$ is a single learned scalar, initialized with $a = 4.6$ (so $w = 0.99$). This preserves direct access to the original embedding.

**2. Smoothed exit loss**: To prevent the gate from abandoning steps (which causes their losses to drift to ~10 nats):

$$\mathcal{L}_{\text{smooth}}^{\text{LM}} = (1 - \epsilon)\sum_{t=0}^{T-1} p_t \mathcal{L}_t + \frac{\epsilon}{T}\sum_{t=0}^{T-1} \mathcal{L}_t$$

With $\epsilon = 0.05$ and $T = 3$, each step retains at least 0.0167 LM weight. Gate regularization and inference still use the original exit probabilities.

**3. Optional step-wise LoRA**: Independent rank-32 LoRA adapters per step were tested but provide little benefit (≤2 points difference), so the fully shared model is used by default.

### 3.4 Baselines

For controlled comparisons, the authors train:

- **From-scratch**: 
  - `dense-ouro-1.4B`: Same 24-layer block, T=1, parameter-matched
  - `dense-ouro-3.9B`: Three stacked blocks (72 layers), matches LoopLM's inference compute at 3 recurrent steps
  
- **Conversion**:
  - Qwen3-1.7B-Base: Parameter-matched
  - Qwen3-4B-Base: Inference-compute-matched reference

All baselines are trained on the same data with the same stages and token budget.

---

## Empirical Validation / Results

### 5.1 LoopLM vs. Controlled Dense Baselines

**Table 2: Benchmark scores after mid-training** (LoopLM at fixed exits; dense-ouro-3.9B matches 3-iteration inference compute):

| Benchmark | dense-ouro-1.4B | LoopLM 1 iter | LoopLM 2 iter | LoopLM 3 iter | LoopLM 4 iter | dense-ouro-3.9B |
|---|---|---|---|---|---|---|
| **GSM8K** | 69.14 ± 1.27 | 60.65 | 76.50 | 81.58 | 83.32 | 82.79 |
| **MATH** | 54.10 ± 0.67 | 44.18 | 60.90 | 63.64 | 63.80 | 64.54 |
| **MATH500** | 55.00 ± 2.22 | 45.00 | 61.40 | 63.60 | 63.20 | 63.60 |
| **DROP (F1)** | 38.59 ± 0.48 | 22.20 | 45.02 | 57.40 | 60.57 | 58.57 |
| **MMLU** | 46.52 ± 0.41 | 40.77 | 52.61 | 55.93 | 56.84 | 58.15 |
| **MMLU-Pro** | 17.48 ± 0.34 | 13.80 | 20.53 | 24.27 | 24.53 | 28.22 |
| **ARC-Challenge** | 46.16 ± 1.46 | 45.99 | 50.60 | 50.51 | 51.02 | 55.03 |
| **ARC-Easy** | 74.75 ± 0.89 | 72.18 | 76.47 | 77.78 | 77.99 | 79.67 |
| **HellaSwag** | 57.28 ± 0.49 | 51.02 | 56.94 | 58.92 | 59.29 | 67.91 |
| **Winogrande** | 59.43 ± 1.38 | 56.67 | 60.22 | 62.90 | 63.93 | 65.19 |
| **CRUXEval-O** | 6.15 ± 0.76 | 6.03 | 18.20 | 17.36 | 16.60 | 18.48 |
| **CRUXEval-I** | 4.61 ± 0.60 | 2.34 | 16.20 | 19.46 | 21.21 | 21.80 |

**Key findings:**
- At equal parameter count, 4-step LoopLM outperforms `dense-ouro-1.4B` on **all 12 benchmarks**
- Largest gains on reasoning tasks: +22 (DROP), +17 (CRUXEval-I), +14 (GSM8K), +10 (MATH, MMLU, CRUXEval-O)
- At equal inference compute (3 iterations), LoopLM comes within one standard error on GSM8K, MATH, MATH500 and within ~1 point on DROP of the 3.9B dense model
- The 4th iteration moves LoopLM ahead on DROP and level on GSM8K
- The remaining gap is in knowledge/commonsense benchmarks (2–9 points), where the larger dense model stores more facts

### 5.2 LoopLM vs. Published Models

Despite using only 0.31T training tokens, LoopLM achieves the **highest DROP score** and **second-highest GSM8K score** in a comparison table of models between 1B–4B parameters. Against Ouro-1.4B (trained on 25× more tokens), LoopLM is higher by +4.4 points on GSM8K and +10.9 points on DROP, but lower on MATH and MATH500.

### 5.3 Looped Qwen Conversion Results

**Table 4: Conversion on curated instruction mixture (60B tokens, evaluated at fixed exit after 2 of 3 steps):**

| Benchmark | Qwen3-1.7B-Base† | Looped Qwen (shared) | Looped Qwen (LoRA-32) | Qwen3-4B-Base† |
|---|---|---|---|---|
| **GSM8K** | 66.64 ± 1.30 | **70.96** ± 1.25 | 69.90 | 82.87 |
| **MATH** | 51.14 ± 0.66 | **53.34** ± 0.66 | 54.46 | 65.92 |
| **MMLU** | 60.62 ± 0.39 | **61.05** ± 0.39 | 60.99 | 71.83 |
| **MMLU-Pro** | 30.76 ± 0.41 | **34.52** ± 0.42 | 36.24 | 49.52 |
| **ARC-C** | 53.84 ± 1.46 | **55.46** ± 1.45 | 54.95 | 64.42 |
| **ARC-E** | 79.97 ± 0.82 | **80.13** ± 0.82 | 80.51 | 86.32 |
| **HellaSwag** | 63.35 ± 0.48 | **64.13** ± 0.48 | 64.29 | 70.88 |
| **Winogrande** | 62.90 ± 1.36 | **63.61** ± 1.35 | 64.01 | 67.80 |

† = continued training on the same mixture for the same number of steps

**Key findings:**
- Looped Qwen improves over the continued dense baseline on **every benchmark** in both data regimes (instruction and web-scale)
- On the curated instruction mixture, gains are statistically clear (exceeding two standard errors): +4.3 (GSM8K), +3.8 (MMLU-Pro), +2.2 (MATH)
- Gains are larger on curated data than web-scale data, showing recurrence benefits depend on training data
- Shared vs. LoRA-32 variants differ by ≤2 points—step-specific parameters are **not necessary**
- Looped Qwen still trails Qwen3-4B-Base (as expected for a model >2× larger)

### 5.4 Normalization Ablation (Appendix A)

The paper identifies that **normalization placement** is critical for recurrent training:

- **Post-norm and sandwich-norm fail catastrophically** (loss plateaus at ~7.4, never recovers)
- The failure is due to gradient vanishing: layers 0–22 receive gradients 8–16 orders of magnitude smaller than the final layer
- Pre-norm, Ouro-norm, and Magic-norm all train normally (losses in 1.98–2.11 range)
- A final normalization (RMS or FRMS) before the LM head is essential for stability: without it, models exhibit 5–11 loss spikes and gradient norms up to 5.2 × 10⁷

### 5.5 Reinforcement Learning Results (Appendix C)

RLTT (a GRPO variant for looped models) was tested in six configurations:

- **GSM8K from pretraining checkpoint** (solve rate ≈ 0.30): **+12.0 points** improvement (35.3 → 47.3)
- **MATH from mid-trained checkpoint** (solve rate ≈ 0.77): **No gain** (−0.5 points)

The key insight: RL helps only when rollouts contain **both successes and failures** (mixed solve rate). When the solve rate is too high (~77%), groups are uniform and the group-relative advantage vanishes. Additionally, a learning rate of 10⁻⁶ fails to move the policy (KL stays near 2×10⁻⁴), while 10⁻⁵ works.

**Useful diagnostic**: KL divergence from reference policy and gradient norm predict whether RL is working. Growing KL + non-negligible gradient norm = policy is moving.

---

## Theoretical and Practical Implications

### Implications for Looped Model Training

1. **Recurrence is the source of gains**: The depth sweep (1→4 iterations) shows that gains come from recurrent computation, not the underlying checkpoint. One iteration trails the dense baseline on every benchmark; two outperform on 11/12; three on all 12.

2. **Training stability recipes transfer**: Learning-rate warmup + larger KL coefficient are sufficient to prevent exit-gate collapse without multi-stage schedules. This removes a major barrier to from-scratch looped training.

3. **Normalization is critical for recurrence**: Post-norm and sandwich-norm fail because shared weights amplify gradient vanishing. This is a unique challenge of recurrent parameter sharing that doesn't appear in standard dense training.

4. **Pretrained checkpoints are convertible**: A single learned scalar (input injection) + smoothed exit loss suffices to convert a dense model to a looped one, at under 1% of pretraining compute. This makes looped models practical to introduce into existing pretrained checkpoints.

### Practical Applications

- **Cost reduction**: The from-scratch recipe uses ~3% of Ouro's training compute (9.9 × 10²¹ vs. 3.4 × 10²³ FLOPs)
- **Compute-efficient inference**: LoopLM at 3 steps uses 36% of the parameters of the matched 3.9B dense model; Looped Qwen at 2 steps uses 22% fewer FLOPs than Qwen3-4B-Base
- **RL guidance**: The paper establishes when RL helps looped models (mixed solve rates, sufficient learning rate) and provides cheap telemetry diagnostics

### Limitations

- The 1.4B LoopLM trails the 3.9B dense model on knowledge/commonsense benchmarks (where parameter count stores facts)
- The conversion does not replace the capacity of substantially larger pretrained models
- The Think-at-Hard reproduction (Appendix J) failed to produce gains, suggesting correctness-directed routing requires more careful design

---

## Conclusion

This work provides two practical paths to building looped language models:

1. **From-scratch training**: A two-stage recipe (210B pretraining + 100B mid-training) with learning-rate warmup and β = 0.15, producing a 1.4B LoopLM that beats a parameter-matched dense twin on all 12 benchmarks and approaches a 3.9B dense model (3× parameters) on math and reading at equal inference compute.

2. **Dense-to-loop conversion**: A single learned input-mixing scalar and a smoothed exit loss convert Qwen3-1.7B-Base into Looped Qwen, which improves on the source checkpoint given the same continued training—with no step-specific parameters required.

Both results are validated against dense baselines trained on the **same data, same stages, and same token budget**, matched in parameters or in inference compute—isolating the gains due to recurrence itself.

### Future Directions

- Whether RLTT can improve mid-trained checkpoints on tasks specifically selected for mixed rollout solve rates
- Online correctness-based depth routing (proposed in Appendix J) that combines Think-at-Hard's supervision with more than two recurrent steps
- Scaling the simplified recipe to larger model sizes and token budgets

---

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