Attention is brilliant at recall because it lets every new token look back at all the tokens before it. That's also its curse: every step re-reads the whole history, so the running cost climbs like the square of the length and the memory (the KV cache) grows with it. On very long inputs you hit a wall. State-space models take the other road — a fixed-size running state updated token by token — and Mamba is the version that finally made them competitive. I built a demo that runs both models live and computes both cost counters as you advance a stream, plus the actual SSM recurrence with a selectivity gate you can flip. Here's the idea.
The quadratic wall vs the fixed state
Attention's per-token cost is the position number t (it re-reads all t earlier tokens), so the total after n tokens is n(n+1)/2 ≈ n²/2, and it must cache all n keys and values. An SSM only ever touches its fixed state of size d, so per-token cost is a constant, the total is d·n (linear), and memory is just d — forever:
// ATTENTION: re-reads the whole history every token
const attnTotal = n => n * (n + 1) / 2; // ~ n^2 / 2 -> the wall
const attnMemory = n => n; // KV cache grows with length
// SSM / MAMBA: one fixed-size state, updated in place
const ssmTotal = (n, d) => d * n; // LINEAR in length
const ssmMemory = d => d; // CONSTANT — streams forever
attnTotal(1024); // 524,800 (quadratic)
ssmTotal(1024, 16); // 16,384 (linear) -> ~32x less work
The demo lets you step or play a token stream and watch attention's red curve peel away from the SSM's straight green line. Push the length up and every doubling roughly quadruples attention's cost but only doubles the SSM's — the gap becomes a chasm, and the memory read-out shows attention's KV cache growing to n entries while the SSM holds steady at d.
The recurrence is one line
The heart of an SSM is a running, decaying summary. Keep a fraction a of the old state, add b times the new input x, read the output out with c. No history is stored — everything the model "knows" is compressed into one number h:
function ssmStep(h, x, a, b){ return a * h + b * x; } // h_t = a*h + b*x
// 'a' near 1 = long memory, near 0 = forget fast
// 'b' = how strongly the new input enters the state
Streaming is just a loop: for each token, update the single state and emit an output. There's no growing array of past tokens — only h lives between steps — which is the whole reason an SSM can process a million-token stream in constant memory.
Why plain SSMs weren't enough — and Mamba's fix
A fixed SSM uses the same a, b for every token, so the state leaks at one fixed rate. It has no way to say "this token matters, hold it" or "new topic, dump the old stuff". A burst of noise decays at exactly the same speed as a crucial fact. That input-blindness is why classic linear SSMs and RNNs lagged behind attention on recall-heavy tasks.
Mamba's twist is selectivity: make the parameters functions of the current token. Now the model gates its own memory — on a "new topic" token drive a→0 to flush the state; on an "important" token drive a→0, b→1 to write the value in cleanly; on an "irrelevant" token drive a→1, b→0 to hold what it has and ignore the input:
function selectiveA(tok, baseA){
if (tok.type === "reset") return 0.04; // "forget the past" -> flush
if (tok.type === "keep") return 0.05; // "write this value" -> clear old
if (tok.type === "noise") return 0.98; // "ignore this" -> hold state
return baseA;
}
The demo runs the real recurrence over a token stream with special tokens — reset, keep, noise — and plots two lines: a selective state and a non-selective one. Watch what happens at the control tokens: the selective line snaps — it flushes on reset and freezes on keep — while the non-selective line just leaks at one fixed rate, unable to tell an important token from noise. That input-dependent gating is the entire idea behind Mamba, and it's what let a fixed-size state finally rival attention.
The one catch: input-dependent parameters break the simple convolution trick that made linear SSMs fast to train. So Mamba uses a hardware-aware parallel scan to compute all the states at once on GPUs during training, then runs as the plain O(1)-memory recurrence at inference.
The lasting trade-off
Attention buys pinpoint recall by re-reading everything, at n² cost and growing memory. An SSM trades that for a fixed-size state: linear time, constant memory, and the ability to stream indefinitely. Recall is attention's edge; efficiency is the SSM's. That's why long-context, streaming, and on-device systems lean on SSMs — and why the frontier increasingly reaches for hybrids that keep a little attention for recall alongside cheap SSM layers.
Advance the stream and watch attention's cost explode while the SSM stays flat:
https://dev48v.infy.uk/ai/days/day43-state-space-models.html
Top comments (0)