Papers · Mamba
AI PAPER

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Introduced a selective state space model that matches Transformer-quality language modeling with linear scaling in sequence length, offering an alternative to quadratic attention.

Authors
Albert Gu, Tri Dao
Institution
Carnegie Mellon University
Published
December NaN, 2023
Citations
2,000
Category
Architecture
Impact
notable
PAPER EXPLAINED

Plain-English Summary

Transformers process sequences with quadratic cost because every token attends to every other token. Mamba proposed an alternative: a state space model where each token updates a running hidden state, processing sequences in linear time regardless of length. The key innovation was making the state update selective — the model learns which information to remember or forget at each step, giving it capabilities comparable to attention.

Mamba matched Transformer performance on language modeling while being significantly faster for long sequences. At 2,048 tokens it was competitive. At longer sequences the efficiency advantage grows dramatically, making it attractive for applications requiring extremely long contexts.

Think of it like reading a book by maintaining a running mental summary rather than flipping back to check every earlier page. The summary gets updated selectively — important details are retained while irrelevant ones fade.

Key Innovation

Selective state spaces make the model’s dynamics input-dependent. Previous state space models used fixed dynamics (the same transition matrix regardless of input), limiting expressiveness. Mamba introduced selection mechanisms that allow the model to decide, based on the current input, what to remember and what to forget. This gives it attention-like ability to focus on relevant context.

The hardware-aware implementation by Tri Dao (of FlashAttention fame) achieves efficient GPU utilization through careful memory management, making the theoretical efficiency gains practical on real hardware.

Impact on the Field

Mamba challenged the assumption that Transformers are the only viable architecture for language modeling. While Transformers remain dominant, Mamba demonstrated that alternatives exist with genuinely different scaling properties. This is particularly relevant for applications requiring extremely long contexts, where quadratic attention becomes prohibitively expensive even with FlashAttention.

The paper sparked renewed interest in recurrent and state-space architectures, leading to hybrid models that combine the best of both approaches.

Models That Built on This

AI21’s Jamba combined Mamba layers with attention layers in a hybrid architecture, getting both efficient long-context processing and the strong in-context learning of attention. Mamba-2 improved the theoretical foundations and efficiency. While pure Mamba models have not displaced Transformers at frontier scale, the hybrid approach has gained traction. NVIDIA’s research team and several startups are exploring Mamba-based architectures for specific long-context applications.