Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Transformers dominate sequence modeling, but attention compares every token with every other one. Training cost grows with the square of the sequence length, and generation needs a KV cache that grows with every token. Many linear-time alternatives exist, and structured state space models (SSMs) such as S4 were the strongest of them: a hidden state is updated at each step as h = Ā·h + B̄·x, and output as y = C·h. They excelled at long continuous signals like audio but lagged on text.

Albert Gu at CMU and Tri Dao at Princeton pinned the weakness on one property. Earlier SSMs are linear time-invariant: the same Δ, A, B and C apply at every position. That lets them run as one long convolution, which is fast, but it also means they cannot decide, based on content, to keep one token and ignore the next. Two synthetic tasks make the point: selective copying (remember the colored tokens scattered among noise) and induction heads (after "Harry Potter ... Harry", predict "Potter").

Mamba's fix, the selection mechanism, makes Δ, B and C functions of the current input. A large Δ lets the input in and largely resets the state; a tiny Δ lets the state pass through unchanged. With N=1 this reduces exactly to a gated RNN. The cost is losing the convolution form, so the authors wrote a hardware-aware scan: load the parameters into fast on-chip SRAM, discretize and run the recurrence there with a parallel scan, write only outputs back to GPU memory, and recompute intermediate states in the backward pass. The Mamba block merges the SSM layer with a gated MLP-like structure and is stacked with no attention at all.

Mamba solved both synthetic tasks, extrapolating induction heads from 256 to a million tokens. On language, it matched a modern Transformer recipe in scaling studies up to 1.3B parameters and beat same-size open models at 2.8B. It also set best results on DNA and audio modeling.

You will meet Mamba-style layers inside production models rather than as standalone Mamba checkpoints: several current open models mix linear-time layers with a few attention layers. Knowing what they buy and give up helps when sizing serving or choosing a model.

No KV cache means bigger batches. A Transformer's memory per sequence grows with context length; an SSM's state is fixed. For long contexts and high concurrency, that is the difference between a handful of sequences per GPU and thousands, and it is where the paper's throughput gain comes from.

Long inputs get cheaper. Linear-time training and constant-time generation steps matter for long logs, genomic data, audio and long documents.

Exact recall is the weak spot. A fixed-size state must compress the past. Later work found Transformers better than state space models at copying long strings from context. If the task is "find the exact clause in this 200-page contract", check recall before switching.

Tooling maturity varies. Serving engines, quantization and fine-tuning recipes are most mature for Transformers; the authors list these "downstream affordances" as an open question.