Mamba & State Space Models
The Motivation
Transformers are powerful because self-attention lets every token look at every other token. But that exactness has a price. For a sequence of length n, attention compares all pairs of positions, so its compute and memory scale as O(n²) in the sequence length. At generation time, a Transformer also keeps a KV cache — the keys and values of every previous token — which grows linearly with the context. For very long sequences (long documents, genomes, audio), both the quadratic attention and the ever-growing cache become expensive.
State Space Models (SSMs) take a different route, borrowed from classical control theory. Instead of re-attending to the whole past, they carry a fixed-size hidden state forward through the sequence with a simple linear recurrence — like an RNN. The state is a running summary of everything seen so far. This gives O(n) sequential inference with constant memory per step, no cache that grows with context.
The tension: attention buys exact all-pairs recall at quadratic cost; an SSM buys linear cost and constant memory by compressing the past into one fixed-size state. The whole story of Mamba is about getting that compression to be smart.
The Core Recurrence
An SSM maps an input sequence to an output sequence through a hidden state h. In its discrete form it is just two lines:
- xₜ — the input at step t (a vector, e.g. a token embedding).
- hₜ — the hidden state, a fixed-size vector that carries information forward. Its size does not depend on sequence length.
- A — the state-transition matrix. It decides how much of the old state survives into the next step — the memory / forgetting behavior.
- B — how the new input is written into the state.
- C — how the state is read out to produce the output yₜ.
Read the recurrence intuitively: keep a fraction of what you knew (A·hₜ₋₁), add what just arrived (B·xₜ), then report a view of the state (C·hₜ). The eigenvalues of A control how fast old information fades.
The clever part: parallel training via convolution
A plain recurrence is sequential — step t needs step t−1 — which is slow to train. But if A, B, C are fixed (the same at every step), the recurrence is linear and time-invariant (LTI), and you can unroll it into a single convolution:
The kernel K̄ = (C·B, C·A·B, C·A²·B, …) can be precomputed and applied with a fast parallel convolution at training time, while the same model runs as a cheap step-by-step recurrence at inference. Structured SSMs like S4 made this practical by giving A a special structure so the long kernel can be computed efficiently and stably.
So a structured SSM is trainable in parallel like a CNN, yet runs as an O(n) constant-memory recurrence like an RNN — the best of both at inference, with a fixed-size state.
Interactive: a State-Space Recurrence
Here is a faithful scalar SSM, hₜ = a·hₜ₋₁ + b·xₜ, yₜ = c·hₜ. Feed it an input sequence and watch the hidden state and output evolve. The key thing to internalize: a alone controls how long the model remembers, and the state stays a single number no matter how long the input — that is the constant-memory recurrence at the heart of every SSM.
What to watch: the hidden state hₜ is a single number that summarizes the entire past. When an impulse arrives, h jumps by b·xₜ, then decays geometrically by a factor of a every step.
Push a near 1 → the state barely fades, so the model remembers far back. Drop a toward 0 → it forgets almost immediately. The state size never grows: one number, no matter how long the sequence — that is the O(n) / constant-memory recurrence.
Memory length τ ≈ 2.8 steps (time to decay to 1/e). a→1 = long memory, a→0 = forgets fast.
When on, the decay a becomes a function of the input: aₜ = a·(1 − σ(s·|xₜ| − 1)). A large input drives the gate up → smaller aₜ → the state resets and writes the new content. A quiet input keeps aₜ ≈ a → the state holds its memory. This is content-based gating, which a fixed (LTI) SSM cannot do.
Try this
- Set a = 0.95 with “Single impulse” — the state lingers.
- Set a = 0.3 — it vanishes in a couple of steps.
- Pick “Burst then big spike”, turn on Selective: the big spike resets the state while the small burst does not.
The impulse response is y = c·b·aᵏ — a clean geometric decay. With a near 1 the response is a long tail (long memory); with small a it is gone almost instantly. In a real S4/Mamba layer the state is a vector and there are many such channels, but every channel obeys exactly this rule.
Mamba's Innovation: Selectivity
Classic structured SSMs (including S4) are linear time-invariant: A, B, C and the discretization timestep Δ are the same for every token. That is exactly what makes the convolution trick work — but it is also a weakness. An LTI system processes every input the same way; it cannot decide, based on content, that this particular token is worth remembering and that one should be ignored. It has no way to do the selective copying or content-based reasoning that attention does naturally.
Mamba (S6) fixes this by making the SSM selective: the parameters B, C, and the timestep Δ become functions of the input (Δ in turn shapes the effective A). The model now chooses, per token, how much to write into the state, how much of the past to keep, and what to read out — content-dependent gating that earlier linear SSMs lacked.
Why this breaks the convolution
Once B, C, Δ depend on the input, the system is time-varying — there is no single fixed kernel, so the elegant LTI convolution no longer applies. Mamba recovers parallel training with a hardware-aware parallel scan: a linear recurrence can still be computed in parallel (associative scan), and Mamba implements it with a kernel that keeps the large state in fast GPU SRAM and avoids materializing it in slow memory. The result is linear scaling in sequence length with a recurrence that is now input-aware.
Important: Mamba is not attention-based. There is no QKᵀ score matrix and no all-pairs comparison. It is a selective (time-varying) recurrence — an evolution of RNN/SSM ideas, not of attention. Try the selective toggle in the demo above: the effective decay aₜ now reacts to each input, so a big spike can reset the state while small inputs leave the memory untouched.
Mamba vs Transformers
| Property | Transformer (self-attention) | Mamba (selective SSM) |
|---|---|---|
| Training compute vs length | O(n²) | O(n) (parallel scan) |
| Per-token inference cost | O(n) — must attend over all prior tokens | O(1) — one fixed-size state update |
| Memory during generation | KV cache grows with context | Constant — fixed-size hidden state |
| Access to the past | Exact, all-pairs recall of every token | Compressed summary in a fixed state |
| Mechanism | Attention (content-based scores) | Selective recurrence (not attention) |
The honest tradeoff: a Transformer can recall any earlier token exactly, which is invaluable for tasks like precise retrieval, but it pays quadratic compute and a growing cache. Mamba runs in linear time with constant memory per step, but it must fit the relevant past into a fixed-size state — selectivity is what lets it decide which past to keep. Hybrid models that interleave attention and SSM blocks aim to get both strengths.
From S4 to Mamba: the lineage
Classical / linear SSMs (LTI)
Fixed A, B, C. Trainable as a convolution, cheap recurrent inference, but content-blind — they treat every token identically.
S4 (structured state spaces)
Gives A a special (diagonal-plus-low-rank / HiPPO-inspired) structure so the long convolution kernel can be computed efficiently and the state captures long-range dependencies. Still LTI.
Mamba / S6 (selective SSM)
Makes B, C, and Δ input-dependent (time-varying). Loses the simple convolution but regains parallelism with a hardware-aware scan. Adds content-based memory control — the missing ingredient.
Common pitfalls & clarifications
- “Mamba is a kind of attention.” No — it has no attention matrix. It is a selective recurrence.
- “SSMs are just RNNs.” They share the fixed-state recurrence, but structured SSMs are linear in the state, which is what enables parallel convolution/scan training and stable long-range memory — things vanilla RNNs struggle with.
- “Selectivity is free.” Making parameters input-dependent breaks the LTI convolution; the parallel scan is what keeps it efficient.
- Constant memory ≠ infinite memory. A fixed-size state must compress the past; very precise long-range recall can still favor attention.
Key Takeaways
- Self-attention gives exact all-pairs recall but costs O(n²) compute and a KV cache that grows with context — expensive for very long sequences.
- An SSM processes a sequence with a linear recurrence hₜ = A·hₜ₋₁ + B·xₜ, yₜ = C·hₜ, giving O(n) inference with a fixed-size hidden state — like an RNN, but trainable in parallel.
- Structured SSMs (S4) are linear time-invariant, so the recurrence unrolls into a convolution that can be trained in parallel and run as a cheap recurrence at inference.
- Mamba's key idea is selectivity: making B, C, and Δ input-dependent so the model chooses what to remember or forget per token — content-based control that LTI SSMs lacked.
- Selectivity makes the SSM time-varying (no fixed kernel); a hardware-aware parallel scan restores efficient parallel training.
- Mamba is not attention-based. It trades attention's exact recall for linear scaling and constant memory per step.