通过循环计算提升推理深度,让Mamba模型更高效解决复杂问题。
Looped State-Space Language Models with Adaptive Exit-State Selection
- 用共享模块循环执行,实现可调控的计算深度
- 在数学推理任务中超越同等参数量的非循环模型
- 支持动态选择输出步数,适合需要深度推理的任务
近期研究表明,许多推理任务受益于更深的计算而非更多独立参数。现有研究主要聚焦于Transformer结构,未验证该原理是否适用于状态空间语言模型。本文探索了循环Mamba和混合Mamba-Transformer架构,通过重复应用共享模块引入显式有限深度的递归计算。在两个受控推理任务(Mano和p-hop归纳)中,循环Mamba始终优于参数匹配的非循环基线,并在若干设置下达到或超过同等有效深度的非循环模型表现。进一步在预训练阶段进行等参数量与等浮点运算量对比,循环模型以更少的独立参数保持下游性能竞争力;尽管在严格等浮点运算约束下,更深的非循环模型仍略胜一筹。最后,将Ouro的两阶段退出门机制引入循环Mamba,实现基于阈值的递归步骤输出选择。由于所有步骤仍被执行,所选退出步代表预测深度而非实际计算节省,实验表明该机制在中等深度下提升下游性能,而真正推理加速需额外状态管理机制。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。