Mamba and State Space Models — Handling Sequences Without Attention

Attention keeps everything and re-reads it on every step, and the price is quadratic cost in sequence length. State space models take the opposite bet: keep one fixed-size note and update it. Starting from a continuous-time linear system, we trace why S4 can secretly become a convolution, what exactly Mamba made 'selective' in order to give that convolution up, and where the whole approach loses to attention.

Efficiently Modeling Long Sequences with Structured State Spaces

Another way to remember

Self-attention is, in principle, remarkably blunt. Every time it processes a new token, it goes back and looks at every previous token again. That is why long-range relationships arrive in a single step, and why it works so well. The price is that doubling the sequence quadruples the number of pairs to compare. Where that quadratic wall comes from is covered in Attention from Scratch.

But consider how people actually do it. When you take minutes in a long meeting, you don't re-read the transcript from the top after every sentence. You keep one page of notes and revise it as new things are said. That page is the same size whether the meeting runs three hours or five.

This "one page of notes" approach is what RNNs used to do. The trouble was that multiplying by the same matrix hundreds of times either erased the past or blew it up, and since you can only advance one step at a time, there was nothing for a GPU to parallelize (see RNNs and LSTMs from Scratch).

State space models (SSMs) rebuild that same idea using the vocabulary of control theory. Mamba then adds one twist that makes the note actually work on data as capricious as language. Let's go down to the equations and see what the note really is.

Intuition: state is compressed history

There are only three characters in this story.

• Input $x_t$: the signal arriving at this step (say, one token's worth of numbers)
• State $h_t$: the note. All of the past, squeezed into a fixed size
• Output $y_t$: what you read back out of the note

And one rule: the new note is the old note, faded a little, plus what just came in. The answer is produced by reading the note.

The decisive property here is that the note's size does not depend on sequence length. Ten thousand tokens in, it is exactly as big as it was at the start. So the cost of one token is constant, and the cost of the whole sequence is linear. That is a fundamentally different way of remembering from attention, which holds on to every past K and V — which is why the KV cache grows in proportion to context length.

Nothing is free, of course. Squeezing into a fixed size means whatever doesn't fit gets thrown away. Exactly where that pays off and where it hurts is the second half of this article.

Mechanism: a continuous-time linear system

SSMs start from a linear system that control engineers have used for the better part of a century. Pretend time $t$ is continuous:

In plain words, Equation (1) says: the rate of change of the note, $h'(t)$, is the current note $h(t)$ reshaped by a matrix $A$, plus the current input $x(t)$ reshaped by a matrix $B$. The output $y(t)$ is the note read out through a matrix $C$. So $A$ governs how the note naturally fades, $B$ decides where in the note an input gets written, and $C$ decides which parts get read.

Strip the symbols away and one sentence is left, which says: how fast the note is changing right now is settled by exactly two things — what the note already holds, and what just walked in. Nothing from the past has to be kept on the side, because the influence of every earlier input has already been folded into $h(t)$. The meeting-minutes analogy from the opening, written down.

A mass on a spring obeys this form. So does the temperature of a room. SSMs bring it to text.

Text, however, is not continuous — it is a discrete run of tokens. So we discretize. Chop time into steps of width $\Delta$ and assume the input is constant within each step (a zero-order hold), and Equation (1) becomes a recurrence:

Equation (2) says the one-step matrix $\bar{A}$ is $A$ scaled by the step width $\Delta$ and then exponentiated; $\bar{B}$ is derived from $\Delta$ in the same way.

Put in words: the differential equation is now behind us. Advancing by one token has become nothing more than "multiply the note by a fixed matrix, then add the new input." Equation (1) is the continuous-time blueprint; Equation (2) is that blueprint turned into something you can run in a loop.

Here $\Delta$ is the dial for how much time one token is worth. Small $\Delta$ and the note barely changes (the input is ignored, the past is preserved); large $\Delta$ and the note gets heavily rewritten. That dial becomes the star of the show once we get to Mamba.

Why the eigenvalues of $A$ set the memory's lifespan

Run Equation (2) forward $k$ steps and whatever entered first has been multiplied by $\bar{A}^k$ — the same matrix, raised to a power. What a matrix power does is governed by its eigenvalues: where the magnitude is below 1, information decays exponentially and disappears; above 1, it explodes. The vanishing and exploding gradients of RNNs show up here wearing a different costume, as the lifespan of a memory.

So an SSM lives or dies by the design of $A$. A randomly initialized $A$ is known to retain almost nothing over long ranges, and the S4 line of work fixed this by importing a theory called HiPPO. HiPPO derives how $A$ should be set so that the past is retained as a polynomial approximation of the incoming signal, and the resulting matrix is used as the initialization. An initialization that "remembers the distant past well," falling out of theory rather than tuning, is one of the more charming aspects of this family.

S4: a recurrence is also a convolution

Here is the central trick of S4 (the Structured State Space Sequence model). Equation (2) looks strictly sequential, but if $\bar{A}, \bar{B}, C$ are the same at every step (time-invariant), you can unroll it into a single expression:

$$y_t = \sum_{k=0}^{t} C\bar{A}^{k}\bar{B}\, x_{t-k}$$

Don't let the summation sign put you off. Spelled out in words, it says: the output at time $t$ is a weighted sum of every input that has arrived so far, and the weight applied to an input from $k$ steps back is the fixed quantity $C\bar{A}^{k}\bar{B}$. The part that matters is that the weight depends only on how far back an input was — never on where in the sequence it happened. So you can build that list of weights once, slide it along the input, and read off every output at once.

That is a convolution of the input sequence with a kernel $\bar{K} = (C\bar{B},\ C\bar{A}\bar{B},\ C\bar{A}^2\bar{B},\ \dots)$. In other words, the same model can be written either as a recurrence or as a convolution.

That duality is what makes it practical. During training the whole input sequence is already available, so you treat it as a convolution, compute it with an FFT in $O(L \log L)$, and — with no step-to-step dependency — keep the GPU saturated. At inference you switch back to the recurrence and advance one token in constant time and constant memory. The RNN's fatal flaw, being unparallelizable, is dodged by wearing a different face during training.

Mamba: making the state selective

S4 still had a weakness. If $\bar{A}, \bar{B}, C$ are shared across all steps, then every token is processed identically. For data with uniform statistics — audio, sensor signals — that's fine. Language is not like that. "The" and a comma and a proper noun all overwrite the note with equal force, and there is no way to say "hold on to this one" or "from here on, drop the previous topic."

Mamba (Gu & Dao, 2023) changes exactly one thing: make $\Delta$, $B$, and $C$ functions of the input $x_t$.

# Pseudocode: derive B, C, Δ from each token (A stays a learned parameter)
B_t     = x_t @ W_B                      # where in the state to write this token
C_t     = x_t @ W_C                      # which part of the state to read
delta_t = softplus(x_t @ W_d + bias)     # how strongly this token should register
A_t     = exp(delta_t * A)               # decay that now varies per timestep
h       = A_t * h + delta_t * B_t * x_t  # update the state
y_t     = C_t @ h

This is the selective state. A large $\Delta_t$ means the token is written in forcefully and the previous state is largely overwritten; a small $\Delta_t$ means the token is essentially ignored and the state is preserved as is. The model can now decide for itself: this one matters, keep it; that one is filler, let it wash through. It is the input and forget gates of an LSTM, reinvented as the step size of a linear system.

The paper demonstrates the effect on two deliberately constructed synthetic tasks: one where specific tokens must be picked out of randomly spaced filler, and an induction-head-style task where a pattern seen once must be reproduced later. Time-invariant S4 struggles on the former; adding selection lets Mamba solve it.

The cost, and how it gets paid (parallel scan)

There is a catch. If $\bar{A}$ differs at every timestep, it can no longer be written as a convolution. The FFT trick that made S4 trainable is gone.

Mamba pays for this with a parallel scan. Accumulating front to back looks inherently sequential, but think about a prefix sum. For $[1,2,3,4]$ you can add neighbors to get $[1+2,\ 3+4]$ and combine the partial results — a tournament bracket, giving depth $\log L$ with $O(L)$ total work. The state update $h_t = \bar{A}_t h_{t-1} + u_t$ can be written in the same associative form, so the same trick applies.

Mamba's implementation goes further and is built around the GPU's memory hierarchy. The state is an array expanded per channel, and a naive implementation would materialize an enormous intermediate in HBM (the GPU's main memory). So discretization, scan, and output multiply are fused into a single kernel, the state lives only in fast SRAM, and the backward pass recomputes instead of storing. Spending extra arithmetic to avoid memory round-trips is precisely the idea behind FlashAttention — and indeed Tri Dao is behind both.

Complexity: living with sequence length

Attention versus SSMs, as a function of sequence length $L$:

The two right-hand columns are what you feel in production. A Transformer revisits the entire past for every token, so each token gets slower as context grows while the KV cache eats memory and drives concurrency down. An SSM's per-token cost and state size are unchanged when context grows tenfold.

One caveat that is routinely misread: the gap between $O(L)$ and $O(L^2)$ only matters once $L$ is genuinely large. On short sequences the constants dominate — implementation efficiency, how fully the GPU is used — and a heavily optimized FlashAttention kernel is often simply faster. Big-O describes the slope, not the intercept.

Where it wins and where it loses

The wins follow from the design. Very long sequences where information is spread broadly across the whole span: genomic sequences, raw audio, sensor time series. Constant per-token inference cost also suits edge devices and long-running streaming. As a language model, the Mamba paper reports scaling that outperforms same-size Transformers.

The losses are just as predictable, because they come from the premise itself: everything must fit in a fixed-size state. Copying a string verbatim from earlier in the input, pulling one specific sentence out of a long document word for word — on these verbatim recall tasks, attention holds a clear advantage, shown both theoretically and empirically (Jelassi et al., 2024). Attention keeps the whole past, so if it's there you can go get it; an SSM cannot reconstruct what its state no longer holds. That is why the gap tends to open up on needle-in-a-haystack style long-context evaluations.

Which is why production has landed on hybrids. Make most layers Mamba and interleave a handful of attention layers at chosen depths. Jamba (AI21, 2024) published exactly this arrangement as an open model, keeping the linear-cost advantage while patching the weakness in verbatim recall.

On the theory side, Mamba-2 (Dao & Gu, 2024) established structured state space duality (SSD): a particular structured form of attention and selective SSMs are two views of the same computation, and you can move between them. Under that lens, Mamba looks less like attention's rival and more like a different coordinate system over the same space.

How this shows up in practice

Who touches it, and when. The person who owns inference cost, the moment they say "we want longer context but the KV cache is eating the GPUs." Teams working with long audio, sensor streams, or logs who have run into the quadratic wall. On-device inference with a hard memory ceiling. Conversely, if the requirement is accurate quotation from internal documents, reach for RAG plus an attention-based model first — that is not a place to pick a pure SSM.

Names you will actually type. The usual implementation is the mamba-ssm package (the Mamba block plus the selective_scan CUDA kernel). There are really only four dials at design time:

• d_state (the $N$ in the equations — the width of the state): more capacity to remember, at the cost of memory and speed
• d_conv: the width of the short 1-D convolution that sits in front of the SSM. It handles local n-grams, and removing it measurably degrades quality
• expand: the channel expansion factor inside the block
• dt_min / dt_max (the initialization range for $\Delta$): the most important one. It sets the effective time constant of memory, and if it is too short for the sequence lengths you care about, long-range dependencies never get learned in the first place

Failure modes that bite.

1. Dependence on the CUDA kernel. The fast selective scan is a custom kernel; if your environment or dtype doesn't line up, it silently falls back to the reference PyTorch implementation and gets orders of magnitude slower. Confirm the fused kernel is actually running before you trust a benchmark.
2. State management is not KV caching. Reusing the KV of a shared prefix (prefix caching) is standard practice with Transformers. Doing the equivalent with an SSM requires snapshotting and restoring the state at that point. If your serving stack has no such mechanism, every branch replays the prompt from scratch and you have thrown away the linear-cost advantage yourself.
3. Choosing the wrong metric. These models break in a specific way: perplexity looks indistinguishable, and then performance collapses on anything demanding verbatim recall. If long-context retrieval or quotation is a requirement, evaluate for it explicitly.
4. Don't frame it as "replacing the Transformer." Most deployments are hybrids. How many layers between attention blocks is a design decision, and it is exactly the thing that gets argued about in review.

Questions you'll get in interviews and design reviews. "Why can't Mamba use the FFT convolution the way S4 does?" — because selection makes $\bar{A}$ vary per timestep, so the system is no longer time-invariant. "Then how is it parallelized?" — with an associative parallel scan that brings the depth down to $\log L$. Answer those two cleanly and you've shown you understand the design philosophy of the whole family.

Summary

• An SSM is a linear system that keeps updating a fixed-size state: linear time and constant memory in sequence length
• S4 exploited time-invariance to wear two faces — a convolution (FFT) while training, a recurrence at inference
• Mamba made $\Delta, B, C$ input-dependent to gain selection, and paid for the lost convolution with a parallel scan and kernel fusion
• It wins on long, uniform sequences and loses on verbatim recall, which is why production settles on hybrids

Next up: structured state space duality — how a selective state can be rewritten as a form of attention — traced through the shape of the matrices.