摘要

扩展推理时计算已成为推动大语言模型(LLM)性能的重要驱动力,这使得推理效率与模型质量一样,成为模型设计的核心关注点。当前基于 Transformer 的模型虽然具备强大的模型质量,但其二次计算复杂度和线性内存需求使得推理成本高昂。这促使了具有线性计算复杂度和常数内存需求的次二次(sub-quadratic)模型的发展。然而,许多近期的线性模型为了算法效率牺牲了模型质量与能力,在状态跟踪等任务上表现不佳。此外,它们在理论上具有线性推理复杂度,但在实践中硬件效率低下。在“推理优先”视角的指导下,我们从状态空间模型(SSM)的观点出发,引入了三项核心方法改进。我们将:(1) 源自 SSM 离散化的更具表达力的循环结构、(2) 支持更丰富状态跟踪的复数值状态更新规则,以及 (3) 在不增加解码延迟的前提下提升模型性能的多输入多输出(MIMO)结构相结合。结合架构优化,我们的 Mamba-3 模型在检索、状态跟踪和下游语言建模任务上取得了显著提升。在 1.5B1.5\text{B}1.5B 参数量级下,Mamba-3 相比次优模型(Gated DeltaNet)平均下游准确率提升了 0.60.60.6 个百分点,而 Mamba-3 的 MIMO 变体进一步提升了 1.21.21.2 个百分点,总计提升 1.81.81.8 个百分点。在不同状态尺寸实验中,Mamba-3 仅使用其前代模型一半的状态尺寸,即达到了与 Mamba-2 相当的困惑度。我们的评估表明,Mamba-3 成功推进了性能-效率的帕累托前沿。

1 引言

测试时计算已成为 LLM 进步的关键驱动力。思维链(chain-of-thought)推理和迭代优化等技术表明,推理时扩展能够解锁新能力 (Snell et al. 2024; Wu et al. 2025)。并行化、智能体(agentic)工作流的快速兴起进一步加剧了对高效推理和部署此类模型的需求 (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 替代方案的性能。

尽管线性模型取得了成功,但在提升其性能(特别是推进模型质量与推理效率之间的帕累托前沿)方面仍有显著进展空间。例如,Mamba-2 旨在通过牺牲部分表达力来提升训练速度和简化架构,从而在推理匹配的模型上表现稍逊。此外,研究已表明它们缺乏某些能力,例如状态跟踪能力较差,甚至在判断比特序列奇偶性等简单任务上表现不佳 (Grazzi, Siems, Zela, et al. 2025; Sarrof, Veitsman, and Hahn 2024)。最后,尽管这些次二次模型因理论上的高效推理而备受推崇并被广泛采用,但其推理算法在硬件上并不高效。具体而言,由于这些算法最初是从训练视角开发的,其解码阶段的算术强度(FLOPs 与内存流量的比率)较低,导致大量硬件计算单元处于空闲状态。

为了从推理优先的范式出发开发性能更强的模型,我们在 Mamba-2 的基础上引入了三项核心方法改进,这些改进深受次二次模型的 SSM 中心视角影响:

  • 指数梯形离散化(Exponential-Trapezoidal Discretization):我们提供了一种离散化时变、选择性 SSM 的简洁技术。通过我们的框架,可以推导出多种新的离散化方法。其中一个实例称为“指数-欧拉法(exponential-Euler)”,从理论上证明了此前缺乏理论依据的 Mamba-1 和 Mamba-2 启发式离散化方法。我们新的“指数梯形”实例是“指数-欧拉法”更具表达力的推广,其循环展开可揭示出对 SSM 输入应用的隐式卷积。结合显式的 BBBCCC 偏置项,Mamba-3 在经验上能够替代语言模型架构中此前被认为对循环模型至关重要的短因果卷积。
  • 复数值状态空间模型(Complex-valued State Space Model):通过将 Mamba-3 的底层 SSM 视为复数值,我们实现了比 Mamba-2 更具表达力的状态更新。这一更新规则专为训练和推理轻量化设计,克服了当前许多线性模型缺乏状态跟踪能力的问题。我们证明了该复数值更新规则等价于数据依赖的旋转位置嵌入(RoPE),并可高效计算 (Su et al. 2023)。经验上,它展现出解决此前线性模型无法处理的合成任务的能力。
  • 多输入多输出(MIMO)SSM:为提升解码时的 FLOP 效率,我们将基于外积的状态更新切换为基于矩阵乘法的状态更新。从 SSM 的信号处理基础来看,这种转换恰好对应于从单输入单输出(SISO)序列动力学到多输入多输出(MIMO)的推广。我们发现 MIMO 特别适合推理,因为额外的表达力允许在内存受限的解码状态更新阶段进行更多计算,而无需增加状态尺寸或牺牲速度。

综上所述,这些改进构成了我们 Mamba-3 层的核心。在方法上,我们注意到这些改进均自然源于 SSM 中心视角,而从线性注意力或测试时回归等其他流行视角来看则并不直观;我们将在第 5 节进一步讨论这些联系。经验上,我们在一系列合成状态跟踪和语言建模任务上验证了新模型的能力:

  • 更高质量:在 1.5B1.5\text{B}1.5B 规模下,Mamba-3 (MIMO) 相比 Transformer 下游语言建模准确率提升 +2.2+2.2+2.2,相比 Mamba-2 提升 +1.9+1.9+1.9,相比 GDN 提升 +1.8+1.8+1.8;而 Mamba-3 (SISO) 相比次优模型 GDN 提升 +0.6+0.6+0.6。此外,在状态尺寸实验中,状态尺寸为 646464 的 Mamba-3 (MIMO) 与状态尺寸为 128128128 的 Mamba-2 困惑度相当,有效以一半的延迟实现了相同的语言建模性能。
  • 新能力:Mamba-3 对 SSM 状态的复数化使其能够解决 Mamba-2 无法处理的合成状态跟踪任务。我们经验证明,高效的类 RoPE 计算能够近乎完美地解决算术任务,而无 RoPE 的 Mamba-3 和 Mamba-2 的表现仅略高于随机猜测。
  • 推理效率:Mamba-3 (MIMO) 提升了硬件利用率。在固定状态尺寸下,其解码 FLOPs 相比 Mamba-2 最多提升 444 倍,同时保持相似的墙钟解码延迟,并同步提升困惑度与下游性能。我们发布了 Mamba-3 的快速训练与推理内核。

Mamba-3 (SISO) 在质量与能力上超越了先前的线性模型,而 Mamba-3 (MIMO) 在匹配 Mamba-2 推理速度的同时,进一步提升了性能。我们的两种 Mamba-3 变体均凭借其强大的建模能力和硬件高效设计,推进了性能-延迟帕累托前沿。

2 预备知识

2.1 符号

标量用纯文本字母表示(如 x,yx, yx,y)。张量(包括向量和矩阵)用粗体字母表示(如 h,C\mathbf{h}, \mathbf{C}h,C)。张量的形状可根据上下文推断。我们将输入序列长度记为 TTT,模型维度记为 DDD,SSM 状态尺寸记为 NNN。对于时间索引,我们使用下标(如 xtx_txt 表示时间 ttt 的输入)。两个张量之间的哈达玛积(逐元素乘积)记为 ⊙\odot。对于向量 v∈Rd\mathbf{v} \in \mathbb{R}^dvRdDiag(v)∈Rd×d\text{Diag}(\mathbf{v}) \in \mathbb{R}^{d \times d}Diag(v)Rd×d 表示以 v\mathbf{v}v 为对角线的对角矩阵。对于跨时间步的标量乘积,我们使用符号 αt:s=∏i=stαi\alpha_{t:s} = \prod_{i=s}^{t} \alpha_iαt:s=i=stαi

2.2 SSM 预备知识

状态空间模型(SSM)通过以下公式描述连续时间线性动力学:
h˙(t)=A(t)h(t)+B(t)x(t)\dot{\mathbf{h}}(t) = \mathbf{A}(t)\mathbf{h}(t) + \mathbf{B}(t)x(t)h˙(t)=A(t)h(t)+B(t)x(t)
y(t)=C(t)⊤h(t)y(t) = \mathbf{C}(t)^\top \mathbf{h}(t)y(t)=C(t)h(t)
其中 h(t)∈RN\mathbf{h}(t) \in \mathbb{R}^Nh(t)RN 是隐藏状态,x(t)∈Rx(t) \in \mathbb{R}x(t)R 是输入,A(t)∈RN×N\mathbf{A}(t) \in \mathbb{R}^{N \times N}A(t)RN×NB(t),C(t)∈RN\mathbf{B}(t), \mathbf{C}(t) \in \mathbb{R}^NB(t),C(t)RN。我们偶尔将 A(t)\mathbf{A}(t)A(t) 称为状态转移,将 B(t)x(t)\mathbf{B}(t)x(t)B(t)x(t) 称为状态输入;这也适用于它们的离散化对应项。对于步长为 Δt\Delta tΔt 的离散序列,Mamba-1 和 Mamba-2 将系统离散化为以下循环:
ht=eΔtAtht−1+ΔtBtxt\mathbf{h}_t = e^{\Delta_t \mathbf{A}_t} \mathbf{h}_{t-1} + \Delta_t \mathbf{B}_t x_tht=eΔtAtht1+ΔtBtxt
yt=Ct⊤hty_t = \mathbf{C}_t^\top \mathbf{h}_tyt=Ctht

Mamba-2 的参数化:Mamba-2 层 (Dao and Gu 2024) 的核心是数据依赖且硬件高效的 SSM。状态转移和状态输入均通过将 Δt∈R>0\Delta_t \in \mathbb{R}_{>0}ΔtR>0B,C∈RN\mathbf{B}, \mathbf{C} \in \mathbb{R}^NB,CRN 从当前 token 投影而变得数据依赖。通过将状态转移参数化为标量乘以单位矩阵(At=AtIN×N\mathbf{A}_t = A_t \mathbf{I}_{N \times N}At=AtIN×N,其中 At∈R<0A_t \in \mathbb{R}_{<0}AtR<0),SSM 循环可利用 GPU 的矩阵乘法 Tensor Core 高效计算。
定义 αt≜eΔtAt∈(0,1)\alpha_t \triangleq e^{\Delta_t A_t} \in (0, 1)αteΔtAt(0,1)γt≜Δt\gamma_t \triangleq \Delta_tγtΔt,更新变为:
ht=αtht−1+γtBtxt,yt=Ct⊤ht(1)\mathbf{h}_t = \alpha_t \mathbf{h}_{t-1} + \gamma_t \mathbf{B}_t x_t, \quad y_t = \mathbf{C}_t^\top \mathbf{h}_t \quad (1)ht=αtht1+γtBtxt,yt=Ctht(1)
数据依赖的状态转移 αt\alpha_tαt 控制层内每个 SSM 的记忆跨度。Δt\Delta_tΔt 同时调节状态转移和状态输入:较大的 Δt\Delta_tΔt 遗忘更快并更强地加权当前 token,而较小的 Δt\Delta_tΔt 则保留隐藏状态,当前 token 的贡献极小。
注 1:在 Mamba-2 中,AtA_tAt 是数据独立的,因为整体离散转移 αt≜eΔtAt\alpha_t \triangleq e^{\Delta_t A_t}αteΔtAt 通过 Δt\Delta_tΔt 变得数据依赖。在 Mamba-3 中,我们经验发现数据依赖的 AtA_tAt 与数据独立的 AtA_tAt 性能相似,为保持一致性选择前者作为默认值,使所有 SSM 参数均数据依赖。

2.3 结构化掩码表示与状态空间对偶性

Mamba-2 表明,一大类 SSM 允许一种矩阵形式,可将时间步循环向量化。通过状态空间对偶性(SSD)框架,循环 SSM 可在并行形式中表示,该形式包含逐元素掩码以建模状态转移衰减。
SSD 为线性循环与可并行化(基于矩阵乘法)计算形式之间的对偶性提供了通用框架:
Y=(L⊙CB⊤)X(2)\mathbf{Y} = (\mathbf{L} \odot \mathbf{C} \mathbf{B}^\top) \mathbf{X} \quad (2)Y=(LCB)X(2)
其中 L∈RT×T\mathbf{L} \in \mathbb{R}^{T \times T}LRT×T 是结构化掩码,B,C∈RT×N\mathbf{B}, \mathbf{C} \in \mathbb{R}^{T \times N}B,CRT×NX∈RT×D\mathbf{X} \in \mathbb{R}^{T \times D}XRT×D 是 SSM 的输入,Y∈RT×D\mathbf{Y} \in \mathbb{R}^{T \times D}YRT×D 是其输出。L\mathbf{L}L 的不同结构产生 SSD 的各种实例。
公式 (2) 也通过将 Q≜C\mathbf{Q} \triangleq \mathbf{C}QCK≜B\mathbf{K} \triangleq \mathbf{B}KBV≜X\mathbf{V} \triangleq \mathbf{X}VX 并将 L\mathbf{L}L 视为数据依赖掩码,建立了循环与注意力之间的一般联系。实际上,SSD 的最简单情况是(因果)线性注意力 (Katharopoulos et al. 2020),其中 L\mathbf{L}L 是因果三角掩码。
Mamba-2 是一种推广形式,其中:
L=[1α11⋮⋱⋱αT⋯2⋯αT1]⋅Diag(γ)(3)\mathbf{L} = \begin{bmatrix} 1 \\ \alpha_1 & 1 \\ \vdots & \ddots & \ddots \\ \alpha_{T \cdots 2} & \cdots & \alpha_T & 1 \end{bmatrix} \cdot \text{Diag}(\gamma) \quad (3)L= 1α1αT21αT1 Diag(γ)(3)
由公式 (1) 中的项 αt,γt\alpha_t, \gamma_tαt,γt 组成。在第 3.1.3 节中,我们将证明 Mamba-3 是 Mamba-2 的推广,具有更具表达力的 L\mathbf{L}L,因此也是 SSD 的一个实例。
在这里插入图片描述

3 方法论

我们引入 Mamba-3,一种具有三项新创新的 SSM:“指数梯形”离散化以实现更具表达力的动力学(第 3.1 节)、用于状态跟踪的复数值状态空间(第 3.2 节),以及提升建模能力和推理时硬件利用率的 MIMO(第 3.3 节)。这些改进解决了当前次二次架构在质量、能力和效率方面的限制。我们在第 3.4 节将它们结合到更新的 Mamba 架构块中。

3.1 指数梯形离散化

结构化 SSM 自然定义为连续时间动力系统,将输入函数 x(t)∈Rx(t) \in \mathbb{R}x(t)R 映射为输出函数 y(t)∈Ry(t) \in \mathbb{R}y(t)Rt>0t>0t>0)。底层连续状态空间系统由状态 h˙(t)\dot{\mathbf{h}}(t)h˙(t) 的一阶常微分方程(ODE)和输出 y(t)y(t)y(t) 的代数方程定义。然而,在序列建模中,数据仅在离散时间步观测,因此需要对 SSM 应用离散化步骤,将其连续时间动力学转换为离散循环。
离散化方法在经典控制理论中已有深入研究,早期深度学习 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。
我们推导的高级直觉源自简单线性 ODE x′(t)=Ax(t)x'(t) = Ax(t)x(t)=Ax(t) 的闭式解 x(t)=eAtx(0)x(t) = e^{At}x(0)x(t)=eAtx(0),其离散化为 xt+1=eΔAxtx_{t+1} = e^{\Delta A}x_txt+1=eΔAxt。在此例中,指数项主导了底层一阶 ODE 的动力学,导致在不严格约束 Δ\DeltaΔ 的情况下使用低阶方法会产生不精确的近似。因此,我们分析指数调整系统 e−Atx(t)e^{-At}x(t)eAtx(t) 的动力学。调整后的系统产生离散循环形式,其中状态转移积分和状态输入积分分别近似——状态转移积分通过右端点近似,即 A(s)≈A(τt)A(s) \approx A(\tau_t)A(s)A(τt) 对所有 s∈[τt−1,τt]s \in [\tau_{t-1}, \tau_t]s[τt1,τt],得到:
h(τt)=exp⁡(∫τt−1τtA(s)ds)h(τt−1)+∫τt−1τtexp⁡(∫ττtA(s)ds)B(τ)x(τ)dτ\mathbf{h}(\tau_t) = \exp\left(\int_{\tau_{t-1}}^{\tau_t} A(s)ds\right)\mathbf{h}(\tau_{t-1}) + \int_{\tau_{t-1}}^{\tau_t} \exp\left(\int_{\tau}^{\tau_t} A(s)ds\right) \mathbf{B}(\tau)x(\tau)d\tauh(τt)=exp(τt1τtA(s)ds)h(τt1)+τt1τtexp(ττtA(s)ds)B(τ)x(τ)dτ
ht≈exp⁡(ΔtAt)ht−1+∫τt−1τtexp⁡((τt−τ)At)B(τ)x(τ)dτ\mathbf{h}_t \approx \exp(\Delta_t A_t)\mathbf{h}_{t-1} + \int_{\tau_{t-1}}^{\tau_t} \exp((\tau_t-\tau)A_t) \mathbf{B}(\tau)x(\tau)d\tauhtexp(ΔtAt)ht1+τt1τtexp((τtτ)At)B(τ)x(τ)dτ
这为状态输入积分的进一步离散化技术奠定了基础。完整推导详见命题 5。
ZOH:经典的零阶保持离散化方法可从上述基础通过右侧积分的特定近似推导得出。通过在区间 [τt−1,τt][\tau_{t-1}, \tau_t][τt1,τt] 上将 At,B(τ),x(τ)A_t, \mathbf{B}(\tau), x(\tau)At,B(τ),x(τ) 视为常数(固定为右端点 τt\tau_tτt 的值),积分结果为 At−1(exp⁡(ΔtAt)−I)BtxtA_t^{-1}(\exp(\Delta_t A_t) - \mathbf{I}) \mathbf{B}_t x_tAt1(exp(ΔtAt)I)Btxt。我们注意到,这正式证明了经典 LTI 系统的 ZOH 公式通过简单替换时变参数即可应用于 LTV。
指数-欧拉法(Mamba-1/-2):虽然 Mamba-1 声称使用上述时变 ZOH 公式,但 Mamba-1 和 Mamba-2 在实际发布实现中使用了额外的近似。该离散化可通过对状态输入积分使用欧拉规则并保持右端点在区间内恒定(图 1)来恢复:
ht≈eΔtAtht−1+(τt−τt−1)e(τt−τt)AtBtxt=eΔtAtht−1+ΔtBtxt(4)\mathbf{h}_t \approx e^{\Delta_t A_t}\mathbf{h}_{t-1} + (\tau_t-\tau_{t-1})e^{(\tau_t-\tau_t)A_t}\mathbf{B}_t x_t = e^{\Delta_t A_t}\mathbf{h}_{t-1} + \Delta_t \mathbf{B}_t x_t \quad (4)hteΔtAtht1+(τtτt1)e(τtτt)AtBtxt=eΔtAtht1+ΔtBtxt(4)
我们将公式 (4) 称为指数-欧拉离散化方法,源于指数积分后接欧拉近似。该推导证明了 Mamba-1/-2 实现中使用的公式。
指数梯形(Mamba-3):然而,欧拉规则仅提供状态输入积分的一阶近似,其局部截断误差量级为 O(Δt2)O(\Delta_t^2)O(Δt2)。相比之下,我们引入了一种广义梯形规则,提供二阶精度的积分近似,优于欧拉规则。具体而言,它使用数据依赖的凸组合近似区间两端点的积分。这推广了经典梯形规则,后者简单平均区间端点(图 1)。
命题 1(指数梯形离散化):通过广义梯形规则近似公式 (16) 中的状态输入积分,可得循环:
ht=eΔtAtht−1+(1−λt)ΔteΔtAtBt−1xt−1+λtΔtBtxt(5)\mathbf{h}_t = e^{\Delta_t A_t}\mathbf{h}_{t-1} + (1-\lambda_t)\Delta_t e^{\Delta_t A_t}\mathbf{B}_{t-1}x_{t-1} + \lambda_t\Delta_t \mathbf{B}_t x_t \quad (5)ht=eΔtAtht1+(1λt)ΔteΔtAtBt1xt1+λtΔtBtxt(5)
≜αtht−1+βtBt−1xt−1+γtBtxt(6)\triangleq \alpha_t \mathbf{h}_{t-1} + \beta_t \mathbf{B}_{t-1}x_{t-1} + \gamma_t \mathbf{B}_t x_t \quad (6)αtht1+βtBt1xt1+γtBtxt(6)
其中 λt∈[0,1]\lambda_t \in [0, 1]λt[0,1] 是数据依赖标量,αt≜eΔtAt\alpha_t \triangleq e^{\Delta_t A_t}αteΔtAtβt≜(1−λt)ΔteΔtAt\beta_t \triangleq (1-\lambda_t)\Delta_t e^{\Delta_t A_t}βt(1λt)ΔteΔtAtγt≜λtΔt\gamma_t \triangleq \lambda_t\Delta_tγtλtΔt
注 2(表达力):指数梯形规则是 (a) 经典梯形规则的推广(当 λt=12\lambda_t = \frac{1}{2}λt=21 时恢复),也是 (b) Mamba-2 欧拉规则的推广(当 λt=1\lambda_t = 1λt=1 时恢复)。
注 3(误差率):这是状态输入积分的二阶离散化,在标准稳定性假设下,若梯形参数满足 λt=12+O(Δt)\lambda_t = \frac{1}{2} + O(\Delta_t)λt=21+O(Δt),其误差量级为 O(Δt3)O(\Delta_t^3)O(Δt3)。然而,我们的消融实验表明不强制此约束经验性能更好。详见附录 A.2 和 A.3。
在这里插入图片描述

3.1.2 指数梯形循环作为隐式卷积

我们的广义指数梯形离散化等价于对 SSM 的状态输入应用大小为 2 的数据依赖卷积。具体而言,普通 SSM 的循环形式首先实例化状态输入 vt=Btxt\mathbf{v}_t = \mathbf{B}_t x_tvt=Btxt,然后计算线性循环 ht=αtht−1+γtvt\mathbf{h}_t = \alpha_t \mathbf{h}_{t-1} + \gamma_t \mathbf{v}_tht=αtht1+γtvt。在公式 (6) 中,我们首先在 vt\mathbf{v}_tvt 上应用宽度为 2 的卷积(由 β,γ\beta, \gammaβ,γ 加权),然后再传入线性循环。
注 4(卷积差异):指数梯形离散化诱导的“卷积”与 Mamba 和 GDN 等序列模型使用的标准短卷积有明显区别。标准短卷积是在核心循环外部独立应用于 xtx_txt(通常还包括 Bt,Ct\mathbf{B}_t, \mathbf{C}_tBt,Ct)的操作,而我们的新离散化可解释为在核心循环内部对状态输入 Btxt\mathbf{B}_t x_tBtxt 应用卷积。

3.1.3 指数梯形循环的并行表示

我们的新循环可作为 SSD 的一种情况实例化,并具有对应于公式 (2) 的并行形式。展开状态循环可得 hT=αT⋯2(γ0α1+β1)B0x0+⋯+γTBTxT\mathbf{h}_T = \alpha_{T \cdots 2}(\gamma_0 \alpha_1 + \beta_1)\mathbf{B}_0 x_0 + \cdots + \gamma_T \mathbf{B}_T x_ThT=αT2(γ0α1+β1)B0x0++γTBTxT,其中 SSM 输出为 yT=αT⋯2(γ0α1+β1)CT⊤B0x0+⋯+γTCT⊤BTxTy_T = \alpha_{T \cdots 2}(\gamma_0 \alpha_1 + \beta_1)\mathbf{C}_T^\top \mathbf{B}_0 x_0 + \cdots + \gamma_T \mathbf{C}_T^\top \mathbf{B}_T x_TyT=αT2(γ0α1+β1)CTB0x0++γTCTBTxT。展开这些行表明,梯形更新诱导的掩码不再是固定的端点平均(如经典梯形规则),而是区间两个端点的数据依赖凸组合。
在 SSD 框架 (2) 下,Mamba-3 对应于掩码 L\mathbf{L}L,其结构是一个 1-半可分离矩阵与一个 2-带矩阵的组合:
L=[γ0(γ0α1+β1)γ1α2(γ0α1+β1)(γ1α2+β2)γ2⋮⋮⋱αT⋯2(γ0α1+β1)⋯γT]=[1α11α2α1α21⋮⋮⋱⋱αT⋯1⋯1][γ0β1γ10β2γ2⋮⋮⋱0⋯γT](7)\mathbf{L} = \begin{bmatrix} \gamma_0 \\ (\gamma_0\alpha_1+\beta_1) & \gamma_1 \\ \alpha_2(\gamma_0\alpha_1+\beta_1) & (\gamma_1\alpha_2+\beta_2) & \gamma_2 \\ \vdots & \vdots & \ddots \\ \alpha_{T \cdots 2}(\gamma_0\alpha_1+\beta_1) & \cdots & \gamma_T \end{bmatrix} = \begin{bmatrix} 1 \\ \alpha_1 & 1 \\ \alpha_2\alpha_1 & \alpha_2 & 1 \\ \vdots & \vdots & \ddots & \ddots \\ \alpha_{T \cdots 1} & \cdots & 1 \end{bmatrix} \begin{bmatrix} \gamma_0 & \beta_1 & \gamma_1 \\ 0 & \beta_2 & \gamma_2 \\ \vdots & \vdots & \ddots \\ 0 & \cdots & \gamma_T \end{bmatrix} \quad (7)L= γ0(γ0α1+β1)α2(γ0α1+β1)αT2(γ0α1+β1)γ1(γ1α2+β2)γ2γT = 1α1α2α1αT11α211 γ000β1β2γ1γ2γT (7)
该并行形式支持训练时基于矩阵乘法的硬件高效 SSM 输出计算。
我们注意到,Mamba-3 的卷积联系也可通过此并行对偶形式看出,其中公式 (7) 中与 2-带矩阵的乘法表示使用权重 β,γ\beta, \gammaβ,γ 进行卷积。在附录 A.1 中,我们使用 SSD 张量收缩机制证明了并行形式等价于具有状态输入卷积的普通 SSM。
注 5:Mamba-3 的结构化掩码可视为对 Mamba-2 的推广,后者 2-带矩阵替换为仅含 γt\gamma_tγt 的对角矩阵 (3)。

3.2 复数值 SSM

现代 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) 表明,将转移矩阵特征值限制为实数、非负会削弱模型在简单状态跟踪任务(主要指可解群 regime (TC0),如奇偶性校验)上的能力,而这些任务单层 LSTM 即可解决。这一局限性在 (Grazzi, Siems, Schrodi, et al. 2024) 定理 1 中被形式化,源于将转移矩阵特征值限制为实数,无法表示“旋转”隐藏状态动力学。例如,考虑二元输入 {0,1}\{0, 1\}{0,1} 的奇偶函数,定义为 ∑txt mod 2\sum_t x_t \bmod 2txtmod2。该任务可使用更新 ht=R(πxt)ht−1\mathbf{h}_t = \mathbf{R}(\pi x_t)\mathbf{h}_{t-1}ht=R(πxt)ht1 完成,其中 R(⋅)\mathbf{R}(\cdot)R() 是 2D 旋转矩阵。此类旋转动力学无法用实数特征值表示。

3.2.1 指数-欧拉离散化的复数值 SSM

为恢复此能力,我们从复数值 SSM (8) 开始,其能够表示状态跟踪动力学。我们证明,在离散化下(命题 5),复数值 SSM 可表述为实数值 SSM,其转移矩阵为由 2×22 \times 22×2 旋转矩阵组成的块对角矩阵(命题 2)。我们进一步证明这等价于在输入和输出投影 B,C\mathbf{B}, \mathbf{C}B,C 上分别应用数据依赖旋转嵌入(命题 3)。最终,Su et al. (2023) 使用的“RoPE 技巧”允许高效实现复数值状态转移矩阵,计算开销相比实数值 SSM 极小。
命题 2(复数到实数 SSM 等价性):考虑复数值 SSM
h˙(t)=Diag(A(t)+iθ(t))h(t)+(B(t)+iB^(t))x(t)(8)\dot{\mathbf{h}}(t) = \text{Diag}(A(t) + i\theta(t))\mathbf{h}(t) + (\mathbf{B}(t) + i\hat{\mathbf{B}}(t))x(t) \quad (8)h˙(t)=Diag(A(t)+iθ(t))h(t)+(B(t)+iB^(t))x(t)(8)
y(t)=Re((C(t)+iC^(t))⊤h(t))y(t) = \text{Re}((\mathbf{C}(t) + i\hat{\mathbf{C}}(t))^\top \mathbf{h}(t))y(t)=Re((C(t)+iC^(t))h(t))
其中 h(t)∈CN/2\mathbf{h}(t) \in \mathbb{C}^{N/2}h(t)CN/2θ(t),B(t),B^(t),C(t),C^(t)∈RN/2\theta(t), \mathbf{B}(t), \hat{\mathbf{B}}(t), \mathbf{C}(t), \hat{\mathbf{C}}(t) \in \mathbb{R}^{N/2}θ(t),B(t),B^(t),C(t),C^(t)RN/2x(t),A(t)∈Rx(t), A(t) \in \mathbb{R}x(t),A(t)R。在指数-欧拉离散化下,该系统等价于实数值 SSM
ht=eΔtAtRtht−1+ΔtBtxt(9)\mathbf{h}_t = e^{\Delta_t A_t} \mathbf{R}_t \mathbf{h}_{t-1} + \Delta_t \mathbf{B}_t x_t \quad (9)ht=eΔtAtRtht1+ΔtBtxt(9)
yt=Ct⊤hty_t = \mathbf{C}_t^\top \mathbf{h}_tyt=Ctht
状态 ht∈RN\mathbf{h}_t \in \mathbb{R}^NhtRN,投影 Bt≜[Bt B^t]∈RN\mathbf{B}_t \triangleq [\mathbf{B}_t \ \hat{\mathbf{B}}_t] \in \mathbb{R}^NBt[Bt B^t]RNCt≜[Ct −C^t]∈RN\mathbf{C}_t \triangleq [\mathbf{C}_t \ -\hat{\mathbf{C}}_t] \in \mathbb{R}^NCt[Ct C^t]RN,转移矩阵 Rt≜Block({R(Δtθt[i])}i=1N/2)∈RN×N\mathbf{R}_t \triangleq \text{Block}(\{\mathbf{R}(\Delta_t \theta_t[i])\}_{i=1}^{N/2}) \in \mathbb{R}^{N \times N}RtBlock({R(Δtθt[i])}i=1N/2)RN×N,其中 R(θ)≜[cos⁡θ−sin⁡θsin⁡θcos⁡θ]\mathbf{R}(\theta) \triangleq \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix}R(θ)[cosθsinθsinθcosθ]。证明见附录 B.1。

命题 3(复数 SSM 与数据依赖 RoPE 等价性):在命题 2 的符号体系下,考虑公式 (9) 定义的实数 SSM 展开 TTT 个时间步。上述 SSM 的输出等价于应用了数据依赖旋转嵌入的普通标量转移矩阵 SSM (4),定义如下:
ht=eΔtAtht−1+(∏i=0tRi⊤)ΔtBtxt,yt=[(∏i=0tRi⊤)Ct]⊤ht(10)\mathbf{h}_t = e^{\Delta_t A_t}\mathbf{h}_{t-1} + \left(\prod_{i=0}^{t} \mathbf{R}_i^\top\right) \Delta_t \mathbf{B}_t x_t, \quad y_t = \left[\left(\prod_{i=0}^{t} \mathbf{R}_i^\top\right) \mathbf{C}_t\right]^\top \mathbf{h}_t \quad (10)ht=eΔtAtht1+(i=0tRi)ΔtBtxt,yt=[(i=0tRi)Ct]ht(10)
其中矩阵乘积表示右乘,例如 ∏i=01Ri=R0R1\prod_{i=0}^{1} \mathbf{R}_i = \mathbf{R}_0 \mathbf{R}_1i=01Ri=R0R1。我们将使用变换后的实数值 SSM 计算复数 SSM 称为“RoPE 技巧”。证明见附录 B.2。
为观察复数 SSM 与 RoPE 嵌入的联系,注意在命题中,数据依赖旋转 Ri\mathbf{R}_iRi 跨时间步聚合并应用于 C,B\mathbf{C}, \mathbf{B}C,B,这在状态空间对偶框架下对应注意力的查询(Q)和键(K)分量。类似地,普通 RoPE (Su et al. 2023) 应用数据独立旋转矩阵,其旋转角度遵循固定频率表 θ[i]=10000−2i/N\theta[i] = 10000^{-2i/N}θ[i]=100002i/N

3.2.2 指数梯形离散化的复数值 SSM

推导指数-欧拉离散化复数 SSM 的循环后,推广到指数梯形离散化是类似的。命题 4 给出了 Mamba-3 带 RoPE 技巧的完整循环。
命题 4(指数梯形离散化的旋转嵌入等价性):使用指数梯形规则离散化复数 SSM(命题 1)得到循环:
ht=αtht−1+βt(∏i=0t−1Ri⊤)Bt−1xt−1+γt(∏i=0tRi⊤)Btxt\mathbf{h}_t = \alpha_t \mathbf{h}_{t-1} + \beta_t \left(\prod_{i=0}^{t-1} \mathbf{R}_i^\top\right) \mathbf{B}_{t-1} x_{t-1} + \gamma_t \left(\prod_{i=0}^{t} \mathbf{R}_i^\top\right) \mathbf{B}_t x_tht=αtht1+βt(i=0t1Ri)Bt1xt1+γt(i=0tRi)Btxt
yt=[(∏i=0tRi⊤)Ct]⊤ht(11)y_t = \left[\left(\prod_{i=0}^{t} \mathbf{R}_i^\top\right) \mathbf{C}_t\right]^\top \mathbf{h}_t \quad (11)yt=[(i=0tRi)Ct]ht(11)
其中 Rt\mathbf{R}_tRt 是命题 2 定义的块对角旋转矩阵。证明见附录 B.3。
我们经验验证,通过数据依赖 RoPE 实现的复数 SSM 能够解决实数值 SSM(无论是否使用标准 RoPE)无法解决的状态跟踪任务(表 5b),支持了理论主张。

3.3 多输入多输出(MIMO)

扩展测试时计算开启了模型能力的新前沿,例如智能体工作流,其中推理占总计算预算的比例不断增加。这使人们对语言模型的推理效率重新聚焦,并推动了采用具有固定隐藏状态尺寸的 SSM 和次二次层,它们提供更低的计算和内存需求。尽管这些新层的墙钟时间低于 Transformer,但其解码严重受限于内存带宽,导致硬件利用率低下。本节,我们利用 SSM 视角引入对 Mamba-3 循环的方法论优化,允许在不增加解码墙钟时间的情况下提升模型 FLOPs,从而在相同解码速度下获得更优模型。
在这里插入图片描述

解码算术强度:为提升硬件效率,需考虑 token 生成的算术强度,定义为 FLOPs 除以给定操作的输入输出字节数。由于 SSM 解码饱和了内存带宽并伴有空闲计算(即受限于内存),我们希望提升其算术强度以有效重叠计算与内存 I/O。具体而言,Mamba 中单次生成的算术强度约为每字节 2.52.52.5 次操作(表 2a),而 NVIDIA H100-SXM5 上 bfloat16 矩阵乘法的算术强度约为每字节 295295295 次操作 (NVIDIA 2022)。因此,SSM 解码远未达到计算受限状态,且目前尚不清楚如何调整 Mamba 现有参数以缓解硬件效率不足。我们注意到此观察普遍适用于其他次二次模型,如因果线性注意力。

从 SISO 到 MIMO:考虑典型 SSM 的单头,头维度为 PPP,涉及堆叠 PPP 份共享相同 αt,Δt\alpha_t, \Delta_tαt,ΔtBt\mathbf{B}_tBt 的 SISO 循环 ht←αtht−1+ΔtBtxt\mathbf{h}_t \leftarrow \alpha_t \mathbf{h}_{t-1} + \Delta_t \mathbf{B}_t x_thtαtht1+ΔtBtxt。得到的广播循环 ht←αtht−1+ΔtBtxt⊤\mathbf{h}_t \leftarrow \alpha_t \mathbf{h}_{t-1} + \Delta_t \mathbf{B}_t x_t^\tophtαtht1+ΔtBtxt 接受向量输入 xt∈RPx_t \in \mathbb{R}^PxtRP 并具有矩阵值状态 ht∈RN×P\mathbf{h}_t \in \mathbb{R}^{N \times P}htRN×P
注意内存流量(输入输出大小)主要由状态 ht\mathbf{h}_tht 主导,而计算主要包括外积 Btxt⊤\mathbf{B}_t x_t^\topBtxt,其 FLOPs 与 NPNPNP 成正比。通过增加后者的维度,将 Bt∈RN→Bt∈RN×R\mathbf{B}_t \in \mathbb{R}^N \rightarrow \mathbf{B}_t \in \mathbb{R}^{N \times R}BtRNBtRN×Rxt∈RP→xt∈RP×Rx_t \in \mathbb{R}^P \rightarrow x_t \in \mathbb{R}^{P \times R}xtRPxtRP×R,内存流量不会显著增加(对小 RRR 而言),而消耗的 FLOPs 增加 RRR 倍(表 2a)。因此,此变换提升了循环的算术强度。此外,算术强度的提升转化为实际收益,因为外积 Btxt⊤\mathbf{B}_t x_t^\topBtxt 变为硬件高效的矩阵乘法(matmul),使用快速 Tensor Core 计算,仅产生微小延迟成本。结果,MIMO 循环比原始 SISO 循环更具表达力,计算 RRR 倍 FLOPs 的同时实际保持了解码速度。
出于类似原因,从状态计算输出 yt←Ct⊤hty_t \leftarrow \mathbf{C}_t^\top \mathbf{h}_tytCtht 通过修改输出投影 Ct∈RN→Ct∈RN×R\mathbf{C}_t \in \mathbb{R}^N \rightarrow \mathbf{C}_t \in \mathbb{R}^{N \times R}CtRNCtRN×R 获得额外秩 RRR。总体而言,此变换等价于将原始单输入单输出(SISO)循环扩展为多输入多输出(MIMO)。

训练 MIMO SSM:尽管 MIMO 公式受推理效率驱动,但 SSM 的训练算法(包括第 3.1、3.2 节的开发)通常为 SISO 模型开发。我们首先观察到 MIMO SSM 可表示为 R2R^2R2 个 SISO SSM,其中 RRR 个共享相同循环的 SISO SSM 为每个 MIMO 输出求和。具体而言,定义 Ct(i)∈RN,Bt(j)∈RN,xt(j)∈R,Δt∈R\mathbf{C}_t^{(i)} \in \mathbb{R}^N, \mathbf{B}_t^{(j)} \in \mathbb{R}^N, x_t^{(j)} \in \mathbb{R}, \Delta_t \in \mathbb{R}Ct(i)RN,Bt(j)RN,xt(j)R,ΔtR,其中 i,j∈{0,…,R−1}i, j \in \{0, \dots, R-1\}i,j{0,,R1},则有:
ht(j)←αtht−1(j)+ΔtBt(j)xt(j)(12)\mathbf{h}_t^{(j)} \leftarrow \alpha_t \mathbf{h}_{t-1}^{(j)} + \Delta_t \mathbf{B}_t^{(j)} x_t^{(j)} \quad (12)ht(j)αtht1(j)+ΔtBt(j)xt(j)(12)
ht=∑j=0R−1ht(j)(13)\mathbf{h}_t = \sum_{j=0}^{R-1} \mathbf{h}_t^{(j)} \quad (13)ht=j=0R1ht(j)(13)
yt(i)←(Ct(i))⊤ht(14)y_t^{(i)} \leftarrow (\mathbf{C}_t^{(i)})^\top \mathbf{h}_t \quad (14)yt(i)(Ct(i))ht(14)
因此,yt(i)=∑jSSM(α,Δ,B(j),C(i),x(j))ty_t^{(i)} = \sum_j \text{SSM}(\alpha, \Delta, \mathbf{B}^{(j)}, \mathbf{C}^{(i)}, x^{(j)})_tyt(i)=jSSM(α,Δ,B(j),C(i),x(j))t,其中 SSM(… )t≜(Ct(i))⊤h(j)\text{SSM}(\dots)_t \triangleq (\mathbf{C}_t^{(i)})^\top \mathbf{h}^{(j)}SSM()t(Ct(i))h(j)。标准 SISO 训练算法可作黑盒直接应用于 MIMO 模型。此观察允许通过并行调用 SISO 算法 R2R^2R2 次训练 MIMO 模型。相比之下,以循环形式计算时,公式 (12)-(14) 可顺序执行,仅产生相对 SISO SSM RRR 倍的开销。

MIMO SSM 的分块算法:许多现代 SISO 循环模型(包括 Mamba-2)使用分块算法计算,序列被分成长度为 CCC 的块。在每个块内应用并行(但渐近较慢)算法,跨块计算循环。分块算法在完全并行和完全顺序算法之间插值。利用此结构,我们可将 MIMO SSM 的训练成本降低至 SISO SSM 的 RRR 倍。此思想也见于 SSD 框架——SSD 在每个块内应用硬件友好的二次算法,跨块使用循环形式,并表明当状态和头维度相当时,将块尺寸设为该维度可产生整体线性时间算法。具体而言,SSD 的块内计算每块产生 (2C2N+2C2P)(2C^2N + 2C^2P)(2C2N+2C2P) FLOPs,总计 TC(2C2N+2C2P)=2TNP(CN+CP)\frac{T}{C}(2C^2N + 2C^2P) = 2TNP(\frac{C}{N} + \frac{C}{P})CT(2C2N+2C2P)=2TNP(NC+PC)(忽略可忽略项)。设 C=P=NC=P=NC=P=N,总 FLOP 数为 8TN28TN^28TN2,对 TTT 呈线性。
SSD 的分块算法可自然推广到 MIMO SSM。此时,状态投影 Bx⊤\mathbf{B}x^\topBx 和状态发射 C⊤h\mathbf{C}^\top \mathbf{h}Ch 的 FLOP 数增加 RRR 倍,而块内分量 C⊤B\mathbf{C}^\top \mathbf{B}CB 的 FLOP 数增加 R2R^2R2 倍。因此,块内计算产生 2(TC(CR)2N+TC(CR)2P)2(TC(CR)^2N + TC(CR)^2P)2(TC(CR)2N+TC(CR)2P) FLOPs,块间计算产生 4TCNP(CR)+2TCNP4\frac{T}{C}NP(CR) + 2\frac{T}{C}NP4CTNP(CR)+2CTNP FLOPs。因此,设 CR=N=PCR = N = PCR=N=P 产生总 FLOP 数 8TRN28TRN^28TRN2,FLOP 数增加 RRR 倍。直观上,设 MIMO 块尺寸为 SISO 块尺寸的 1R\frac{1}{R}R1,即 CMIMO←1RCSISOC_{\text{MIMO}} \leftarrow \frac{1}{R} C_{\text{SISO}}CMIMOR1CSISO,保持 SISO 块内 FLOP 数同时将块数增加 RRR 倍,从而产生整体 RRR 倍 FLOP 增加而非 R2R^2R2 倍,同时保持算法硬件友好。
实际训练速度取决于内核实现策略细节、架构选择(如如何实例化 MIMO 参数)和问题维度,但应不超过 RRR 倍慢。我们发布的 Triton Mamba-3 SISO 内核与 Triton Mamba-2 内核大致相当,MIMO 内核在 R=4R=4R=4 时仅产生 2×2\times2× 减速,因为计算延迟可与内存移动并行化。表 6 基准测试了各种内核的预填充速度,等价于训练内核的前向传播。

MIMO 实例化:在 MIMO 参数化的多种选择中,Mamba-3 的方法实现了一种平衡,保留了其 SISO 对应物的状态尺寸和 SSM 数量,同时避免了参数计数的过度增长。将 SISO SSM 朴素转换为秩 RRR 的 MIMO SSM 将导致参数增加 RRR 倍,因为所有建模 SSM 输入的投影(B,C,x\mathbf{B}, \mathbf{C}, xB,C,x)都会增加。块级组件(如门控 zzz 和输出 yyy 投影)也会受影响。这种参数激增在更大规模下将不可行。为抵消此问题,我们进行以下修改:Mamba 的多值注意力(MVA)头结构导致 B,C\mathbf{B}, \mathbf{C}B,C 跨头共享,因此这些组件的投影可直接转换以纳入新 MIMO 秩 RRR,仅使整个层参数计数从 DNDNDN 轻微增加至 DNRDN RDNR。然而,SSM 输入 xtx_txt、输出 yty_tyt 和门控 ztz_tzt 是每头独有的,因此主导参数计数。此处,直接调整投影将使每头参数计数从 DPDPDP 增加至 DPRDPRDPR。相反,我们保留原始 SISO 投影,并使用可学习的、数据独立的向量逐元素缩放投影输出的每个维度至大小 RRR,使每头产生 DP+PRDP + PRDP+PR 参数。这将乘性增加缓解为更合理的加性参数增加。附录 C 详细说明了参数化,我们论文中的所有 MIMO 变体均通过缩减 MLP 宽度与 SISO 对应物进行参数匹配。
注 6:为简化起见,本节所有讨论均针对较简单的 2 项循环(如指数-欧拉离散化产生的循环);推广到 3 项指数梯形循环是类似的。
在这里插入图片描述

3.4 Mamba-3 架构

整体架构遵循 Llama (Grattafiori et al. 2024),交替使用 Mamba-3 和 SwiGLU 块,并采用预归一化。Mamba-3 块保留了其前代的整体布局,同时引入了几项关键修改。
更新 SSM 循环:SSD 层被替换为命题 4 定义的更具表达力的复数值指数梯形 SSM。Mamba-3 默认采用 SISO SSM 以与其他类 SISO 模型进行公平比较,但其 MIMO 变体可作为更强的基线 Mamba-3 替代品进行训练和部署(表 3)。我们的 SSM A\mathbf{A}A 是复数,包含实部和虚部,均由数据依赖投影生成。结合图 2,这被划分为实数值 AAA 和虚数值 Θ\ThetaΘ;前者如 Mamba-2 一样传入 SSD 黑盒,后者通过 RoPE 技巧计算。
BC/QK 归一化:在 B,C\mathbf{B}, \mathbf{C}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\mathbf{B}, \mathbf{C}B,C 偏置:类似于 Yu and Erichson (2025) 证明在 Mamba-1 的块级变体中为 B\mathbf{B}B 添加通道特定偏置可实现通用逼近能力,Mamba-3 在 BCNorm 后将可学习的、头特定的、通道级偏置纳入 B\mathbf{B}BC\mathbf{C}C 组件。我们假设这些偏置也在模型中诱导了类似卷积的行为。具体而言,为 B\mathbf{B}BC\mathbf{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) 的 100B100\text{B}100B token 进行预训练,使用 Llama-3.1 tokenizer (Grattafiori et al. 2024),上下文长度为 2K2\text{K}2K,采用相同的标准训练协议。训练和评估细节见附录 D。
在所有四个模型规模上,Mamba-3 在各种下游任务上均优于流行基线(表 3)。我们强调,Mamba-3 未使用在许多高性能线性模型中被经验识别为重要组件的外部短卷积 (Allen-Zhu 2025; Gu and Dao 2024; S. Yang, Kautz, and Hatamizadeh 2025)。

4.1.1 MIMO

我们旨在通过调查 MIMO 模型(秩 R=4R=4R=4)的语言建模能力进一步验证 MIMO 的收益,训练设置与前述相同。为确保总参数计数与基于 SISO 的模型相当,我们降低了 MIMO 模型中 MLP 层的内部维度以补偿 MIMO 投影的增加。例如,在 1.5B1.5\text{B}1.5B 参数模型中,MLP 内部维度仅减少 6.6%6.6\%6.6%,从 409640964096 降至 382438243824。详见附录 C。
在验证困惑度和语言评估任务套件上(表 3),我们观察到从 SISO 切换到 MIMO 为 Mamba-3 模型带来了显著提升。具体而言,我们在 1.5B1.5\text{B}1.5B 模型上实现了 0.110.110.11 的显著困惑度收益,图 3 展示了验证损失的下移。在语言评估方面,相比 SISO 我们在大多数任务上看到收益,平均提升 1.21.21.2 个百分点。
在这里插入图片描述

4.1.2 检索能力

除标准语言建模外,线性模型的一个重要衡量标准是检索能力——模型从序列早期回忆信息的能力 (A. Arora et al. 2025; S. Arora, Eyuboglu, et al. 2025)。与可自由通过增长 KV 缓存重访过去上下文的注意力模型不同,线性模型必须将上下文压缩到固定尺寸状态中。这种权衡反映在 Transformer 基线显著更强的检索分数上。为在此视角下评估 Mamba-3,表 4 将其与基线在真实世界和合成“大海捞针”(NIAH)任务 (Hsieh et al. 2024) 上进行比较,使用第 4.1 节预训练的 1.5B1.5\text{B}1.5B 模型。我们将任务序列长度限制为 2K2\text{K}2K token 以匹配训练设置,并对真实世界任务采用填空格式以镜像下一 token 预测目标,遵循 S. Arora, Eyuboglu, et al. (2025) 和 S. Arora, Timalsina, et al. (2024)。
Mamba-3 在真实世界关联回忆和问答(TQA, SQuAD)上具有竞争力,但在从半结构化或非结构化数据中提取信息(SWDE, FDA)时表现挣扎。然而,在合成 NIAH 任务上,Mamba-3 在大多数情况下超越或匹配基线,并显著展现出比其 Mamba-2 前代更优的分布外检索能力。
使用混合模型提升检索:由于固定状态尺寸天然存在检索弱点,我们预测线性层将主要用于混合架构中,通过二次自注意力层缓解此缺点。为评估 Mamba-3 在此架构范式中的表现,我们以 5:15:15:1 的线性层与 NoPE 自注意力 (B. Yang et al. 2025) 交错方式训练同规模混合模型。如先前工作所示 (Waleffe et al. 2024),混合模型优于 Transformer 基线。我们发现向 Mamba-3 层重新引入预输出投影 RMSNorm(预门控,表 4 中分组 RMSNorm)提升了 NIAH 任务的长度泛化检索能力,代价是上下文内真实世界检索任务略有下降,当与自注意力混合时,作为线性序列混合主干极具竞争力。然而,理想归一化类型(分组 vs 默认)及其位置(预 vs 后门控)仍不明确,因为存在相互竞争权衡(附录 E,表 9),我们发现混合模型及其确切特征和动态是复杂且常违反直觉的,近期工作如 Cabannes et al. (2025) 也呼应了这一点。
在这里插入图片描述

4.2 SSM 方法论消融

表 5a 消融了 Mamba-3 引入的核心 SSM 组件变化,主要是引入 BC 偏置和指数梯形离散化。我们报告了 440M440\text{M}440M 规模模型在 Chinchilla 最优 token 数上的预训练测试困惑度。我们发现偏置和指数梯形 SSM 协同良好,使当前许多线性模型使用的短卷积变得冗余。
我们经验证明 Mamba-3 中的数据依赖 RoPE 实现了状态跟踪。遵循 Grazzi, Siems, Zela, et al. (2025),我们在乔姆斯基层级(Chomsky hierarchy)任务上评估——奇偶性、模算术(无括号)和模算术(有括号)——并在表 5b 报告缩放准确率。Mamba-3 解决了奇偶性和模算术(无括号),并大幅缩小了模算术(有括号)的准确率差距。相比之下,无 RoPE 的 Mamba-3、带标准 RoPE 的 Mamba-3 和 Mamba-2 均未能学习这些任务。我们使用启用状态跟踪的 GDN 变体,观察到 Mamba-3 具有竞争力——匹配奇偶性并在两个模算术任务上接近其性能。实验设置见附录 D。
在这里插入图片描述

4.3 推理效率与性能权衡

由于 dstated_{\text{state}}dstate 决定了本文考虑次二次模型的解码运行时间(第 3.3 节),我们将其用作推理速度的代理。通过绘制验证困惑度(模型性能代理)关于 dstated_{\text{state}}dstate 的图,我们旨在全面描绘次二次模型如何在性能与推理速度之间权衡。
图 3 展示了本文考虑 Mamba 模型的帕累托前沿。对于每个数据点,我们在 Fineweb-Edu 数据集上训练 440M440\text{M}440M 参数模型至 2×2\times2× Chinchilla 最优 token,模型配置 dstate∈{16,32,64,128}d_{\text{state}} \in \{16, 32, 64, 128\}dstate{16,32,64,128}。如预期,我们观察到验证损失与 dstated_{\text{state}}dstate 之间的负相关。此外,从 Mamba-2 到 Mamba-3 帕累托前沿总体下移,表明模型更强:在此设置下,2×2\times2× 更小状态尺寸的 Mamba-3 实现了比其 Mamba-2 对应物更好的预训练困惑度,从而以相同质量获得更快模型,或以相同速度获得更好模型。
从 Mamba-3 的 SISO 变体移动到 MIMO 变体(设 MIMO 秩 R=4R=4R=4 并降低 MLP 内部维度以匹配 SISO 参数)观察到进一步下移。
我们在附录 E 图 6 中将比较扩展至 GDN 基线,同样显示 Mamba-3 优于 GDN。
在这里插入图片描述

4.4 快速 Mamba-3 内核

我们以提供实际设置中快速推理的优化内核补充 Mamba-3 的方法论进步。我们为 Mamba-3 实现了新系列推理内核——使用 Triton 用于前向(预填充)路径,使用 CuTe DSL 用于解码——并在表 6 中将其每 token 解码延迟与发布的 Mamba-2 和 GDN Triton 内核进行比较。评估在单张 H100 上测量批量 128128128 的单次解码步骤,涵盖 FP32 和 BF16 数据类型;模型为 1.5B1.5\text{B}1.5B 参数,模型维度 204820482048,状态维度 ∈{64,128}\in \{64, 128\}{64,128}。在所有配置中,SISO 在基线中实现最低延迟。MIMO 凭借更高算术强度,增加了解码 FLOPs 而未显著增加解码运行时间。我们的基准测试表明 CuTe DSL 解码实现具有竞争力,且 Mamba-3 的附加组件(指数梯形更新、复数值状态和 MIMO 投影)是轻量级的。这支持了我们整体的推理优先视角:Mamba-3 允许简单、低延迟实现,同时提供强大的经验性能。
表 7 基准测试了不同解码序列长度的端到端延迟及相同序列长度的预填充时间。解码时间与表 6 一致,Mamba-3 (SISO) 最快;Mamba-3 (MIMO) 与 Mamba-2 相当;所有线性方法在序列长度增长时均快于优化注意力。我们还看到 MIMO 在预填充时产生适度开销,如第 3.3 节所述。基准测试详情见附录 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 层,以促进高效计算和长上下文任务的稳定学习。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,导致在如 S5S_5S5 上的置换组合等任务上失败,除非扩展原始组件。类似地,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) 是单输入单输出 LTI 系统,输入的每一维分配了独立的 SSM。此类 SISO 模型具有比经典 RNN 大得多的循环状态, necessitating 更复杂的数学 machinery 以高效计算它们。为简化模型,S5 (Smith, Warrington, and Linderman 2023) 和 LRU (Orvieto et al. 2023) 用直接应用于整个向量化输入的多输入多输出 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 收敛的视角(第 5.1 节),每个框架的解释和动机略有不同,可导致不同的设计空间和扩展。特别是,线性注意力和测试时训练更密切相关,或许可在关联记忆框架下归类,明确旨在通过“键值”存储记忆输入数据;要么通过规范 KV 方法(即二次注意力)的近似(在 LA 中),要么通过最小化软优化目标(在 TTT 中)。另一方面,状态空间模型有不同的谱系,反映在术语(如 A,B,C,XA, B, C, XA,B,C,X 而非 Q,K,VQ, K, VQ,K,V)及其自然扩展中。值得注意的是,Mamba-3 的方法论改进均与 SSM 视角 specifically 相关,在关联记忆框架中动机较弱。

  1. 指数梯形离散化:SSM 视角包含控制系统的连续 ODE 离散化;我们的指数梯形离散化源于改进的离散化方法。由于关联记忆方法不使用离散化,如何在替代视角下解释如指数梯形的 3 项循环并不明显。
  2. 复数值状态转移:复数值 SSM 长期以来是动力系统的支柱,将复数值视为选择性 SSM 的扩展是自然的。另一方面,关联记忆框架将 AAA 状态转移解释为目标函数的系数,例如对应优化目标中 L2 正则化(或权重衰减)项的权重 (K. A. Wang, Shi, and Fox 2025)。然而,复数值作为回归目标系数无意义;因此,Mamba-3 在此框架内不明显可解释。
  3. 多输入多输出:MIMO 是状态空间模型文献中的经典概念,在关联记忆(线性注意力或测试时训练)框架中不自然出现。然而,我们注意到本文引入的 MIMO 公式不直接 tied 到 SSM 理论——而是受计算视角驱动——我们的技术也可适配其他现代循环模型。
    线性时间序列模型的开发持续取得蓬勃发展,此处的讨论仅捕捉了其中一部分。我们预期随着这些模型的持续演进,将出现日益增长的统一框架、改进的理解和新推广。

6 结论与未来工作

我们引入 Mamba-3,一种在先前 SSM 基础上具有多项方法论改进的状态空间模型:通过指数梯形离散化实现更强大的循环;通过复数值状态转移提升表达力;以及通过 MIMO 公式提升推理效率和建模能力。Mamba-3 的基础 SISO 版本在独立使用和交错混合架构中均提供了强大的语言建模结果,并推进了先前线性序列模型在性能-效率权衡上的帕累托前沿。MIMO 版本以较慢训练为代价换取更强建模能力,同时保持与 Mamba-2 相当的推理效率。
综上所述,Mamba-3 中的技术展示了源于状态空间模型视角的简单且理论驱动的改进,并为高效序列模型开辟了新的方向和设计原则。

致谢
我们衷心感谢 Schmidt Sciences AI2050 研究员项目、Google ML 与系统青年教员奖、Google 研究学者项目、普林斯顿语言与智能(PLI)、Together AI 和 Cartesia AI 的支持。KL 受 NSF GRFP 资助(Grant DGE2140739)。我们也感谢 Sukjun Hwang 和 Gaurav Ghosal 的有益反馈与讨论。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐