PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Benhao Huang, Chufan Shi, Junlin Chen, Shicheng Wen, Zhengzhong Liu, Eric Xing, Xuezhe Ma

arXiv:2610.23036v1cs.LG

TL;DR

Ouro- and Huginn-style looped models bundle recurrence placement, input injection, and recurrent-state organization, leaving their effects insufficiently isolated. The paper separates these choices in matched dense-model experiments and evaluates the complete architectures in MoE models; Huginn outperforms Ouro in both settings.

  • Problem

    Ouro- and Huginn-style architectures vary recurrence placement, input injection, and recurrent-state organization together, and their differences have not been systematically isolated.

    Source: §1 · p. 3

  • Method

    The paper separates these three design axes through controlled Ouro-to-Huginn comparisons under matched parameter scale, logical depth, and token budgets, then evaluates the complete architectures in MoE models.

    Source: §1 · p. 3

  • Results

    8 of 10 benchmarks: Huginn-style MoE outperforms Ouro, while dense Huginn outperforms Ouro across all ten benchmarks and MoE Huginn has more balanced expert utilization.

    Source: §1 · p. 3, §4.5 · pp. 16, 18

  • Takeaways & Limitations

    In MoE, Huginn approaches and on some tasks surpasses its feedforward counterpart despite matched memory and FLOP constraints.

    Source: §6 · p. 19

    The MoE study compares complete configurations that change the iteration envelope, input injection, and state initialization together, so their individual contributions remain unattributed.

    Source: §5 · p. 18

Abstract

from arXiv · show

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.

1 Introduction

Looped models reuse Transformer blocks across logical depth, but Ouro and Huginn combine several architectural choices that have not been systematically isolated. The paper separates those choices and reports gains for the Huginn design in dense and MoE models.

1 Introduction

  • Looped models trade additional computation for reduced memory usage as hardware compute throughput grows faster than memory capacity and bandwidth.Source: §1 · p. 2
  • Matched parameter scale, logical depth, and token budgets support controlled comparisons of Ouro-to-Huginn architectural changes.The models are pretrained from scratch and evaluated using a common protocol.Source: §1 · p. 3
  • The framework separates three choices often bundled together: recurrence placement, input reintroduction, and recurrent-state organization.Source: §1 · p. 3
  • The complete Huginn design outperforms Ouro across all ten dense benchmarks, and Huginn again outperforms Ouro in MoE models.Huginn MoE models also distribute expert load more evenly.Source: §1 · p. 3

2 Preliminary

Ouro and Huginn differ in how they arrange recurrence, inject inputs, and organize recurrent state. The section introduces notation for these choices and a common framework for comparing them.

2 Preliminary

  • Feedforward models use distinct blocks, Ouro reuses a full Transformer stack, and Huginn places a tied iterative body between untied prelude and coda stacks.Source: Figure 3 · p. 4

2.1 Notations

  • A tied iterative model applies recurrent updates zt+1 = Rθ(˜zt) over iterations, with a prelude, recurrent body, coda, and input-write map defined in the notation.The prelude and coda are applied once, while the recurrent core is repeatedly applied with shared parameters.Source: §2.1 · p. 4

2.2 Comparing Ouro and Huginn

  • Ouro sets the prelude, input-write, and coda maps to identity, initializes z0 from token embeddings, and reuses the same stack for T iterations.Source: §2.2 · p. 4
  • Huginn retains nontrivial prelude and coda maps and writes the fixed prelude representation into the recurrent stream at each iteration.Source: §2.2 · p. 4
  • The framework separates iteration envelope, input injection, and latent-state organization; Table 1 presents instantiations under a 112-layer budget, with each row isolating one choice.Source: §2.2 · p. 5, Table 1 · p. 5

3 Controlled Ablations of Looped-Model Design

The controlled ablations compare loop designs at matched model scale and logical depth, separating recurrence placement, input injection, and recurrent-state organization. Across ten benchmarks, the results distinguish the sandwich envelope’s reasoning gains, input injection’s context-related gains, and the mixed returns of additional state choices.

3.4.1 Q1: Does Moving Recurrence into a Sandwich Envelope Help?

  • At 730M parameters and 336B tokens, the sandwich envelope raises MATH500 by 12.00 points and DROP by 2.61 points, with reasoning gains persisting across training budgets.The comparison places recurrence in a tied middle core between untied prelude and coda layers; MATH500, BBH-CoT, and DROP favor this envelope across budgets.Source: Table 2 · p. 7, §3.4.1 · p. 7

3.4.2 Q2: Does Input Injection Help?

  • With their respective input-conditioning mechanisms, the middle-loop model outperforms full-stack Ouro on eight of ten benchmarks, trailing on ARC-Challenge and HellaSwag.The comparison is between prelude-state injection in the middle loop and raw-token injection in Ouro.Source: §3.4.2 · pp. 8, 10, Table 3 · p. 9
  • Input injection improves context- and specification-dependent tasks but can impair quantitative reasoning, including MATH500 and GSM8K in the middle-loop model.Prelude-state injection improves MMLU, BBH-CoT, DROP, HumanEval+, and MBPP+; it reduces MATH500 by 3.60 points and GSM8K by 2.51 points.Source: §3.4.2 · pp. 9–10

3.4.3 Q3: Do Separate Recurrent Latent States Help?

  • Random initialization provides no consistent cross-task advantage: it improves two benchmarks by more than one point but reduces four by more than one point.The evaluation does not test multi-start path independence or extrapolation to recurrent depths beyond those used during training.Source: §3.4.3 · p. 12
  • A shared-module H/L state hierarchy also provides no consistent aggregate benefit and substantially degrades MATH500 under matched recurrent compute and training controls.The conclusion is limited to the shared-module setting and does not rule out separately parameterized H and L modules.Source: §3.4.3 · p. 13

3.5 Summary

  • Along the construction path, the sandwich envelope and persistent input injection account for the broadest gains, while a separate random-initialized state adds no consistent benefit.The envelope’s clearest gains are on mathematical and multi-step reasoning; injection adds gains especially on BBH-CoT and code generation, while random state initialization reverses some prior gains.Source: §3.5 · p. 14

4 Transfer to MoE Models

In MoE models, Huginn outperforms Ouro on most benchmarks, routes experts more evenly, and narrows the performance gap to a matched-depth feedforward reference.

4.1 Configurations for MoE Comparison

  • Both recurrent models use 28 routed layers, 25 experts, top-2 routing, 8.0B resident parameters, and 793.9M active parameters per physical pass; the 112-layer reference has 32.0B resident parameters.Unrolling the recurrent models yields approximately 3.2B active parameter applications per token; the reference is matched in logical depth and active compute.Source: §4.1 · p. 15

4.2 Huginn-Style MoE versus Ouro

  • Huginn MoE beats Ouro on eight of ten benchmarks, with largest gains on GSM8K (+4.70) and MATH500 (+3.60).The table reports benchmark performance after 500B training tokens.Source: §4.2 · p. 15, Table 5 · p. 15
  • Huginn’s gains over Ouro are largest on GSM8K (+4.70) and MATH500 (+3.60), with additional leads on MMLU (+2.78), DROP (+1.95), and HumanEval+ (+1.22).Ouro leads on ARC-Challenge by 0.86 points and BBH-CoT by 0.20; Huginn’s leads on HellaSwag and TriviaQA are 0.28 and 0.07 points.Source: §4.2 · p. 15

4.3 Routing Diagnostics

  • At 500B tokens, Huginn’s load-balancing loss is 1.571 versus 1.652 for feedforward and 1.899 for Ouro; freezing loop-1 routes reduces accuracy on all six evaluated tasks.Lower load-balancing loss indicates more even expert utilization; the route-freezing intervention retains iteration-specific mixture weights.Source: §4.3 · p. 16, Figure 10 · p. 16, Figure 12 · p. 17

4.4 Huginn-MoE versus Feedforward MoE

  • The feedforward MoE leads on seven of ten benchmarks, while Huginn leads on DROP and GSM8K, matches it on MATH500, and reduces the mean gap from 4.96 to 1.71 points.Across MATH500, GSM8K, HumanEval+, and MBPP+, Huginn averages 61.48 versus 62.79 for the feedforward reference.Source: §4.4 · p. 16, Figure 13 · p. 18

4.5 Summary

  • The section’s summary finds that Huginn combines more balanced expert utilization and meaningful iteration-specific expert selection with a narrowed feedforward gap.The feedforward baseline still leads overall, but Huginn outperforms it on DROP and GSM8K and matches it on MATH500.Source: §4.5 · pp. 16, 18

5 Future Work

The authors identify model-scale and replication limits, along with unresolved attribution in the MoE transfer study.

5 Future Work

  • The experiments cover finite-budget behavior at one practical scale, so they do not establish a scaling law; models with hundreds of billions of parameters exceed the current compute budget.The authors also call for matched comparisons across model sizes and training budgets, plus multi-seed replications to quantify uncertainty.Source: §5 · p. 18
  • The MoE comparison changes iteration envelope, input injection, and state initialization together, leaving their separate contributions unattributed.The authors propose matched MoE ablations that vary these axes separately.Source: §5 · p. 18

6 Conclusions

The paper isolates three design axes in looped models and reports gains for the prelude–loop–coda topology. Its MoE transfer study retains Huginn’s advantages, while broader claims remain bounded by open questions about scale, seeds, and design-space coverage.

6 Conclusions

  • Under matched parameter scale, logical depth, and token budget, the prelude–loop–coda sandwich improves mathematical reasoning, context extraction, and other reasoning-intensive tasks.Input injection primarily benefits knowledge retrieval and context access, while contributing little to quantitative reasoning.Source: §6 · p. 19
  • The MoE study finds Huginn’s advantages over Ouro persist, with more balanced expert utilization and performance approaching or surpassing its feedforward counterpart on some tasks.The models are trained from scratch under the same data budget and matched memory and FLOP constraints.Source: §6 · p. 19
  • The looped-model design space remains underexplored, and meaningful comparisons require matched parameter scale, training budget, inference cost, and evaluation protocol.Open questions include interactions with conditional sparsity and architecture–system co-design.Source: §6 · p. 19

A Model and Training Configuration

The appendix specifies the shared Transformer block structure, points to the MoE geometry accounting, and lists optimization settings.

A Model and Training Configuration

  • All models use a decoder-only PreNorm block with hidden size 1536, 24 query heads, 6 key–value heads, and a dense SwiGLU MLP of width 4352.MoE variants retain the block structure and replace only the dense feed-forward branch with 25 width-2432 SwiGLU experts using sigmoid top-2 routing and no shared experts.Source: A Model and Training Configuration · p. 24
  • Table 6 gives the MoE models’ geometry and non-embedding parameter accounting.Source: Table 6 · p. 24
  • Training uses TxT360, sequence length 8192, AdamW, 200 warmup steps, cosine decay, and weight decay 0.1; the best learning rate is selected from 4×10^-4, 6×10^-4, and 8×10^-4.Each run uses an effective batch of 512 sequences, or 4,194,304 tokens per optimizer step.Source: A Model and Training Configuration · p. 24

B Evaluation Protocol

The evaluation protocol specifies task-specific metrics, prompts, generation limits, and stopping rules. Likelihood tasks score answer choices without generation; generation tasks use greedy single-sample decoding.

B Evaluation Protocol

  • Likelihood tasks score all answer choices without generation, while generation tasks use greedy single-sample decoding, fixed demonstrations, and no chat template.Evaluation batch size and cache use vary only for throughput.Source: B Evaluation Protocol · p. 24
  • Task-specific demonstrations and stopping rules govern chain-of-thought, short-answer, mathematics, and code generation evaluations.ARC-C reports answer-token-normalized choice accuracy, HellaSwag raw highest-likelihood choice accuracy, and BBH-CoT uses the three-shot protocol.Source: B Evaluation Protocol · p. 25
  • Prompts longer than the input cap are left-truncated; in MMLU, the 2,048-token cap can remove early demonstrations while preserving the current question and choices.The two cap values specify maximum input and generated-token lengths; “–” marks likelihood tasks without decoded output.Source: B Evaluation Protocol · pp. 24–25
  • HumanEval+ and MBPP+ use EvalPlus extra tests; HumanEval+ uses 512-token traces, with matched 1,024-token checks producing the same pass@1 values.Most MATH500 rows use standard Minerva prompts with EOS stopping, while three Q3 rows use a pinned prompt variant and next-Problem: stopping.Source: B Evaluation Protocol · p. 25

C MoE Routing Probe

A routing probe measures expert choices across recurrent iterations, then tests their performance relevance by freezing route identities after iteration 1.

C MoE Routing Probe

  • The probe uses a 500B-token Huginn-style MoE checkpoint with 12 recurrent layers, eight iterations, 25 experts per layer, and top-2 routing.Its 48-input probe set spans six tasks; each adjacent-iteration transition has a nonzero expert-switch rate, averaging 49.4%.Source: C MoE Routing Probe · p. 25
  • Figure 11 heatmaps fix one recurrent layer per panel, with experts in columns, iterations in rows, and cells showing top-2 selection fractions.Cell fractions aggregate valid token decisions.Source: C MoE Routing Probe · p. 25
  • Freezing expert identities for iterations 2–8 lowers accuracy on all six evaluated tasks, with drops of 11.2 points on MMLU and 9.6 on ARC-Easy.BoolQ drops 8.7 points, ARC-Challenge 8.3, PIQA 1.5, and COPA 1.0.Source: C MoE Routing Probe · p. 25

D Related Work on Recurrent Depth and Sparse Routing

The related work is organized around how prior research informs this paper’s comparisons: attribution, controlled mechanisms, and compute-matched sparse models.

D Related Work on Recurrent Depth and Sparse Routing

D.1 Complete Recipes Leave Architectural Attribution Unresolved

Prior work establishes full-stack and middle-loop recurrence as viable recipes, but endpoint comparisons leave architectural attribution unresolved. The review also situates controlled recurrent-state studies and sparse-model comparisons.

D.1 Complete Recipes Leave Architectural Attribution Unresolved

  • Huginn’s middle-loop and Ouro’s full-stack recurrence establish viable pretraining recipes, but complete-recipe comparisons cannot isolate which architectural choices drive differences.The paper’s controlled dense path holds resident scale, logical block executions, data, checkpoint tokens, and evaluation protocol fixed while varying the iteration envelope, input write, and state organization.Source: §D.1 · p. 26

D.2 Input and State Design Shape Recurrent Dynamics

  • Prior work links input injection and state channels to recurrent dynamics, while the paper’s input-write ablations compare write versus no write within each family.Raw-token and contextualized writes remain coupled to different recurrence envelopes, so those comparisons do not isolate write type from envelope.Source: §D.2 · p. 26

D.3 Sparse Routing Adds Conditional Capacity to Recurrent Depth

  • LoopMoE leads a matched vanilla MoE on eight of nine benchmarks at 200B tokens, while persistence of its 9B gap through full training remains unresolved.This paper tests whether Huginn–Ouro ordering transfers to MoE, using resident- and active-matched models trained through 500B tokens and a compute-matched 112-layer feedforward reference; component-level MoE ablations remain open.Source: §D.3 · p. 27

Related papers

Loading 2610.23036v1…