Mamba-3: Improved Sequence Modeling using State Space Principles
Mamba-3:利用状态空间原理改进序列建模
https://arxiv.org/pdf/2603.15569
![]()
摘要
扩展推理时计算已成为大型语言模型(LLM)性能的重要驱动因素,使得推理效率成为与模型质量并列的模型设计核心焦点。虽然当前基于 Transformer 的模型提供了强大的模型质量,但其二次方计算和线性内存使得推理成本高昂。这激发了具有降低的线性计算和恒定内存需求的次二次方模型的发展。然而,许多最近的线性模型为了算法效率牺牲了模型质量和能力,在状态追踪等任务上失败。此外,它们理论上线性的推理在实践中仍然硬件效率低下。在推理优先视角的指导下,我们引入了三个受线性模型的状态空间模型(SSM)视角启发的核心方法改进。我们结合了:(1) 源自 SSM 离散化的更具表现力的递归,(2) 能够实现更丰富状态追踪的复数值状态更新规则,以及 (3) 在不增加解码延迟的情况下提高模型性能的多输入多输出(MIMO)公式。结合架构改进,我们的 Mamba-3 模型在检索、状态追踪和下游语言建模任务上取得了显著增益。在 1.5B 规模上,与下一个最佳模型(Gated DeltaNet)相比,Mamba-3 将平均下游准确率提高了 0.6 个百分点,Mamba-3 的 MIMO 变体进一步将准确率提高了 1.2 个百分点,总增益为 1.8 个百分点。在状态大小实验中,尽管使用其前身一半的状态大小,Mamba-3 实现了与 Mamba-2 相当的困惑度。我们的评估证明了 Mamba-3 推进性能 - 效率帕累托前沿的能力。
1 引言
测试时计算已成为大型语言模型进步的关键驱动因素,思维链推理和迭代细化等技术表明,推理时扩展可以解锁新能力(Snell et al. 2024; Wu et al. 2025)。并行、智能体工作流的迅速兴起仅加剧了对此类模型高效推理和部署的需求(Anthropic 2026; OpenAI 2026)。这种范式转变使得推理效率(Kwon et al. 2023; Li et al. 2024)至关重要,因为 AI 系统的实际影响现在关键取决于它们在部署期间执行大规模推理的能力。模型架构设计在确定推理效率方面起着根本作用,因为架构选择直接决定了生成期间的计算和内存需求。虽然基于 Transformer 的模型(Vaswani et al. 2017)是当前的行业标准,但它们从根本上受限于通过 KV 缓存线性增加的内存需求和通过自注意力机制二次方增加的计算需求。这些缺点激发了最近关于次二次方模型的工作路线,例如状态空间模型(SSM)和线性注意力,它们在保持恒定内存和线性计算的同时,实现了与其 Transformer 对应模型相当或更好的性能。这些模型已进入主流,诸如 Mamba-2(Dao and Gu 2024)和 Gated DeltaNet(GDN)(Schlag, Irie, and Schmidhuber 2021; S. Yang, B. Wang, Y. Zhang, et al. 2025)等层最近被纳入大规模混合模型中,这些模型以高得多的效率匹配纯 Transformer 替代方案的性能(Kimi Team et al. 2025; NVIDIA et al. 2025; Tencent Hunyuan Team et al. 2025; A. Yang et al. 2025)。
尽管线性模型取得了成功,但在提高其性能方面仍有重大进展空间,特别是在推进模型质量和推理效率之间的帕累托前沿方面。例如,Mamba-2 旨在提高相对于 Mamba-1(Gu and Dao 2024)的训练速度和简单性,但牺牲了一些表现力,因此在推理匹配的模型上表现更差。此外,它们已被证明缺乏某些能力,例如糟糕的状态追踪能力,即简单地确定比特序列的奇偶性(Grazzi, Siems, Zela, et al. 2025; Sarrof, Veitsman, and Hahn 2024)。最后,尽管这些次二次方模型因其理论上高效的推理而受到重视并因此被广泛采用,但它们的推理算法在硬件上效率不高。特别是,因为这些算法是从训练视角开发的,它们的解码阶段具有低算术强度(FLOPs 与内存流量的比率),导致大部分硬件保持空闲。
为了从推理优先范式开发性能更高的模型,我们在 Mamba-2 之上引入了三个核心方法变更,受次二次方模型的以 SSM 为中心的视角影响。
指数 - 梯形离散化。我们提供了一种离散化时变、选择性 SSM 的简单技术。通过我们的框架,我们可以推导几种新的离散化方法。我们的一个实例化,称为“指数 - 欧拉”,形式化了 Mamba-1 和 Mamba-2 的启发式离散化,此前缺乏理论依据。我们新的“指数 - 梯形”实例化是“指数 - 欧拉”的更具表现力的泛化,其中递归可以展开以揭示应用于 SSM 输入的隐式卷积。结合显式 B , C偏置项,Mamba-3 可以在经验上替换语言模型架构中的短因果卷积,此前假设这对于递归模型是必不可少的。
复数值状态空间模型。通过将 Mamba-3 的底层 SSM 视为复数值,我们实现了比 Mamba-2 更具表现力的状态更新。这种更新规则的变更旨在使训练和推理轻量级,克服了许多当前线性模型中缺乏状态追踪能力的问题。我们表明,我们的复数值更新规则等价于数据依赖的旋转嵌入,并且可以高效计算(Su et al. 2023),并在经验上证明了其解决先前线性模型能力之外的合成任务的能力。
多输入多输出(MIMO)SSM。为了提高解码期间的 FLOP 效率,我们从基于外积的状态更新切换到基于矩阵乘法的状态更新。从 SSM 的信号处理基础视角来看,这种转换恰好与从单输入单输出(SISO)序列动力学到多输入多输出(MIMO)的泛化一致。在这里,我们发现 MIMO 特别适合推理,因为额外的表现力使得在解码期间内存受限的状态更新期间能够进行更多计算,而不增加状态大小和损害速度。
综上所述,这些改进构成了我们Mamba-3层的核心。在方法论上,我们注意到这些都自然地从以 SSM 为中心的视角产生,但从现代线性层的其他流行视角(如线性注意力或测试时回归)来看并非显而易见;我们在第 5 节进一步讨论这些联系。在经验上,我们在一系列合成状态追踪和语言建模任务上验证了我们新模型的能力和功能。
- 更好的质量。在 1.5B 规模上,Mamba-3 (MIMO) 相比 Transformer 将下游语言建模准确率提高了+2.2,相比 Mamba-2 提高了+1.9 个百分点,相比 GDN 提高了,而 Mamba-3 (SISO) 相比下一个最佳模型 GDN 提高了+0.6 个百分点。此外,在状态大小实验中,状态大小为 64 的 Mamba-3 (MIMO) 与状态大小为 128 的 Mamba-2 的困惑度相匹配,有效地以一半的延迟实现了相同的语言建模性能
- 新能力。Mamba-3 对 SSM 状态的复数化使其能够解决 Mamba-2 无法解决的合成状态追踪任务。我们在经验上证明,高效的类 RoPE 计算能够近乎完美地解决算术任务,而没有 RoPE 的 Mamba-3 和 Mamba-2 的表现并不比随机猜测更好。
- 推理效率。Mamba-3 (MIMO) 提高了硬件利用率。在固定状态大小下,相比 Mamba-2,它将解码 FLOPs 提高了高达4×,同时保持相似的挂钟解码延迟,并同时改善困惑度和下游性能。我们发布了 Mamba-3 的快速训练和推理内核。¹
Mamba-3 (SISO) 相比之前的线性模型提高了质量和能力,而 Mamba-3 (MIMO) 相比 Mamba-3 (SISO) 和其他强基线进一步提升了性能,同时保持了与 Mamba-2 匹配的推理速度。我们的两种 Mamba-3 变体通过其强大的建模能力和硬件高效设计,推进了性能 - 延迟帕累托前沿。
2 预备知识
2.1 符号
![]()
2.2 SSM 预备知识
![]()
![]()
2.3 结构化掩码表示与状态空间对偶性
Mamba-2 表明,一大类 SSM 允许采用矩阵形式,该形式将时间步递归向量化。通过状态空间对偶性(SSD)框架,递归 SSM 可以在并行形式内表示,该形式结合了元素级掩码来模拟状态转移衰减。
SSD 为线性递归与可并行化(基于矩阵乘法)计算形式之间的对偶性提供了一个通用框架。
![]()
![]()
![]()
3 方法论
我们介绍了 Mamba-3,一种具有三个新创新的状态空间模型:用于更具表现力动态的“指数 - 梯形”离散化(第 3.1 节),用于状态追踪的复数值状态空间(第 3.2 节),以及用于提高建模能力和推理时硬件利用率的多输入多输出(MIMO)(第 3.3 节)。这些进展解决了当前次二次方架构在质量、能力和效率方面的局限性。我们在第 3.4 节将这些结合到一个更新后的 Mamba 架构块中。
3.1 指数 - 梯形离散化
![]()
离散化方法在经典控制理论中已得到充分研究,在早期深度学习中的 SSM 工作中使用了几种规范公式(Gu, Goel, and Ré 2022; Gu, Gupta, et al. 2022; Smith, Warrington, and Linderman 2023)。这些机制传统上被陈述并应用于线性时不变(LTI)系统,它们的推导不直接适用于线性时变(LTV)系统。此外,虽然 Mamba-1 在没有证明的情况下将零阶保持(ZOH)方法适应于 LTV 系统,但与选择性 SSM 相关的复杂性促使使用额外的启发式近似,这种近似缺乏理论依据,并且不对应任何已建立的离散化技术。在以下小节中,我们通过我们的离散化框架形式化了当前 LTV SSM 中使用的先前启发式方法,并利用它提出了一种更具表现力的离散化方案。
![]()
3.1.1 指数调整离散化概览
我们引入了一种简单的推导,它导出了针对 LTV(线性时变)状态空间模型的一类新离散化方法。该方法可以通过多种方式实例化;我们展示了其中一种实例化产生了 Mamba-1/2 中使用的启发式方法,从而在理论上证明了它的合理性(指数 - 欧拉)。我们还介绍了一种用于 Mamba-3 的更强大的离散化方法(指数 - 梯形)。
![]()
![]()
![]()
![]()
我们将公式 (4) 称为指数 - 欧拉离散化方法,源于指数积分后接欧拉近似。这一推导证明了 Mamba-1/-2 实现中使用的公式的合理性。
指数 - 梯形(Mamba-3)。然而,欧拉法则仅提供状态输入积分的一阶近似,其局部截断误差按
缩放。相比之下,我们引入了一种广义梯形法则,它提供了积分的二阶精确近似,比欧拉法则具有更高的精度。具体而言,它用依赖于数据的、两个区间端点的凸组合来近似积分。这一泛化扩展了经典的梯形法则(Süli and Mayers 2003),后者简单地平均区间端点(图 1)。
命题 1(指数 - 梯形离散化)。通过广义梯形法则近似公式 (16) 中的状态输入积分,得到递归关系,
![]()
![]()
![]()
![]()
这种并行形式使得针对 SSM 输出的硬件高效、以矩阵乘法(matmul)为中心的训练计算成为可能。
我们注意到,Mamba-3 的卷积连接也可以通过这种并行对偶形式看出,其中公式 (7) 中乘以 2-带状矩阵代表了权重为 β , γ的卷积。在附录 A.1 中,我们利用 SSD 张量收缩机制证明了该并行形式等价于带有状态 - 输入卷积的基础 SSM。
备注 5。Mamba-3 的结构化掩码可以被视为对 Mamba-2 的泛化,后者不使用 2-带状矩阵,而是仅拥有包含 γ t
的对角矩阵 (3)。
3.2 复数值状态空间模型
现代 SSM 的设计以效率为核心目标,这是由扩展到更大模型和更长序列的需求所驱动的。例如,连续的架构逐渐简化了状态转移矩阵:S4 (Gu, Goel, and Ré 2022) 使用了复数值正规加低秩(NPLR)矩阵,Mamba (Gu and Dao 2024) 将其简化为实数对角线,而 Mamba-2 (Dao and Gu 2024) 进一步将其简化为单个缩放单位矩阵。尽管这些简化在很大程度上保持了语言建模性能,但最近的研究 (Grazzi, Siems, Zela, et al. 2025; Merrill, Petty, and Sabharwal 2025; Sarrof, Veitsman, and Hahn 2024) 表明,限制于实数、非负特征值转移会降低模型在简单状态追踪任务上的能力——这里主要指可解群机制(TC⁰),例如奇偶性校验——这可以通过单层 LSTM 解决。这一局限性在 (Grazzi, Siems, Schrodi, et al. 2024) 的定理 1 中被形式化,它源于将转移矩阵的特征值限制在实数上,这无法表示“旋转”隐藏状态动态。例如,考虑定义在二进制输入 { 0 , 1 }
上的奇偶函数,定义为
。此任务可以使用更新公式执行:
,其中 R ( ⋅ )
是一个 2-D 旋转矩阵。这种旋转动态无法用实特征值表示。
3.2.1 具有指数 - 欧拉离散化的复数 SSM
为了恢复这种能力,我们从复数 SSM (8) 开始,它能够表示状态追踪动态。我们表明,在离散化下(命题 5),复数 SSM 可以表述为具有由 2 × 2 旋转矩阵组成的块对角转移矩阵的实数 SSM(命题 2)。随后我们表明,这等同于分别在输入和输出投影 B , C
上应用数据依赖旋转嵌入。这一结果建立了复数 SSM 与数据依赖 RoPE 嵌入之间的理论联系(命题 3)。最后,Su et al. (2023) 中使用的“RoPE 技巧”允许高效实现复数值状态转移矩阵,与实数值 SSM 相比,其计算开销极小。
![]()
![]()
![]()
![]()
![]()
3.2.2 具有指数 - 梯形离散化的复数 SSM
在推导了具有指数 - 欧拉离散化的复数 SSM 的递归后,推广到指数 - 梯形离散化是相似的。命题 4 提供了 Mamba-3 带有 RoPE 技巧的完整递归。
命题 4(具有指数 - 梯形离散化的旋转嵌入等价性)。用指数 - 梯形法则(命题 1)离散化复数 SSM 产生递归
![]()
![]()
证明见附录 B.3。 我们实证验证了,我们通过数据依赖 RoPE 实现的复数 SSM,能够解决带有和不带有标准 RoPE 的实数值 SSM 无法解决的状态追踪任务(表 5b),支持了理论主张。
3.3 多输入,多输出
扩展测试时计算开启了模型能力的新前沿,例如智能体工作流,其中推理占据了整体计算预算越来越大的份额。这使得人们对语言模型的推理效率重新关注,并促进了 SSM 和次二次方层的采用,这些层具有固定大小的隐藏状态,因此提供更低的计算和内存需求。尽管这些新层与 Transformer 相比具有更低的挂钟时间,但它们的解码严重受限于内存,导致硬件利用率低。在本节中,我们利用 SSM 视角引入了一种对 Mamba-3 递归的方法论改进,允许在不增加解码挂钟时间的情况下增加模型 FLOPs,从而在相同解码速度下获得更好的模型。
解码算术强度。为了提高硬件效率,我们需要考虑 token 生成的算术强度,定义为给定操作的 FLOPs 除以输入 - 输出字节数。由于 SSM 解码用空闲计算饱和了内存带宽(即受限于内存),我们希望增加其算术强度,以有效地将计算与内存 I/O 重叠。更具体地说,Mamba 中单次生成的算术强度约为每字节 2.5 次操作(表 2a),而对于 NVIDIA H100-SXM5,bfloat16 矩阵乘法的算术强度约为每字节 295 次操作(NVIDIA 2022)。因此,SSM 解码远未达到计算受限状态,而且此外,尚不清楚如何调整 Mamba 中的现有参数来缓解硬件效率的缺乏。我们注意到,这一观察结果通常适用于其他次二次方模型,例如因果线性注意力。
![]()
![]()
![]()
![]()
![]()
![]()
![]()
![]()
![]()
这缓解了参数的倍数增长,使其变为更合理的加法参数数量增长。附录 C 详细说明了参数化细节,且我们论文中的所有 MIMO 变体均通过降低 MLP 宽度,使其参数数量与对应的 SISO 模型相匹配。
备注 6。为简单起见,本节的所有讨论均针对较简单的 2 项递归(例如由指数 - 欧拉离散化产生的递归);推广到 3 项指数 - 梯形递归的情况是类似的。
3.4 Mamba-3 架构
整体架构遵循 Llama(Grattafiori et al. 2024),交替使用 Mamba-3 和 SwiGLU 块,并采用预归一化(pre-norm)。Mamba-3 块保留了其前身的整体布局,同时引入了几个关键修改。
更新的 SSM 递归。SSD 层被替换为命题 4 中定义的更具表现力的复数值指数 - 梯形 SSM。Mamba-3 默认采用 SISO SSM,以便与其他类 SISO 模型进行公平比较,但其 MIMO 变体可以训练和部署为基线 Mamba-3 的更强替代方案(表 3)。我们的 SSM A A 是复数的,包含由数据依赖投影产生的实部和虚部。结合图 2,这被划分为实数值 A A 和虚数值 Θ Θ;前者像 Mamba-2 中一样传入 SSD 黑盒,而后者通过 RoPE 技巧计算。
![]()
![]()
BC / QK 归一化。在 B , C
投影之后添加了 RMS 归一化,这模仿了现代 Transformer(Henry et al. 2020; Wortsman et al. 2023)和其他近期线性模型(Hu et al. 2025; S. Yang, Kautz, and Hatamizadeh 2025)中常用的 QKNorm。我们将其称为 BC 归一化(BCNorm)或 QK 归一化(QKNorm),两者可互换使用。我们发现 BCNorm 也能稳定大规模运行,从而使得在我们的纯 Mamba-3 模型中移除了门后 RMSNorm 层(该层在 Mamba-2 中引入用于稳定性)。然而,在混合模型中,移除的 RMSNorm 层对于长上下文外推至关重要(表 4)。
B , C偏置。类似于 Yu and Erichson (2025),其证明了在 Mamba-1 的分块变体中向 B 添加通道特定偏置赋予了通用逼近能力,Mamba-3 在 BCNorm 之后将可学习的、头特定的、通道维度的偏置整合到了 B 和 C 分量中。
我们假设这些偏置也在模型中诱导了类似卷积的行为。具体而言,向 B B 和 C C添加偏置将数据独立组件引入 SSM 中,使其功能更类似于卷积。关于偏置参数化的消融实验位于附录 F。
数据独立偏置参数的组合,加上指数 - 梯形离散化(其本身在状态输入上诱导了卷积),在经验上能够免除 Mamba-2 和大多数现代递归模型中存在的短因果卷积及其伴随的激活函数(第 4.2 节)。
4 实证验证
我们通过 Mamba-3 模型在一系列合成和现实世界任务上,实证验证了我们以 SSM 为中心的方法论变更。第 4.1 节在语言建模和基于检索的任务上评估 Mamba-3。第 4.2 节消融了我们新 SSM 组件(如离散化和复数转移)的效果。第 4.3 节探讨了 Mamba-3 家族的推理效率以及 MIMO Mamba-3 相比 SISO 变体在固定推理计算下的优势,第 4.4 节基准测试了我们 Mamba-3 训练和推理内核的性能。
4.1 语言建模
所有模型均使用 FineWeb-Edu 数据集(Penedo et al. 2024)的 1000 亿 token 进行预训练,使用 Llama-3.1 分词器(Grattafiori et al. 2024),上下文长度为 2K,采用相同的标准训练协议。训练和评估细节可在附录 D 中找到。
在所有四个模型规模上,Mamba-3 在各种下游任务上优于流行的基线(表 3)。我们强调,Mamba-3 不使用外部短卷积,该卷积已被经验性地确定为许多高性能线性模型中的重要组件(Allen-Zhu 2025; Gu and Dao 2024; S. Yang, Kautz, and Hatamizadeh 2025)。
4.1.1 MIMO
我们旨在通过在相同设置下训练秩 R = 4
的 MIMO 模型来调查其语言建模能力,从而进一步验证 MIMO 的收益。为了确保总参数数量与基于 SISO 的模型相当,我们降低了 MIMO 模型中 MLP 层的内部维度,以补偿由于 MIMO 投影导致的增加。例如,在 1.5B 参数模型中,MLP 内部维度仅减少了 6.6%,从 4096 减少到 3824。详见附录 C。
在验证困惑度和我们的语言评估任务套件(表 3)上,我们看到当 Mamba-3 模型从 SISO 转向 MIMO 时有显著增益。即,我们在 1.5B 模型上实现了 0.11 的显著困惑度增益,图 3 说明了我们验证损失的向下偏移。在语言评估方面,与 SISO 相比,我们在大多数任务上看到增益,导致相比 SISO 平均增益 1.2 个百分点。
4.1.2 检索能力
除了标准语言建模外,线性模型的一个重要衡量标准是它们的检索能力——它们能从序列早期回忆信息的程度如何(A. Arora et al. 2025; S. Arora, Eyuboglu, et al. 2025)。与注意力模型不同,注意力模型可以通过增长的 KV 缓存自由回顾过去上下文,线性模型必须将上下文压缩到固定大小的状态中。这种权衡反映在 Transformer 基线显著更强的检索分数上。为了在此视角下评估 Mamba-3,表 4 使用我们第 4.1 节中的预训练 1.5B 模型,在现实世界和合成大海捞针(NIAH)任务(Hsieh et al. 2024)上将其与基线进行比较。我们将任务序列长度限制为 2K token 以匹配训练设置,并遵循 S. Arora, Eyuboglu, et al. (2025) 和 S. Arora, Timalsina, et al. (2024),对我们的现实世界任务采用完形填空风格格式,以反映下一个 token 预测目标。
![]()
Mamba-3 在现实世界关联回忆和问答(TQA, SQuAD)上具有竞争力,但在从半结构化或非结构化数据中提取信息时表现挣扎(SWDE, FDA)。然而,在合成 NIAH 任务上,Mamba-3 在大多数情况下超过或匹配基线,并且值得注意的是,相比其前身 Mamba-2,表现出明显更好的分布外检索能力。
使用混合模型改进检索。由于固定状态大小的自然基于检索的弱点,我们预测线性层将主要用在混合架构中,该架构通过二次方自注意力层减轻这一缺点。为了评估 Mamba-3 在此架构范式内的表现,我们以交错方式训练相同规模的混合模型,线性层与 NoPE 自注意力(B. Yang et al. 2025)的比例为 5:1。正如先前工作(Waleffe et al. 2024)所见,混合模型优于 Transformer 基线。我们发现,将预输出投影 RMSNorm(表 4 中的 pre-gate, grouped RMSNorm)重新引入 Mamba-3 层,提高了长度泛化检索能力,代价是轻微的上下文内现实世界检索任务,并且当与自注意力混合时,作为线性序列混合骨干网络具有高度竞争力。然而,由于相互竞争的权衡(附录 E,表 9),理想的归一化类型(分组 vs 默认)及其放置位置(门前 vs 门后)仍不清楚,因为我们发现混合模型及其确切特征和动态是复杂且通常反直觉的,最近的工作如 Cabannes et al. (2025) 也呼应了这一点。
4.3 推理效率与性能的权衡
![]()
![]()
4.4 快速 Mamba-3 内核
我们通过优化的内核补充了 Mamba-3 的方法论进展,这些内核在实际场景中提供快速推理。我们为 Mamba-3 实现了一系列新的推理内核——使用前向(预填充)路径的 Triton 和解码的 CuTe DSL——并在表 6 中将它们的每 token 解码延迟与 Mamba-2 和 GDN 发布的 Triton 内核进行比较。该评估在单个 H100 上以批量大小 128 测量单个解码步骤,针对 FP32 和 BF16 数据类型;模型为 1.5B 参数,模型维度为 2048,状态维度 ∈ {64, 128}。在所有配置中,SISO 在基线中实现了最低的延迟。MIMO 凭借其更高的算术强度,在不显著增加解码运行时间的情况下增加了解码 FLOPs。我们的基准测试表明,我们的 CuTe DSL 解码实现具有竞争力,且 Mamba-3 的额外组件(指数 - 梯形更新、复数值状态和 MIMO 投影)是轻量级的。这支持了我们整体的推理优先视角:Mamba-3 允许简单、低延迟的实现,同时提供强大的实证性能。
表 7 基准测试了不同解码序列长度下的端到端延迟以及相同序列长度的预填充时间。解码时间与表 6 一致,其中 Mamba-3 (SISO) 最快;Mamba-3 (MIMO) 与 Mamba-2 相当;并且随着序列长度增长,所有线性方法都比优化注意力更快。我们还看到,正如第 3.3 节所讨论的,MIMO 为预填充带来了适度的开销。基准测试的详细信息见附录 G。
5 相关工作
5.1 线性时间序列混合器
越来越多的工作寻求用线性运行时间的替代方案替换基于二次方 softmax 的注意力机制(Bahdanau, Cho, and Bengio 2014; Vaswani et al. 2017)。主要方法可以分为三大框架:线性注意力、测试时训练和状态空间模型。
许多新兴的线性注意力(LA)模型旨在通过核特征图近似 softmax 注意力(Choromanski et al. 2022; Katharopoulos et al. 2020),而最近的模型已经抛弃了特征图,转而使用查询和键之间的原始点积,并由衰减或掩码调节(Yutao Sun et al. 2023; S. Yang, B. Wang, Shen, et al. 2024)。最近,用键值对调节状态内存的快速权重编程器 Schlag, Irie, and Schmidhuber (2021) 也被归入“线性注意力”这一统称之下。S. Yang, Kautz, and Hatamizadeh (2025) 和 S. Yang, B. Wang, Y. Zhang, et al. (2025) 源于这一工作路线,并通过用 delta 规则递归替换加法内存更新来增强传统线性注意力。这进一步激发了一系列工作,以提高基于 delta 规则构建的线性模型的效率和能力(Hu et al. 2025; Kimi Team et al. 2025)。
一条并行的测试时训练(TTT)或测试时回归(TTR)工作线将序列建模视为推理期间的在线学习任务。在这里,递归状态代表过去输入的压缩摘要,递归步骤更新状态以记忆新信息(Yu Sun et al. 2025; Tandon et al. 2025; T. Zhang et al. 2025)。等价地,这些方法可以被视为全局回归目标的优化,递归状态更新代表迭代优化过程,例如梯度下降的变体(K. A. Wang, Shi, and Fox 2025)。
结构化状态空间模型(SSM)是受经典信号处理和动态系统启发的现代递归模型的另一种视角。早期版本的 SSM 如 S4(Gu, Goel, and Ré 2022; Gupta, Gu, and Berant 2022; Smith, Warrington, and Linderman 2023)使用具有结构化状态转移矩阵(例如对角或低秩加对角)的线性时不变(LTI)层,以促进长上下文任务的高效计算和稳定学习(Gu, Goel, and Ré 2022; Gupta, Gu, and Berant 2022; Smith, Warrington, and Linderman 2023)。Mamba-1(Gu and Dao 2024)中向 SSM 引入时变、输入依赖的选择性,减少了自注意力和线性模型在信息密集模态(尤其是语言建模)上的差距。随后,Mamba-2(Dao and Gu 2024)通过结构化状态空间对偶性(SSD)形式化了 SSM 与(线性)注意力之间的联系,我们在此工作中以此为基础。
5.2 状态追踪与复数状态空间模型
表现力与状态追踪。最近的工作刻画了递归、恒定内存混合器可以维持的状态类型,揭示了先前基于 SSM 的模型中的算法缺陷。Merrill, Petty, and Sabharwal (2025) 表明,在有限精度下,实际 SSM 坍缩为 TC0,导致在诸如 上的排列组合等任务上失败,除非扩展原语。类似地,Yu and Erichson (2025) 证明了单层 Mamba 不是通用逼近器。已经提出了几种修改来提高表现力。例如,同一项工作表明,块偏置变体仅需细微更改(通过块分解或通道特定偏置)即可恢复通用逼近属性。允许负特征值或非三角转移使得线性 RNN——包括对角和 Householder/DeltaNet 形式——能够捕捉奇偶性,并在温和假设下捕捉正则语言(Grazzi, Siems, Zela, et al. 2025)。复数值参数化提供了另一条增强表现力的途径。
复数状态空间模型。Mamba 之前的结构化 SSM 经常是复数值的,根植于传统 SSM 理论。它们通常也在视觉和音频等领域表现出色,这些领域具有明确的基于频率的信息内容,而不是语言。虽然一些模型如 H3(Fu et al. 2023)、RetNet(Yutao Sun et al. 2023)和 Megalodon(Ma et al. 2024)在针对语言建模时保留了复数值 SSM,但它们仍然明显逊于 Transformer。
此外,由于这些模型是 LTI 的,并且使用与现代选择性 SSM(如 Mamba)非常不同的算法计算(特别是卷积或显式递归),它们通常不使用 RoPE 技巧来处理复数部分。一个例外是 RetNet,它引入了一种介于线性注意力和 Mamba-2 之间的模型,使用常数标量衰减(不同于 LA 中的无衰减和 Mamba-2 中的数据依赖衰减),并通过 RoPE 实现额外的常数复数相位。
一般来说,经验发现复数对语言建模没有帮助,因此在 Mamba-1 及其后继版本中被淘汰,包括线性注意力和测试时训练方面的并行工作线。Mamba-3 代表了第一个具有复数值状态转移的现代递归模型,引入它是为了增加表现力和状态追踪能力的特定目的。通过结合 RoPE 技巧,据我们所知,这代表了基于理论动机的数据依赖 RoPE 的首次使用。
5.3 多输入,多输出
S4(Gu, Goel, and Ré 2022)是一个单输入单输出(SISO)LTI 系统,其中输入的每个维度都被分配了自己独立的 SSM。此类 SISO 模型具有比经典 RNN 大得多的递归状态,并且需要更复杂的数学机制来高效计算它们。为了简化模型,S5(Smith, Warrington, and Linderman 2023)和 LRU(Orvieto et al. 2023)用直接应用于整个向量化输入的多输入多输出(MIMO)SSM 替换了一组 SISO SSM。这一改变降低了有效状态容量,但通过直接使用并行扫描计算递归,启用了一条替代计算路径。虽然这种状态容量和建模性能之间的权衡在 LTI 模型中不太明显,但 Mamba-1 (S6)(Gu and Dao 2024)和 Mamba-2(Dao and Gu 2024)由于大状态大小在时变设置中的重要性,回到了 SISO 系统。与增加的状态大小相关的计算瓶颈通过 Mamba-1 的硬件感知并行扫描算法和 Mamba-2 的基于矩阵乘法的算法得到解决。
将 MIMO 引入 Mamba-3 与先前的工作显著不同。与旨在简化训练算法但以略微降低表现力为代价的先前 MIMO 模型不同,Mamba-3 的 MIMO 结构旨在提高建模能力同时保持推理效率。因此,其状态扩展保持在 Mamba-1/-2 水平,以维持建模能力,同时权衡额外的训练计算。
5.4 状态空间模型视角
虽然现代递归模型有几个 largely converge 的不同观点(第 5.1 节),但每个框架都有略微不同的解释和动机,可能导致不同的设计空间和扩展。特别是,线性注意力和测试时训练联系更紧密,或许可以在关联记忆框架下归为一类,该框架明确旨在通过“键值”存储记忆输入数据;要么通过 LA 中对规范 KV 方法(即二次方注意力)的近似,要么通过最小化 TTT 中的软优化目标。另一方面,状态空间模型有不同的渊源,这反映在术语(例如, A , B , C , X 而不是 Q , K , V
)及其自然扩展中。值得注意的是,Mamba-3 的方法论改进都特别与 SSM 观点相关,并且较少由关联记忆框架驱动。
- 指数 - 梯形离散化。SSM 观点需要对 governing 系统的连续 ODE 进行离散化;我们的指数 - 梯形离散化源于改进的离散化方法。由于关联记忆方法不使用离散化,如何在替代观点下解释诸如指数 - 梯形之类的 3 项递归并不明显。
- 复数值状态转移。复数 SSM 长期以来一直是动态系统的主要内容,将复数值视为选择性 SSM 的扩展是很自然的。另一方面,关联记忆框架将 A A 状态转移解释为目标函数的系数,例如对应于优化目标中 L2 正则化(或权重衰减)项的权重(K. A. Wang, Shi, and Fox 2025)。然而,作为回归目标的系数,复数值是没有意义的;因此,Mamba-3 在这些框架内并不明显可解释。
- 多输入,多输出。MIMO 是状态空间模型文献中的经典概念,并不自然出现在关联记忆(线性注意力或测试时训练)框架中。然而,我们确实注意到,本文介绍的 MIMO 公式并不直接局限于 SSM 理论——而是由计算视角驱动——并且我们的技术也可以适应其他现代递归模型。
线性时间序列模型的开发继续取得蓬勃进展,这里的讨论仅捕捉了其中的一部分。随着这些模型的不断发展,我们预计将出现一个不断增长的统一框架空间、改进的理解和新的泛化。
6 结论与未来工作
我们介绍了 Mamba-3,这是一种状态空间模型,相比先前的 SSM 具有几项方法论改进:通过指数 - 梯形离散化实现更强大的递归;通过复数值状态转移提高表现力;以及通过 MIMO 公式实现更高的推理效率和建模能力。Mamba-3 的基础 SISO 版本提供了强大的语言建模结果,无论是独立使用还是在交错混合架构中,并且在性能 - 效率权衡的帕累托前沿上超越了先前的线性序列模型。MIMO 版本以较慢的训练速度换取了更强的建模能力,同时与 Mamba-2 相比保持了具有竞争力的推理效率。总的来说,Mamba-3 中的技术展示了从状态空间模型视角出发的简单且有理论依据的改进,并为高效序列模型开辟了新的方向和设计原则。
原文链接:https://arxiv.org/pdf/2603.15569
特别声明:以上内容(如有图片或视频亦包括在内)为自媒体平台“网易号”用户上传并发布,本平台仅提供信息存储服务。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.