🤖 AI

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

本文表明，将循环架构应用于基于 Mamba 的状态空间模型，通过增加计算深度增强了推理能力和参数效率，同时引入了自适应退出状态选择以优化预测深度，尽管文中指出实际的推理时间节省仍需要进一步的状态处理机制。

原作者： Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

发布于 2026-07-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你拥有一个名叫 Mamba 的天才、超级聪明的机器人厨师。这位厨师以速度极快且内存效率极高而闻名，无需一整套庞大且各种各样的工具，就能做出大餐。但有一个问题：当面对非常棘手的食谱——比如复杂的数学谜题或“寻找隐藏食材”的游戏时，Mamba 有时会卡住。它需要更深层的思考，但它没有足够的“层数”来进行深度思考，除非建造一个全新的、巨大的厨房（那将耗资巨大）。

于是，来自立教大学和东京大学的研究人员提出了一个简单而大胆的问题：如果我们不建造更大的厨房，而是让 Mroba 反复烹饪同一道菜，在每一次循环中都精炼其风味呢？

“循环”厨房：一个厨师，多次尝试

在 AI 世界中，大多数模型就像工厂的流水线。一段数据（比如句子中的一个词）沿着流水线移动，经过 24 个不同的工作站（层），然后完成任务。如果你想让机器人思考得更深入，通常的做法就是增加更多的工作站。

作者们尝试了另一种方法。他们构建了一个循环 Mamba（Looped Mamba）。想象一下，Mamba 是一个单一的、超级有才华的工作站。他们没有将数据发送到一条由 24 个不同工作站组成的长线上，而是让数据连续回到这个同一个工作站 24 次。

神奇之处： 机器人在每一次循环中都使用完全相同的“大脑”（参数）。这就像一个学生为了理解一段课文，把同一个段落反复读四遍，而不是去读四个不同的、令人困惑的段落。
结果： 在被称为 Mano（模运算）和 p-hop induction（寻找隐藏模式）的复杂逻辑谜题上，这种“重读”策略效果惊人。一个只有 4 层但循环了 6 次的模型（4 ⊗ 6）击败了一个拥有 24 个独特层的标准模型。这表明，对于推理任务而言，思考的深度比工具的多样性更重要。

“出口门控”：选择最佳版本

但是等等，如果机器人不停地重读这段文字直到永远呢？那是浪费时间。研究人员添加了一个聪明的特性，叫做出口门控（Exit Gate）。

把这个门控想象成一个智能选择器。随着 Mamba 处理数据，门控会询问：“答案变得清晰了吗？”

如果在 2 次循环后答案已经很清晰了，门控会说：“停！我们使用这次循环的结果。”
如果问题很难，门控会说：“继续！”并选择第 3 次或第 4 次循环的结果。

这被称为自适应退出状态选择（Adaptive Exit-State Selection）。论文发现，对于较小的模型（约 1.4 亿参数），这个门控是一个游戏规则改变者。它允许模型为每个特定的词选择“恰到好处”的思考时间，通常比强制机器人始终使用最大循环次数的结果表现得更好。

然而，这里有一个转折。研究人员非常谨慎地指出，这目前还不能真正节省电脑的电量或时间。 尽管门控选择了一个早期的结果，但机器人在后台仍然物理上运行了所有的循环。这是因为 Mamba 的记忆状态是连续的：第 3 次循环的结果取决于第 2 次循环留下的状态，而第 2 次又取决于第 1 次。你不能在不破坏后续所有词的链条的情况下，直接从中间某个循环中“跳出”。因此，门控选择的是预测深度（即使用哪个版本的答案），但*实际运行的时间（wall-clock computation）*保持不变。要真正节省时间，我们需要一种处理机器人记忆状态的新方法，作者表示这是未来的工作方向。

现实检查：哪些没奏效（以及哪些仍是未知数）

这篇论文对它没有解决的问题进行了坦诚的说明。

规模并不总是意味着更好（针对此技巧）： 当他们将他们的循环 Mamba 与一个运行成本相当（具有相同的计算量或“iso-FLOPs”）的标准非循环模型进行比较时，标准模型在“模型有多困惑”（困惑度/perplexity）这一基本测试中仍然胜出。循环模型在推理任务上表现出色，但在仅仅预测句子中下一个词的任务上，标准的深层模型仍然略胜一筹。
规模很重要： “出口门控”在 1.4 亿参数的模型上表现得非常完美。但当他们尝试将其应用于更大的 3.7 亿参数模型时，门控并没有带来太大的改进。较大的模型似乎更倾向于直接运行所有的循环。作者认为，较大的模型可能更依赖于最后一次深层的思考过程。
它不是解决一切的万能药： 研究人员在合成谜题（如数学树和模式寻找）上测试了这一点。虽然这些证明了概念的可行性，但他们承认，我们还不确定这种“循环”技巧在处理像写小说或解决复杂法律案件这类开放式的、现实世界的推理任务时，是否也能表现得同样出色。

核心结论

该论文表明，通过反复使用同一个聪明的“大脑”，可以在不构建更大、更昂贵的大脑的情况下，让 AI 变得更聪明、更擅长推理。 这是一种“参数高效型”的扩展维度。

已证实： 在特定的逻辑谜题上，循环一个小型的 Mamba 模型可以击败同等规模的标准模型，并能媲美规模大得多的模型。
暗示： 这种方法可能是未来提高 AI 效率的关键，特别是对于那些需要深度、循序渐进思考的任务。
尚未解决： 我们仍需弄清楚如何让“出口门控”真正节省计算时间（而不仅仅是选择输出），以及这种方法是否适用于驱动当今最强大 AI 的那些数十亿参数级的巨型模型。

简而言之，作者们不仅是建造了一个更大的机器人，而是通过让一个小机器人通过第二次（或第三次、第四次）观察问题，教会了它如何进行深度思考。目前来看，这种“再看一眼”的方法看起来非常有前景。

技术摘要：具有自适应退出状态选择的循环状态空间语言模型

问题陈述

大型语言模型（LLMs）的最新进展主要由模型规模、训练数据和计算密集型推理驱动。然而，这种缩放范式引入了显著的部署挑战，包括内存消耗增加，以及由于 Transformer 中自注意力机制维护庞大的键值（KV）缓存而导致的推理成本快速增长。虽然状态空间模型（SSMs），特别是 Mamba-2，提供了一种具有良好长上下文扩展性的高效替代方案，但目前尚不清楚“计算深度优于独立参数”这一原则（近期在循环 Transformer 中得到验证）是否同样适用于 SSMs。此外，现有的关于循环架构的研究几乎完全集中在 Transformer 主干网络上，这在理解循环参数共享如何在 Mamba 和混合 Mamba-Transformer 架构中发挥作用方面留下了空白。

方法论

作者提出了 Looped Mamba 和 Looped Hybrid Mamba–Transformer 架构。其核心机制是通过显式的有限深度循环，共享单个 Mamba（或混合）块的参数。

架构

符号表示： 模型表示为 $N \otimes R$ ，其中 $N$ 是独立参数化层的数量， $R$ 是循环步数。有效深度为 $NR$。
Looped Mamba： 重复应用一个共享的 $N$ $N$ 层 Mamba 主干 $M_\theta$ $M_{θ}$ 。
- $h^{(r)} = M_\theta(h^{(r-1)})$ ，其中 $r = 1, \dots, R$ 。
- 最终输出由第 $R$ 步的隐藏状态导出。
Looped Hybrid： 将循环扩展到包含 Mamba-2 层和一小部分因果自注意力层的主干网络。循环应用于整个混合块，通过状态空间更新和基于注意力的交互进行隐藏表示的迭代细化。
训练协议：
1. 预训练： 使用最终循环目标（仅在第 $R$ 步计算损失）进行优化。
2. 退出门阶段 1： 引入熵正则化的自适应计算目标（改编自 Ouro），利用来自所有循环步的输出，学习用于退出深度的标记级分布。
3. 退出门阶段 2： 冻结主干和头部，仅训练退出门，以鼓励仅在下一步能改善预测损失时才继续计算。

自适应退出状态选择

论文将 Ouro 的两阶段退出门适配到了 Mamba。在推理时，阈值 $q$ 根据累积退出概率来控制每个标记的退出步 $z_i$ 的选择。

语义说明： 作者澄清，在 SSMs 中，由于状态连续性的约束，所有 $R$ 步的状态链必须被保留。因此，“退出”是指选择哪一个循环步的隐藏状态传递给语言建模头；它本身并不会在没有额外的状态处理机制的情况下减少墙钟计算时间或执行步数。

核心贡献

受控实证研究： 作者在匹配的 等参数（相同 $N$ ，不同 $R$ ）和 等 FLOPs（相同有效深度 $NR$）协议下，将 Looped Mamba 和混合模型与非循环 Mamba 及 Transformer 基准进行了严格对比。
推理任务评估： 该研究在合成推理任务（Mano 和 $p$ -hop induction）以及语言模型预训练的下游基准测试中验证了该方法。
自适应退出机制： 引入了用于循环状态空间模型的阈值控制自适应退出状态选择机制，证明了其提高中间深度下游性能的能力。

实验结果

合成推理任务

Mano 任务（模运算）： Looped Mamba 一致优于参数匹配的非循环基准模型。在最难的设置（ $L=24$ ）中，一个 $4 \otimes 6$ 模型达到了 99.35% 的准确率，而 $4 \otimes 1$ 基准模型仅为 7.50%。在等 FLOPs 条件下（有效深度为 24）， $4 \otimes 6$ 模型显著优于 $24 \otimes 1$ 非循环基准模型（99.35% 对比 37.45%），这表明循环参数共享比增加非共享层提供了更强的组合计算归纳偏置。
$p$ -hop Induction（递归检索）： 浅层非循环 Mamba 模型表现挣扎（例如 $1 \otimes 1$ 约为 38-41%），而增加循环显著提高了性能。 $1 \otimes 12$ 模型达到了接近完美的准确率（~99.89%），表明对于递归检索而言，递归深度比不同的参数量更为关键。

语言建模与下游基准

等参数： 在固定参数预算下，循环深度提升了 Mamba-2 和 LLaMA-2 主干网络的验证困惑度（PPL）。
等 FLOPs： 较深的非循环模型（ $24 \otimes 1$ ）在验证困惑度方面保持了对循环模型的优势。然而，循环模型在下游基准测试中仍具竞争力，有时甚至优于更深的非循环基座。
自适应退出：
- 140M 规模： 自适应退出门提升了下游性能。对于 140M Mamba-2 模型，阶段 2 退出门的平均得分为 38.87，平均选择的退出步为 3.00，优于 fixed-loop4 基准（37.79）。
- 370M 规模： 最佳下游性能仍集中在 fixed-loop4 设置附近。退出门变体并未超越峰值 fixed-loop 得分，这表明较大的模型可能更依赖于最终的循环步，或者需要更好的门控校准。
- 权衡： 退出门提供了一个可控的性能与选择深度之间的边界，允许模型有效地利用中间循环状态，而不是强制执行固定的最大深度。

意义与主张

论文声称 循环深度为 Mamba 和混合语言模型提供了一个极具前景的、参数高效的缩放轴。主要发现包括：

深度 vs. 参数： 许多推理任务受益于更大的计算深度，而非额外的独立参数。循环重用 Mamba 层可以在不增加参数量的情况下，显著提高组合计算和递归计算能力。
混合架构的可行性： 循环深度原则也有效地适用于混合 Mamba-Transformer 架构，即注意力层与 Mamba 层共同进行循环。
自适应选择： 虽然目前的 SSM 实现要求执行完整的状态过程，但自适应退出状态选择允许选择最优的预测深度，从而提高在中间深度下的下游性能（特别是在 140M 规模下）。

作者对其主张保持了谦逊，指出实验仅限于高达 370M 参数的模型，并且实际的推理时间计算节省需要未来的工作来研究如何完全跳过循环步的状态处理机制。结果是以受控实证证据的形式呈现，而非对大规模性能的完整表征。