State Space Models: From S4 to Selection

The state space model comes from a part of engineering that predates deep learning by decades. It describes a system with an internal state evolving continuously in response to an input.

The form is two equations. The rate of change of the state equals a matrix A applied to the current state, plus a matrix B applied to the current input. The output equals a matrix C applied to the state.

A governs how the state evolves on its own, B how input enters it, C how the output is read from it. The state is a compressed summary of everything the system has seen, and its size is fixed by design.

To use this on token sequences it has to be discretised, converting the continuous dynamics into a step-by-step update with a step size usually written as delta. After discretisation the equations become a recurrence: the new state is a matrix applied to the old state plus a matrix applied to the new input, and the output is read from the state.

Which is exactly the shape the previous lesson arrived at from the other direction.

The property that made this line viable is that A, B and C do not depend on the input. They are fixed parameters, the same at every position, which makes the system linear and time-invariant.

That has a consequence worth unfolding. Expand the recurrence: the state at step t is a sum over all previous inputs, each multiplied by A raised to some power times B. Since those coefficients depend only on how far back the input is, and not on what it was or where it occurred, the whole output is the input sequence convolved with one fixed kernel.

A convolution is not sequential. The kernel can be precomputed, and the convolution evaluated for every position simultaneously, using the fast Fourier transform to do it in n log n time.

So the same model has two forms. A recurrence for generation, stepping one token at a time with constant memory. And a convolution for training, computing every position at once.

That is the trilemma resolved again, by a different route, and it is why state space models became interesting rather than remaining a control theory curiosity.

The early attempts did not work, and the reason is a detail that sounds like a footnote and is not.

With A initialised randomly, the model failed on long sequences. Repeatedly applying a random matrix either shrinks the state toward zero, erasing the past, or grows it without bound. Neither retains useful history.

HiPPO supplied the fix. It is a construction for A derived from a specific question: what matrix makes the state an optimal compression of the input history, in the sense of maintaining the coefficients of a polynomial approximation to everything seen so far?

That question has a closed-form answer, and initialising A to it gives the state a principled job. It is not an arbitrary summary; it is an online approximation of the entire past, updated exactly as each element arrives.

The empirical effect was decisive. The same architecture with random initialisation performs poorly on long-range tasks and with HiPPO initialisation performs well.

Structured state space models, S4, are essentially this plus a parameterisation of A that keeps the convolution kernel cheap to compute, since a general matrix power series would be prohibitive.

S4 worked well on signals: audio, time series, long-range synthetic benchmarks. On language it lagged transformers, and the reason is the same property that made it fast.

Because A, B and C are fixed, the system processes every token identically. Its behaviour depends on position but never on content. The same update is applied to the word the and to a rare technical term that determines the rest of the document.

For a physical signal that is appropriate, since the dynamics genuinely are the same at every instant. For language it is badly wrong. Some tokens should overwrite the state, some should be absorbed and forgotten, and some should be held for a long time. Which is which depends entirely on what the token is.

A time-invariant system cannot make that distinction. It has no mechanism to look at a token and decide.

Attention makes exactly that content-dependent decision, in its similarity computation, which is a good explanation of the gap.

So the requirement is clear: make the model's behaviour depend on its input. The difficulty is that time-invariance is what the convolution rested on.

Mamba's change is small to state and consequential in effect. Make the parameters functions of the current input.

The step size delta, and the input and output matrices B and C, become learned projections of the current token rather than fixed values. Delta is the important one: it controls how much the state advances per step, so making it input-dependent lets the model decide, per token, whether to let the state evolve substantially or barely at all.

That is a content-dependent forget gate, arrived at from control theory rather than from recurrent networks, and it is the same conclusion the previous lesson reached about decay.

What it buys is selectivity. The model can absorb a filler word with almost no state change, flush the state at a document boundary, and hold a salient entity across a long span.

What it costs is the convolution. With parameters varying by position, the coefficients relating input t to output s no longer depend only on their separation. There is no single kernel, so there is no convolution, and the fast Fourier transform route is gone.

The model is now a sequential recurrence, which is precisely the problem that sank recurrent networks.

The rescue comes from a classical result in parallel computing. Sequential dependence does not always mean sequential execution.

A prefix sum looks strictly sequential: each partial sum needs the one before it. It is not. Because addition is associative, partial results can be combined in a tree, giving the whole sequence in logarithmic depth rather than linear. That is a parallel scan, and it generalises to any associative operator.

The selective state space recurrence is a linear update, and composing two such updates yields another of the same form. The operator is associative, so the recurrence can be evaluated by a scan.

That restores parallel training. Not as a single convolution, but in logarithmic depth rather than linear, which is the property that matters.

The implementation is where the practical work sits. A naive scan writes the full state for every position to main memory, and the state is large. Mamba's hardware-aware version keeps the state in on-chip memory, fuses discretisation, scan and output projection into one kernel so intermediates never leave the chip, and recomputes intermediate states during the backward pass rather than storing them.

The algorithm was known. Making it faster than the memory traffic it avoids was the contribution.

Each step in this line bought something and gave something up, and laying them in order makes the design space legible.

The continuous system is the starting point, with no notion of discrete tokens.

Discretisation makes it a recurrence over positions, giving constant-memory inference and sequential training.

Time-invariance turns the recurrence into a convolution, recovering parallel training, at the cost of behaviour that cannot depend on content.

HiPPO initialisation makes the state a principled compression of history, which is what made long-range performance work at all.

Selection makes the parameters input-dependent, buying content-dependent memory and losing the convolution.

The parallel scan recovers parallel training from the associativity of the update, and hardware-aware kernels make it fast in practice.

Read as a whole, the line converges on the same object the previous lesson reached: a linear recurrence with an input-dependent forget gate, trainable in parallel because its update composes associatively. The two lines were never actually separate.

The most satisfying result in this area is that the two lines this path has followed describe one object.

Mamba-2 introduces structured state space duality. The observation is that a selective state space model can be written as a matrix transformation of the input sequence, output equals M times input, where M is a structured matrix determined by the model's parameters.

The structure in question is semiseparable: a matrix whose blocks have low rank, which admits a subquadratic parameterisation and linear-time algorithms.

With the transition matrix restricted to a scalar times the identity, that matrix becomes a form of attention, with the softmax replaced by cumulative gating and the causal mask replaced by data-dependent decay.

So a state space model with a suitably structured transition and a form of linear attention are the same computation viewed from two angles. Two research communities using different vocabularies had converged on one mechanism.

The payoff is practical rather than aesthetic. The duality gives a chunked algorithm: quadratic attention-style matrix multiplication within a chunk, which uses tensor cores efficiently, and linear recurrence between chunks, which keeps memory bounded. Best of both, chosen at the block level.

A Mamba block is not a state space model in isolation, and the surrounding machinery is doing real work.

The input is projected up to a wider inner dimension, typically twice the model dimension, giving the state room to work in.

A short causal convolution, on the order of four positions wide, is applied before the recurrence. This handles very local patterns cheaply, so the state does not have to spend capacity on immediate neighbours.

The selective state space computation runs on that, with delta, B and C projected from the input.

A gating branch runs in parallel, projecting the input through a nonlinearity and multiplying the recurrence's output elementwise. This is the same gating idea that gated linear units brought to feed-forward blocks, and it supplies the nonlinearity the linear recurrence lacks.

Finally a projection back down to the model dimension.

The block replaces both the attention and the feed-forward sublayer of a transformer, rather than only attention, so a like-for-like comparison uses roughly twice as many Mamba blocks as transformer layers at equal parameter count.

Three results carry forward from the state space route.

A discretised linear dynamical system is a recurrence, and time-invariance turns that recurrence into a convolution, which is a second independent way to get parallel training with a constant-size state.

The state's usefulness depends on how it is parameterised. HiPPO made the difference between a state that forgets or explodes and one that maintains a principled compression of history, which is a reminder that initialisation is a modelling decision here rather than a detail.

And content-dependence is what language needs. Selection buys it by making the parameters functions of the input, at the cost of the convolution, with the parallel scan recovering trainability from associativity.

The duality result then collapses the distinction this path has been maintaining. Linear attention with decay and a selective state space model are the same computation, and the chunked algorithm that follows is what makes either fast on real hardware.

What none of this addresses is the capacity limit. A fixed state still holds a fixed amount, and the last lesson is about what that costs and what the field did about it.