Looped State-Space Language Models with Adaptive Exit-State Selection

2026-07-11

AI Summary

This work investigates whether enhancing computational depth—rather than merely increasing parameter count—can effectively improve reasoning capabilities in state space language models and address the challenge of selecting output states within recurrent structures. To this end, the authors propose a recurrent Mamba architecture combined with a Mamba-Transformer hybrid design, incorporating limited recursive computation through shared modules and introducing a two-stage adaptive exit gating mechanism. This study marks the first integration of recurrent computation into state space models. Under equal-parameter or equal-FLOPs conditions, the proposed approach outperforms non-recurrent baselines on Mano and p-hop inductive reasoning tasks while maintaining competitive performance on downstream tasks with fewer parameters, thereby demonstrating that increased computational depth is more effective than simple parameter scaling.

Abstract

Recent work on looped language models suggests that many reasoning problems benefit from greater computational depth rather than from additional independent parameters. Existing studies, however, focus almost exclusively on Transformer backbones, leaving open whether this principle also applies to state-space language models. We investigate Looped Mamba and Looped Hybrid Mamba-Transformer architectures, which repeatedly apply a shared Mamba (or hybrid) block to introduce explicit finite-depth recurrent computation. On two controlled reasoning tasks-Mano (modular-arithmetic manipulation) and p-hop induction-Looped Mamba consistently outperforms parameter-matched non-looped baselines and, in several settings, matches or exceeds non-looped models of equal effective depth. We then extend the study to language model pre-training under matched iso-parameter and iso-FLOPs protocols, which jointly disentangle the effects of parameter sharing and effective depth: looped models remain competitive on downstream benchmarks with substantially fewer distinct parameters, although deeper non-looped models retain an advantage in validation perplexity under strict iso-FLOPs comparisons. Finally, we adapt Ouro's two-stage exit gate to Looped Mamba for threshold-controlled selection among recurrent-step outputs. Since all recurrent steps are still executed, the selected exit step represents prediction depth rather than reduced wall-clock computation. At the scales studied, adaptive exit-state selection improves downstream performance at intermediate depths, while actual inference-time savings require additional state-handling mechanisms.

Problem

• looped language models
• state-space models
• computational depth
• parameter efficiency
• reasoning tasks

Innovation

• Looped State-Space Models
• Adaptive Exit-State Selection
• Mamba Architecture
• Computational Depth
• Parameter Sharing