Mamba: Linear-Time Sequence Modeling with Selective State Spaces

By 2023, virtually every foundation model — GPT, Claude, Gemini, BERT — was built on the Transformer. Its self-attention mechanism connects every token to every other token in a single step, which is why it models language so well.

But this power comes at a steep price:

• Quadratic compute. Self-attention compares every pair of tokens: doubling the sequence length quadruples the work. A 100K-token document costs 10,000× more than a 1K-token paragraph.

• Linear KV cache. During inference, the model stores a Key and Value vector for every past token. As context grows, so does this cache, consuming GPU memory and slowing down generation.

Many alternatives tried to fix this — linear attention, global convolutions, structured state space models like S4 — but all shared the same blind spot: they treated every input identically, regardless of whether it was a crucial fact or meaningless noise.

Before understanding what Mamba adds, we need the foundation it builds on: structured state space models (S4).

Think of a state space model as a pipeline with hidden memory. An input signal $x(t)$ enters, updates a hidden state $h(t)$, and produces an output $y(t)$. The hidden state is like the water level in a tank — it accumulates input over time and the output reads from it.

Four matrices control this pipeline: $\boldsymbol{A}$ governs how the state evolves on its own (the leak rate of the tank), $\boldsymbol{B}$ controls how new input flows in, $\boldsymbol{C}$ controls what we read out, and $\Delta$ sets the step size (discretization). Crucially, in prior SSMs all four are fixed — the same for every input. This is called Linear Time Invariance (LTI).

Gu and Dao identified a fundamental tension: sequence modeling is about compressing context into a finite state. Attention avoids compression entirely by storing every token (the KV cache) — effective but expensive. Recurrent models compress into a fixed-size state — efficient but lossy.

The key question is: what gets compressed? In LTI models, the compression is content-blind: the same matrix $\boldsymbol{A}$ governs every step, so the model cannot choose to remember "Harry" more than "um". This is why LTI models fail on tasks requiring content-based reasoning, like the Selective Copying task: copying specific tokens (not all tokens) from scattered positions.

The solution is selection: make the SSM parameters $\Delta$, $\boldsymbol{B}$, and $\boldsymbol{C}$ depend on the current input. Now the model can:

• Remember a relevant token by opening the gate ($\Delta$ large → state resets and absorbs the new input)
• Ignore noise by closing the gate ($\Delta$ small → state persists, input slides off)
• Control what enters the state ($\boldsymbol{B}$ input-dependent) and what exits ($\boldsymbol{C}$ input-dependent)

The change from S4 to Mamba's S6 is remarkably simple. Instead of parameters being fixed constants, three of them become functions of the input:

• $\boldsymbol{B}_t = \text{Linear}_N(x_t)$ — the input projection is now input-dependent
• $\boldsymbol{C}_t = \text{Linear}_N(x_t)$ — the output projection is now input-dependent
• $\Delta_t = \text{softplus}(\text{Linear}(x_t))$ — the step size is now input-dependent

The matrix $\boldsymbol{A}$ stays fixed (it only matters through its interaction with $\Delta$ via discretization, so selectivity in $\Delta$ already provides selectivity in $\boldsymbol{A}$).

This is the entire innovation at the mathematical level. But it has a profound consequence: the model is no longer time-invariant, which means the convolution shortcut no longer works.

The selection mechanism is not as alien as it might seem — it has deep roots in classical RNN gating. The paper proves (Theorem 1) that when the state dimension is 1, the selective SSM reduces exactly to a gated RNN:

$g_t = \sigma(\text{Linear}(x_t))$, then $h_t = (1 - g_t)\, h_{t-1} + g_t\, x_t$

This is the familiar LSTM-style gate: $g_t$ near 1 means "forget the old state, absorb the new input"; $g_t$ near 0 means "keep the old state, ignore the current input."

Mamba generalizes this in two ways: the state dimension $N$ is much larger (typically 16), giving the model a richer hidden representation; and the gating comes from principled discretization of a continuous system rather than heuristic design.

Making SSM parameters input-dependent breaks the convolution shortcut. The naive recurrence is sequential and requires materializing a state tensor of shape $(B, L, D, N)$ — far larger than the input. On a GPU, this means constant trips between slow global memory (HBM) and fast local memory (SRAM), which is the real bottleneck for non-matmul operations.

Mamba solves this with three classical techniques fused together:

• Kernel fusion. Load the small parameters $(Δ, A, B, C)$ from HBM to SRAM once, perform discretization and the full scan and the output multiplication all in SRAM, then write only the final output back to HBM. This reduces memory reads by a factor of $N$ (the state dimension).

• Parallel scan. The recurrence $h_t = \bar A_t h_{t-1} + \bar B_t x_t$ is associative, so it can be parallelized with a work-efficient prefix sum (Blelloch scan), avoiding the sequential bottleneck.

• Recomputation. During the backward pass, instead of storing the full intermediate states, recompute them on the fly from the small inputs already in SRAM. This keeps memory usage equal to FlashAttention.

Most SSM architectures (H3, Hyena) alternate between an SSM block and a standard MLP block, similar to how Transformers alternate attention and MLP. Mamba simplifies this by merging both into a single homogeneous block, inspired by the Gated Attention Unit.

Each Mamba block has two parallel branches:

• Branch 1 (main): Linear projection → 1D convolution → SiLU activation → Selective SSM
• Branch 2 (gate): Linear projection → SiLU activation

The outputs of both branches are multiplied element-wise, then projected back to the model dimension. The 1D convolution is a short local convolution (kernel size 4) that provides local context before the SSM processes global context.

Two Mamba blocks have the same parameter count as one Transformer attention+MLP pair (about $12D^2$ parameters), making them directly comparable.

Mamba was evaluated across three modalities — language, DNA, and audio — and set new records for sub-quadratic models on all of them:

• Language: Mamba-3B outperformed Transformers of the same size (Pythia-3B) by 4 points on common-sense reasoning and even exceeded Pythia-7B — a model more than twice its size. Against strong baselines using the LLaMA recipe (Transformer++), Mamba matched perplexity at all scales from 125M to 1.3B parameters.

• DNA: Mamba was the only model whose perplexity improved continuously with longer context, all the way up to million-length sequences. On a classification task distinguishing five great ape species (which share 99% of their DNA), Mamba significantly outperformed HyenaDNA.

• Audio: On speech generation (SC09), a small 6M-parameter Mamba outperformed much larger GAN and diffusion models. A parameter-matched 24M model cut the FID score to 0.67 — the best autoregressive result to date.

• Speed: Mamba achieves 5× higher inference throughput than Transformers of similar size, because it has no KV cache — just a fixed-size hidden state. This means a Mamba-6.9B at inference is faster than a Transformer-1.3B.

Mamba showed that making recurrence content-aware — a conceptually simple change — is enough to close the gap between linear-time models and Transformers. The principle of selection extends beyond SSMs: any recurrent architecture benefits from input-dependent dynamics. This is now an active frontier of research, with Mamba-3, state-space-attention hybrids, and vision/audio Mamba variants building on this foundation.