Overview
This research investigates looped state-space language models, specifically Looped Mamba and Looped Hybrid Mamba-Transformer architectures. The study explores whether the principle that reasoning problems benefit from greater computational depth, rather than additional independent parameters, applies to state-space language models, as previously suggested for Transformer backbones. The investigation includes performance comparisons on controlled reasoning tasks and language model pre-training protocols, as well as an adaptation of an adaptive exit-state selection mechanism.
Research Context
Prior work on looped language models has indicated that many reasoning tasks benefit more from increased computational depth than from a larger number of independent parameters. However, these earlier studies primarily focused on Transformer-based architectures. A gap existed regarding whether this observation extends to state-space language models. This research aims to address that gap by examining Looped Mamba and Looped Hybrid Mamba-Transformer models, which incorporate explicit finite-depth recurrent computation through repeated application of a shared block.
Approach
The study employed Looped Mamba and Looped Hybrid Mamba-Transformer architectures. These architectures involve the recurrent application of a shared Mamba or hybrid block. Two controlled reasoning tasks, Mano (modular-arithmetic manipulation) and p-hop induction, were used for initial evaluations. The researchers performed comparisons against parameter-matched non-looped baselines and, in some settings, against non-looped models of equal effective depth.
Further investigation extended to language model pre-training. This pre-training phase utilized matched iso-parameter and iso-FLOPs protocols. These protocols were designed to disentangle the effects of parameter sharing from effective depth. Finally, the researchers adapted Ouro's two-stage exit gate to Looped Mamba. This adaptation was intended for threshold-controlled selection among recurrent-step outputs. It is noted that, with this adaptation, all recurrent steps are still executed, meaning the selected exit step indicates prediction depth, not reduced wall-clock computation.
Findings
- On the Mano and p-hop induction controlled reasoning tasks, Looped Mamba consistently outperformed parameter-matched non-looped baselines.
- In several settings on these reasoning tasks, Looped Mamba matched or exceeded the performance of non-looped models with equal effective depth.
- During language model pre-training under matched iso-parameter protocols, looped models remained competitive on downstream benchmarks while using substantially fewer distinct parameters.
- Under strict iso-FLOPs comparisons in language model pre-training, deeper non-looped models maintained an advantage in validation perplexity compared to looped models.
- The adaptation of Ouro's two-stage exit gate to Looped Mamba improved downstream performance at intermediate depths.
Key Limitations Mentioned by Researchers
The study notes that for adaptive exit-state selection, at the scales investigated, actual inference-time savings would require additional state-handling mechanisms. The current setup for exit-state selection represents prediction depth because all recurrent steps are executed.