循环状态空间语言模型采用自适应退出状态选择

背景

近年来，循环语言模型（Looped Language Models）的研究表明，许多推理问题受益于更大的计算深度，而非额外的独立参数。此前的工作几乎完全集中在Transformer主干上，而这一原则是否同样适用于状态空间语言模型（State-Space Language Models）尚不清楚。状态空间模型（如Mamba）以其线性复杂度和高效的序列建模能力著称，但能否通过循环计算来提升推理能力仍是一个开放问题。

核心内容

本文研究了两种循环架构：Looped Mamba 和 Looped Hybrid Mamba-Transformer。它们通过重复应用共享的Mamba（或混合）模块，引入显式的有限深度循环计算。具体来说，模型将同一个参数化块多次应用，使得每个输入序列经历多步迭代变换，从而在参数总量不变的前提下增加推理时的计算深度。

实验设计：研究在两个受控推理任务上进行评估：

• Mano（模算术操作）：涉及模运算的符号推理。
• p-hop induction：一种需要多步归纳跳跃的推理任务。

结果：

• Looped Mamba 在这两个任务上持续优于参数匹配的非循环基线（即相同参数量但非循环的Mamba模型）。
• 在多种设置下，Looped Mamba 的表现达到或超过了同等有效深度的非循环模型（即深度相同但参数独立的模型）。

预训练实验：研究进一步扩展到语言模型预训练，采用了**参数匹配（iso-parameter）和计算量匹配（iso-FLOPs）**两种对比协议。这两个协议共同分离了参数共享与有效深度的影响：

• 在参数匹配条件下，循环模型以更少的独立参数在下游基准上保持竞争力。
• 在严格计算量匹配（iso-FLOPs）的比较中，更深非循环模型在验证困惑度上仍保持优势。

自适应退出状态选择：本文借鉴了Ouro的两阶段退出门控（exit gate），并将其适配到Looped Mamba，用于对循环步输出进行阈值控制的选择。由于所有循环步仍然会被执行，选中的退出步代表的是预测深度，而非减少实际计算量。在研究的规模下，自适应退出状态选择在中深度上提升了下游性能，而要实现实际推理时间节省，则需要额外的状态处理机制（例如提前退出但需处理中间状态）。

关键要点

• Looped Mamba 和 Looped Hybrid Mamba-Transformer 将共享模块重复应用，在不增加参数的前提下增加推理深度，与Transformer上的循环思想类似。
• 在受控推理任务（Mano和p-hop induction）中，循环架构一致优于参数匹配的非循环基线，且有时能达到与更深非循环模型相当的性能。
• 预训练对比表明：在参数匹配下，循环模型能以更少独立参数在下游任务上保持竞争力；但在计算量匹配下，非循环的深模型仍更优。
• 自适应退出状态选择（基于Ouro的两阶段门控）可以在中间深度改善下游表现，但由于所有循环步都被执行，实际推理加速需要额外的状态处理（如提前终止并有效利用中间状态）。

意义与影响

本研究首次系统性地将循环计算原则从Transformer扩展到状态空间模型（Mamba），验证了“通过共享模块增加计算深度”这一范式在非Transformer架构中同样有效。这为构建参数高效、计算可扩展的语言模型提供了新思路。

具体而言：

• 对于资源受限场景，Looped Mamba提供了用更少独立参数达到类似性能的路径，降低了存储和通信开销。
• 自适应退出机制表明，模型可以在不同推理步上动态选择输出，从而可能根据任务复杂度自适应调整深度，但实际加速仍需进一步优化。
• 本研究也指出了当前循环状态空间模型的局限：在严格计算量匹配下，更深的非循环模型在困惑度上仍占优，暗示循环共享带来的参数效率有时以表达能力为代价，需要更深入的理论分析。

未来工作可探索更高效的退出机制、状态压缩技术，以及将循环Mamba扩展到更大规模模型和更多实际任务。