Towards Looped Models Done Right, Part I: Topology, Input Injection, Recurrent-State Design

AI summary

Overview

Research area: Machine learning — efficient language-model architecture, specifically "looped" (recurrent-depth) Transformers that reuse the same physical blocks across multiple logical layers.

Technical level: Advanced. The paper assumes familiarity with Transformer blocks, Mixture-of-Experts routing, pretraining recipes, and benchmark suites.

Scope (one sentence): This is Part I of a study that disentangles three design axes of looped language models — where recurrence is placed, how the input is re-injected, and how the recurrent state is organized — through controlled from-scratch pretraining at matched scale, then tests whether the resulting principles hold in Mixture-of-Experts models.

What This Paper Is About

Looped language models save memory by applying the same Transformer blocks repeatedly instead of storing a separate block for every layer. Two architectural families dominate — Ouro (recurrent full stack, initialized from token embeddings, no injection) and Huginn (recurrence sandwiched between untied prelude and coda layers, with the prelude representation re-injected at every iteration and a separately initialized recurrent state) — but prior work bundles these choices together, so nobody knows which ones actually cause the performance difference.

The goal is to isolate each choice one at a time, measure its effect on ten benchmarks under matched parameters, logical depth, and token budgets, and then check whether the winning recipe survives the move to sparse MoE models.

Key Contributions

1. A common notation and three-axis framework for looped models. The authors write a general tied iterative model as a prelude map, a state initializer, a per-step input write, a tied recurrent body, and a coda, then express Ouro and Huginn as special cases of it. This splits the design space into iteration envelope, input injection, and latent-state organization.

2. A controlled Ouro-to-Huginn transformation at matched scale. All variants are trained from scratch at 730M stored parameters, an unrolled scale of 2.9B parameter applications, and 112 logical block executions (either R²⁸₄ = 28×4 or P₈R¹²₈C₈ = 8 + 12×8 + 8), across four token budgets up to 336B tokens.

3. Stepwise ablations that attribute gains to specific axes (Q1 envelope, Q2 injection, Q3 latent state), showing the sandwich envelope and input injection carry most of the benefit while random initialization and separate H/L states do not.

4. A transfer study to MoE models with routing diagnostics, including a causal intervention that freezes expert identities across loop iterations to test whether iteration-specific routing matters.

Main Findings

• The sandwich envelope helps reasoning, not knowledge. Moving recurrence into an 8-layer prelude / 12-layer tied core repeated 8 times / 8-layer coda sandwich raised MATH500 by 12.00 points, GSM8K by 7.05, DROP by 2.61, BBH-CoT by 2.47, and MBPP+ by 2.12 relative to full-stack Ouro at 460 TPP; ARC-Challenge fell 0.94 and HumanEval+ fell 1.22, and knowledge-heavy tasks showed no consistent gain.

• Input injection mainly helps context- and specification-dependent tasks. With prelude-state injection in the middle loop, BBH-CoT rose 6.63, HumanEval+ 5.49, MBPP+ 4.23, MMLU 2.53, and DROP 1.39; raw-token injection into full-stack Ouro produced the same pattern of gains (1.80 BBH-CoT, 2.44 HumanEval+, 4.50 MBPP+, 1.79 MMLU, 0.91 DROP).

• Injection can hurt quantitative reasoning. Prelude-state injection lowered MATH500 by 3.60 and GSM8K by 2.51 in the middle loop; raw-token injection lowered MATH500 by 1.60 in full-stack Ouro while slightly improving GSM8K by 0.53. The authors state the mechanism is not isolated.

• The combined injected middle loop still beats the combined injected Ouro. With both models' respective injection mechanisms enabled, the prelude–loop–coda model wins on eight of ten benchmarks, trailing only on ARC-Challenge and HellaSwag.

• Random state initialization is not a free win. Versus z₀ = e, random initialization gained 3.34 on ARC-Challenge and 1.22 on GSM8K, but lost more than one point on MMLU, MATH500, HumanEval+, and MBPP+, with the other four within one point; direct initialization was better on six of ten benchmarks. The authors frame random init as a task- and objective-dependent inductive bias.

• A shared-module H/L hierarchy gives no consistent benefit. Using H = 2 high-level cycles and L = 3 low-level updates (eight total recurrent-body applications), the H/L variant gained 2.23 on ARC-Challenge, 1.80 on DROP, and 1.22 on HumanEval+, but lost 7.20 on MATH500, 2.64 on MBPP+, and 2.19 on GSM8K. The authors explicitly note this does not rule out benefits from separately parameterized H and L modules as in HRM-Text.

• The full Huginn design beats Ouro everywhere in the dense setting. At the 730M-parameter and 336B-token scale, the complete Huginn architecture outperforms Ouro across all ten benchmarks.

• The advantage transfers to MoE. At 8.0B resident / 793.9M active parameters per physical pass (about 3.2B active parameter applications per token when unrolled) and 500B training tokens, Huginn beats Ouro on eight of ten benchmarks; five gains exceed one point, led by GSM8K (+4.70), MATH500 (+3.60), MMLU (+2.78), DROP (+1.95), and HumanEval+ (+1.22). Ouro leads on ARC-Challenge (0.86) and BBH-CoT (0.20).

• Huginn routes experts more evenly. Normalized load-balancing loss at 500B tokens is 1.571 for Huginn, 1.652 for the 112-layer feedforward MoE, and 1.899 for Ouro (uniform routing = 1).

• Iteration-specific expert selection matters causally. Forcing iterations 2–8 to reuse the loop-1 top-2 expert identities, while keeping iteration-specific mixture weights, reduced accuracy on all six evaluated tasks (MMLU, ARC-Easy, BoolQ, ARC-Challenge, PIQA, COPA).

• MoE narrows the gap to a feedforward model of matched active compute. The feedforward D¹¹²_MoE with 32.0B resident parameters leads on seven of ten benchmarks, but Huginn leads on DROP and GSM8K and matches it on MATH500 (both 53.40); over MATH500, GSM8K, HumanEval+, and MBPP+, Huginn averages 61.48 versus 62.79. The mean gap falls from 4.96 points in the dense setting to 1.71 points in the MoE setting, smaller on eight of nine directly comparable benchmarks and at all four training checkpoints.

• Memory efficiency under matched FLOPs. Under matched training and inference FLOPs, the Huginn-style model uses 75% fewer resident parameters than the 112-layer feedforward MoE baseline while outperforming it on DROP and GSM8K, matching it on MATH500, and remaining competitive elsewhere.

Methodology in Plain English

The authors treat architecture as a set of switchable knobs and change one knob at a time.

1. Fix the budget. Every looped model stores 28 Transformer blocks' worth of weights (730M parameters) but executes 112 blocks' worth of computation, so comparisons are fair against both a 28-block feedforward baseline (matched memory) and a 112-block feedforward baseline (matched compute).

2. Q1 — move the loop. Compare Ouro's "run all 28 blocks four times" against a sandwich: 8 untied layers, then a 12-layer block run 8 times, then 8 untied layers.

3. Q2 — add a write. Give each iteration a direct path back to the input using a learned diagonal gate that blends the current state with a projected input. In the sandwich, the write source is the prelude's contextualized representation; in Ouro it must be the raw token embeddings.

4. Q3 — change the state. Compare initializing the loop state from the prelude output, from Gaussian noise, and a two-timescale H/L scheme where one shared body alternates three low-level refinements with one high-level update.

5. Train and evaluate identically. All models use the TxT360 dataset, sequence length 8192, a 4,194,304-token global batch, AdamW with a swept learning rate, 200 warmup steps and cosine decay, and the same ten benchmarks; checkpoints are taken at 58, 115, 230, and 460 tokens per parameter.

6. Repeat with experts. Replace each dense feed-forward branch with a top-2 mixture of 25 experts of width 2432 (attention and normalization unchanged), retrain the two endpoint architectures plus a 112-layer feedforward MoE reference, and log routing statistics.

Why This Matters

Impact on research. Looped models are a plausible answer to the widening gap between fast compute and slow memory, but the literature has been split into Ouro-like and Huginn-like camps without evidence about which components matter. This paper converts an architectural taste argument into a measurable decomposition, and its negative results (random initialization, H/L states) are as useful as its positive ones — they tell the field which complexity to drop.

Real-world applications:

• On-device and edge deployment, where resident parameter count, not FLOPs, is the binding constraint — the 75%-fewer-resident-parameters result speaks directly to this.
• Long-context document and specification processing, since input injection produced its largest gains on tasks requiring recovery of prompt-provided evidence, instructions, and worked examples.
• Code assistants, where HumanEval+ and MBPP+ improved by 5.49 and 4.23 points in the middle-loop injection ablation.
• Cost-constrained serving of reasoning and math workloads, where the MoE variant led the feedforward reference on GSM8K and DROP and tied on MATH500 at a fraction of the stored weights.

Industry relevance. The recipe is pretraining-scale and reproducible from scratch, and the authors release code at https://github.com/ifm-ai/xllm-loop. Any team weighing memory budget against quality for a fixed compute envelope gets a concrete design checklist plus an honest accounting of the trade-off between context-following and quantitative reasoning.

Future Directions

• Refine the injection mechanism to stop hurting math. The authors show the MATH500 and GSM8K regressions from prelude-state injection but state they do not isolate the mechanism, suggesting write strength, injection location, and the interaction between the fixed prelude state and evolving state as candidates.

• Test separately parameterized H and L modules. The paper notes that a closer HRM-Text-style variant with distinct modules for the two update streams "is still under evaluation" and is not part of the present comparison, so the hierarchy's value in that setting is unresolved.

• Extend the diagnostics beyond finite unrolling. The authors caution that their evaluations do not directly measure multi-start path independence or extrapolation to recurrent depths beyond those used in training.

• Push the MoE comparison further. The 1.71-point average gap to the feedforward MoE remains, and the paper reports MoE benefits under matched active compute; whether the residual gap closes with different expert counts, routing schemes, or longer training is left open. This paper is labeled Part I, with the truncated content ending at the start of a future-work section, so the authors' own follow-up agenda is not fully reported here.

Target Audience

Architecture researchers and pretraining engineers working on efficient or recurrent-depth language models; practitioners deciding between dense and MoE looped designs under a fixed memory or FLOP budget; and anyone who wants a rigorous template for "one-knob-at-a-time" architecture ablations at scale. Readers seeking results on inference-time recurrent-depth extrapolation beyond the trained number of loops, or on separately parameterized hierarchical modules, will not find them here.

Authors’ abstract

Looped language models reuse physical Transformer blocks across logical depth, but established architectures entangle recurrence placement, input injection, and recurrent state organization. We isolate these axes in models trained from scratch at matched parameter scale, logical depth, and token budgets, and then test whether the resulting design principles transfer to resident- and active-parameter-matched Mixture-of-Experts (MoE) models. The prelude-recurrent block-coda sandwich architecture consistently improves performance on context extraction and reasoning-intensive tasks. Input injection primarily enhances knowledge retrieval and context utilization, whereas random initialization and separate H/L recurrent states provide mixed or negative returns. Taken together, these results favor the Huginn-style design: at the 730M-parameter and 336B-token scale, the complete Huginn architecture outperforms Ouro across all ten benchmarks. The advantages of Huginn over Ouro transfer to MoE models. At the 8BA0.8B (32BA3B when unrolled) scale trained on 500B tokens, Huginn outperforms Ouro on nearly all benchmarks spanning knowledge, reasoning, and coding, while also exhibiting more balanced expert routing. Under matched training and inference FLOPs, the Huginn-style model uses 75% fewer resident parameters than a 112-layer feedforward MoE baseline. Despite its smaller physical parameters, it outperforms the feedforward baseline on DROP and GSM8K, matches it on MATH500, and remains competitive on the other benchmarks.