RESEARCH PAPER / 2023

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

Carnegie Mellon University +2 · landmark

Affiliations: Carnegie Mellon University, Princeton University, Stanford University

Abstract

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.

The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.

Research summary

Mamba introduces input-dependent selective state space models with a hardware-aware parallel scan algorithm, achieving linear-time scaling and Transformer-matching performance across multiple modalities. By fundamentally rethinking sequence modeling through dynamic information routing and architectural simplification, the paper provides a highly efficient, scalable alternative to attention that has rapidly reshaped the landscape of foundation model design and spurred widespread adoption across academia and industry.

Methodology

The paper introduces Selective State Space Models (SSMs), a principled evolution of structured SSMs that addresses the critical limitation of time-invariance in prior sequence models. By parameterizing the discretization step ($\Delta$) and projection matrices ($B$, $C$) as functions of the input token, the model gains dynamic, content-dependent routing capabilities, effectively simulating attention-like selectivity without quadratic complexity. The authors further design a hardware-aware parallel scan algorithm that maintains $O(N)$ training complexity while maximizing GPU utilization, alongside a recurrent inference mode that enables fast autoregressive generation. The architectural choice to entirely remove attention and MLP blocks is bold and theoretically elegant, forcing the SSM to handle both long-range dependency modeling and non-linear transformations. The mathematical formulation is rigorous, clearly bridging continuous-time dynamical systems with discrete token processing, and the algorithmic contributions are well-grounded in modern GPU architecture constraints.

Results & evaluation

The empirical evaluation is extensive and spans multiple modalities, including language (Pile, SlimPajama), audio (Speech Commands, LibriSpeech), and genomics (Nucleotide Transformer). The Mamba-3B model demonstrates strong scaling behavior, outperforming same-sized Transformers and matching 2x-sized counterparts on both pretraining loss and downstream benchmarks. Throughput and memory benchmarks convincingly show 5x faster inference and linear scaling up to million-token sequences. The evaluation covers autoregressive generation, zero-shot transfer, and fine-tuning, providing a holistic performance profile. However, comparisons against highly optimized modern Transformer variants (e.g., FlashAttention-2, GQA, sliding-window attention) could be more thorough, and some downstream tasks show marginal rather than decisive gains, suggesting that architectural efficiency does not universally translate to superior representation quality across all tasks.

Reproducibility

Reproducibility is strong. The paper provides explicit mathematical derivations, clear algorithmic pseudocode for the selective scan, and detailed architectural hyperparameters. The open-source release includes optimized CUDA kernels for the parallel scan, pre-trained weights, and training scripts, which have enabled rapid community verification and adoption. While reproducing exact training throughput requires specific GPU hardware and kernel compilation, the core methodology and baseline implementations are fully accessible and well-documented.

Limitations

The selective mechanism inherently introduces sequential dependencies during inference, which, despite being highly optimized, still limits the massive parallelism achievable by bidirectional or fully parallelizable architectures. Removing MLP blocks may constrain the model's capacity for highly localized, non-linear feature transformations, potentially necessitating increased depth to compensate. The paper primarily focuses on autoregressive and unidirectional tasks; bidirectional encoder-only applications (e.g., masked language modeling) are underexplored. Additionally, while long-context scaling is linear, empirical retrieval capabilities on extreme-length "needle-in-a-haystack" tasks still lag behind specialized attention-based retrieval mechanisms.

Broader impact

Mamba fundamentally challenges the Transformer's dominance in sequence modeling by demonstrating that linear-time, state-space-based architectures can achieve competitive performance across diverse modalities. Its computational efficiency lowers the barrier to training and deploying long-context models, with significant implications for resource-constrained environments, edge deployment, and sustainable AI. The work has already catalyzed a new research direction in hybrid architectures and efficient sequence modeling. As with any foundational architecture, widespread adoption raises standard concerns regarding accelerated deployment of highly capable models, alignment challenges, and the environmental footprint of large-scale pretraining, though the efficiency gains partially mitigate computational costs. Mamba introduces input-dependent selective state space models with a hardware-aware parallel scan algorithm, achieving linear-time scaling and Transformer-matching performance across multiple modalities. By fundamentally rethinking sequence modeling through dynamic information routing and architectural simplification, the paper provides a highly efficient, scalable alternative to attention that has rapidly reshaped the landscape of foundation model design and spurred widespread adoption across academia and industry.

Curator note

Gu & Dao; SSM alternative to Transformer; linear scaling in sequence length

Related later work

Author of this paper? Send a correction.