Towards Looped Models Done Right: Part I — Summary

Summary (Overview)

• Disentangles three architectural axes in looped language models—iteration envelope, input injection, and latent-state organization—that prior work (Ouro and Huginn families) varied simultaneously, preventing clean attribution of design choices.
• Controlled ablations at 730M parameters / 336B tokens show that the prelude–recurrent block–coda "sandwich" envelope provides the clearest gains (e.g., +12.00 on MATH500, +7.05 on GSM8K), while input injection primarily boosts context-dependent tasks (BBH-CoT +6.63, HumanEval+ +5.49) but can hurt quantitative reasoning (MATH500 −3.60).
• Random initialization and shared H/L recurrent states provide mixed or negative returns, suggesting these are task-dependent inductive biases rather than universally beneficial design choices.
• Huginn-style advantages transfer to MoE models: at 8B resident / 0.8B active parameters (500B tokens), Huginn outperforms Ouro on 8/10 benchmarks, achieves more balanced expert routing, and iteration-specific expert selection is causally important for performance.
• Apples-to-apples comparison: Huginn MoE uses 75% fewer resident parameters than a 112-layer feedforward MoE baseline, yet surpasses it on DROP and GSM8K, matches it on MATH500, and narrows the mean gap from 4.96 (dense) to 1.71 points (MoE).

---

Introduction and Theoretical Foundation

Background and Motivation

Looped language models reuse the same Transformer blocks across logical depth, trading additional computation for reduced memory usage. This paradigm is motivated by the observation that compute throughput grows faster than memory capacity/bandwidth in modern hardware. Two dominant architectural lineages have emerged:

• Ouro-style: Recurrently applies a tied full stack initialized from token embeddings.
• Huginn-style: Places recurrence between untied prelude and coda layers, repeatedly injects a fixed prelude representation, and uses a separately initialized recurrent state.

Common Iterative-Model Framework

The paper introduces a unifying notation for tied iterative models:

$$\mathbf {e} = P _ {\theta} (\mathbf {x} _ {0}),\tag{1}$$ $$\mathbf {z} _ {0} = \phi_ {\theta} (\mathbf {e}, \pmb {\xi}),\tag{2}$$ $$\tilde {\mathbf {z}} _ {t} = W _ {\theta} (\mathbf {z} _ {t}, \mathbf {e}),\tag{3}$$ $$\mathbf {z} _ {t + 1} = R _ {\theta} (\tilde {\mathbf {z}} _ {t}), \qquad t = 0, \ldots , T - 1,\tag{4}$$ $$\mathbf {h} = C _ {\theta} (\mathbf {z} _ {T}).\tag{5}$$

where $P_\theta$ is the prelude, $R_\theta$ the tied iterative body, $C_\theta$ the coda, $\phi_\theta$ initializes the recurrent state, and $W_\theta$ is the per-step input write.

Three Separable Design Axes

1. Iteration envelope: Placement and sharing pattern of $P_\theta, R_\theta, C_\theta$
2. Input injection: Whether $W_\theta$ receives a persistent input condition (e.g., token embeddings vs. prelude-encoded $\mathbf{e}$)
3. Latent-state organization: How $\phi_\theta$ initializes the loop state and the structure of states (e.g., fast/slow latents)

---

Methodology

Model Configurations

All looped models store 28 Transformer blocks (730M parameters) but reuse subsets across iterations, yielding 112 logical block executions (2.9B unrolled parameter applications):

• Ouro-style: $R_{28}^{4}$ — tied 28-block stack executed 4 times
• Huginn-style: $P_8 R_{12}^8 C_8$ — 8-layer prelude, 8-step loop over 12-layer tied core, 8-layer coda
• Feedforward baselines: $D_{28}$ (730M, resident-matched) and $D_{112}$ (2.9B, depth-matched)

Transformer block: PreNorm decoder-only backbone, 24 query heads, 6 KV heads, head dim 64, RoPE base $10^6$, hidden dim 1536, SwiGLU MLP width 4352.

Training Recipe

• Data: TxT360 dataset (deduplicated Common Crawl + 14 curated sources)
• Sequence length: 8192; batch: 512 sequences (4,194,304 tokens/step)
• Optimizer: AdamW, LR ∈ {4×10⁻⁴, 6×10⁻⁴, 8×10⁻⁴}, 200 warmup steps, cosine decay, weight decay 0.1
• Budgets: 42B, 84B, 168B, 336B tokens (58–460 TPP for 730M models)

Evaluation Benchmarks

---

Empirical Validation / Results

Q1: Sandwich Envelope (Iteration Envelope)

Table 2: Q1 envelope comparison at 460 TPP (730M parameters; 336B tokens)

Key findings: The sandwich envelope improves instance-conditioned, multi-step reasoning (MATH500, BBH-CoT, DROP) but does not consistently improve knowledge-intensive or requirement-constrained tasks (code generation).

Q2: Input Injection

The input write uses a diagonal operator:

$$D (\mathbf {z} _ {t}, \mathbf {v}) = \boldsymbol {\alpha} \odot \mathbf {z} _ {t} + \boldsymbol {\delta} \odot \mathbf {W} _ {\mathrm{in}} \mathbf {v}, \qquad \boldsymbol {\delta} = \mathrm{softplus} (\mathbf {b} _ {\delta}), \qquad \boldsymbol {\alpha} = \exp \{- \boldsymbol {\delta} \odot \exp (\mathbf {a}) \}.$$

Key findings (middle loop, prelude write vs. no write):

• Gains: BBH-CoT +6.63, HumanEval+ +5.49, MBPP+ +4.23, MMLU +2.53, DROP +1.39
• Losses: MATH500 −3.60, GSM8K −2.51
• Despite the math trade-off, the complete injected middle-loop design outperforms injected full-stack Ouro on 8/10 benchmarks.

Q3: Recurrent Latent States

Table 4: Matched recurrent-state comparison at 460 TPP

Key findings: Random initialization produces 2 gains and 4 losses > 1 point; H/L hierarchy produces 3 gains and 3 losses > 1 point, with the largest degradation on MATH500 (−7.20). Neither provides consistent benefit.

MoE Transfer Results

Table 5: Benchmark performance of MoE models (8B resident / 793.9M active, 500B tokens)

Routing diagnostics:

• Load-balancing loss at 500B tokens: Huginn 1.571 < Feedforward 1.652 < Ouro 1.899 (lower = more balanced)
• Causal intervention: Forcing iterations 2–8 to reuse iteration-1 expert identities (retaining iteration-specific weights) reduces accuracy on all six evaluated tasks, proving iteration-specific expert selection matters.

Feedforward comparison: Huginn uses 75% fewer resident parameters (8B vs. 32B) than $D_{112}^{MoE}$ yet surpasses it on DROP and GSM8K, matches on MATH500, with a mean gap of only 1.71 points (vs. 4.96 in dense).

---

Theoretical and Practical Implications

Architectural Design Principles

1. Localized recurrence is beneficial: Placing recurrence in the middle of the network (sandwich envelope) allows the recurrent core to focus on iterative state refinement while untied prelude/coda specialize in input encoding and output decoding—consistent with theories that repeated middle-layer computation supports multi-step reasoning.

2. Input injection is a double-edged sword: Persistent injection keeps prompt-relevant context accessible across iterations (benefiting context-dependent tasks) but may "pull" the state back toward input-conditioned representations, interfering with progressive state transformation needed for quantitative reasoning.

3. State initialization is task-dependent: Random initialization should be viewed as an inductive bias, not a universally beneficial choice—its value likely depends on whether tasks require path independence or algorithmic equilibrium behavior.

Practical Implications

• Memory efficiency without performance sacrifice: Looped MoE models can match or exceed much larger feedforward models on several benchmarks while storing 75% fewer parameters—highly relevant for deployment on memory-constrained hardware.
• MoE narrows the loop-to-feedforward gap: Sparse routing appears to compensate for the capacity limitations of tied weights, making looped models more competitive in the MoE regime.
• Better expert utilization: Huginn's design yields more balanced routing, potentially improving hardware utilization and training efficiency.

---

Conclusion

Main Takeaways

1. The prelude–loop–coda sandwich envelope is the primary driver of Huginn's advantage over Ouro, especially for mathematical and multi-step reasoning.
2. Persistent input injection provides complementary gains on context-dependent tasks but requires careful design to avoid harming quantitative reasoning.
3. Random initialization and H/L hierarchical states provide no consistent benefit in matched, large-scale, finitely unrolled settings.
4. All conclusions transfer to MoE models, where Huginn also achieves better expert utilization and narrows the gap to feedforward baselines.

Future Directions

• Scaling laws: Repeat controlled comparisons across model sizes and training budgets to establish scaling behavior.
• Multi-seed replications: Quantify seed-to-seed uncertainty of reported differences.
• MoE component attribution: Vary the three axes separately within MoE models (current study compares complete configurations only).
• Architecture–system co-design: Explore how loop architectures interact with conditional sparsity and hardware to advance the performance–efficiency Pareto frontier.
• HRM-Text-style variants: Evaluate separately parameterized H/L modules (the shared-module H/L finding does not rule out benefits from distinct modules).