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

# 循环状态空间语言模型与自适应退出状态选择
来源：https://arxiv.org/html/2607.10110
Zhenxuan Yu
Rikkyo University
Yutaka Matsuo
The University of Tokyo
Takeshi Kojima²²footnotemark:2
The University of Tokyo
Yusuke Iwasawa²²footnotemark:2
The University of Tokyo

###### 摘要

关于循环语言模型的最新研究表明，许多推理问题受益于更大的计算深度，而非更多独立参数。然而，现有研究几乎只关注Transformer骨干网络，尚不清楚这一原理是否也适用于状态空间语言模型。我们研究了循环Mamba和循环混合Mamba–Transformer架构，它们重复使用共享的Mamba（或混合）模块，引入显式有限深度循环计算。在两种受控推理任务（Mano（模运算操作）和p-hop归纳）上，循环Mamba始终优于参数匹配的非循环基线，并且在若干设置中，匹配或超越了同等有效深度的非循环模型。然后，我们将研究扩展到在匹配的iso-parameter和iso-FLOPs协议下的语言模型预训练，以共同解耦参数共享和有效深度的影响：循环模型在下游基准上使用大幅减少的不同参数仍具有竞争力，尽管在严格的iso-FLOPs比较下，更深层的非循环模型在验证困惑度上保持优势。最后，我们将Ouro的两阶段退出门适配到循环Mamba，用于在循环步输出之间进行阈值控制的选择。由于所有循环步仍会执行，所选的退出步代表预测深度而非减少的墙钟计算时间。在所研究的规模中，自适应退出状态选择在中间深度提升了下游性能，而实际的推理时间节省则需要额外的状态处理机制。

¹¹代码将公开发布。

循环状态空间语言模型与自适应退出状态选择
Zhenxuan Yu††感谢: 邮箱:[email protected]
Rikkyo University
Yutaka Matsuo††感谢: 邮箱:{t.kojima, matsuo, iwasawa}@weblab.t.u-tokyo.ac.jp
The University of Tokyo
Takeshi Kojima²²footnotemark:2
The University of Tokyo
Yusuke Iwasawa²²footnotemark:2
The University of Tokyo

## 1 引言

大型语言模型（LLMs）在语言理解、文本生成和复杂推理方面取得了显著进步（DeepSeek-AI et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib1)；Yang et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib6)）。然而，这些进展主要依赖于扩大模型规模、训练数据以及计算密集型的训练和测试时推理（Hoffmann et al., 2022 (https://arxiv.org/html/2607.10110#bib.bib7)；Snell et al., 2024 (https://arxiv.org/html/2607.10110#bib.bib2)）。这种扩展范式给实际部署带来了重大挑战：模型权重和中间激活增加了内存消耗，而基于Transformer的自注意力机制需要在长上下文推理期间维护大量的键值缓存，导致推理成本随上下文长度迅速增长（Beltagy et al., 2020 (https://arxiv.org/html/2607.10110#bib.bib5)）。因此，在固定参数、FLOP或内存预算下提升模型能力，已成为高效LLM架构设计的核心问题。

状态空间模型（SSMs），特别是Mamba-2（Gu and Dao, 2024 (https://arxiv.org/html/2607.10110#bib.bib3)；Dao and Gu, 2024 (https://arxiv.org/html/2607.10110#bib.bib11)），为注意力机制提供了一种高效的替代方案，具有良好的长上下文扩展性，而混合Mamba-Transformer设计（Lieber et al., 2024 (https://arxiv.org/html/2607.10110#bib.bib4)；Waleffe et al., 2024 (https://arxiv.org/html/2607.10110#bib.bib8)；Blakeman et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib12)）进一步成为现代LLMs的实用骨干网络。

另一方面，循环Transformer在深度上重用共享模块以获得有效计算深度，且无需额外参数，近期工作表明这有利于推理（Dehghani et al., 2019 (https://arxiv.org/html/2607.10110#bib.bib9)；Saunshi et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib13)）。

我们的核心思想是跨层共享Mamba或混合模块的参数，并通过显式有限深度循环重复应用它们，构建更深的有效计算路径，而不增加不同可训练参数的数量。这种设计使我们能够解耦参数数量、计算预算和有效深度的影响。我们在需要组合计算和递归检索的受控合成推理任务（Mano和p-hop归纳）上验证了循环Mamba，然后扩展到在受控的iso-parameter和iso-FLOPs协议下的语言模型预训练，共同分离参数共享、额外计算和有效深度的贡献。我们进一步评估下游推理基准，并将循环Mamba与循环Transformer和循环混合Mamba-Transformer进行比较。我们还将Ouro (Zhu et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib14)) 的训练后退出门机制适配到循环Mamba，实现阈值控制的自适应退出状态选择。实验表明，在可比的选定预测深度下，带退出门的变体相比固定深度推理获得了更高的下游性能，同时仅引入极少的额外参数。

本文的主要贡献如下：首先，我们在匹配的iso-parameter和iso-FLOPs协议下，针对基于Mamba和混合Mamba-Transformer的语言模型，呈现了显式有限深度循环的受控实证研究，两种协议共同解耦了参数共享、计算预算和有效深度的影响。其次，我们在受控推理任务（Mano, p-hop归纳）以及包含下游基准的语言模型预训练上展示了实验结果，表明循环深度在固定参数预算下提升了组合和递归计算能力。第三，我们为循环状态空间语言模型引入了一种自适应退出状态选择机制，将Ouro的训练后退出门适配到Mamba和混合骨干网络，并展示了阈值控制的自适应退出状态选择能够在可比的选定预测深度下提高下游性能。

## 2 相关工作

#### 状态空间模型与Mamba。
状态空间模型（SSMs），特别是Mamba和Mamba-2，为纯Transformer架构的效率瓶颈提供了一个有前景的替代方案（Gu and Dao, 2024 (https://arxiv.org/html/2607.10110#bib.bib3)）。先前的工作表明，Mamba风格的模型在小到中等规模的语言建模中能够匹配或超越Transformer，而Mamba-2则通过结构化状态空间对偶性进一步建立了SSMs与注意力之间的联系，并引入了更高效的核心层设计（Dao and Gu, 2024 (https://arxiv.org/html/2607.10110#bib.bib11)）。由于与标准注意力相比具有有利的序列建模和内存效率特性，基于Mamba的架构为高效的长序列语言建模和推理提供了一条有吸引力的路径。同时，混合Mamba-Transformer (Lieber et al., 2024 (https://arxiv.org/html/2607.10110#bib.bib4)；Waleffe et al., 2024 (https://arxiv.org/html/2607.10110#bib.bib8)) 架构开始出现在最近的LLM系统中；例如，Nemotron 3 结合了Mamba、Transformer和混合专家组件，以提高吞吐量并支持长上下文建模（Blakeman et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib12)）。这些发展表明，Mamba和混合Mamba-Transformer设计正成为现代高效LLM架构的重要组成部分。

#### 递归Mamba模型。
一项密切相关的最新研究也在递归推理设置中探索了Mamba-2。Wang and Reid (2026 (https://arxiv.org/html/2607.10110#bib.bib17)) 将Tiny Recursive Model (TRM) (Jolicoeur-Martineau, 2025 (https://arxiv.org/html/2607.10110#bib.bib16)) 中的注意力模块替换为Mamba-2注意力混合模块，并在ARC-AGI上评估了所得模型。他们的结果表明，Mamba-2混合算子能够保持并在某些情况下提升小型任务专用模型中的递归推理能力。然而，他们的重点是一个大约7M参数的小型TRM风格架构，专门针对ARC-AGI风格的抽象推理进行训练。相比之下，我们研究的是50M–370M规模的通用预训练Mamba语言模型，并在标准下游基准上进行评估。

#### 隐式与自迭代语言模型。
Schöne et al. (2025 (https://arxiv.org/html/2607.10110#bib.bib15)) 的密切相关工作研究了向固定点自迭代的隐式SSM；它们的迭代充当固定点求解器，使用幻影梯度训练，而我们研究的是对整个Mamba（或混合）骨干网络进行显式的、梯度跟踪的有限深度循环，从而能够进行受控的iso-parameter和iso-FLOPs比较。

#### 其他循环模型。
循环或递归深度架构在序列建模中有着悠久的历史。Universal Transformers 在序列表示上递归应用相同的变换，结合了Transformer的并行性和全局感受野以及递归归纳偏置；他们还引入了自适应计算机制以实现动态停止（Dehghani et al., 2019 (https://arxiv.org/html/2607.10110#bib.bib9)）。最近，Saunshi et al. (2025 (https://arxiv.org/html/2607.10110#bib.bib13)) 研究了用于推理的循环Transformer，并表明一个\(k\)层Transformer循环\(L\)次几乎可以匹配，有时甚至超越一个非循环的\(kL\)层模型在合成和下游推理任务上的表现，同时使用远少于的不同参数。这些结果表明，许多推理任务可能更依赖于有效深度而非独立参数。然而，这些研究主要关注Transformer骨干网络。相比之下，我们研究相同的循环深度原理是否适用于Mamba和混合Mamba-Transformer架构。基于这一工作线，Ouro (Zhu et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib14)) 进一步引入了一个熵正则化退出门，学习循环步上的令牌级分布，用于循环Transformer的自适应退出状态选择；我们将此机制适配到循环Mamba和混合Mamba-Transformer骨干网络，并研究了其选定深度与性能的权衡。

#### 关于循环次二次模型的并行工作。
与本文并行的工作LT2 (Deng et al., 2026 (https://arxiv.org/html/2607.10110#bib.bib19)) 研究了具有线性和稀疏注意力的循环架构，并在0.6B–1.3B规模上预训练了循环Mamba-2模型，观察到纯Mamba-2的循环模型在下游任务上落后于基于注意力的循环模型，并可能出现训练不稳定性。值得注意的是，LT2 有意避免在大规模上使用自适应计算。我们的研究在三个维度上是互补的：(i) 我们进行了严格受控的iso-parameter *和* iso-FLOPs比较，以解耦参数共享与有效深度，包括在合成组合推理任务上；(ii) 我们研究了混合Mamba–Transformer骨干网络，其中注意力层与Mamba层一起循环；(iii) 我们为循环状态空间模型引入了一种自适应退出状态选择机制，这是LT2未涉及的。

## 3 提出方法

### 3.1 整体架构

#### 循环Mamba模型。
我们使用循环模型 (Saunshi et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib13)) 的符号研究标准循环模型构造与Mamba骨干网络。对于任何序列到序列函数\(f\)，令 \(f^{(R)} = \underbrace{f \circ f \circ \cdots \circ f}_{R \ \mathrm{times}}\) 表示通过对\(f\)递归应用\(R\)步得到的函数。在本文中，\(f\)对应一个\(N\)层Mamba骨干网络\(M_{\theta}\)，我们将得到的循环Mamba模型记为\(N \otimes R\)，其中\(N\)是独立参数化的Mamba层数，\(R\)是循环步数。因此有效深度为\(NR\)。由于同一\(N\)层骨干网络在循环步中重复使用，参数预算随\(N\)扩展，而FLOP预算随\(NR\)扩展。因此，\(N \otimes 1\)对应标准非循环模型，\(N \otimes R\)与\(N \otimes 1\)是iso-parameter的，并且与非循环模型\(NR \otimes 1\)是iso-FLOPs的。

给定输入令牌\(x\)，我们计算

\[
\begin{aligned}
h^{(0)} &= \mathrm{emb}(x), \\
h^{(r)} &= M_{\theta}(h^{(r-1)}), \quad r = 1, \ldots, R.
\end{aligned}
\]

并从 \(p_{\theta}^{(R)}(\cdot \mid x_{\leq t}) = \mathrm{softmax}\!\left(\mathrm{lmhead}(h_{t}^{(R)})\right)\) 得到最终语言模型分布。

在预训练期间，我们仅优化长度为\(T\)的序列的最后一个循环步：

\[
\mathcal{L}_{\mathrm{final}} = \mathbb{E}_{x}\left[\sum_{t=1}^{T-1} -\log p_{\theta}^{(R)}(x_{t+1} \mid x_{\leq t})\right].
\]

#### 循环混合Mamba-Transformer模型。
我们还将相同的循环构造扩展到混合Mamba-Transformer骨干网络。在此设置中，循环骨干网络\(M_{\theta}\)仍然是一个\(N\)层参数化模块，但其一小部分层实现为因果自注意力，其余层为Mamba-2层。完整的注意力配置见附录E.4 (https://arxiv.org/html/2607.10110#A5.SS4)。对于140M混合模型，我们将索引为\(\{5,10,15,20\}\)的层替换为因果注意力层；对于370M混合模型，我们将索引为\(\{6,18,30,42\}\)的层替换为因果注意力层。重要的是，循环应用于整个混合骨干网络，而非仅Mamba层：

\[
h^{(r)} = M_{\theta}^{\mathrm{hybrid}}(h^{(r-1)}), \qquad r = 1, \ldots, R.
\]

这一选择遵循标准循环模型公式，其中循环应用于完整的序列到序列模块，而非单个子模块。它还允许在每次循环步中重新访问Mamba状态空间更新和基于注意力的令牌交互。因此，每次循环执行另一轮Mamba和注意力组件之间的交互，实现隐藏表示的迭代细化。这种设计使我们能够评估当少量注意力层插入到原本基于Mamba的骨干网络中时，循环参数共享是否仍然有效。

#### 三阶段训练流程。
我们分三个阶段训练循环Mamba模型。首先，我们在没有退出门的情况下，使用上述最终循环目标预训练基础循环模型。其次，我们继续使用从Ouro (Zhu et al., 2025 (https://arxiv.org/html/2607.10110#bib.bib14)) 改编的熵正则化自适应计算目标进行训练，联合使用所有循环步的输出，学习一个关于退出深度的令牌级分布。第三，我们冻结Mamba骨干网络和语言建模头，仅训练退出门参数，并使用聚焦的自适应目标。