State Space Models: Why Linear Recurrences Rival Attention for Long Sequences
Attention made long context practical, but it made it expensive. Every new token looks back at every earlier token, so the cost of a sequence grows quadratically with its length. State space models take a different bargain: compress the past into a fixed-size state, update that state with a linear recurrence, and let the state carry exactly the information the future needs. Done well, that bargain gives near-linear scaling, constant-time generation per token, and surprisingly strong results on very long sequences.
This article explains where state space models come from, why the naive recurrence fails, and how the modern versions - structured state spaces and selective scans - fix it. The running example is a long sequence task: predicting the next value of a signal after ten thousand steps, where attention is slow and a plain recurrent network forgets.
From differential equations to sequences
A continuous-time linear state space model describes a hidden state that evolves under an input signal:
Here is an -dimensional state, is the input, is the output, and , , and are matrices. The state is a memory: it summarises the entire input history, filtered through the dynamics in .
Real data arrives in discrete steps, so the model is discretised. With a step size , the zero-order hold discretisation gives:
and the discrete recurrence:
That recurrence is the whole point. Processing a new token costs a fixed matrix-vector product, independent of sequence length. Generation needs only the current state, not a growing cache of keys and values. Training, meanwhile, can be parallelised: because the recurrence is linear, it can be rewritten as a convolution over the whole sequence, which runs efficiently on modern hardware.
So why did this old idea from control theory not take over sooner? Because a generic matrix gives a generic memory, and generic memories are bad at remembering the right things over long horizons.
The memory problem: what should the state remember?
Suppose the input is a long stream and the task needs a fact from a thousand steps ago. A linear recurrence multiplies the state by at every step. If the eigenvalues of are much smaller than one, old information decays exponentially and is gone. If they are larger than one, the state explodes. Only a carefully structured keeps information alive at many timescales at once.
The structured state space approach, introduced in the S4 model, solves this by initialising to approximate the HiPPO operator. HiPPO answers a precise question: given the history of a signal up to time , what are the coefficients of the best polynomial approximation to that history, under a measure that weights recent time more heavily? The state then literally stores a compressed, multi-scale summary of the past, rather than a fading echo of it.
In practice, S4 parametrises as a diagonal plus low-rank matrix. That structure keeps the HiPPO-like memory properties while making the convolution kernel computable in near-linear time. The result was a model that, for the first time, handled raw sequences tens of thousands of steps long - audio waveforms, long documents, genomic sequences - where attention-based models were either too slow or too memory hungry.
But S4 had a limitation that mattered for language: its dynamics were time-invariant. The same , and applied at every step, whatever the input. The model could not decide, based on content, to remember this token and ignore that one. Attention can do exactly that: it selects. A fixed linear filter cannot.
Selectivity: letting the input choose what to remember
The Mamba architecture removes the time-invariance constraint. Its parameters become functions of the input at each step:
where , and are computed from itself, typically through small learned projections. The step size also becomes input-dependent: a large step forgets the past quickly, a small step holds on to it.
Intuition helps here. Think of the state as a notebook with limited pages. A time-invariant model writes every incoming note in the same handwriting, on the same schedule, whether the note matters or not. A selective model decides, per note, whether to write it large, write it small, or skip it, and whether to erase an old page to make room. That content-based gating is what lets selective state space models do selective copying and induction-like tasks that defeated earlier recurrent models.
The price is that the convolution trick no longer applies: with time-varying dynamics, the recurrence cannot be rewritten as a fixed convolution. Mamba recovers efficiency with a hardware-aware selective scan, an associative scan that computes the recurrence in parallel across the sequence while keeping the state in fast on-chip memory and never materialising the full intermediate sequence. The algorithm, not just the maths, is part of the model.
A practical example: long-context copying
Consider a synthetic but revealing task: a sequence of key-value pairs, followed by a query key, where the model must output the matching value. Attention solves this by attending directly to the matching pair. A time-invariant state space model struggles, because it compresses all pairs into the same fixed summary and cannot single one out.
A selective model behaves differently. When it sees a pair, it can open the gate and write it into the state; when it sees filler tokens, it can close the gate and leave the state almost untouched. At query time, the output projection can read out the stored value. In code, the core recurrence for a single channel looks like this:
# h: state vector, x_t: input at step t
# A, B, C, delta: input-dependent parameters at step t
h = A_bar * h + B_bar * x_t
y_t = C @ h
The full model stacks many such channels, mixes them with projections and a gating nonlinearity, and scans them in parallel. The lesson for practitioners is not the two lines above, but the design pattern: a fixed-size state plus content-dependent updates can implement a form of attention-like selection without a quadratic attention matrix.
Where state space models win, and where they do not
State space models win on three axes:
- Length. Cost per token during generation is constant, and memory is the state size, not the sequence length. Sequences of a hundred thousand steps or more become feasible.
- Throughput. The scan formulation processes long sequences with high hardware utilisation, and inference avoids the growing key-value cache that slows long-context attention.
- Continuous and irregular data. Because the model starts from a differential equation, it handles irregularly sampled time series, audio and sensor streams naturally, where token-based attention needs extra machinery.
They do not win everywhere. Attention still has an advantage when the task requires many arbitrary lookups into a long context, such as retrieving several scattered facts and combining them, because a compressed state is a bottleneck by design. A state of dimension cannot losslessly store an arbitrary sequence of length much larger than ; it can only store a useful summary. Hybrid architectures, which interleave state space layers with occasional attention layers, are therefore a sensible default: use the recurrence for the bulk of the sequence and attention where precise recall matters.
Training stability is another practical concern. The discretised matrix must stay stable, initialisation matters more than in a transformer, and selective scans need careful implementation to match the speed claims. A naive loop over time steps in Python will be slow; the published speedups come from fused kernels, not from the recurrence alone.
Limitations and common pitfalls
- The state is a bottleneck, not a database. Do not expect a fixed-size state to memorise a long document verbatim. Judge state space models on summarised, filtered or predicted information, not on exact recall of arbitrary details.
- Selectivity is learned, not guaranteed. The gating that makes Mamba work must be learned from data. On tasks with little training signal for selection, a selective model can behave much like a fixed one.
- Long does not mean infinite. Performance still degrades as sequences grow far beyond the training length, just more gracefully than a model whose cost explodes. Always test at the deployment length, not only at the training length.
- Benchmarks can mislead. A model that wins on long-range synthetic benchmarks may not win on your data, where the relevant horizon may be short and attention's precise lookup may matter more. Measure on the actual task, at the actual length.
The takeaway
State space models reframe sequence modelling as a filtering problem: maintain a state, update it linearly, and let the input decide what enters and what fades. Structured initialisation gives the state a good multi-scale memory, and input-dependent selectivity gives it the ability to choose. The result is not a universal replacement for attention, but a genuinely different trade-off: linear scaling and constant-time generation, paid for with a compressed rather than complete memory of the past. For long sequences, that trade is often the right one.
Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial
Top comments (0)