Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Reading guide

Problem and motivation

Foundation models overwhelmingly rely on the Transformer architecture and self-attention. While dense attention routes information effectively across a sequence context window, it incurs two fundamental limitations: quadratic computational cost $O(L^2)$ with respect to sequence length $L$, and the requirement of storing past keys and values (the KV cache), leading to linear-time inference per step and memory bottlenecks.

Prior subquadratic approaches—such as linear attention, recurrent architectures, and structured state space models (SSMs) like S4—scale linearly or pseudo-linearly in time. However, structured SSMs have historically failed to match Transformer quality on discrete, information-dense modalities such as text. The authors identify the root cause: prior SSMs are linear time-invariant (LTI). Because their transition dynamics do not change depending on input content, they cannot selectively compress context, filter out irrelevant noise, or modulate retention across variable spacing.

Key ideas

• Selection Mechanism as Context Compression: Sequence modeling requires compressing history into a state. Prior LTI SSMs use fixed transition parameters, forcing them to treat all timesteps uniformly. By parameterizing SSM matrices based on the current input token, the model gains the content-awareness needed to filter or remember information selectively.
• Overcoming the Non-Convolving Bottleneck: Making parameters input-dependent breaks the mathematical equivalence between linear recurrences and global convolutions. Without convolutions, naive recurrence would require materializing an expanded latent state of size $(B,L,D,N)$ in GPU high-bandwidth memory (HBM). Mamba solves this by fusing discretization, parallel associative scan, and output projection into fast GPU on-chip SRAM, keeping HBM traffic minimal.
• Unified Architecture: Mamba replaces interleaved Multi-Head Attention (MHA) and MLP blocks with a single homogeneous block that combines the gated design of H3 with the SwiGLU-style gated MLP.

Method

Continuous state space models map an input $x(t)\in \mathbb{R}$ to $y(t)\in \mathbb{R}$ via an implicit latent state $h(t)\in {\mathbb{R}}^N$: $$h^{\prime }(t)=\mathbf{A}h(t)+\mathbf{B}x(t),y(t)=\mathbf{C}h(t)$$

Using zero-order hold (ZOH) discretization with step size $\Delta$, the continuous system transforms into discrete recurrence parameters: $$\overline{\mathbf{A}}=\exp (\Delta \mathbf{A}),\overline{\mathbf{B}}=(\Delta \mathbf{A}{)}^{-1}(\exp (\Delta \mathbf{A})-\mathbf{I})\cdot \Delta \mathbf{B}$$ $$h_t=\overline{\mathbf{A}}{h}_{t-1}+\overline{\mathbf{B}}{x}_t,y_t=\mathbf{C}{h}_t$$

The Selection Mechanism (S6)

Instead of keeping $(\Delta ,\mathbf{B},\mathbf{C})$ constant over time, Mamba makes them input-dependent:

• ${\mathbf{B}}_t=s_B(x_t)={\mathsf{Linear}}_N(x_t)$
• ${\mathbf{C}}_t=s_C(x_t)={\mathsf{Linear}}_N(x_t)$
• ${\Delta }_t={\tau }_{\Delta }(\mathsf{Parameter}+s_{\Delta }(x_t))$, with ${\tau }_{\Delta }=\mathsf{softplus}$ and $s_{\Delta }(x)={\mathsf{Broadcast}}_D({\mathsf{Linear}}_1(x))$

This generalizes classical RNN gating: when $N=1,\mathbf{A}=-1,\mathbf{B}=1$, the recurrence simplifies directly to a gated update $h_t=(1-g_t)h_{t-1}+g_t{x}_t$, where $g_t=\sigma (\mathsf{Linear}(x_t))$.

Hardware-Aware State Expansion

To maintain high throughput without materializing the $(B,L,D,N)$ state in HBM:

1. Parameters $(\Delta ,\mathbf{A},\mathbf{B},\mathbf{C})$ are loaded from HBM directly into on-chip SRAM.
2. Discretization into $(\overline{\mathbf{A}},\overline{\mathbf{B}})$ and a work-efficient parallel associative scan are computed entirely inside SRAM.
3. The latent state is contracted with $\mathbf{C}$ to shape $(B,L,D)$ before writing back to HBM.
4. Intermediate states are recomputed in SRAM during the backward pass instead of being saved in HBM.

Results

• Synthetic Extrapolation: On the Induction Heads task, Mamba trained on sequence length 256 extrapolates perfectly (100% accuracy) up to length $2^{20}=1,048,576$ ($4000\times$ sequence length), whereas attention baselines fail or OOM past $2^{14}$. On Selective Copying, Mamba reaches 99.8% accuracy compared to 57.0% for H3 and 18.3% for S4.
• Language Modeling: Evaluated under Chinchilla scaling up to 1.3B parameters on the Pile, Mamba matches or outperforms Transformer++ (rotary embeddings, SwiGLU, RMSNorm). On downstream zero-shot tasks (e.g., LAMBADA, PIQA, HellaSwag, WinoGrande, ARC), Mamba-2.8B attains an average accuracy of 63.3%, outperforming Pythia-2.8B (59.1%), Pythia-6.9B (61.7%), and matching or exceeding 6B–7B parameter models.
• Genomics and Audio: On Great Apes DNA classification at sequence length 1M ($2^{20}$), Mamba achieves 71.67% (1.4M params) and 81.31% (7M params) accuracy compared to 54.87% for HyenaDNA. On SC09 speech generation, Mamba reduces FID from 1.99 (SaShiMi) to 0.94 (6.1M params) and 0.67 (24.3M params).
• Efficiency: The fused scan kernel executes 20–40× faster than standard PyTorch scan implementations and runs up to 7× faster than FlashAttention-2 at 32K context. Inference throughput is up to 5× higher than comparable Transformers.

Limitations and open questions

• Continuous Signal Performance: While input-dependent selection improves discrete reasoning (text, DNA), ablations on raw audio waveform pretraining show that purely linear time-invariant (LTI) models (like standard S4) perform better on uniformly sampled, continuous data.
• Model Scale: Evaluated models are capped around 3B parameters and 300B pretraining tokens; empirical performance at 7B–70B+ scales and beyond requires further investigation.
• Downstream Ecosystem: Transformers possess extensive tooling for prompt engineering, in-context manipulation, KV-cache quantization, and adapter tuning. Understanding how SSMs adapt across these interaction modes remains an active exploration.