Mamba-3: Improved Sequence Modeling using State Space Principles

作者: Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang, Aviv Bick, J. Zico Kolter, Tri Dao, Albert Gu
机构: Carnegie Mellon University, Princeton University, Together AI, Cartesia AI

A1 主要贡献

核心问题与研究目标: 尽管基于Transformer的模型在模型质量上表现出色，但其在推理过程中存在二次方计算复杂度和线性内存增长的问题，导致推理成本高昂。这催生了亚二次方模型（如状态空间模型SSMs）的发展，它们具有恒定的内存和线性计算需求。然而，许多现有的线性模型为了追求算法效率而牺牲了模型质量和能力，例如在状态追踪等任务上表现不佳。此外，这些模型理论上的线性推理在实践中硬件效率并不高，因为其解码阶段的算术强度低，导致硬件资源闲置。因此，本文的研究目标是从“推理优先”的视角出发，改进现有的线性模型（特别是Mamba-2），在提升模型质量和能力的同时，提高其硬件推理效率，从而推动性能-效率帕累托前沿。

创新点: 本文在Mamba-2的基础上，受状态空间模型（SSM）观点的启发，引入了三项核心方法论改进，共同构成了Mamba-3模型：
1. 指数-梯形离散化 (Exponential-Trapezoidal Discretization): 提出了一种新的、更具表达能力的SSM离散化方法。该方法源于对时变、选择性SSM的离散化框架，该框架不仅为Mamba-1/2中缺乏理论依据的启发式离散化（本文称之为“指数-欧拉”方法）提供了形式化解释，还引入了“指数-梯形”这一更精确的实例化。这种新离散化方法在SSM的循环中隐含地应用了一个卷积操作，结合显式的B、C偏置项，使得Mamba-3能够替代语言模型架构中通常认为必不可少的短因果卷积。

复值状态空间模型 (Complex-valued State Space Model): 将Mamba-3的底层SSM视为复值系统，从而实现比Mamba-2更具表达力的状态更新规则。这种轻量级的更新规则旨在克服当前许多线性模型在状态追踪能力上的缺陷。研究表明，这种复值更新等价于一种数据依赖的旋转位置嵌入（RoPE），并且可以高效计算，从而解决了先前线性模型无法完成的合成任务。
多输入多输出SSM (Multi-Input, Multi-Output, MIMO): 为了提高解码过程中的浮点运算（FLOP）效率，将状态更新从基于外积的方式转变为基于矩阵乘法的方式。从信号处理的角度看，这相当于将单输入单输出（SISO）的序列动态推广到多输入多输出（MIMO）。MIMO结构特别适用于推理，因为它在不增加状态大小（从而不影响解码速度）的情况下，通过增加内存密集型状态更新阶段的计算量来提升模型的表达能力。

主要成果:
- 更优的质量: 在1.5B规模下，Mamba-3（MIMO）在下游语言建模任务上的平均准确率比Transformer高2.2个百分点，比Mamba-2高1.9个百分点，比当时最优模型GDN高1.8个百分点。Mamba-3（SISO）也比GDN高0.6个百分点。在状态大小实验中，状态大小为64的Mamba-3（MIMO）与状态大小为128的Mamba-2达到了相当的困惑度，即用一半的延迟实现了同等性能。
- 全新的能力: Mamba-3的复值状态使其能够解决Mamba-2无法解决的合成状态追踪任务，如算术任务。
- 更高的推理效率: 在固定状态大小下，Mamba-3（MIMO）将解码FLOPs提高了多达4倍，同时保持与Mamba-2相似的壁钟解码延迟，并提升了困惑度和下游任务性能。

A3 背景知识

2.1 符号表示

基本定义。标量用普通字体字母（如 $a, b$）表示。张量，包括向量和矩阵，用粗体字母（如 $\mathbf{h}, \mathbf{C}$）表示。张量的形状可以从上下文中推断。输入序列长度表示为 $L$，模型维度为 $D$，SSM状态大小为 $N$。对于时间索引，我们使用下标（如 $x_t$ 表示时间 $t$ 的输入）。两个张量之间的Hadamard积（逐元素乘积）用 $\odot$ 表示。对于向量 $\mathbf{v} \in \mathbb{R}^N$，我们用 $\text{Diag}(\mathbf{v}) \in \mathbb{R}^{N \times N}$ 表示对角线为向量 $\mathbf{v}$ 的对角矩阵。对于跨时间步的标量乘积，我们使用符号 $\alpha_{t \cdots s} = \alpha_t \times \cdots \times \alpha_s = \prod_{i=s}^t \alpha_i$。

2.2 SSM 基础

连续时间线性动力学。状态空间模型（SSMs）通过以下方程组描述连续时间的线性动力学：

$$\dot{h}(t) = A(t) h(t) + B(t) x(t), \qquad y(t) = C(t)^\top h(t),$$

其中，$h(t) \in \mathbb{R}^N$ 是隐藏状态，$x(t) \in \mathbb{R}$ 是输入，$A(t) \in \mathbb{R}^{N \times N}$, $B(t), C(t) \in \mathbb{R}^N$。我们有时会将 $A(t)$ 称为状态转移矩阵，将 $B(t)x(t)$ 称为状态-输入；这也适用于它们的离散化对应物。对于步长为 $\Delta_t$ 的离散序列，Mamba-1和Mamba-2将该系统离散化为以下循环形式：

$$h_{t}=e^{\Delta_{t}\boldsymbol{A}_{t}}h_{t-1}+\Delta_{t}\boldsymbol{B}_{t}x_{t},$$ $$y_t = C_t^\top h_t.$$

Mamba-2的参数化。Mamba-2层【11，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. 2024. https://arxiv.org/abs/2405.21060】的核心是一个数据依赖且硬件高效的SSM。状态转移和状态-输入都通过从当前词元投影出 $\Delta_t \in \mathbb{R}_{>0}$ 和 $B, C \in \mathbb{R}^N$ 来实现数据依赖。通过将状态转移矩阵 $A_t$ 参数化为标量乘以单位矩阵（$A_t = \lambda_t I_{N \times N}$，其中 $\lambda_t \in \mathbb{R}_{<0}$），SSM的循环可以用GPU的矩阵乘法张量核心高效计算。定义 $\alpha_t \triangleq e^{\Delta_t \lambda_t} \in (0, 1)$ 和 $\gamma_t \triangleq \Delta_t$，更新过程变为：

$$ \boldsymbol{h}_{t} = \alpha_{t} \boldsymbol{h}_{t-1} + \gamma_{t} \boldsymbol{B}_{t} x_{t}, \quad y_{t} = \boldsymbol{C}_{t}^{\top} \boldsymbol{h}_{t}. $$

数据依赖的状态转移 $\alpha_t$ 控制了层内每个SSM的记忆范围。特别是 $\Delta_t$ 同时调节状态转移和状态-输入：较大的 $\Delta_t$ 会导致更快的遗忘并更强地加权当前词元，而较小的 $\Delta_t$ 则会保留隐藏状态，并使当前词元的贡献最小化。

备注1。在Mamba-2中，$\lambda_t$ 是数据无关的，因为总的离散转移 $\alpha_t \triangleq e^{\Delta_t \lambda_t}$ 通过 $\Delta_t$ 已经是数据依赖的。在Mamba-3中，我们根据经验发现数据依赖的 $\lambda_t$ 与数据无关的 $\lambda_t$ 性能相似，为了保持一致性，我们选择前者作为默认设置，这样所有的SSM参数都是数据依赖的。

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

状态空间对偶性（SSD）框架。Mamba-2证明了一大类SSM可以表示为一种矩阵形式，该形式将时间步的循环向量化。通过状态空间对偶性（SSD）框架，循环SSM可以被表示为一种并行形式，该形式包含一个逐元素的掩码来模拟状态转移的衰减。SSD为线性循环和可并行化（基于矩阵乘法）的计算形式之间提供了一个通用的对偶框架：

$$\boldsymbol{Y} = (\boldsymbol{L} \odot \boldsymbol{C}\boldsymbol{B}^\top)\boldsymbol{X}$$

其中 $L \in \mathbb{R}^{T \times T}$ 是一个结构化掩码，$B, C \in \mathbb{R}^{T \times N}$，$X \in \mathbb{R}^{T \times D}$ 是SSM的输入，$Y \in \mathbb{R}^{T \times D}$ 是其输出。对 $L$ 施加不同的结构会产生SSD的各种实例化。

与注意力机制的联系。方程（2）也揭示了循环与注意力之间的普遍联系，只需设置 $Q \triangleq C$, $K \triangleq B$, $V \triangleq X$，并将 $L$ 视为一个数据依赖的掩码。实际上，SSD最简单的情况是（因果）线性注意力【26，Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. 2020. https://arxiv.org/abs/2006.16236】，其中 $L$ 是因果三角掩码。

Mamba-2的并行形式。Mamba-2是线性注意力的一个泛化，其掩码 $L$ 由方程（1）中的项 $\alpha_t, \gamma_t$ 构成：

$$\begin{aligned} L = \begin{bmatrix} 1 & & & \\ \alpha_{1} & 1 & & \\ \vdots & & \ddots & \\ \alpha_{T \dots 1} & \dots & \alpha_{T} & 1 \end{bmatrix} \cdot \text{Diag}(\gamma) \end{aligned}$$

在第3.1.3节中，我们将展示Mamba-3是Mamba-2的一个泛化，其 $L$ 更具表达力，因此也是SSD的一个实例。

图1：左图：由指数-梯形规则（第3.1节）导出的结构化掩码是衰减掩码和双带卷积掩码的乘积。右图：欧拉法（保持终点）与梯形法（平均端点）的积分近似对比。

A2 方法细节

我们引入了Mamba-3，这是一个具有三项创新的状态空间模型：“指数-梯形”离散化以实现更具表达力的动态（第3.1节），复值状态空间以实现状态追踪（第3.2节），以及多输入多输出（MIMO）以提高建模能力和推理时硬件利用率（第3.3节）。这些进展解决了当前亚二次方架构在质量、能力和效率方面的局限性。我们在第3.4节中将这些改进整合到一个更新的Mamba架构块中。

3.1 指数-梯形离散化

SSM离散化背景。结构化SSM天然地被定义为连续时间动力系统，它将输入函数 $u(\tau) \in \mathbb{R}$ 映射到输出函数 $y(\tau) \in \mathbb{R}$，其中时间 $\tau > 0$。底层的连续状态空间系统由状态 $h'(\tau)$ 的一阶常微分方程（ODE）和输出 $y(\tau)$ 的代数方程定义。然而，在序列建模中，数据仅在离散时间步上被观察到，这需要对SSM应用离散化步骤，将其连续时间动态转换为离散循环。

传统离散化方法的局限性。离散化方法在经典控制理论中已有深入研究，并且一些规范公式已在早期的深度学习SSM工作中被使用【19，Efficiently Modeling Long Sequences with Structured State Spaces. 2022. https://arxiv.org/abs/2111.00396；20 ，On the Parameterization and Initialization of Diagonal State Space Models. 2022. https://arxiv.org/abs/2206.11893；44 ，Simplified State Space Layers for Sequence Modeling. 2023. https://arxiv.org/abs/2208.04933】。这些机制传统上是针对线性时不变（LTI）系统陈述和应用的，其推导过程不直接适用于线性时变（LTV）系统。此外，虽然Mamba-1在没有证明的情况下将零阶保持（ZOH）方法应用于LTV系统，但与选择性SSM相关的复杂性促使研究人员使用了一种额外的启发式近似，这种近似缺乏理论依据，并且不对应于任何已建立的离散化技术。在接下来的小节中，我们通过我们的离散化框架将当前LTV SSM中使用的先前启发式方法形式化，并利用该框架提出一个更具表达力的离散化方案。

表1：规范的线性时不变离散化方法（顶部）和从我们的指数调整框架中派生的自定义线性时变离散化方法（底部），以及它们在深度学习中使用的结构化SSM中的出现情况。我们的理论将先前的Mamba离散化形式化为指数-欧拉法，并用更具表达力的指数-梯形法对其进行了扩展。广义离散化框架将连续SSM $h'(\tau) = A(\tau)h(\tau) + B(\tau)u(\tau)$ 转换为离散循环 $h_t = \alpha_t h_{t-1} + \beta_{t-1}B_{t-1}x_{t-1} + \gamma_t B_t x_t$，其中不同的离散化方法会产生不同的 $\alpha_t, \beta_t, \gamma_t$ 公式。

3.1.1 指数调整离散化概述

一种新的LTV-SSM离散化推导。我们引入一个简单的推导，它能导出一类新的用于LTV状态空间模型的离散化方法。该方法可以有多种实例化；我们展示了其中一个实例化（指数-欧拉）恰好是Mamba-1/2中使用的启发式方法，从而在理论上对其进行了证明。我们还引入了一个更强大的离-散化方法（指数-梯形），用于Mamba-3。

指数调整系统的直觉。我们推导的高层直觉源于简单线性ODE $u'(t) = Au(t)$ 的闭式解 $u(t) = e^{At}u(0)$，其离散化形式为 $u_{t+1} = e^{\Delta t} u_t$。在这个例子中，指数项主导了底层一阶ODE的动态，导致在不显著限制 $\Delta$ 的情况下使用低阶方法会产生不精确的近似。因此，我们分析了指数调整系统 $e^{-At}u(t)$ 的动态。调整后的系统产生一个离散循环形式，其中状态转移和状态-输入积分被分开近似——状态转移积分通过右手近似法进行近似，即对于所有 $\tau \in [\tau_{t-1}, \tau_t]$，$A(\tau) \triangleq A(\tau_t)$，得到：

$$\begin{aligned} \begin{aligned} \boldsymbol{h}(\tau_t) &= \underbrace{\exp \left(\int_{\tau_{t-1}}^{\tau_t} A(s) d s\right) \boldsymbol{h}(\tau_{t-1})}_{\text{via right-hand approximation}} + \underbrace{\int_{\tau_{t-1}}^{\tau_t} \exp \left(\int_{\tau}^{\tau_t} A(s) d s\right) \boldsymbol{B}(\tau) x(\tau) d \tau}_{\text{via different discretization schemes}}, \\ \boldsymbol{h}_t &\approx \exp \left(\Delta_t A_t\right) \boldsymbol{h}_{t-1} + \int_{\tau_{t-1}}^{\tau_t} \exp \left(\left(\tau_t-\tau\right) A_t\right) \boldsymbol{B}(\tau) x(\tau) d \tau, \end{aligned} \end{aligned}$$

这为状态-输入积分的进一步离散化技术奠定了基础。完整的推导在命题5中详细说明。

零阶保持（ZOH）。经典的零阶保持离散化方法可以从上述基础出发，通过对右侧积分进行特定近似推导得出。通过将 $A_t$, $B(\tau)$, $u(\tau)$ 在区间 $[\tau_{t-1}, \tau_t]$ 上视为常数，且其值固定为右端点 $\tau_t$ 处的值，积分结果为 $A_t^{-1}(\exp(\Delta_t A_t) - I) B_t u_t$。我们注意到，这正式证明了LTI系统的经典ZOH公式可以通过简单地将参数 $A, B, \Delta$ 替换为它们随时间变化的对应项来应用于LTV系统。

指数-欧拉法 (Mamba-1/-2)。虽然Mamba-1声称使用了上述时变ZOH公式，但Mamba-1和Mamba-2在发布的实现中实际上使用了额外的近似。这种离散化方法可以通过使用欧拉法则【47，An Introduction to Numerical Analysis. Cambridge University Press, 2003】来近似状态-输入积分，并在整个区间内保持（右）端点恒定（如图1所示）来恢复：

$$\begin{aligned} \begin{aligned} \boldsymbol{h}_t & \approx e^{\Delta_t A_t} \boldsymbol{h}_{t-1}+\left(\tau_t-\tau_{t-1}\right) e^{\left(\tau_t-\tau_t\right) A_t} \boldsymbol{B}_t x_t \\ & =e^{\Delta_t A_t} \boldsymbol{h}_{t-1}+\Delta_t \boldsymbol{B}_t x_t . \end{aligned} \end{aligned}$$

我们将方程（4）称为指数-欧拉离散化方法，源于指数积分后跟欧拉近似。这个推导为Mamba-1/-2实现中使用的公式提供了理论依据。

指数-梯形法 (Mamba-3)。然而，欧拉法则只提供了状态-输入积分的一阶近似，其局部截断误差与 $O(\Delta_t^2)$ 成正比。相比之下，我们引入了一种广义梯形法则，它提供了积分的二阶精确近似，比欧拉法则具有更高的精度。具体来说，它用区间两端点的数据依赖的凸组合来近似积分。这种泛化扩展了经典的梯形法则【47，An Introduction to Numerical Analysis. Cambridge University Press, 2003】，后者简单地对区间端点进行平均（图1）。

命题1 (指数-梯形离散化)。通过广义梯形法则近似方程(16)中的状态-输入积分，得到以下循环：

$$\begin{aligned} \begin{aligned} \boldsymbol{h}_t &= e^{\Delta_t A_t} \boldsymbol{h}_{t-1} + (1-\lambda_t) \Delta_t e^{\Delta_t A_t} \boldsymbol{B}_{t-1} x_{t-1} + \lambda_t \Delta_t \boldsymbol{B}_t x_t, \\ &=: \alpha_t \boldsymbol{h}_{t-1} + \beta_t \boldsymbol{B}_{t-1} x_{t-1} + \gamma_t \boldsymbol{B}_t x_t, \end{aligned} \end{aligned}$$

其中 $\lambda_t \in [0, 1]$ 是一个数据依赖的标量，$\alpha_t \triangleq e^{\Delta_t A_t}$，$\beta_t \triangleq (1-\lambda_t)\Delta_t e^{\Delta_t A_t}$，$\gamma_t \triangleq \lambda_t \Delta_t$。

备注2 (表达能力)。指数-梯形法是 (a) 经典梯形法（当 $\lambda_t = \frac{1}{2}$ 时恢复）和 (b) Mamba-2的欧拉法（当 $\lambda_t = 1$ 时恢复）的泛化。

备注3 (误差率)。这是状态-输入积分的二阶离散化方法，在标准的稳定性假设下，其误差与 $O(\Delta_t^3)$ 成正比，前提是梯形参数满足 $\lambda_t = \frac{1}{2} + O(\Delta_t)$。然而，我们的消融实验表明，不强制执行此约束对经验性能更好。详见附录A.2和A.3。

新颖性。据我们所知，我们新的离散化框架及其两个实例——指数-欧拉和指数-梯形——对于深度学习中使用的结构化SSM来说都是新颖的。表1比较和总结了状态空间模型的规范和常用离散化方案。

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

等价于卷积。我们的广义指数-梯形离散化等价于在SSM的状态-输入上应用一个大小为2的数据依赖卷积。具体来说，一个常规的循环形式SSM会先实例化状态-输入 $v_t = B_t x_t$，然后计算线性循环 $h_t = \alpha_t h_{t-1} + \gamma_t v_t$。而在方程（6）中，我们首先在 $v_t$ 上应用一个由 $\beta, \gamma$ 加权的宽度为2的卷积，然后再将其输入到线性循环中。

备注4 (卷积差异)。由指数-梯形离散化引入的“卷积”与Mamba和GDN等序列模型使用的标准短卷积有明显区别。标准的短卷积是在核心循环之外对 $x_t$（通常还有 $B_t, C_t$）应用的独立操作，而我们的新离散化可以解释为在核心循环内部对状态-输入 $B_t x_t$ 进行的卷积。

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

作为SSD的实例。我们的新循环可以被实例化为SSD的一个特例，并具有与方程（2）对应的并行形式。从 $h_0 = \gamma_0 B_0 x_0$ 开始展开状态循环得到 $h_t = \alpha_{t \cdots 2} (\gamma_0 \alpha_1 + \beta_1) B_0 x_0 + \cdots + \gamma_t B_t x_t$，其中SSM的输出是 $y_t = \alpha_{t \cdots 2} (\gamma_0 \alpha_1 + \beta_1) C_t^\top B_0 x_0 + \cdots + \gamma_t C_t^\top B_t x_t$。展开这些行表明，由梯形更新导出的掩码不再是端点的固定平均（如经典梯形法则），而是两个区间端点的数据依赖的凸组合。

掩码结构。在SSD框架（2）下，其并行形式为 $Y = (L \odot CB^\top)X$，Mamba-3对应的掩码 $L$ 的结构是一个1-半可分矩阵与一个2-带状矩阵的组合：

$$\begin{aligned} \boldsymbol{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 & & & \ddots \\ \alpha_{T\dots2}(\gamma_{0}\alpha_{1} + \beta_{1}) & & \dots & \gamma_{T} \end{bmatrix} = \begin{bmatrix} 1 & & & \\ \alpha_{1} & 1 & & \\ \alpha_{2}\alpha_{1} & \alpha_{2} & 1 & \\ \vdots & & & \ddots \\ \alpha_{T\dots1} & & \dots & 1 \end{bmatrix} \begin{bmatrix} \gamma_{0} & & & \\ \beta_{1} & \gamma_{1} & & \\ 0 & \beta_{2} & \gamma_{2} & \\ \vdots & & & \ddots \\ 0 & & \dots & \gamma_{T} \end{bmatrix}. \end{aligned}$$

这种并行公式使得SSM输出的计算能够以硬件高效的、以矩阵乘法为中心的方式进行训练。

与卷积的并行视角联系。我们注意到，Mamba-3的卷积连接也可以通过这种并行的对偶形式看到，其中乘以方程（7）中的2-带状矩阵表示与权重 $\beta, \gamma$ 进行卷积。在附录A.1中，我们使用SSD张量收缩机制来证明并行形式等价于一个带有状态-输入卷积的普通SSM。

备注5。Mamba-3的结构化掩码可以看作是Mamba-2的泛化，后者不是2-带状矩阵，而是一个只有 $\gamma_t$ 的对角矩阵（3）。

3.2 复值 SSMs

现代SSM的简化与局限。现代SSM的设计以效率为中心目标，其动机是需要扩展到更大的模型和更长的序列。例如，连续的架构逐步简化了状态转移矩阵：S4【19，Efficiently Modeling Long Sequences with Structured State Spaces. 2022. https://arxiv.org/abs/2111.00396】使用了复值的正规加低秩（NPLR）矩阵，Mamba【18 ，Mamba: Linear-Time Sequence Modeling with Selective State Spaces. 2024. https://arxiv.org/abs/2312.00752】将其简化为实数对角矩阵，而Mamba-2【11 ，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. 2024. https://arxiv.org/abs/2405.21060】进一步简化为单个缩放的单位矩阵。尽管这些简化在很大程度上保持了语言建模的性能，但最近的研究【17 ，Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. 2025. https://arxiv.org/abs/2411.12537；32 ，The Illusion of State in State-Space Models. 2025. https://arxiv.org/abs/2404.08819；42 ，The Expressive Capacity of State Space Models: A Formal Language Perspective. 2024. https://arxiv.org/abs/2405.17394】表明，将转移矩阵的特征值限制为实数且非负，会降低模型在简单状态追踪任务上的能力——这里主要指可解群范畴（TC0），如奇偶校验——而这些任务可以由单层LSTM解决。这种局限性，在【16 ，Is Mamba Capable of In-Context Learning? 2024. https://arxiv.org/abs/2402.03170】的定理1中被形式化，源于将转移矩阵的特征值限制为实数，这无法表示“旋转”的隐藏状态动态。例如，考虑对二进制输入$\{0, 1\}$的奇偶校验函数，定义为 $\sum_t x_t \pmod 2$。这个任务可以使用更新规则：$h_t = R(\theta x_t) h_{t-1}$ 来执行，其中 $R(\cdot)$ 是一个二维旋转矩阵。这种旋转动态无法用实数特征值来表示。

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

恢复状态追踪能力。为了恢复这种能力，我们从复值SSM（8）开始，它能够表示状态追踪动态。我们证明，在离散化（命题5）下，复值SSM可以被表述为具有由2x2旋转矩阵组成的块对角转移矩阵的实值SSM（命题2）。然后我们证明这等价于在输入和输出投影B和C上分别应用数据依赖的旋转嵌入。这个结果建立了复值SSM和数据依赖的RoPE嵌入之间的理论联系（命题3）。最后，Su等人【46，RoFormer: Enhanced Transformer with Rotary Position Embedding. 2023. https://arxiv.org/abs/2104.09864】使用的“RoPE技巧”允许高效实现复值状态转移矩阵，与实值SSM相比，计算开销极小。

命题2 (复值到实值SSM的等价性)。考虑一个复值SSM：

$$\begin{aligned} \begin{aligned} \dot{\boldsymbol{h}}(t) & =\operatorname{Diag}(A(t)+i \boldsymbol{\theta}(t)) \boldsymbol{h}(t)+(B(t)+i \hat{B}(t)) x(t), \\ y(t) & =\operatorname{Re}\left((C(t)+i \hat{C}(t))^{\top} \boldsymbol{h}(t)\right), \end{aligned} \end{aligned}$$

其中 $h(t) \in \mathbb{C}^{N/2}$，$\theta(t), B(t), \hat{B}(t), C(t), \hat{C}(t) \in \mathbb{R}^{N/2}$，以及 $u(t), y(t) \in \mathbb{R}$。在指数-欧拉离散化下，该系统等价于一个实值SSM：

$$\begin{aligned} \begin{aligned} \boldsymbol{h}_t &= e^{\Delta_t A_t} \boldsymbol{R}_t h_{t-1} + \Delta_t B_t x_t, \\ y_t &= \boldsymbol{C}_t^\top h_t, \end{aligned} \end{aligned}$$

其状态为 $h_t \in \mathbb{R}^N$，投影为：

$$\begin{aligned} \boldsymbol{B}_{t} := \begin{bmatrix} B_{t} \\ \hat{B}_{t} \end{bmatrix} \in \mathbb{R}^{N}, \quad \boldsymbol{C}_{t} := \begin{bmatrix} C_{t} \\ -\hat{C}_{t} \end{bmatrix} \in \mathbb{R}^{N}, \end{aligned}$$

以及一个转移矩阵：

$$\begin{aligned} \boldsymbol{R}_t := Block\left(\{R(\Delta_t \boldsymbol{\theta}_t[i])\}_{i=1}^{N/2}\right) \in \mathbb{R}^{N \times N}, \quad R(\theta) := \begin{bmatrix} \cos(\theta) & -\sin(\theta) \\ \sin(\theta) & \cos(\theta) \end{bmatrix} . \end{aligned}$$

证明在附录B.1中给出。命题2表明，状态维度为 $N/2$ 的离散化复值SSM，等价于一个状态维度加倍（$N$）的实值SSM，其转移矩阵是一个标量衰减的、由2x2数据依赖旋转矩阵（$e^{\Delta_t A_t} R_t$）组成的块对角矩阵。

命题3 (复值SSM与数据依赖RoPE的等价性)。在命题2建立的符号体系下，考虑在方程（9）中定义的实值SSM展开 $t$ 个时间步。该SSM的输出等价于一个基于普通标量转移矩阵的SSM（4），但在SSM的B和C分量上应用了数据依赖的旋转嵌入，定义如下：

$$h_t = e^{\Delta_t A_t} h_{t-1} + \left( \prod_{i=0}^t R_i^\top \right) \Delta_t B_t x_t, \qquad y_t = \left[ \left( \prod_{i=0}^t R_i^\top \right) C_t \right]^\top h_t$$

其中矩阵乘积表示右矩阵乘法，例如，$\prod_{i=0}^t R_i^\top = R_0^\top R_1^\top$。我们将使用变换后的实值SSM来计算复值SSM的方法称为“RoPE技巧”。证明在附录B.2中给出。

与RoPE嵌入的联系。要观察复值SSM与RoPE嵌入的联系，请注意在上述命题中，数据依赖的旋转 $R_i$ 随时间步累积并应用于 $C, B$，根据状态空间对偶性框架，这对应于注意力机制的查询（Q）和键（K）分量（第2.3节）。类似地，标准的RoPE【46，RoFormer: Enhanced Transformer with Rotary Position Embedding. 2023. https://arxiv.org/abs/2104.09864】应用的是数据无关的旋转矩阵，其旋转角度遵循一个固定的频率调度 $\theta[i] = 10000^{-2i/D}$。

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

推广到指数-梯形法。在推导出具有指数-欧拉离散化的复值SSM的循环之后，推广到指数-梯形离散化是相似的。命题4为Mamba-3提供了带有RoPE技巧的完整循环公式。

命题4 (使用指数-梯形离散化的旋转嵌入等价性)。用指数-梯形规则（命题1）离散化一个复值SSM，得到以下循环：

$$\begin{aligned} \begin{aligned} \boldsymbol{h}_t &= \alpha_t \boldsymbol{h}_{t-1} + \beta_t \left( \prod_{i=0}^{t-1} R_i^\top \right) B_{t-1} x_{t-1} + \gamma_t \left( \prod_{i=0}^t R_i^\top \right) B_t x_t, \\ y_t &= \left[ \left( \prod_{i=0}^t R_i^\top \right) C_t \right]^\top h_t. \end{aligned} \end{aligned}$$

这里，$R_i$ 是在命题2中定义的块对角旋转矩阵。证明在附录B.3中。我们通过经验验证，我们通过数据依赖的RoPE实现的复值SSM能够解决实值SSM（无论有无标准RoPE）无法解决的状态追踪任务（表5b），这支持了理论主张。

3.3 多输入，多输出

推理效率的挑战。扩展测试时计算已在模型能力方面开辟了新前沿，例如在智能体工作流中，推理占据了总计算预算越来越大的份额。这使得语言模型的推理效率重新成为焦点，并推动了SSM和亚二次方层的采用，这些层具有固定大小的隐藏状态，因此提供了更低的计算和内存需求。尽管这些新层与Transformer相比具有更短的壁钟时间，但它们的解码过程严重受内存限制，导致硬件利用率低。在本节中，我们使用SSM的视角引入对Mamba-3循环的方法论改进，允许在不增加解码壁钟时间的情况下增加模型FLOPs，从而在解码速度相同的情况下获得更好的模型。

表2：(a) SISO 和 (b) MIMO 的算术强度。批次和头维度可以消掉。MIMO的算术强度随秩R线性增加，从而在解码等内存密集型阶段实现更好的硬件利用率。这里N是状态大小（扩展因子），P是头维度。对于Mamba-3，通常R ≪ N, P。
(a) SISO (2字节数据)。

图片描述
(b) MIMO (2字节数据)。

解码算术强度。为了提高硬件效率，我们需要考虑词元生成的算术强度，其定义为给定操作的FLOPs除以输入输出字节数。由于SSM解码会使内存带宽饱和而计算单元闲置（即受内存限制），我们希望增加其算术强度，以便有效地将计算与内存I/O重叠。具体来说，Mamba中单次生成的算术强度约为每字节2.5次操作（表2a），而NVIDIA H100-SXM5【34，NVIDIA H100 Tensor Core GPU White Paper. Tech. rep. NVIDIA, 2022】上bfloat16矩阵乘法的算术强度约为每字节295次操作。因此，SSM解码远未达到计算密集型状态，而且尚不清楚如何调整Mamba中的现有参数来缓解硬件效率不足的问题。我们注意到，这一观察普遍适用于其他亚二次方模型，如因果线性注意力。

从SISO到MIMO。考虑一个典型的SSM的单个头，其头维度为 $P$，它涉及到将SISO循环 $h_t \leftarrow \alpha_t h_{t-1} + \Delta_t B_t x_t$ 与 $N$ 个共享相同 $\alpha_t, \Delta_t$ 和 $B_t$ 的副本堆叠起来。由此产生的广播循环 $h_t \leftarrow \alpha_t h_{t-1} + \Delta_t B_t x_t^\top$ 接受向量输入 $x_t \in \mathbb{R}^P$ 并具有矩阵值的状态 $h_t \in \mathbb{R}^{N \times P}$。

提高算术强度。请注意，内存流量（输入和输出大小）主要由状态 $h_t$ 决定，而计算主要包括外积 $B_t x_t^\top$，其FLOPs与 $N \cdot P$ 成正比。通过增加后者的维度，将 $B_t \in \mathbb{R}^N \rightarrow B_t \in \mathbb{R}^{N \times R}$ 和 $x_t \in \mathbb{R}^P \rightarrow x_t \in \mathbb{R}^{P \times R}$，内存流量不会显著增加（对于小的 $R$），而消耗的FLOPs增加了 $R$ 倍（表2b）。因此，这种变换增加了循环的算术强度。此外，算术强度的增加转化为实际收益，因为外积 $B_t x_t^\top$ 变成了硬件高效的矩阵-矩阵乘积（matmul），它使用快速的张量核心进行计算，仅产生边际的延迟成本。结果是，MIMO循环比原始SISO循环更具表现力，计算量增加了 $R$ 倍，而实际上保持了解码速度。

MIMO整体变换。出于类似的原因，从状态计算输出 $y_t \leftarrow C_t^\top h_t$ 通过将输出投影修改为 $C_t \in \mathbb{R}^N \rightarrow C_t \in \mathbb{R}^{N \times R}$ 获得了一个额外的秩 $R$。总体而言，这种变换等价于将原始的单输入单输出（SISO）循环扩展为多输入多输出（MIMO）。

训练MIMO SSMs。虽然MIMO公式的动机是推理效率，但SSM的训练算法（包括我们在第3.1节、第3.2节中的发展）通常是为SISO模型开发的。我们从观察到MIMO SSMs可以表示为 $R^2$ 个SISO SSMs开始，其中 $R$ 个共享相同循环的SISO SSMs对每个MIMO输出进行求和。具体来说，定义 $C_t^{(i)} \in \mathbb{R}^N, B_t^{(j)} \in \mathbb{R}^N, x_t^{(j)} \in \mathbb{R}, \Delta_t \in \mathbb{R}$，其中 $i, j \in \{0, \dots, R-1\}$，那么我们有：

$$ \boldsymbol{h}_{t}^{(j)} \leftarrow \alpha_{t} \boldsymbol{h}_{t-1}^{(j)}+\Delta_{t} \boldsymbol{B}_{t}^{(j)} \boldsymbol{x}_{t}^{(j)} $$ $$h_{t}=\sum_{j=0}^{R-1} h_{t}^{(j)}$$ $$y_t^{(i)} \leftarrow (C_t^{(i)})^\top h_t$$

因此，$y_t^{(i)} = \sum_j \text{SSM}_{\alpha, \Delta, B^{(j)}, C^{(i)}, x^{(j)}}_t$，其中 $\text{SSM}_{\alpha, \Delta, B^{(j)}, C^{(i)}, x^{(j)}}_t := (C_t^{(i)})^\top h_t^{(j)}$ 且 $h_t^{(j)}$ 来自(12)。此外，对标准SISO-based SSM模型的改进可以直接应用于MIMO模型，因为底层的SISO训练算法可以作为黑盒使用。这一观察使得MIMO模型可以通过并行调用SISO算法 $R^2$ 次来训练。相比之下，当以循环形式计算时，方程(12)、(13)和(14)可以顺序执行，相对于SISO SSMs仅产生 $R$ 倍的开销（回想一下关于MIMO解码FLOPs的讨论）。

MIMO SSMs的分块算法。许多现代SISO循环模型，包括Mamba-2，都是使用分块算法计算的，其中序列被分成长度为 $C$ 的块。在每个块内，应用并行（但渐近较慢）算法，而在块之间计算循环。分块算法在完全并行和完全顺序两种极端算法之间进行插值。通过利用这种结构，我们可以将MIMO SSMs的训练成本降低到SISO SSMs的 $R$ 倍。这个想法也出现在SSD框架中——SSD在每个块内应用硬件友好的二次算法，而在块之间使用循环形式，并表明当状态和头维度相当时，将块大小设置为这个维度会得到一个总体的线性时间算法。具体来说，SSD的块内计算每个块产生 $\approx 2T D^2 + 2T N^2$ FLOPs，总计为 $\frac{L}{C} (2C D^2 + 2C N^2) = 2LD(D+N)$。块间计算每个块产生 $4N D C + 2N C$ FLOPs，总计为 $\frac{L}{C} (4N D C + 2N C) = 4LND + \frac{L}{C} 2NC$（忽略可忽略的项）。设置 $C=D=N$，总FLOPs计数为 $8LD^2$，这在 $L$ 上是线性的。

MIMO的分块算法泛化。SSD的分块算法可以自然地推广到MIMO SSMs。在这种情况下，状态投影 $Bx^\top$ 和状态发射 $C^\top h$ 的FLOPs计数增加 $R$ 倍，而块内分量 $C^\top B$ 的FLOPs计数增加 $R^2$ 倍。因此，块内计算产生 $2 \cdot \frac{L}{C} (RC)^2 D + \frac{L}{C} (RC)^2 N$ FLOPs，块间计算产生 $4 \cdot \frac{L}{C} C (RD) N + 2 \cdot \frac{L}{C} C N$ FLOPs。因此，设置 $RC = D = N$ 得到的总FLOPs计数为 $8 L R D^2$，是 $R$ 倍的FLOPs增长。直观地，将MIMO的块大小设置为SISO块大小的 $R$ 倍，即 $C_{\text{MIMO}} \leftarrow R C_{\text{SISO}}$，可以保持SISO的块内FLOPs计数，同时将块数增加 $R$ 倍，从而导致总FLOPs计数增加 $R$ 倍而不是 $R^2$ 倍，同时保持算法对硬件友好。

实际训练速度。算法在实践中的训练速度取决于内核实现策略的细节、架构选择（如MIMO参数如何实例化）和问题维度，但应该不会比SISO慢 $R^2$ 倍。我们发布的Triton Mamba-3 SISO内核与Triton Mamba-2内核大致相当，而当 $R=4$ 时，MIMO内核仅导致2倍的减速，因为计算延迟可以与内存移动并行。表6对各种内核的预填充（prefill）速度进行了基准测试，这等同于训练内核的前向传递。

MIMO实例化。在MIMO参数化的各种选择中，Mamba-3的方法实现了平衡，保留了其SISO对应物的状态大小和SSM数量，同时避免了参数数量的过度增长。将SISO SSM朴素地转换为秩为 $R$ 的MIMO SSM会导致参数增加 $R$ 倍，因为所有模拟SSM输入的投影 $B, C, x$ 都会增加。块级组件，如门控 $z$（为简单起见至今被忽略）和输出 $y$ 的投影也会受到影响。这种参数数量的涌入在较大模型规模下是难以处理的。为了解决这个问题，我们做了如下改变。Mamba的多值注意力（MVA）头结构导致 $B, C$ 在头之间共享，所以这些组件的投影可以直接转换为包含新的MIMO秩 $R$，整个层的参数数量仅从 $2D^2$ 轻微增加到 $2D^2R$。然而，SSM输入 $x_t$，输出 $y_t$ 和门控 $z_t$ 对每个头都是唯一的，因此主导了参数数量。在这里，直接调整投影会将每个头的参数数量从 $3D^2$ 增加到 $3D^2R$。相反，我们保留原始的SISO投影，并用一个可学习的、数据无关的向量将投影输出的每个维度逐元素缩放到大小 $R$，每个头产生 $3D^2 + 3DR$ 个参数。

图2：对比Mamba-2和Mamba-3架构：关键更新包括指数-梯形离散化、数据依赖的RoPE嵌入、MIMO投影、QK归一化和可学习的偏置。

参数增长控制。这缓解了乘法增长，使其成为更合理的加法参数增长。附录C详细说明了参数化，我们论文中所有的MIMO变体都通过减少MLP宽度来与其SISO对应物进行参数匹配。

备注6。为简单起见，本节中的所有讨论都是针对更简单的两项循环，例如由指数-欧拉离散化产生的循环；推广到三项指数-梯形循环是相似的。

3.4 Mamba-3 架构

整体架构。整体架构遵循Llama【15，The Llama 3 Herd of Models. 2024. https://arxiv.org/abs/2407.21783】，交替使用Mamba-3和SwiGLU块，并采用前置归一化（pre-norm）。Mamba-3块保留了其前身的总体布局，同时引入了几个关键修改。

更新的SSM循环。SSD层被命题4中定义的更具表达力的复值指数-梯形SSM所取代。Mamba-3默认使用SISO SSM，以便与其他类SISO模型进行公平比较，但其MIMO变体可以作为基线Mamba-3的更强替代品进行训练和部署（表3）。我们的SSM A是复值的，实部和虚部都由数据依赖的投影产生。如图2所示，这被划分为实值 $\Delta$ 和虚值 $\Theta$；前者像在Mamba-2中一样被传递到SSD黑盒中，而后者则通过RoPE技巧计算。

BC / QK归一化。在B、C投影之后添加了RMS归一化，这与现代Transformer【22，Query-Key Normalization for Transformers. 2020. https://arxiv.org/abs/2010.04245；56 ，Small-scale proxies for large-scale Transformer training instabilities. 2023. https://arxiv.org/abs/2309.14322】和其他最近的线性模型【24 ，Comba: Improving Bilinear RNNs with Closed-loop Control. 2025. https://arxiv.org/abs/2506.02475；60 ，Gated Delta Networks: Improving Mamba2 with Delta Rule. 2025. https://arxiv.org/abs/2412.06464】中常用的QKNorm类似。我们交替使用BC归一化（BCNorm）或QK归一化（QKNorm）这两个术语。我们发现BCNorm也能够稳定大规模运行，导致在我们的纯Mamba-3模型中移除了后门控RMSNorm层（在Mamba-2中为稳定性而引入）。然而，在混合模型中，被移除的RMSNorm层对于长上下文外推至关重要（表4）。

B, C偏置。与Yu和Erichson【63，Block-Biased Mamba for Long-Range Sequence Processing. 2025. https://arxiv.org/abs/2505.09022】类似，该工作证明了在Mamba-1的分块变体中向B添加通道特定的偏置赋予了其通用逼近能力，Mamba-3在BCNorm之后将可学习的、特定于头的、通道维度的偏置合并到B和C组件中。

偏置的作用。我们假设这些偏置也在模型中引入了类似卷积的行为。具体来说，向B和C添加偏置会向SSM中引入数据无关的组件，这些组件的功能更类似于卷积。关于偏置参数化的消融实验位于附录F。数据无关偏置参数与指数-梯形离散化（其本身在状态-输入上引入了卷积）的结合，根据经验能够消除Mamba-2和大多数现代循环模型中存在的短因果卷积及其伴随的激活函数（第4.2节）。

A4 实验结果

我们通过一系列合成和真实世界的任务，对我们以SSM为中心的方法论变革在Mamba-3模型上进行了经验验证。第4.1节评估了Mamba-3在语言建模和基于检索的任务上的表现。第4.2节对我们新的SSM组件（如离散化和复值转换）的效果进行了消融研究。第4.3节探讨了Mamba-3家族的推理效率以及在固定推理计算下MIMO Mamba-3相对于SISO变体的优势。第4.4节对我们的Mamba-3训练和推理内核的性能进行了基准测试。

4.1 语言建模

预训练设置。所有模型都使用FineWeb-Edu数据集【39，The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale. 2024. https://arxiv.org/abs/2406.17557】的100B个词元进行预训练，采用Llama-3.1分词器【15 ，The Llama 3 Herd of Models. 2024. https://arxiv.org/abs/2407.21783】，上下文长度为2K，并遵循相同的标准训练协议。训练和评估的详细信息可在附录D中找到。

下游任务表现。在所有四个模型规模上，Mamba-3在各种下游任务上的表现均优于流行的基线模型（表3）。我们强调，Mamba-3并未使用外部的短卷积，而该组件在许多高性能线性模型【1，“Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers”. 2025. https://ssrn.com/abstract=5240330；18 ，Mamba: Linear-Time Sequence Modeling with Selective State Spaces. 2024. https://arxiv.org/abs/2312.00752；60 ，Gated Delta Networks: Improving Mamba2 with Delta Rule. 2025. https://arxiv.org/abs/2412.06464】中被经验性地确定为重要组成部分。

表3：在100B FineWeb-Edu词元上训练的模型的下游语言建模评估。每个规模的最佳结果加粗显示，次佳结果加下划线，不包括Mamba-3 MIMO变体。所有模型均采用相同的训练流程。Mamba-3 SISO在每个模型规模上都优于Mamba-2和其他模型，而秩R=4的MIMO变体进一步提升了建模能力。

4.1.1 MIMO

MIMO的增益验证。我们旨在通过在相同设置下训练秩 $R=4$ 的MIMO模型，进一步验证MIMO带来的增益。为确保总参数数量与基于SISO的模型相当，我们减小了MIMO模型中MLP层的内部维度，以补偿因MIMO投影增加的参数。例如，在1.5B参数模型中，MLP内部维度仅减少了6.6%，从4096降至3824。更多细节见附录C。

MIMO的性能提升。在验证集困惑度（validation perplexity）和我们的语言评估任务套件上（表3），我们观察到从SISO转向MIMO的Mamba-3模型有显著的提升。具体来说，我们在1.5B模型上实现了0.11的显著困惑度增益，图3展示了我们验证损失的下降趋势。在语言评估方面，与SISO相比，我们在大多数任务上都看到了提升，平均增益比SISO高出1.2个百分点。

4.1.2 检索能力

线性模型的检索挑战。除了标准的语言建模，衡量线性模型的一个重要标准是其检索能力——即它们从序列早期回忆信息的能力如何【3，Mechanistic evaluation of Transformers and state space models. 2025. https://arxiv.org/abs/2505.15105；4 ，Simple linear attention language models balance the recall-throughput tradeoff. 2025. https://arxiv.org/abs/2402.18668】。与注意力模型可以通过不断增长的KV缓存自由地回顾过去上下文不同，线性模型必须将上下文压缩到一个固定大小的状态中。这种权衡反映在Transformer基线模型显著更强的检索分数上。为了从这个角度评估Mamba-3，表4比较了它与基线模型在真实世界和合成的“大海捞针”（NIAH）任务【23 ，RULER: What’s the Real Context Size of Your Long-Context Language Models? 2024. https://arxiv.org/abs/2404.06654】上的表现，使用的是我们从第4.1节预训练的1.5B模型。我们将任务序列长度限制为2K词元以匹配训练设置，并为我们的真实世界任务采用完形填空（cloze-style）格式，以模仿下一个词元预测目标，这遵循了Arora等人【4 ，Simple linear attention language models balance the recall-throughput tradeoff. 2025. https://arxiv.org/abs/2402.18668；5 ，Just read twice: closing the recall gap for recurrent language models. 2024. https://arxiv.org/abs/2407.05483】的做法。

Mamba-3的检索表现。Mamba-3在真实世界的关联回忆和问答（TQA, SQuAD）任务上具有竞争力，但在从半结构化或非结构化数据（SWDE, FDA）中提取信息时表现不佳。然而，在合成的NIAH任务上，Mamba-3在大多数情况下超过或与基线持平，并显著展示出比其前身Mamba-2更好的分布外检索能力。

通过混合模型提升检索能力。由于固定状态大小模型在检索方面的天然弱点，我们预测线性层将主要用于混合架构中，通过二次自注意力层来缓解这一缺点。为了评估Mamba-3在这种架构范式中的表现，我们以5:1的线***与NoPE自注意力【59，Rope to Nope and Back Again: A New Hybrid Attention Strategy. 2025. https://arxiv.org/abs/2501.18795】层的交错方式，在相同规模下训练了我们的混合模型。正如先前工作【54 ，An Empirical Study of Mamba-based Language Models. 2024. https://arxiv.org/abs/2406.07887】所见，混合模型优于Transformer基线。我们发现，将前置输出投影RMSNorm（在表4中为pre-gate, grouped RMSNorm）重新引入Mamba-3层，可以提高长泛化检索能力，但代价是轻微牺牲了上下文内的真实世界检索任务性能，并且在与自注意力混合时，作为线性序列混合骨干网络具有很强的竞争力。然而，理想的归一化类型（分组vs默认）及其位置（门控前vs门控后）仍不清楚，因为存在相互竞争的权衡（附录E，表9），我们发现混合模型及其确切特性和动态是复杂且常常不直观的，这一点在最近的工作如Cabannes等人【8，Short window attention enables long-term memorization. 2025. https://arxiv.org/abs/2509.24552】中也得到了呼应。

表4：通过真实世界和合成检索任务混合评估的检索能力。真实世界检索任务使用原始数据集的完形填空变体，并截断至2K长度。Mamba-3在关联回忆、问答和大海捞针（NIAH）任务上的长度泛化能力强，但在半结构化和非结构化数据的信息提取方面表现不佳。Transformer基线使用RoPE，这可能解释了其长度泛化问题，而混合模型使用NoPE（无位置嵌入）。我们发现，在Mamba-3 SISO混合模型中添加一个门控前、分组的RMSNorm可以提高NIAH任务的长度泛化能力，但代价是真实世界检索性能略有下降。

4.2 SSM 方法论消融实验

核心组件消融。表5a对Mamba-3引入的核心SSM组件变化进行了消融实验，主要是BC偏置和指数-梯形离散化的引入。我们报告了在440M规模的模型上，以Chinchilla最优词元数训练后的预训练测试困惑度。我们发现偏置和指数-梯形SSM协同作用良好，并使得许多当前线性模型使用的短卷积变得多余。

状态追踪能力验证。我们通过经验证明，Mamba-3中的数据依赖RoPE能够实现状态追踪。遵循Grazzi等人【17，Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. 2025. https://arxiv.org/abs/2411.12537】的方法，我们在Chomsky层级结构的任务上进行评估——奇偶校验（Parity）、模块化算术（无括号）和模块化算术（有括号）——并在表5b中报告了缩放后的准确率。Mamba-3解决了奇偶校验和模块化算术（无括号）任务，并在模块化算术（有括号）任务上几乎弥合了准确率差距。相比之下，没有RoPE的Mamba-3、使用标准RoPE【46 ，RoFormer: Enhanced Transformer with Rotary Position Embedding. 2023. https://arxiv.org/abs/2104.09864】的Mamba-3以及Mamba-2都未能学会这些任务。我们使用了GDN的具备状态追踪能力的变体，并观察到Mamba-3具有竞争力——在奇偶校验上与之匹敌，并在两个模块化算术任务上接近其性能。实验设置详见附录D 。

图片描述
(a) 440M规模下的组件消融实验。我们的BC偏置和指数-梯形离散化的组合使得普遍使用的短卷积成为可选项。

表5: 左图：Mamba-3 SISO核心建模组件的消融实验，结果基于数据集的测试集。右图：形式语言评估（缩放准确率，%）。越高越好。SISO模型在短序列上训练，在长序列上评估以测试长度泛化能力。对于GDN，我们报告了特征值范围为[-1, 1]的变体。

图片描述
(b) 形式语言任务上的性能比较。结果显示，与Mamba-2不同，Mamba-3具备源于数据依赖RoPE嵌入的状态追踪能力。

4.3 推理效率与性能的权衡

状态大小与推理速度。由于 $d_{\text{state}}$ 决定了本文所考虑的亚二次方模型的解码运行时间（第3.3节），我们用它作为推理速度的代理指标。通过将验证困惑度（模型性能的代理指标）作为 $d_{\text{state}}$ 的函数进行绘图，我们旨在构建一个关于亚二次方模型如何在性能与推理速度之间进行权衡的整体图景。

帕累托前沿分析。图3展示了本文所考虑的Mamba模型的这样一个帕累托前沿。对于每个数据点，我们在Fineweb-Edu数据集上将一个440M参数的模型训练到2倍Chinchilla最优词元数，其中模型配置的 $d_{\text{state}}$ 为{16, 32, 64, 128}。正如预期的那样，我们观察到验证损失与 $d_{\text{state}}$ 之间存在负相关关系。此外，从Mamba-2到Mamba-3，帕累托前沿普遍下移，表明模型更强：在这种设置下，状态大小减半的Mamba-3比其Mamba-2对应物获得了更好的预训练困惑度，从而在相同质量下模型更快，或在相同速度下模型更好。

MIMO的进一步提升。当从Mamba-3的SISO变体转向MIMO变体时，观察到帕累托前沿进一步下移（其中我们设置MIMO秩 $R=4$ 并减少MLP内部维度以与SISO变体参数匹配）。我们在附录E的图6中将比较扩展到包括GDN基线，该图也显示Mamba-3与GDN相比具有优势。

图3：不同Mamba变体中状态大小（推理速度代理）与预训练困惑度（性能代理）的探索。与先前的循环SISO模型相比，Mamba-3改善了帕累托前沿，而引入MIMO通过更好的建模性能在不增加状态大小的情况下进一步推动了前沿。

4.4 快速的 Mamba-3 内核

优化的推理内核。我们通过优化的内核来补充Mamba-3的方法论进步，这些内核在实际设置中提供了快速的推理。我们为Mamba-3实现了一系列新的推理内核——前向（prefill）路径使用Triton，解码路径使用CuTe DSL——并在表6中将它们的单词元解码延迟与已发布的Mamba-2和GDN的Triton内核进行比较。评估在单块H100上以批次大小128测量了单步解码，数据类型为FP32和BF16；模型为1.5B参数，模型维度为2048，状态维度∈{64, 128}。

性能基准测试。在所有配置中，SISO在基线中实现了最低的延迟。MIMO具有更高的算术强度，增加了解码FLOPs而没有显著增加解码运行时间。我们的基准测试表明，我们的CuTe DSL解码实现具有竞争力，并且Mamba-3的附加组件（指数-梯形更新、复值状态和MIMO投影）是轻量级的。这支持了我们总体的推理优先视角：Mamba-3允许简单、低延迟的实现，同时提供了强大的经验性能。

表6：不同模型、精度和dstate值下的内核延迟（毫秒）比较。Mamba-3相比Mamba-2引入的开销极小，并具有高效的实际实现。我们的Mamba-3 SISO内核在常用的bf16, dstate=128设置下比参考的Mamba-2和GDN内核更快。Mamba-3 MIMO (R=4) 相比SISO增加的成本很小。

端到端延迟。表7对不同解码序列长度的端到端延迟和相同序列长度的预填充时间进行了基准测试。解码时间与表6一致，其中Mamba-3 (SISO) 最快；Mamba-3 (MIMO) 与Mamba-2相当；并且随着序列长度的增长，所有线性方法都比优化的注意力机制快。我们还看到MIMO对预填充产生了一定的开销，如第3.3节所述。基准测试的细节在附录G中。

表7：不同序列长度下的预填充（Prefill）和预填充+解码（Prefill+Decode）延迟。Mamba-3在其前向传递中增加的开销极小，并保持了有竞争力的解码延迟。详情见附录G。

A4 实验环境

数据集: 使用 FineWeb-Edu 数据集【39，The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale. 2024. https://arxiv.org/abs/2406.17557】中的 100B 个词元进行预训练。下游评估使用了 LAMBADA【38，The LAMBADA dataset: Word prediction requiring a broad discourse context. 2016. https://arxiv.org/abs/1606.06031】、HellaSwag【64，HellaSwag: Can a Machine Really Finish Your Sentence? 2019. https://arxiv.org/abs/1905.07830】、PIQA【7，PIQA: Reasoning about Physical Commonsense in Natural Language. 2019. https://arxiv.org/abs/1911.11641】、Arc-Easy/Challenge【10，Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge. 2018. https://arxiv.org/abs/1803.05457】、WinoGrande【41，WinoGrande: An Adversarial Winograd Schema Challenge at Scale. 2019. https://arxiv.org/abs/1907.10641 】和 OpenBookQA【33，Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering. 2018. https://arxiv.org/abs/1809.02789】等。检索任务使用了 SWDE, SQuAD【40，“Know What You Don’t Know: Unanswerable Questions for SQuAD”. 2018.】，FDA, TriviaQA【25，TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension. 2017. https://arxiv.org/abs/1705.03551 】, NQ【28，“Natural Questions: A Benchmark for Question Answering Research”. 2019.】，和 DROP【12，DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs. 2019. https://arxiv.org/abs/1903.00161】的完形填空变体。
模型架构:
- Mamba家族: 标准扩展因子为2，状态维度 $d_{\text{state}}=128$，头维度为64。
- Transformer基线: 遵循Dao和Gu (2024)的设置【11，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. 2024. https://arxiv.org/abs/2405.21060】。
- GDN基线: 遵循Yang等人 (2025)的设置【60，Gated Delta Networks: Improving Mamba2 with Delta Rule. 2025. https://arxiv.org/abs/2412.06464】，其中 $d, d_{\text{key}} = 128$, $d_{\text{val}} = 256$。
- MIMO变体: 秩 $R=4$，并通过减小MLP层的隐藏维度来匹配SISO模型的参数量。
硬件配置: 内核延迟基准测试在单块 NVIDIA H100-SXM 80GB GPU 上进行。
软件配置:
- 分词器: Llama-3.1分词器【15，The Llama 3 Herd of Models. 2024. https://arxiv.org/abs/2407.21783】。
- 训练: 所有模型均使用 bfloat16 精度进行训练。
- 评估: 使用 LM Evaluation Harness【14，The Language Model Evaluation Harness. Version v0.4.3. 2024. doi: 10.5281/zenodo.12608602】。
- 自定义内核: Mamba-3的prefill内核使用 Triton（SISO）和 TileLang（MIMO），解码内核使用 CuTe DSL。
- Transformer推理引擎: 使用 vLLM (v0.11.0)。

A7 补充细节

5.1 线性时间序列混合器

发展背景。越来越多的工作旨在用线性运行时间的替代方案取代二次方的softmax注意力机制【6，Neural Machine Translation by Jointly Learning to Align and Translate. 2014. https://arxiv.org/abs/1409.0473；53 ，“Attention is all you need”. 2017】。主要方法可分为三大框架：线性注意力、测试时训练和状态空间模型。

线性注意力（LA）。许多早期的线性注意力（LA）模型旨在通过核特征图来近似softmax注意力【9，Rethinking Attention with Performers. 2022. https://arxiv.org/abs/2009.14794；26 ，Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. 2020. https://arxiv.org/abs/2006.16236】，而近期的模型则摒弃了特征图，转而使用查询和键之间的原始点积，并通过衰减或掩码进行调节【49 ，Retentive Network: A Successor to Transformer for Large Language Models. 2023. https://arxiv.org/abs/2307.08621；61 ，Gated Linear Attention Transformers with Hardware-Efficient Training. 2024. https://arxiv.org/abs/2312.06635】。最近，使用键值对调节状态记忆的快速权重编程器【43 ，Linear Transformers Are Secretly Fast Weight Programmers. 2021. https://arxiv.org/abs/2102.11174】也归入“线性注意力”这一总称下。Yang等人【60 ，Gated Delta Networks: Improving Mamba2 with Delta Rule. 2025. https://arxiv.org/abs/2412.06464；62 ，Parallelizing Linear Transformers with the Delta Rule over Sequence Length. 2025. https://arxiv.org/abs/2406.06484】的工作源于这一系列，他们通过用delta规则循环替换加性记忆更新来增强传统线性注意力。这进一步催生了大量基于delta规则改进线性模型效率和能力的工作【24 ，Comba: Improving Bilinear RNNs with Closed-loop Control. 2025. https://arxiv.org/abs/2506.02475；27 ，Kimi Linear: An Expressive, Efficient Attention Architecture. 2025. https://arxiv.org/abs/2510.26692】。

测试时训练（TTT）或测试时回归（TTR）。另一条平行的研究路线将序列建模视为推理过程中的在线学习任务。在这里，循环状态代表了过去输入的压缩摘要，而循环步骤则更新状态以记忆新信息【48，Learning to (Learn at Test Time): RNNs with Expressive Hidden States. 2025. https://arxiv.org/abs/2407.04620；50 ，End-to-End Test-Time Training for Long Context. 2025. https://arxiv.org/abs/2512.23675；65 ，Test-Time Training Done Right. 2025. https://arxiv.org/abs/2505.23884】。等价地，这些方法可以被看作是优化一个全局回归目标，而循环状态更新则代表了梯度下降等迭代优化过程的变体【55 ，Test-time regression: a unifying framework for designing sequence models with associative memory. 2025. https://arxiv.org/abs/2501.12352】。

结构化状态空间模型（SSM）。这是现代循环模型的另一种观点，灵感来源于经典信号处理和动力系统。早期版本的SSM，如S4【19，Efficiently Modeling Long Sequences with Structured State Spaces. 2022. https://arxiv.org/abs/2111.00396；21 ，Diagonal State Spaces are as Effective as Structured State Spaces. 2022. https://arxiv.org/abs/2203.14343；44 ，Simplified State Space Layers for Sequence Modeling. 2023. https://arxiv.org/abs/2208.04933】，使用线性时不变（LTI）层和结构化的状态转移矩阵（例如对角矩阵或对角加低秩矩阵）来促进高效计算和长上下文任务的稳定学习【19 ，Efficiently Modeling Long Sequences with Structured State Spaces. 2022. https://arxiv.org/abs/2111.00396；21 ，Diagonal State Spaces are as Effective as Structured State Spaces. 2022. https://arxiv.org/abs/2203.14343；44 ，Simplified State Space Layers for Sequence Modeling. 2023. https://arxiv.org/abs/2208.04933】。在Mamba-1【18 ，Mamba: Linear-Time Sequence Modeling with Selective State Spaces. 2024. https://arxiv.org/abs/2312.00752】中引入时变、输入依赖的选择性机制，缩小了自注意力与线性模型在信息密集模态（特别是语言建模）上的差距。随后，Mamba-2【11 ，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. 2024. https://arxiv.org/abs/2405.21060】通过我们在这项工作中构建的结构化状态空间对偶性（SSD），形式化了SSM与（线性）注意力之间的联系。

5.2 状态追踪与复值状态空间模型

表达能力与状态追踪。近期的工作描述了循环、常数内存混合器可以维持的状态类型，揭示了先前基于SSM的模型的算法缺陷。Merrill等人【32，The Illusion of State in State-Space Models. 2025. https://arxiv.org/abs/2404.08819】指出，在有限精度下，实际的SSM会退化到TC0，导致在诸如对S5进行置换组合等任务上失败，除非扩展其基本操作。同样，Yu和Erichson【63 ，Block-Biased Mamba for Long-Range Sequence Processing. 2025. https://arxiv.org/abs/2505.09022】证明了单层Mamba不是通用逼近器。为提高表达能力，已提出了几种修改。例如，同一项工作表明，通过分块分解或通道特定偏置的微小改动，分块偏置变体重新获得了通用逼近属性。允许负特征值或非三角转移矩阵使得线性RNN——包括对角和Householder/DeltaNet形式——能够捕捉奇偶性，并在温和假设下处理常规语言【17 ，Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. 2025. https://arxiv.org/abs/2411.12537】。复值参数化为增强表达能力提供了另一条途径。

复值状态空间模型。Mamba之前的结构化SSM通常是复值的，这根植于传统的SSM理论。它们也通常在视觉和音频等具有明确基于频率信息内容的领域表现出色，而不是语言领域。虽然一些模型如H3【13，Hungry Hungry Hippos: Towards Language Modeling with State Space Models. 2023. https://arxiv.org/abs/2212.14052】、RetNet【49 ，Retentive Network: A Successor to Transformer for Large Language Models. 2023. https://arxiv.org/abs/2307.08621】和Megalodon【31 ，Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length. 2024. https://arxiv.org/abs/2404.08801】在针对语言建模时保留了复值SSM，但它们的性能仍然明显不如Transformer 。

与现代SSM的区别。此外，由于这些模型是LTI的，并且使用与现代选择性SSM（如Mamba）截然不同的算法（特别是卷积或显式循环）进行计算，它们通常不使用RoPE技巧来处理复数部分。一个例外是RetNet，它引入了一个介于线性注意力和Mamba-2之间的模型，该模型使用恒定的标量衰减（与LA中无衰减和Mamba-2中数据依赖衰减相对）以及通过RoPE实现的额外恒定复相位。

Mamba-3中的复值。总的来说，经验上发现复数对语言建模没有帮助，因此在Mamba-1及其后继者中被逐步淘汰，包括线性注意力和测试时训练的并行工作线。Mamba-3代表了第一个具有复值状态转移的现代循环模型，其引入是为了增加表达能力和状态追踪能力的特定目的。通过结合RoPE技巧，据我们所知，这是首次在理论动机下使用数据依赖的RoPE。

5.3 多输入，多输出

SISO与MIMO的演变。S4【19，Efficiently Modeling Long Sequences with Structured State Spaces. 2022. https://arxiv.org/abs/2111.00396】是一个单输入单输出（SISO）的LTI系统，其中输入的每个维度都被分配了其独立的SSM。这类SISO模型的循环状态比经典RNN大得多，需要更复杂的数学工具来高效计算。为了简化模型，S5【44 ，Simplified State Space Layers for Sequence Modeling. 2023. https://arxiv.org/abs/2208.04933】和LRU【37 ，Resurrecting Recurrent Neural Networks for Long Sequences. 2023. https://arxiv.org/abs/2303.06349】用一个直接作用于整个向量化输入的多输入多输出（MIMO）SSM取代了SISO SSM集合。这一改变降低了有效状态容量，但通过直接用并行扫描计算循环，开辟了另一条计算路径。虽然状态容量和建模性能之间的这种权衡在LTI模型中不太明显，但Mamba-1 (S6)【18，Mamba: Linear-Time Sequence Modeling with Selective State Spaces. 2024. https://arxiv.org/abs/2312.00752】和Mamba-2【11 ，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. 2024. https://arxiv.org/abs/2405.21060】由于在时变设置下大状态尺寸的重要性而回归到SISO系统。与增加状态尺寸相关的计算瓶颈通过Mamba-1的硬件感知并行扫描算法和Mamba-2的基于矩阵乘法的算法得到解决。

Mamba-3中MIMO的动机。Mamba-3引入MIMO与先前的工作有显著不同。与之前旨在以牺牲轻微表达能力为代价简化训练算法的MIMO模型不同，Mamba-3的MIMO结构旨在增加建模能力同时保持推理效率。因此，其状态扩展保持在Mamba-1/-2的水平以维持建模能力，同时以额外的训练计算为代价。

5.4 状态空间模型视角

不同框架的趋同与差异。尽管现代循环模型有几种很大程度上趋同的不同观点（第5.1节），但每个框架都有略微不同的解释和动机，这可能导致不同的设计空间和扩展。特别是，线性注意力和测试时训练关系更紧密，或许可以归为一个关联记忆的框架，该框架明确旨在通过“键值”存储来记忆输入数据；无论是在LA中通过对经典KV方法的近似（即二次注意力），还是在TTT中通过最小化软优化目标。另一方面，状态空间模型有不同的血统，这反映在术语（例如，$A, B, C, \Delta$ 而不是 $Q, K, V$）和其自然扩展上。值得注意的是，Mamba-3中的方法论改进都与SSM视角特定相关，而从关联记忆框架的动机较少。

SSM视角的独特性。

指数-梯形离散化。SSM视角需要对支配系统的连续ODE进行离散化；我们的指数-梯形离散化源于一种改进的离散化方法。由于关联记忆方法不使用离散化，因此不清楚如何从其他视角解释像指数-梯形这样的三项循环。
复值状态转移。复值SSM长期以来一直是动力系统的主力，将复值视为选择性SSM的扩展是自然的。另一方面，关联记忆框架将 $A$ 状态转移解释为目标函数的系数，例如对应于优化目标中L2正则化（或权重衰减）项的权重【55，Test-time regression: a unifying framework for designing sequence models with associative memory. 2025. https://arxiv.org/abs/2501.12352】。然而，复值作为回归目标的系数是无意义的；因此，Mamba-3在这些框架内不易解释。
多输入，多输出。MIMO是状态空间模型文献中的一个经典概念，在关联记忆（线性注意力或测试时训练）框架中并不自然出现。然而，我们确实注意到，本文中介绍的MIMO公式并不直接与SSM理论相关联——而是从计算角度出发——我们的技术也可以适应于其他现代循环模型。

未来展望。线性时间序列模型的开发仍在蓬勃发展，这里的讨论只涵盖了其中的一部分。我们预计随着这些模型的不断发展，将会出现更多统一的框架、更深入的理解和新的泛化。

A5 结论

我们引入了Mamba-3，一个在先前SSM基础上进行了几项方法论改进的状态空间模型：通过指数-梯形离散化实现了更强大的循环；通过复值状态转移提高了表达能力；以及通过MIMO公式实现了更高的推理效率和建模能力。Mamba-3的基础SISO版本在独立使用和交错混合架构中都表现出强大的语言建模结果，并在性能-效率权衡上推动了现有线性序列模型的帕累托前沿。MIMO版本以较慢的训练速度换取了更强的建模能力，同时与Mamba-2相比保持了有竞争力的推理效率。总而言之，Mamba-3中的技术展示了从状态空间模型视角出发的简单且有理论动机的改进，并为高效序列模型开辟了新的方向和设计原则。

A6 附录

A 指数-梯形离散化

命题5 (常数变易法 (Tenenbaum and Pollard 1985))。考虑线性SSM

$$\dot{h}(t)=A(t) h(t)+B(t) x(t),$$

其中 $h(t) \in \mathbb{R}^N$，$A(t) \in \mathbb{R}$ 是一个标量衰减，以及 $B(t)x(t) \in \mathbb{R}^N$。对于 $\Delta_t$ 离散化的时间网格 $\tau_t = \tau_{t-1} + \Delta_t$，隐藏状态满足方程(15)，然后可以近似为方程(16)，误差为 $O(\Delta_t^2)$。对状态-输入的剩余积分的近似可以有不同的误差界，具体取决于所使用的方法：一个例子可以在附录A.2中找到。

$$\begin{aligned} \begin{aligned} \boldsymbol{h}(\tau_t) &= \exp \left( \int_{\tau_{t-1}}^{\tau_t} A(s) \, ds \right) \boldsymbol{h}(\tau_{t-1}) + \int_{\tau_{t-1}}^{\tau_t} \exp \left( \int_{\tau}^{\tau_t} A(s) \, ds \right) \boldsymbol{B}(\tau) x(\tau) \, d\tau, \\ \boldsymbol{h}_t &\approx e^{\Delta_t A_t} \boldsymbol{h}_{t-1} + \int_{\tau_{t-1}}^{\tau_t} e^{(\tau_t - \tau) A_t} \boldsymbol{B}(\tau) x(\tau) \, d\tau. \end{aligned} \end{aligned}$$

证明。从初始的线性SSM开始，应用一个积分因子 $z(t) \triangleq e^{\int_0^t -A(s)ds}$ 以方便积分。

$$z(t)\dot{h}(t)=z(t)A(t)h(t)+z(t)B(t)x(t)$$

考虑到 $z'(t) = -A(t)z(t)$；通过重新整理项并在时间网格 $[\tau_{t-1}, \tau_t]$ 之间积分

$$\int_{\tau_{t-1}}^{\tau_t} \frac{d}{d\tau} (z(\tau)h(\tau)) d\tau = \int_{\tau_{t-1}}^{\tau_t} z(\tau)B(\tau)x(\tau)d\tau$$

得到

$$z(\tau_t)h(\tau_t) - z(\tau_{t-1})h(\tau_{t-1}) = \int_{\tau_{t-1}}^{\tau_t} z(\tau)B(\tau)x(\tau)d\tau,$$

这可以整理成一个更熟悉的形式

$$h(\tau_{t})=z(\tau_{t})^{-1}z(\tau_{t-1})h(\tau_{t-1})+\int_{\tau_{t-1}}^{\tau_{t}}z(\tau_{t})^{-1}z(\tau)B(\tau)x(\tau)d\tau.$$

代入积分因子 $z(t)$ 对应于

$$ \boldsymbol{h}(\tau_t) = \exp \left( \int_{\tau_{t-1}}^{\tau_t} A(s)ds \right) \boldsymbol{h}(\tau_{t-1}) + \int_{\tau_{t-1}}^{\tau_t} \exp \left( \int_{\tau}^{\tau_t} A(s)ds \right) B(\tau)x(\tau)d\tau. $$

我们用右手假设来近似状态转移积分，其中对于 $\forall \tau \in [\tau_{t-1}, \tau_t], A(\tau) \triangleq A(\tau_t)$，我们称之为 $A_t$，

$$ \boldsymbol{h}_{t} \approx \underbrace{\exp \left(\Delta_{t} A_{t}\right) \boldsymbol{h}_{t-1}}_{\text {right-hand approximation }}+\underbrace{\int_{\tau_{t-1}}^{\tau_{t}} \exp \left(\left(\tau_{t}-\tau\right) A_{t}\right) \boldsymbol{B}(\tau) x(\tau) d \tau}_{\text {to be approximated }} $$

这会产生一个阶数为 $O(\Delta_t^2)$ 的局部截断误差。因此，我们已经近似了调整后的底层ODE的指数动态，并将状态-输入积分留给任何一系列方法来近似。□

A.1 指数-梯形离散化的掩码矩阵

证明。当观察张量收缩形式时，我们根据Mamba-2论文将维度称为 $T = (\text{Time}), N = (\text{State}), P = (\text{Head})$。通过这种对掩码的分解，我们可以将 $L = \text{contract}(TJ, JS \to TS)(L_1, L_2)$。原始的收缩可以看作是：

$$\text{contract}(TN, SN, TS, SP \rightarrow TP)(C, B, L, X)$$

现在我们可以将其视为：

$$\text{contract}(TN, SN, TJ, JS, SP \rightarrow TP)(C, B, L_1, L_2, X)$$

这可以分解为以下步骤：

$$\begin{aligned} \begin{aligned} Z &= \text{contract}(SN, SP \rightarrow SNP)(B, X) \\ Z' &= \text{contract}(JS, SNP \rightarrow JNP)(L_2, Z) \\ H &= \text{contract}(TJ, JNP \rightarrow TNP)(L_1, Z') \\ Y &= \text{contract}(TN, TNP \rightarrow TP)(C, H) \end{aligned} \end{aligned}$$

我们可以将这一步：contract(JS, SNP → JNP)(L2, Z) 视为在状态-输入（B, X外积）上应用一个大小为2的卷积，然后在传统SSD的 $L=L_1$ 矩阵衰减之前进行。□

A.2 指数-梯形离散化误差率

标准假设。我们假设：$A(t), B(t), x(t)$ 在每个时间步上是有界的且为 $C^3$，因此 $g(t)$ 有三个有界导数；映射 $h \mapsto A(t)h + B(t)x(t)$ 在 $h$ 上是Lipschitz的，这对于线性系统是成立的；$\lambda_t$ 位于一个有界区间内，以使更新是零稳定的。

证明。令 $g(\tau) \triangleq e^{(\tau_t - \tau)A_t} B(\tau)x(\tau)$ 表示命题5中第二项的被积函数。由于 $A(t), B(t), x(t)$ 在 $[\tau_{t-1}, \tau_t]$ 上是 $C^3$ 的，函数 $g$ 有三个有界导数。在 $\tau_{t-1}$ 处对 $g$ 进行二阶泰勒展开得到：

$$\int_{t_{k-1}}^{t_{k}} g(\tau) d \tau=\Delta_{t} g\left(t_{k-1}\right)+\frac{\Delta_{t}^{2}}{2} g^{\prime}\left(t_{k-1}\right)+\frac{\Delta_{t}^{3}}{6} g^{\prime \prime}\left(t_{k-1}\right)+O\left(\Delta_{t}^{4}\right).$$

回想一下，这个积分的梯形近似由下式给出：

$$Q_{\lambda}=\Delta_{t}\left[\left(1-\lambda_{t}\right) g\left(t_{k-1}\right)+\lambda_{t} g\left(t_{k}\right)\right].$$

使用泰勒展开式展开 $g(\tau_t)$：$g(\tau_t) = g(\tau_{t-1}) + \Delta_t g'(\tau_{t-1}) + \frac{\Delta_t^2}{2} g''(\tau_{t-1}) + O(\Delta_t^3)$。将此代入 $Q_\lambda$，

$$\begin{aligned} \begin{aligned} Q_{\lambda} & =\Delta_{t}\left[\left(1-\lambda_{t}\right) g\left(t_{k-1}\right)+\lambda_{t} g\left(t_{k}\right)\right] \\ & =\Delta_{t} g\left(t_{k-1}\right)+\lambda_{t} \Delta_{t}^{2} g^{\prime}\left(t_{k-1}\right)+\lambda_{t} \frac{\Delta_{t}^{3}}{2} g^{\prime \prime}\left(t_{k-1}\right)+O\left(\Delta_{t}^{4}\right) . \end{aligned} \end{aligned}$$

因此，误差由下式给出：

$$\int_{t_{k-1}}^{t_k} g(\tau) d \tau-Q_\lambda=\left(\frac{1}{2}-\lambda_t\right) \Delta_t^2 g^{\prime}\left(t_{k-1}\right)+\left(\frac{1}{6}-\frac{\lambda_t}{2}\right) \Delta_t^3 g^{\prime \prime}\left(t_{k-1}\right)+O\left(\Delta_t^4\right) .$$

在假设 $\lambda_t = \frac{1}{2} + c_t \Delta_t$ 的情况下，其中 $c_t = O(1)$，那么 $\frac{1}{2} - \lambda_t = -c_t \Delta_t = O(\Delta_t)$，因此 $\Delta_t^2$ 项是 $O(\Delta_t^3)$。因此，

$$\int_{t_{k-1}}^{t_k} g(\tau) d \tau-Q_\lambda=O\left(\Delta_t^3\right),$$

这产生了一个 $O(\Delta_t^3)$ 的局部截断误差。

A.3 指数-梯形参数化

设置。所有运行都使用Mamba-3 (SISO) 440M模型，在Chinchilla规模下训练，其他架构和优化超参数与表3中相同。

默认参数化。默认模型使用数据依赖的门控 $\lambda_t = \sigma(u_t)$，其中 $u_t$ 是当前输入词元的学习投影。在表8中，我们尝试了 $\lambda_t$ 的不同参数化，发现默认参数化在经验上表现最好。因此，我们选择了更简单的默认参数化，它不强制 $\lambda_t = \frac{1}{2} + O(\Delta_t)$。

B 复值SSM证明

B.1 命题2的证明

命题2 (复值到实值SSM的等价性)。考虑一个复值SSM：

$$\begin{aligned} \begin{aligned} \dot{\boldsymbol{h}}(t) & = \operatorname{Diag}(A(t) + i \boldsymbol{\theta}(t)) \boldsymbol{h}(t) + (B(t) + i \hat{B}(t)) x(t), \\ y(t) & = \operatorname{Re}\left((C(t) + i \hat{C}(t))^{\top} \boldsymbol{h}(t)\right), \end{aligned} \end{aligned}$$

$$\begin{aligned} \begin{aligned} \boldsymbol{h}_t &= e^{\Delta_t A_t} \boldsymbol{R}_t \boldsymbol{h}_{t-1} + \Delta_t \boldsymbol{B}_t \boldsymbol{x}_t, \\ y_t &= \boldsymbol{C}_t^{\top} \boldsymbol{h}_t, \end{aligned} \end{aligned}$$

其状态为 $h_t \in \mathbb{R}^N$，投影为：

$$\begin{aligned} \boldsymbol{B}_t := \begin{bmatrix} B_t \\ \hat{B}_t \end{bmatrix} \in \mathbb{R}^N, \quad \boldsymbol{C}_t := \begin{bmatrix} C_t \\ -\hat{C}_t \end{bmatrix} \in \mathbb{R}^N, \end{aligned}$$

以及一个转移矩阵：

$$\begin{aligned} \boldsymbol{R}_t := \operatorname{Block}\left(\left\{R\left(\Delta_t \boldsymbol{\theta}_t[i]\right)\right\}_{i=1}^{N / 2}\right) \in \mathbb{R}^{N \times N}, \quad R(\theta):=\left[\begin{array}{cc} \cos (\theta) & -\sin (\theta) \\ \sin (\theta) & \cos (\theta) \end{array}\right] . \end{aligned}$$

证明。我们首先给出 $N=2$ 的推导；对于一般的偶数 $N$，通过将坐标成对分组，可以得到块对角结构。令 $h_k + i\hat{h}_k$ 表示复化隐藏状态，参数为 $A(t) + i\theta(t)$ 和 $B(t) + i\hat{B}(t)$ 分别用于转移和输入。根据常数变易公式（命题5），在一个步长 $[\tau_{t-1}, \tau_t]$ 上应用零阶保持和欧拉法则得到：

$$h_k + i \hat{h}_k = e^{\Delta_t (A_t + i \theta_t)} (h_{k-1} + i \hat{h}_{k-1}) + \Delta_t (B_t + i \hat{B}_t) x_t.$$

展开指数项，

$$e^{\Delta_t(A_t+i\theta_t)} = e^{\Delta_t A_t}\left(\cos(\Delta_t\theta_t) + i\sin(\Delta_t\theta_t)\right),$$

所以在实坐标下 $h_t = \begin{bmatrix} h_t \\ \hat{h}_t \end{bmatrix} \in \mathbb{R}^2$ 的循环变为：

$$\begin{aligned} \boldsymbol{h}_t = e^{\Delta_t A_t} \underbrace{\begin{bmatrix} \cos(\Delta_t \theta_t) & -\sin(\Delta_t \theta_t) \\ \sin(\Delta_t \theta_t) & \cos(\Delta_t \theta_t) \end{bmatrix}}_{R(\Delta_t \theta_t)} \boldsymbol{h}_{t-1} + \Delta_t \begin{bmatrix} B_t \\ \hat{B}_t \end{bmatrix} x_t. \end{aligned}$$

对 $N/2$ 对这样的坐标进行堆叠，得到块对角转移矩阵：

$$\begin{aligned} \boldsymbol{h}_{t}=e^{\Delta_{t} A_{t}} \operatorname{Block}\left(\left\{R\left(\Delta_{t} \theta_{t}[i]\right)\right\}_{i=1}^{N / 2}\right) \boldsymbol{h}_{t-1}+\Delta_{t}\left[\begin{array}{l} B_{t} \\ \hat{B}_{t} \end{array}\right] x_{t} . \end{aligned}$$

对于输出，

$$\begin{aligned} y_t = \text{Re}\left( (\boldsymbol{C}_t + i\hat{\boldsymbol{C}}_t)^\top (h_t + i\hat{h}_t) \right) = \begin{bmatrix} \boldsymbol{C}_t \\ -\hat{\boldsymbol{C}}_t \end{bmatrix}^\top \boldsymbol{h}_t, \end{aligned}$$

这定义了命题中的实投影 $C_t \in \mathbb{R}^N$。这证明了复值SSM与带有旋转的实块对角系统之间的等价性。□

B.2 命题3的证明

$$h_t = e^{\Delta_t A_t} h_{t-1} + \left( \prod_{i=0}^t R_i^\top \right) \Delta_t B_t x_t, \qquad y_t = \left[ \left( \prod_{i=0}^t R_i^\top \right) C_t \right]^\top h_t$$

其中矩阵乘积表示右矩阵乘法，例如，$\prod_{i=0}^1 R_i^\top = R_0^\top R_1^\top$。我们将使用变换后的实值SSM来计算复值SSM的方法称为“RoPE技巧”。

证明。考虑SSM

$$h_t = e^{\Delta_t A_t} R_t h_{t-1} + \Delta_t B_t x_t, \quad y_t = C_t^\top h_t,$$

其中（如命题3中）$A_t \in \mathbb{R}$ 是一个标量（因此 $e^{\Delta_t A_t}$ 是一个标量并与旋转通勤），并且 $R_t$ 是块对角正交/酉矩阵，因此 $R_t^{-1} = R_t^\top$，并且矩阵 $R_s, R_i$ 通勤，即 $R_s R_i = R_i R_s$。

展开循环，约定空积为单位矩阵，

$$h_t = \sum_{i=0}^t \left( \prod_{s=i+1}^t e^{\Delta_s A_s} R_s \right) \Delta_i B_i x_i.$$

因此

$$y_t = C_t^\top h_t = \sum_{i=0}^t C_t^\top \left( \prod_{s=i+1}^t e^{\Delta_s A_s} \boldsymbol{R}_s \right) \Delta_i \boldsymbol{B}_i x_i.$$

利用其酉性质，

$$\prod_{s=i+1}^t R_s = \left(\prod_{s=0}^t R_s\right) \left(\prod_{s=0}^i R_s\right)^{-1} = \left(\prod_{s=0}^t R_s\right) \left(\prod_{s=0}^i R_s^\top\right).$$

由于 $e^{\Delta_s A_s}$ 是标量，它们与旋转通勤；因此

$$\begin{aligned} \begin{aligned} y_t & = \sum_{i=0}^t C_t^\top \left( \prod_{s=0}^t R_s \right) \left( \prod_{s=i+1}^t e^{\Delta_s A_s} \right) \left( \prod_{s=0}^i R_s^\top \right) \Delta_i B_i x_i \\ & = \left[ \left( \prod_{s=0}^t R_s^\top \right) C_t \right]^\top \sum_{i=0}^t \left( \prod_{s=i+1}^t e^{\Delta_s A_s} \right) \left( \prod_{s=0}^i R_s^\top \right) \Delta_i B_i x_i \end{aligned} \end{aligned}$$

定义旋转后的参数 $\bar{C}_t := (\prod_{s=0}^t R_s^\top) C_t$ 和 $\bar{B}_i := (\prod_{s=0}^i R_s^\top) B_i$。那么，

$$y_t = \overline{\boldsymbol{C}}_t^{\top} \sum_{i=0}^t \left( \prod_{s=i+1}^t e^{\Delta_s A_s} \right) \Delta_i \overline{\boldsymbol{B}}_i x_i.$$

等价地，引入旋转后的状态 $\tilde{h}_t := (\prod_{s=0}^t R_s^\top) h_t$，

$$\tilde{\boldsymbol{h}}_t=e^{\Delta_t A_t} \tilde{\boldsymbol{h}}_{t-1}+\Delta_t \overline{\boldsymbol{B}}_t x_t, \quad y_t=\overline{\boldsymbol{C}}_t^{\top} \tilde{\boldsymbol{h}}_t,$$

B.3 命题4的证明

命题4 (使用指数-梯形离散化的旋转嵌入等价性)。用指数-梯形规则（命题1）离散化一个复值SSM，得到以下循环：

$$\begin{aligned} \begin{aligned} \boldsymbol{h}_{t} & =\alpha_{t} \boldsymbol{h}_{t-1}+\beta_{t}\left(\prod_{i=0}^{t-1} \boldsymbol{R}_{i}^{\top}\right) \boldsymbol{B}_{t-1} x_{t-1}+\gamma_{t}\left(\prod_{i=0}^{t} \boldsymbol{R}_{i}^{\top}\right) \boldsymbol{B}_{t} x_{t}, \\ y_{t} & =\left[\left(\prod_{i=0}^{t} \boldsymbol{R}_{i}^{\top}\right) \boldsymbol{C}_{t}\right]^{\top} \boldsymbol{h}_{t} . \end{aligned} \end{aligned}$$

这里，$R_i$ 是在命题2中定义的块对角旋转矩阵。

证明。我们从复值SSM（如命题2中）开始

$$\begin{aligned} \begin{aligned} \dot{\boldsymbol{h}}(t) & = \operatorname{Diag}\left(A(t)+i \boldsymbol{\theta}(t)\right) \boldsymbol{h}(t)+\left(B(t)+i \hat{B}(t)\right) x(t), \\ y(t) & = \operatorname{Re}\left((C(t)+i \hat{C}(t))^{\top} \boldsymbol{h}(t)\right), \end{aligned} \end{aligned}$$

其中 $A(t) \in \mathbb{R}$ 是一个标量，$\theta(t), B(t), \hat{B}(t), C(t), \hat{C}(t) \in \mathbb{R}^{N/2}$。

回忆命题5，

$$ \boldsymbol{h}_t \approx e^{\Delta_t(A_t + i \boldsymbol{\theta}_t)} \boldsymbol{h}_{t-1} + \int_{\tau_{t-1}}^{\tau_t} e^{(\tau_t - \tau)(A_t + i \boldsymbol{\theta}_t)} (B(\tau) + i \hat{B}(\tau)) x(\tau) d\tau. $$

对上述积分应用命题1，我们得到

$$h_t = e^{\Delta_t (A_t + i \theta_t)} h_{t-1} + \beta_t e^{i \Delta_t \theta_t} (B_{t-1} + i \hat{B}_{t-1}) x_{t-1} + \gamma_t (B_t + i \hat{B}_t) x_t,$$

其中

$$\alpha_t := e^{\Delta_t A_t}, \quad \beta_t := (1 - \lambda_t)\Delta_t e^{\Delta_t A_t}, \quad \gamma_t := \lambda_t \Delta_t,$$

由于 $e^{\Delta_t(A_t+i\theta_t)} = e^{i\Delta_t\theta_t}e^{\Delta_t A_t}$ 且如命题2所示，乘以 $e^{i\Delta_t\theta_t}$ 在实坐标中是一个块对角旋转，我们得到实数 $N$ 维循环

$$\begin{aligned} \begin{aligned} \boldsymbol{h}_t &= \alpha_t \boldsymbol{R}_t \boldsymbol{h}_{t-1} + \beta_t \boldsymbol{R}_t \boldsymbol{B}_{t-1} x_{t-1} + \gamma_t \boldsymbol{B}_t x_t, \\ y_t &= \boldsymbol{C}_t^\top \boldsymbol{h}_t, \end{aligned} \end{aligned}$$

其中 $R_t \triangleq \text{Block}\{R(\Delta_t \theta_t[i])\}_{i=1}^{N/2}$，其中 $R(\theta) \triangleq \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix}$，并且投影 $B_t \triangleq \begin{bmatrix} B_t \\ \hat{B}_t \end{bmatrix}$，$C_t \triangleq \begin{bmatrix} C_t \\ -\hat{C}_t \end{bmatrix}$。注意 $R_t$ 是正交的，所以 $R_t^{-1} = R_t^\top$，并且 $R_s, R_i$ 通勤，即 $R_s R_i = R_i R_s$。

我们定义如下。

$$\tilde{\boldsymbol{h}}_t := \left( \prod_{s=0}^t \boldsymbol{R}_s^\top \right) \boldsymbol{h}_t, \quad \overline{\boldsymbol{B}}_t := \left( \prod_{s=0}^t \boldsymbol{R}_s^\top \right) \boldsymbol{B}_t, \quad \overline{\boldsymbol{C}}_t := \left( \prod_{s=0}^t \boldsymbol{R}_s^\top \right) \boldsymbol{C}_t.$$

用 $\prod_{s=0}^t R_s^\top$ 左乘方程(25)并使用 $R_t^\top R_t = I$，

$$\begin{aligned} \begin{aligned} \tilde{\boldsymbol{h}}_t &= \alpha_t \tilde{\boldsymbol{h}}_{t-1} + \beta_t \overline{\boldsymbol{B}}_{t-1} x_{t-1} + \gamma_t \overline{\boldsymbol{B}}_t x_t, \\ y_t &= \overline{\boldsymbol{C}}_t^{\top} \tilde{\boldsymbol{h}}_t. \end{aligned} \end{aligned}$$

这是一个普通的标量转移SSM，其数据依赖的旋转嵌入通过 $R_s^\top$ 的累积乘积被吸收到 $B, C$ 中。□

C Mamba-3的MIMO

带MIMO的Mamba。对于给定的批次、头和序列位置 $t$，考虑输入 $U_t \in \mathbb{R}^D$。也用 $P, R \in \mathbb{N}$ 分别表示头维度和MIMO秩。我们首先通过一组以张量收缩符号定义的投影来获得SSM参数，如下所示：

$$\begin{aligned} \begin{aligned} \boldsymbol{B}_{t} & =\operatorname{contract}(D N R, D \rightarrow N R)\left(\boldsymbol{W}_{B}, \boldsymbol{U}_{t}\right) & \boldsymbol{C}_{t} & =\operatorname{contract}(D N R, D \rightarrow N R)\left(\boldsymbol{W}_{C}, \boldsymbol{U}_{t}\right), \\ \boldsymbol{X}_{t}^{\prime} & =\operatorname{contract}(P D, D \rightarrow P)\left(\boldsymbol{W}_{X^{\prime}}, \boldsymbol{U}_{t}\right) & \boldsymbol{X}_{t} & =\operatorname{contract}(P R, P \rightarrow P R)\left(\boldsymbol{W}_{X}, \boldsymbol{X}_{t}^{\prime}\right), \end{aligned} \end{aligned}$$

其中 $W_B, W_C, W_X', W_X$ 是模型参数。此外，我们以与 $X_t$ 相同的方式，使用权重 $W_Z'$ 和 $W_Z$ 获得残差门项 $Z_t$。这种参数化用于防止参数数量增加 $R$ 倍。

MIMO SSM的更新与输出。状态更新和SSM输出然后通过以下MIMO SSM计算：

$$H_{t}=a_{t}H_{t-1}+B_{t}X_{t}^{\top}\in\mathbb{R}^{N\times P},\quad Y_{t}=H_{t}^{\top}C_{t}\in\mathbb{R}^{P\times R}.$$

中间输出 $Y_t'$ 是通过残差函数 $\mathcal{G}$ 获得的，$Y_t' \leftarrow \mathcal{G}(Y_t, Z_t)$，在我们的情况下 $\mathcal{G}(Y_t, Z_t) := Y_t \odot \text{SiLU}(Z_t)$。最后，层的输出 $O_t \in \mathbb{R}^D$ 通过以下下投影计算：

$$ \boldsymbol{O}_{t}^{\prime}=\operatorname{contract}(P R, P R \rightarrow P)\left(\boldsymbol{W}_{O^{\prime}}, \boldsymbol{Y}_{t}^{\prime}\right) \qquad \boldsymbol{O}_{t}=\operatorname{contract}(P D, P \rightarrow D)\left(\boldsymbol{W}_{O}, \boldsymbol{O}_{t}^{\prime}\right). $$

这个公式通过提供一个轻量级的参数化来增强现有的Mamba-3架构，将每个头内的独立SISO SSM集合转换为一组MIMO SSM。

MIMO参数匹配。Mamba-3的MIMO变体与其SISO对应物相比会产生额外的参数。因此，我们减少MLP层的隐藏维度以与SISO变体进行参数匹配，如下所示：

D 实验细节

语言建模。我们的预训练过程遵循Dao和Gu (2024)【11，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. 2024. https://arxiv.org/abs/2405.21060】的D.2节。每个规模的所有模型都遵循相同的过程，并使用bfloat16进行训练。Mamba系列模型的训练采用了标准的扩展因子2、状态维度128和头维度64。Transformer基线遵循Dao 和Gu (2024)【11，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. 2024. https://arxiv.org/abs/2405.21060】，GDN基线遵循Yang等人 (2025)【60，Gated Delta Networks: Improving Mamba2 with Delta Rule. 2025. https://arxiv.org/abs/2412.06464】，其中 $d, d_{\text{dim}} = 128, d_{\text{dim}} = 256$。所有模型都使用Llama-3.1分词器【15，The Llama 3 Herd of Models. 2024. https://arxiv.org/abs/2407.21783】。

评估。我们利用LM Evaluation Harness【14，The Language Model Evaluation Harness. Version v0.4.3. 2024. doi: 10.5281/zenodo.12608602】来测试我们预训练模型在LAMBADA (OpenAI版本)【38，The LAMBADA dataset: Word prediction requiring a broad discourse context. 2016. https://arxiv.org/abs/1606.06031】、HellaSwag【64 ，HellaSwag: Can a Machine Really Finish Your Sentence? 2019. https://arxiv.org/abs/1905.07830】、PIQA【7 ，PIQA: Reasoning about Physical Commonsense in Natural Language. 2019. https://arxiv.org/abs/1911.11641】、Arc-Easy/Arc-Challenge【10 ，Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge. 2018. https://arxiv.org/abs/1803.05457】、WinoGrande【41 ，WinoGrande: An Adversarial Winograd Schema Challenge at Scale. 2019. https://arxiv.org/abs/1907.10641】和OpenBookQA【33 ，Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering. 2018. https://arxiv.org/abs/1809.02789】上的零样本语言建模能力。

真实世界和合成检索。对于我们的真实世界检索任务，我们在一个常见的套件上进行评估，包括SWDE【4，Simple linear attention language models balance the recall-throughput tradeoff. 2025. https://arxiv.org/abs/2402.18668】、SQuAD【40 ，“Know What You Don’t Know: Unanswerable Questions for SQuAD”. 2018.】、FDA【4，Simple linear attention language models balance the recall-throughput tradeoff. 2025. https://arxiv.org/abs/2402.18668】、TriviaQA【25 ，TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension. 2017. https://arxiv.org/abs/1705.03551】、NQ【28 ，“Natural Questions: A Benchmark for Question Answering Research”. 2019.】和DROP【12，DROP: A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs. 2019. https://arxiv.org/abs/1903.00161】。我们使用了Arora等人【4 ，Simple linear attention language models balance the recall-throughput tradeoff. 2025. https://arxiv.org/abs/2402.18668；5 ，Just read twice: closing the recall gap for recurrent language models. 2024. https://arxiv.org/abs/2407.05483】提供的上述任务的完形填空格式版本，因为原始数据集是问答格式，对于仅预训练的模型来说具有挑战性。所有任务都被截断以匹配训练上下文长度。合成的NIAH任务【23 ，RULER: What’s the Real Context Size of Your Long-Context Language Models? 2024. https://arxiv.org/abs/2404.06654】也使用LM Evaluation Harness运行。

状态追踪合成任务。训练遵循一个序列长度课程，将最小长度设置为3，并将最大长度从40逐步增加到160。最终模型在256长度上进行评估。每个课程运行104步，批次大小为256。我们对Parity任务使用单层模型，对模块化算术任务使用三层模型。状态大小选择为64，我们对 $d_{\text{model}} \in \{32, 64\}$ 和8个在$10^{-4}$到$10^{-2}$之间对数间隔的学习率进行扫描，并报告最佳验证准确率。

E 额外的实验结果

表9：在1.5B规模的预训练混合Mamba-3 SISO模型上，对可选的归一化类型（分组vs默认）和位置（门控前vs门控后）进行的消融实验。所有模型都有BCNorm。平均而言，没有额外归一化在上下文内检索性能上表现最强，而门控前、分组的RMS在合成检索上表现最佳，尤其是在比其训练上下文更长的长度上。

图4：预训练的1.5B模型在不同上下文长度下，在留出的FineWeb-Edu测试集上的性能。Mamba-3表现出强大的长度外推能力，而Mamba-2在更长的上下文中表现不佳。

图5：与Mamba-2和Gated DeltaNet等强基线相比，Mamba-3表现出更好的预训练性能。这些是我们完全预训练的1.5B模型在FineWeb-Edu上的验证困惑度。

与Gated DeltaNet的状态大小比较。我们还在图6中将Mamba变体的状态大小使用效率与Gated DeltaNet基线进行了比较。我们强调直接比较GDN与Mamba风格模型存在困难，因为它们的头结构不同（GDN是多头，而Mamba是多值）。我们的实验将GDN的 $d_{\text{value}}$ 固定为2，并相应地减小头维度以改变相对总状态大小。与图3类似，我们训练440M模型到2倍Chinchilla词元数（40倍词元-参数比），并对Mamba模型的 $d_{\text{state}} = \{32, 64, 128\}$ 和GDN的 $d_{\text{head dim}} = \{32, 64, 128\}$ 进行扫描。我们对所有模型进行了参数匹配。

图6：状态大小（推理速度代理）与预训练困惑度（性能代理）的探索。Mamba-3和Mamba-3 MIMO继续设定帕累托前沿。

F 架构消融实验

实验设置。我们在本节中探讨我们的模型架构消融实验。所有模型都在440M规模下训练至Chinchilla最优词元数（20倍词元与参数比），实验过程与我们的预训练模型相同，如附录D所述，除非另有说明。

B, C偏置参数化。Mamba-3模型的独立 $B$ 和 $C$ 偏置是特定于头和通道的，并在QK-Norm之后添加到B和C中。虽然最终Mamba-3模型中的偏置是可训练的、数据无关的参数，并初始化为全一，但我们在表10a中探讨了各种偏置参数化。我们发现，只要偏置为正，我们的模型对偏置的初始化不是很敏感。我们选择全一初始化是因其简单性。

B, C偏置的影响。我们还在表10b中探讨了移除 $B$ 或 $C$ 偏置对性能的影响（使用时，偏置以我们的默认参数化进行初始化）。与Yu和Erichson (2025)【63，Block-Biased Mamba for Long-Range Sequence Processing. 2025. https://arxiv.org/abs/2505.09022】发现单独使用 $B$ 偏置能够提高Mamba-1的性能不同，我们的实验发现仅有 $B$ 偏置会轻微损害性能，而 $B$ 和 $C$ 偏置具有协同效应。

(a) $B$ 和 $C$ 偏置的参数化对模型性能的影响，通过预训练困惑度衡量。我们发现我们默认的全一初始化（第一行）提供了最佳性能，但只要偏置为正，性能就不敏感。

图片描述
(b) 同时对 $B$ 和 $C$ 应用偏置可获得最佳性能。仅应用 $B$ 偏置（Block-Biased (Yu and Erichson 2025)【63，Block-Biased Mamba for Long-Range Sequence Processing. 2025. https://arxiv.org/abs/2505.09022 】 Mamba-3变体）与无偏置基线相比没有显著增益。

表10: Mamba-3中 $B, C$ 偏置初始化（左）和存在性（右）的消融实验。

G 推理内核延迟分析

G.1 内核实现和融合结构

内核细节。在表6中，我们详细说明了我们延迟分析中使用的内核的DSL（Triton, TileLang, CuTe, PyTorch）和融合级别。对于Mamba-2和Gated DeltaNet (GDN)，我们直接使用相应作者公开发布的Triton内核。对于Mamba-3，我们实现了具有可比融合结构的新推理内核：前向SISO使用与旋转位置嵌入融合的Triton内核，前向MIMO使用具有相同融合级别的TileLang内核，而解码路径使用与门控和MIMO投影融合的CuTe内核。

符号说明。在表11和12中，我们缩写 IP = 输入投影，Conv = 1D卷积，Gate = 门控，OP = 输出投影。颜色表示实现后端（Torch, Triton, TileLang, CuTe）。

G.2 扩展的预填充和预填充+解码延迟测量

模型。我们对Mamba-3 1.5B (SISO)、Mamba-2 1.5B、Gated DeltaNet 1.5B以及一个通过vLLM引擎(v0.11.0)实现的强大Transformer基线（Llama-3.2 1B）进行了基准测试。所有循环模型都在1.5B规模下训练，具有 $d_{\text{model}} = 2048$ 和24层。对于Mamba变体，我们设置状态大小为128，头维度为64；对于GDN，我们使用QK头维度为128。

设置。预填充的序列长度在 $L \in \{512, 1024, 2048, 4096, 16384\}$ 上进行扫描，并解码相同数量的词元。对于所有序列长度，我们使用128的批次大小。为了报告vLLM在序列长度16384的数字，我们在相同序列长度下以批次大小16测量性能。然后我们将结果乘以8以近似批次大小128的性能，因为直接在此设置下测量会超出GPU内存。这提供了一个合理的估计，因为每个批次由GPU上的每个SM独立处理，所以我们预计Transformer模型的性能会随批次大小线性扩展。对于循环模型，当输入和输出张量的大小在序列长度16384时超出GPU内存时，我们利用一种状态传递方法，将序列分两半处理，同时在段之间传播隐藏状态，以避免一次性实例化整个序列。我们使用单块H100-SXM 80GB GPU，并报告三次重复的壁钟时间（秒）。

观察结果。我们观察到（i）Mamba-3增加了最小的前向传递成本，表明指数-梯形更新、复状态追踪和MIMO参数化仍然是轻量级的；（ii）解码延迟在循环模型中具有竞争力；以及（iii）循环混合器比vLLM Llama-3.2-1B随上下文长度的增长更平缓，后者由于KV-cache开销随 $L$ 增长得快得多。