Mamba: How Selective State Spaces Achieve 5x Faster Inference Than Transformers

Every Transformer-based model pays a cost that grows with the square of the sequence length. For a 4,096-token context the attention mechanism computes roughly 16 million pairwise score terms. At 128,000 tokens the same computation requires approximately 16 billion terms, along with a KV cache that grows linearly with context length and saturates GPU memory. Researchers explored linear attention, gated convolutions, and recurrent networks as lower-complexity alternatives for years, but none matched Transformer quality on language tasks well enough to displace attention in production models. Mamba, published December 1, 2023 (arXiv:2312.00752) by Albert Gu of Carnegie Mellon University and Tri Dao of Princeton, changed that. By introducing a selective state space mechanism with input-dependent parameters and a hardware-aware parallel training algorithm, Mamba achieved linear time inference while outperforming same-size Transformers on language modeling benchmarks. Within 18 months of publication, the architecture had appeared in Jamba from AI21 Labs, Falcon Mamba 7B from Technology Innovation Institute, and NVIDIA's Nemotron-H family.

The paper and its two authors

The Mamba paper, titled 'Mamba: Linear-Time Sequence Modeling with Selective State Spaces,' was submitted to arXiv on December 1, 2023, with a revised version posted on May 31, 2024. The paper has exactly two authors. Albert Gu is the researcher behind the S4 (Structured State Space Sequence) model and the HiPPO (High-order Polynomial Projection Operators) framework, both foundational to the SSM lineage. Tri Dao is best known for Flash Attention (arXiv:2205.14135), the hardware-aware attention kernel now used in virtually every major training stack. The combination brought Gu's expertise in structured state space design and Dao's record of GPU kernel optimization into a single architecture. The code and pretrained weights were released at github.com/state-spaces/mamba at the time of publication.

State space models before Mamba: HiPPO and S4

State space models have roots in classical control theory, where a hidden state vector evolves according to a linear differential equation driven by an input signal. Gu's HiPPO paper (NeurIPS 2020, arXiv:2008.07669) introduced a framework for compressing long sequences into polynomial projections of history using structured matrices. S4 (ICLR 2022, arXiv:2111.00396) built on HiPPO, parameterizing the state transition matrix as a diagonal plus low-rank (DPLR) structure and training with a frequency-domain convolution. S4 matched Transformers on several Long Range Arena tasks with sequences up to 16,000 steps while scaling linearly with length, but it underperformed Transformers on language modeling. The reason was structural: SSMs before Mamba used fixed, input-independent transition matrices. Every token was processed the same way regardless of content, so the model had no mechanism to decide which tokens deserved attention and which did not.

The selective mechanism: input-dependent SSM parameters

Mamba's central innovation is making the SSM parameters functions of the current input token. In a standard SSM, the matrices A (state transition), B (input projection), and C (output projection) are fixed at training time and applied identically to every position. In Mamba, the B and C matrices and a discretization step size delta are computed from the current input using learned linear projections. This means the model can choose, for each token, how strongly to update the hidden state with new information and how much of the existing state to retain. A token carrying important semantic content can gate the state aggressively; an irrelevant token can pass through with minimal state update. The paper describes this property as selective propagation and forgetting, and identifies it as the missing capability that prevented earlier SSMs from competing with Transformers on content-based language tasks.

The selectivity change removes the ability to use efficient frequency-domain convolutions, which S4 depended on during training. A selective SSM cannot be precomputed as a fixed convolution kernel because the kernel changes with every input. Mamba resolves this with a hardware-aware parallel algorithm based on the parallel prefix scan, a classical parallel computing primitive. By operating primarily in SRAM rather than high-bandwidth memory and fusing the scan with surrounding operations into a single GPU kernel, Mamba avoids the memory transfer bottleneck that otherwise makes recurrent computation on GPUs slow. The result is a model that trains with parallel scans and runs inference as a true recurrent model, requiring no growing KV cache and consuming constant memory per generation step regardless of sequence length.

• B, C, and delta are computed per token from the input, making the state update content-dependent.
• Fixed SSMs (S4, S4D, Hyena) treat all tokens equally; selective SSMs gate information based on token content.
• Training uses parallel prefix scans (O(L log L) work, highly parallelizable on GPU).
• Inference is purely recurrent: O(1) memory per step, no KV cache, constant throughput regardless of context length.
• No attention layers and no MLP blocks appear in the base Mamba architecture; the SSM layer handles both mixing and gating.

What the numbers show: throughput and quality

The paper evaluates Mamba at four scales: 130 million, 370 million, 790 million, and 2.8 billion parameters. On The Pile language modeling benchmark, Mamba-1.4B matches Pythia Transformer baselines at the same scale, and Mamba-2.8B matches Pythia-6.9B, a model more than twice its size. The abstract states that Mamba-3B outperforms Transformers of the same size and matches Transformers twice its size in both pretraining perplexity and downstream evaluation. On inference throughput measured on an NVIDIA A100 GPU at batch size 1, Mamba reports 5 times higher throughput than a same-size Transformer. Because Mamba's per-step recurrent cost is constant, while Transformer inference cost grows with KV cache size, the throughput advantage widens as sequences grow longer. The paper also evaluates Mamba on audio (SC09 speech modeling) and genomics (DNA sequence prediction), where it outperforms prior SSMs and reaches parity with Transformers on those modalities.

Adoption: Jamba, Falcon Mamba, Nemotron-H, and the hybrid era

Three months after the Mamba paper, AI21 Labs released Jamba (arXiv:2403.19887, March 2024), the first large-scale hybrid Transformer-Mamba-MoE model. Jamba interleaved attention and Mamba layers at a 1:7 ratio (one attention layer per seven Mamba layers), with mixture-of-experts blocks added every two layers. The design supported context lengths up to 256,000 tokens while maintaining a smaller KV cache than a pure Transformer at that length, since only the attention layers cache keys and values. AI21 Labs scaled the design in August 2024 with Jamba 1.5 (arXiv:2408.12570): 398 billion total parameters with 94 billion active, 72 layers interleaving Mamba and attention, 16 MoE experts, and 256K-token context support.

Technology Innovation Institute released Falcon Mamba 7B in 2024, a pure Mamba model with 7.27 billion parameters trained on 5.8 trillion tokens and no attention layers whatsoever. Falcon Mamba 7B outperformed LLaMA 3.1-8B, Mistral 7B, and Falcon2-11B on MMLU, GSM8K, and ARC benchmarks while maintaining constant memory and throughput at any sequence length. NVIDIA's Nemotron-H family (arXiv:2504.03624, April 2025) replaced 92% of attention layers with Mamba2 blocks across 8B, 47B, and 56B models, achieving up to 3 times faster throughput than LLaMA-3.1 and Qwen-2.5 at the same parameter count. IBM's Bamba-9B, built on Mamba2, delivered 2 times throughput over comparable Transformer models while matching LLaMA-3.1-8B accuracy on 7 times less training data. PromptingIndex covers Mamba alongside Flash Attention, grouped query attention, and speculative decoding. Mamba occupies a distinct position among them: rather than optimizing the Transformer attention mechanism, it replaces attention entirely, and its adoption across five major labs in under two years confirms that selective state spaces have earned a permanent role in the model-builder's toolkit.