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:

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):

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):

† = 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