Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5$\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.
核心贡献 · Key contributions
提出选择性状态空间模型(SSM),支持输入依赖的过滤与遗忘,克服了先前 LTI SSM 的关键弱点。 Proposes selective state space models (SSMs) that allow input-dependent filtering and forgetting, overcoming key weakness of prior LTI SSMs.
设计了面向硬件的并行递归算法,实现序列长度线性时间的高效训练与推理。 Designs a hardware-aware parallel algorithm for the recurrent mode, enabling efficient training and inference with linear scaling in sequence length.
引入 Mamba 架构,不含注意力或 MLP 块,推理吞吐量比 Transformer 高 5 倍。 Introduces Mamba, a simplified architecture without attention or MLP blocks, achieving 5x higher inference throughput than Transformers.
在语言、音频和基因组学上达到最先进性能,Mamba-3B 超越同等规模 Transformer,媲美两倍规模 Transformer。 Demonstrates state-of-the-art performance across language, audio, and genomics, with Mamba-3B outperforming Transformers of the same size and matching those twice its size.
表明选择性 SSM 能有效建模长上下文,性能在百万级序列长度上持续提升。 Shows that selective SSMs enable effective long-context modeling, with performance improving up to million-length sequences.
开源模型代码和预训练检查点,促进进一步研究。 Provides open-source model code and pre-trained checkpoints to facilitate further research.
局限 · Limitations
选择机制可能在 LTI SSM 擅长的连续信号数据(如音频、视频)上降低性能。 Selection mechanism may reduce performance on continuous signal data where LTI SSMs excel, such as audio and video.
Mamba 未在更大规模(如 7B+参数)上验证,可扩展性不确定,而 Transformer 在该规模占主导地位。 Mamba has not been validated at larger scales (e.g., 7B+ parameters) where Transformers are dominant, leaving scalability uncertain.
硬件感知算法依赖特定 GPU 内存层次(如 HBM 和 SRAM),可能限制推广到其他硬件。 The hardware-aware algorithm relies on specific GPU memory hierarchy (e.g., HBM and SRAM), potentially limiting generalizability to other hardware.
Mamba 缺乏基于 Transformer 模型的丰富生态系统,如微调、提示和指令调优,可能限制实际采用。 Mamba lacks the rich ecosystem of Transformer-based models, such as fine-tuning, prompting, and instruction tuning, which may limit practical adoption.
输入依赖的选择机制可能引入训练不稳定,或需要比 LTI 模型更仔细的超参数调优。 The input-dependent selection mechanism may introduce training instability or require more careful hyperparameter tuning compared to LTI models.
论文章节 · Sections(共 33)
Mamba: 具有选择性状态空间的线性时间序列建模Mamba: Linear-Time Sequence Modeling with Selective State Spaces
1 引言1 Introduction
2 状态空间模型2 State Space Models
3 选择性状态空间模型3 Selective State Space Models
3.1 动机:作为一种压缩手段的选择3.1 Motivation: Selection as a Means of Compression
3.2 使用选择机制改进 SSM3.2 Improving SSMs with Selection
3.3 选择性 SSM 的高效实现3.3 Efficient Implementation of Selective SSMs
3.3.1 先前模型的动机3.3.1 Motivation of Prior Models
3.3.2 选择性扫描概述:硬件感知状态扩展3.3.2 Overview of Selective Scan: Hardware-Aware State Expansion
3.4 简化的 SSM 架构3.4 A Simplified SSM Architecture
3.5 选择机制的特性3.5 Properties of Selection Mechanisms
3.5.1 与门控机制的联系3.5.1 Connection to Gating Mechanisms
3.5.2 选择机制的解释3.5.2 Interpretation of Selection Mechanisms