Towards Looped Models Done Right, Part I: Topology, Input Injection, Recurrent-State Design
AuthorsBenhao Huang, Chufan Shi, Junlin Chen, Shicheng Wen, Zhengzhong Liu, Eric Xing, Xuezhe Ma
Affiliations[
The paper finds that a carefully designed looped Transformer can outperform comparable models while using far fewer resident parameters.
Key results
Parameter scale for the controlled dense looped-model experiments.
TxT360 training tokens used for the headline dense comparison.
Points gained over full-stack Ouro at the matched dense-model scale.
Points gained over Ouro MoE after 500B training tokens.
Normalized loss at 500B tokens, versus 1.899 for Ouro MoE.
Fewer resident parameters than the compute-matched 112-layer feedforward MoE.
What the paper found
This study separates three looped-model design choices: where recurrence sits, whether the input is reinjected, and how recurrent state is initialized or organized. In controlled dense experiments on TxT360, Huginn’s prelude–loop–coda sandwich improves reasoning over full-stack Ouro: at 730M parameters and 336B training tokens, it raises MATH500 by 12.00 points and DROP by 2.61 points. Reinjecting the prelude representation helps context- and specification-dependent tasks, including coding, but can hurt quantitative reasoning; random initialization and shared high/low states provide no consistent gains. The advantage transfers to mixture-of-experts models: after 500B tokens, Huginn beats Ouro on eight of ten benchmarks, including a 4.70-point gain on GSM8K, and has a lower expert load-balancing loss, 1.571 versus 1.899. Against a compute-matched 112-layer feedforward MoE, Huginn uses 75% fewer resident parameters, surpasses it on DROP and GSM8K, and matches it on MATH500. Freezing expert identities across loop iterations reduces accuracy on all six tested tasks, indicating that iteration-specific routing contributes to performance.
Original 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.
Read the original paperMore in Transformers
Browse all 48 papers →A Conditional No-Hair Theorem for Transformers: How Language Models Forget Their Prompt
Shuyang Xiang
The paper finds that a transformer’s direct memory of its prompt fades as it generates text, while the text itself can carry that memory forward.
DAGFormer: Global Topology Prediction for Transformer Cross-Layer Connections
Xiaocong Yang, Yuren Hao, ChengXiang Zhai
DAGFormer lets a Transformer learn how information should flow across its layers, improving language-model performance across several evaluations.
Decoding the Functional Roles of Register and High-Norm Patch Tokens in Vision Transformers
Neel Varma, Andrew Rufail, Dipika Khullar, Vasu Sharma
This study shows that ViT register tokens carry crucial high-level semantics, while high-norm patch tokens mainly encode lower-level visual structure and contribute little to overall representations.