Mamba changes its state-space model for every token it reads

The small-state bargain

Attention keeps everything. Every token a transformer has read stays available in the key-value cache, and any later token can look back at any earlier one exactly. That is why attention is expressive, and it is also why it costs what it costs: work and memory that grow with the square of the sequence, because every position may consult every other.

A state-space model refuses that deal. It carries a single hidden state of fixed size, updates it one token at a time, and reads its output from that state alone. Nothing about the past survives except what the state has kept. This is the property that makes the model cheap, and Gu and Dao name the tension it creates precisely: the whole design problem is what a fixed-size state is allowed to forget.¹

A large state can hold more of the past but costs more to store and move; a small state is fast but must discard. The question the Mamba paper sets itself is not how to make the state bigger. It is how to make a small state keep the right things.

The recurrence that treats every token alike

The structured state-space models Mamba builds on, S4 and its relatives, earn their speed with a restriction. S4 fixes the SSM matrices once and applies the same ones at every step of the sequence.² Its contribution was making those fixed matrices trainable and stable at long range: a low-rank correction to the transition that lets it be diagonalized and computed through a Cauchy kernel. The mechanics matter less here than the consequence of holding the matrices constant.

Because Ā, B̄, and C do not depend on the step, the recurrence has a second face. Unroll it, and the output at every position is the same fixed weighting of the inputs that came before. That weighting is a convolution kernel, one filter, computed once and slid across the whole sequence. The paper names this property and its payoff in a single stroke.

The convolution is why S4 trains quickly. A single kernel over the whole sequence is one large parallel operation, the kind a GPU is built for, with no step-by-step dependency to serialize. But a fixed filter applies the same response to every position regardless of what the position contains. The same update fires for a token that matters and a token that does not. Gu and Dao state the cost of that bluntness directly.

The paper makes the defect concrete with two tasks a fixed filter fails. Selective Copying asks the model to reproduce specific tokens while the spacing between them is randomized, so no fixed stride can latch onto them; the paper notes the randomized spacing is exactly what removes the shortcut an LTI model would otherwise exploit. The Induction Heads task asks the model to complete a pattern it has seen once before in the same sequence, the mechanism widely thought to underlie in-context learning.¹ Both require the model to act on what a token is, and a filter that ignores content cannot.

Making the parameters depend on the token

The fix is smaller than the problem sounds. If constant dynamics are what blind the model, let the dynamics vary with the input. Mamba computes three of the SSM parameters, the step size Δ and the matrices B and C, as functions of the current token instead of fixing them. Each becomes a linear projection of the input, recomputed at every position.

Here is what that buys, in terms of the state itself. The step size Δ controls how much of the current token is written into the state and how much of the old state is carried forward. When Δ is large at a position, the recurrence turns sharply toward that token: the state absorbs it and lets older content fade. When Δ is small, the token barely moves the state, which glides past almost unchanged. The paper reads Δ in exactly this register, as the SSM's version of a gate.¹ B decides how the chosen token enters the state; C decides what the state exposes as output. Making all three depend on the token turns the recurrence into a controller that, position by position, decides what to admit into its fixed memory and what to spend it on.

The model is no longer applying one response to every token. It reads each token, and from it computes how that token should move the state. The two synthetic tasks that defeat a fixed filter fall to this: with input-dependent dynamics the model can hold the tokens Selective Copying asks it to keep and skip the filler between them, and it can form the induction pattern it previously could not.¹

The gain is not free, and the paper is candid about the bill. The moment the parameters vary with position, the single fixed kernel is gone. There is no longer one filter to slide across the sequence, because the filter is different at every step.

The scan that never writes the state down

Without the convolution, training has to run the recurrence step by step, and a naive step-by-step loop is where SSMs go to die on a GPU. The state Mamba carries is not one vector but a large tensor: a separate hidden state for every channel of every sequence in the batch. Written out across all positions, that is an array of shape (batch, length, channels, state), far too big to keep in the fast memory near the compute units. Materialize it in the GPU's main memory and the recurrence spends all its time moving that tensor back and forth, doing tiny arithmetic between long waits. The work is not the bottleneck. The traffic is.

The answer is a hardware-aware selective scan, and it is worth being concrete about what runs where. A GPU has a large, slow pool of high-bandwidth memory (HBM) and a small, fast scratchpad near the arithmetic units (SRAM). The scan keeps the expensive tensor out of HBM entirely.

Two moves make this pay. First, the recurrence is run as a parallel scan rather than a serial loop. A linear recurrence is associative: the combined effect of two adjacent steps is itself a step of the same form, so the sequence can be reduced in a logarithmic number of parallel passes instead of one pass of length L. The dependency that looked inherently serial is not. Second, the intermediate states are simply not saved.

The move is a fused kernel: the released implementation carries the selective scan as a CUDA operator, selective_scan_cuda, wired in through ops/selective_scan_interface.py, so the discretization, recurrence, and output reduction happen in one pass without ever spilling the state.³ The convolution S4 relied on gave parallelism for free; selection gives it up, and the scan buys it back by respecting the memory hierarchy instead of the algebra.

What selection wins

With the scan carrying the cost, selection turns into a general language model rather than a synthetic-task trick. Mamba at 3 billion parameters matches transformers of roughly twice its size across pretraining and downstream evaluation, while its cost in sequence length grows linearly rather than quadratically.¹ For a practitioner the linear scaling is the structural win: a model whose per-token cost does not climb as the context lengthens.

The throughput figure deserves care, because it is easy to overstate. The abstract's headline of 5× is an inference number with a specific cause. A transformer at generation time holds a key-value cache that grows with every token, and that cache is what limits how many sequences it can serve at once. Mamba has no such cache; its state is fixed. The paper attributes the speedup precisely to that, larger serving batches, not to a faster training step or a faster single stream.¹ It is a serving advantage that follows from the fixed state, and the same fixed state is about to become the limitation.

Where a compressed state runs out

Selection decides what a fixed state keeps. It cannot make the state hold more than it holds. That distinction is where three independent groups, none of them authors of Mamba, drew the honest boundary of the mechanism, and their results are not a footnote to the idea but part of understanding it.

The sharpest is copying. Jelassi and colleagues at Harvard studied the task of reproducing a string from the context, and proved the gap is structural rather than a matter of training.

This is the exact inverse of the compression framing the paper opened with. Selection chooses what to keep; it does not raise the ceiling on how much can be kept. On a task that needs every token back, a fixed state runs out, and no amount of clever gating changes the size of the container.⁴

Merrill and colleagues at NYU and the Allen Institute struck from a different angle, at the intuition that a recurrence must, by carrying a state forward, be doing genuine sequential computation. They show it need not. An SSM cannot express any function outside the circuit class TC⁰, the same ceiling that bounds transformers, which means it cannot solve problems like permutation composition that require tracking state through a long chain of updates.⁵

Arora and colleagues at Stanford give the practical shape of the same limit and, usefully, a partial way around it. Because memory is finite, a recurrent model cannot recall and use everything in a long context, which makes its in-context learning brittle. But much of the lost recall is recoverable at the prompt: their Just-Read-Twice method repeats the context before the query so the model sees what it will be asked about, and reports an average gain of 11.0 ± 1.3 points.⁶ The limit is a property of fixed memory, and it can be managed by feeding the model its context in a friendlier order rather than by enlarging the state.

These do not contradict the paper. A naive reading of Mamba would expect selection to solve copying, since it solves Selective Copying on the synthetic benchmark, and it does not. The resolution is the one the whole reconstruction points at: selection improves which content survives compression; it does not remove the compression bound. The two synthetic tasks were about choosing what to keep, and selection is good at that. Copying and recall are about keeping everything, and no fixed state can.

When to reach for a selective SSM

The tidy framing of Mamba as an attention-free alternative did not survive its own authors. In their follow-up, Dao and Gu show that selective SSMs and attention are two readings of the same structured matrix computation, a duality rather than a rivalry, and use it to build a core layer that runs several times faster than Mamba's original scan while staying competitive on language modeling.⁷ The people with the most reason to sell selection as a free win are the ones who reframed it as a shared structure with the model it was meant to replace. That is the strongest evidence that the right way to hold it is as a tradeoff.

So the decision is not which architecture is better in general. It is which end of the small-state bargain a task lives on. Where the work is throughput, streaming, or very long context and the model can afford to compress what it has seen, a selective SSM gives linear cost and cache-free serving that attention cannot. Where the work is exact recall, copying, or state tracking through a long chain, the fixed state is the wrong tool, and the answer is a hybrid, an attention layer where the recall must be exact, or a prompt that hands the model its context twice.