Mamba-3: Improved Sequence Modeling using State Space Principles

📄 arXiv: 2603.15569v1 📥 PDF

作者: Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang, Aviv Bick, J. Zico Kolter, Tri Dao, Albert Gu

分类: cs.LG

发布日期: 2026-03-16

备注: ICLR 2026


💡 一句话要点

Mamba-3:利用状态空间模型原理改进序列建模,提升推理效率与模型质量。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 序列建模 状态空间模型 推理效率 复值状态 MIMO结构 语言模型 亚二次模型

📋 核心要点

  1. 现有Transformer模型推理成本高昂,亚二次模型虽降低计算复杂度,但常牺牲模型质量,且硬件效率仍有提升空间。
  2. Mamba-3通过改进状态空间模型(SSM)的离散化、引入复值状态更新和MIMO结构,提升模型表达能力和推理效率。
  3. 实验表明,Mamba-3在检索、状态跟踪和语言建模任务中显著提升,在相同困惑度下,状态大小仅为Mamba-2的一半。

📝 摘要(中文)

大规模语言模型(LLM)的性能提升日益依赖于推理时的计算量,因此模型设计除了关注模型质量外,推理效率也至关重要。虽然当前基于Transformer的模型具有强大的模型质量,但其二次计算复杂度和线性内存需求使得推理成本高昂。这促使了亚二次模型的开发,这些模型降低了线性计算复杂度并实现了恒定的内存需求。然而,许多最新的线性模型为了算法效率而牺牲了模型质量和能力,在状态跟踪等任务上表现不佳。此外,它们在理论上的线性推理在实践中仍然缺乏硬件效率。受状态空间模型(SSM)观点的启发,我们从推理优先的角度出发,引入了三个核心方法改进:(1) 从SSM离散化导出的更具表现力的递归;(2) 启用更丰富的状态跟踪的复值状态更新规则;(3) 用于在不增加解码延迟的情况下获得更好模型性能的多输入多输出(MIMO)公式。结合架构改进,我们的Mamba-3模型在检索、状态跟踪和下游语言建模任务中取得了显著的提升。在15亿参数规模下,Mamba-3相比于次优模型(Gated DeltaNet)平均下游精度提高了0.6个百分点,而Mamba-3的MIMO变体进一步将精度提高了1.2个百分点,总计提升了1.8个百分点。在不同状态大小的实验中,Mamba-3在仅使用其前身一半的状态大小的情况下,实现了与Mamba-2相当的困惑度。我们的评估表明,Mamba-3能够提升性能-效率的帕累托前沿。

🔬 方法详解

问题定义:现有Transformer模型在推理时计算复杂度高,内存需求大,导致推理成本高昂。虽然一些亚二次模型试图解决这个问题,但往往以牺牲模型质量为代价,并且在实际硬件上的推理效率仍然不足。论文旨在设计一种既能保持模型质量,又能提高推理效率的序列建模方法。

核心思路:论文的核心思路是借鉴状态空间模型(SSM)的优点,并对其进行改进,以实现高效且高质量的序列建模。通过更具表达力的递归、复值状态更新规则和MIMO结构,提升模型在状态跟踪等任务上的能力,同时保持较低的计算复杂度。

技术框架:Mamba-3的技术框架基于状态空间模型(SSM),主要包含以下几个关键模块:1) SSM离散化:采用改进的离散化方法,导出更具表达力的递归公式。2) 复值状态更新:引入复数域的状态更新规则,增强模型的状态跟踪能力。3) MIMO结构:采用多输入多输出(MIMO)结构,在不增加解码延迟的情况下,提升模型性能。此外,还包括一些架构上的改进,以进一步优化模型性能。

关键创新:Mamba-3的关键创新在于以下三个方面:1) 更具表达力的递归公式:通过改进SSM的离散化方法,获得了更强大的序列建模能力。2) 复值状态更新规则:利用复数域的特性,增强了模型对状态的跟踪和记忆能力。3) MIMO结构:通过多输入多输出的设计,提高了模型的并行性和效率,同时提升了模型性能。

关键设计:Mamba-3的关键设计包括:1) SSM离散化的具体公式和参数设置;2) 复值状态更新规则的实现细节,例如复数的初始化和更新方式;3) MIMO结构的具体实现,包括输入和输出的维度、连接方式等;4) 架构上的其他改进,例如激活函数、归一化方法等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Mamba-3在1.5B参数规模下,相比Gated DeltaNet,平均下游精度提高了0.6个百分点,MIMO变体进一步提升了1.2个百分点,总计提升1.8个百分点。在相同困惑度下,Mamba-3仅使用Mamba-2一半的状态大小,显著提升了性能-效率。

🎯 应用场景

Mamba-3具有广泛的应用前景,包括自然语言处理、语音识别、计算机视觉等领域。其高效的推理能力使其特别适用于资源受限的场景,例如移动设备、嵌入式系统等。未来,Mamba-3有望成为新一代高效序列建模的基础模型,推动相关领域的发展。

📄 摘要(原文)

Scaling inference-time compute has emerged as an important driver of LLM performance, making inference efficiency a central focus of model design alongside model quality. While the current Transformer-based models deliver strong model quality, their quadratic compute and linear memory make inference expensive. This has spurred the development of sub-quadratic models with reduced linear compute and constant memory requirements. However, many recent linear models trade off model quality and capability for algorithmic efficiency, failing on tasks such as state tracking. Moreover, their theoretically linear inference remains hardware-inefficient in practice. Guided by an inference-first perspective, we introduce three core methodological improvements inspired by the state space model (SSM) viewpoint of linear models. We combine: (1) a more expressive recurrence derived from SSM discretization, (2) a complex-valued state update rule that enables richer state tracking, and (3) a multi-input, multi-output (MIMO) formulation for better model performance without increasing decode latency. Together with architectural refinements, our Mamba-3 model achieves significant gains across retrieval, state-tracking, and downstream language modeling tasks. At the 1.5B scale, Mamba-3 improves average downstream accuracy by 0.6 percentage points compared to the next best model (Gated DeltaNet), with Mamba-3's MIMO variant further improving accuracy by another 1.2 points for a total 1.8 point gain. Across state-size experiments, Mamba-3 achieves comparable perplexity to Mamba-2 despite using half of its predecessor's state size. Our evaluations demonstrate Mamba-3's ability to advance the performance-efficiency Pareto frontier.