For a decade the deep learning canon was a tidy taxonomy of four primitives: the MLP, the CNN, the RNN, and (since 2017) the Transformer. Each one is defined by how a single layer mixes information across positions. MLPs mix everything-to-everything but ignore structure. CNNs mix only nearby positions. RNNs mix positions sequentially through a hidden state. Transformers mix everything-to-everything with content-dependent weights.
This lesson is about a fifth primitive that has matured enough to sit on that shelf: the State-Space Model (SSM) and its modern selective variant, Mamba. We are deliberately not looking at this from the "NLP application" angle that you may have seen in the NLP track. We are looking at it as an architecture — a primitive that is parallel to the others — and asking the only question that matters: what does this primitive let you do that the others cannot?
Learning Objectives
After this lesson, you will be able to:
Explain why O(N^2) attention is unaffordable at very long context and what an SSM offers as a replacement
Derive the continuous SSM equations and discretize them via zero-order hold
Explain how an SSM unrolls into a convolution and why that enables O(N log N) FFT-parallel training
Describe the HiPPO matrix and why it gives provably-good long-range memory
Identify the selectivity problem in vanilla SSMs and how Mamba's input-dependent B, C, and step size fix it
Explain why selectivity breaks the convolutional view and what the selective parallel scan replaces it with
Compare Transformer, Mamba, S4, and RNN on training cost, inference cost, memory, and long-range ability
Describe when Mamba beats a Transformer and when a Transformer still wins
The Transformer is brilliant at one job and terrible at one job. The job it is brilliant at is content-based mixing — for every output position, every input position can have a different weight, chosen by the input itself. That is why Transformers handle ambiguous reference, syntax, and structured retrieval far better than RNNs ever did.
The job the Transformer is terrible at is scaling cost with sequence length. A single self-attention layer at sequence length $N$ with hidden dimension $d$ does roughly $N^2 \cdot d$ FLOPs and stores an $N \times N$ attention matrix in memory. Doubling the context quadruples the work. At $N = 100,000$, an attention matrix in fp16 already costs about $20,\text{GB}$ per layer per head before activations — the architecture's whole defining feature becomes its bottleneck.
At inference the picture is worse, not better. For autoregressive generation the Transformer keeps a KV-cache of size $O(N \cdot d \cdot L)$ that grows linearly per token. Every new token attends back over the entire cache, so each generated token costs $O(N \cdot d)$. Generating $M$ tokens at context $N$ costs $O(M \cdot N \cdot d)$ — and the cache itself dominates GPU memory long before the model weights do.
RNNs are the opposite trade. A vanilla RNN does $O(N)$ work and keeps an $O(d)$ state, regardless of how long the sequence is. The catch is that you cannot train them in parallel (each step depends on the previous one) and the gradient propagated through many time steps tends to vanish or explode, so they cannot actually learn long-range dependencies even though they have the capacity to represent them.
The SSM family was built around a single audacious goal: keep the RNN-like cost profile, and find a mathematical structure that (a) trains in parallel and (b) actually preserves long-range memory.
What Do You Think?
A Transformer at sequence length N with hidden dimension d. What is the FLOP count of a single self-attention layer?
The SSM does not start its life as a discrete neural network. It starts as a piece of classical control theory — a model of a physical system whose hidden state $x(t)$ evolves continuously in time, driven by an input signal $u(t)$, and emits an output $y(t)$.
dtdx(t)=Ax(t)+Bu(t),y(t)=Cx(t)+Du(t)
The hidden state lives in $\mathbb{R}^N$ (think $N$ in the tens or hundreds, not the thousands), the input $u(t)$ is a single scalar channel, and $y(t)$ is a single scalar output. To turn this into a deep learning layer with many channels, you simply run many independent copies — one SSM per channel — with their own learned $A$, $B$, $C$, $D$. The continuous-time view (sometimes called the continuous-timeStochastic CalculusStochastic calculus extends derivatives to random processes — Brownian motion, Itô integrals, and stochastic differential equations. The math diffusion models, flow matching, and score-based generative models all live in.Learn more → view, since the same machinery appears in stochastic differential equations) is what lets us reason about the system before we have committed to a sampling rate.
A neural network does not see continuous signals. It sees a sequence $u_1, u_2, \ldots, u_N$ sampled at discrete steps. We must therefore decide how to discretize the continuous SSM — that is, work out the discrete recurrence that an actual computer will run.
The cleanest way to discretize an ODE driven by a piecewise-constant input is the zero-order hold (ZOH): assume the input is constant over each step of size $\Delta$, then integrate the ODE exactly over that step. The closed-form result is a matrix exponential.
Aˉ=exp(ΔA),Bˉ=(ΔA)−1(exp(ΔA)−I)ΔB
With $\bar A$ and $\bar B$ in hand the discrete recurrence is the familiar one-line update:
xk=Aˉxk−1+Bˉuk,yk=Cxk+Duk
Two things to notice. First, $\Delta$ is not a fixed hyperparameter — in S4 it is learned per channel, because different channels want to remember on different timescales (a high-frequency edge detector wants a small $\Delta$; a paragraph-level summarizer wants a large $\Delta$). Second, the entire recurrence is linear: there is no nonlinearity inside the state update. All the expressiveness has to come from the structure of $A$ and from stacking these layers with elementwise nonlinearities in between (gates, MLPs, residuals).
Quick check
In a learned SSM layer, what role does the step size Delta play?
Here is the trick that makes the SSM more than a fancy linear RNN. Because the recurrence is linear and time-invariant (the same $\bar A, \bar B, C$ are applied at every step), you can unroll it in closed form.
Starting from $x_0 = 0$:
x1=Bˉ
Reading off the output $y_k = C x_k + D u_k$:
yk=i=0∑kCAˉk−iBˉui+Duk
The vector $K = (C\bar B,; C\bar A\bar B,; C\bar A^2 \bar B,; \ldots)$ is called the SSM convolution kernel. The output $y$ is simply $K \ast u$ (causal convolution) plus the residual $Du$. This is the single most consequential identity in the SSM literature: a linear recurrence has been re-expressed as a single, long, fixed convolution.
Why is this such a big deal? Because convolutions can be evaluated in parallel. In particular, a length-$N$ convolution can be computed via the FFT in $O(N \log N)$ — and the FFT, unlike a sequential recurrence, has no data dependency between time steps. So during training (when the whole sequence is known up front) the SSM evaluates in $O(N \log N)$ with full GPU parallelism. During inference (when tokens arrive one at a time) you switch back to the recurrent form and pay $O(1)$ per token with $O(N)$ state, just like an RNN.
You get the best of both worlds: CNN-style parallel training and RNN-style streaming inference, from the same set of weights.
A vanilla random initialization of $A$ does not work. Even with the convolutional speedup, an SSM with a generic $A$ matrix has the same long-range memory problems as a vanilla RNN: powers of $\bar A$ either explode or shrink to zero, and the kernel either blows up or forgets everything past a few hundred steps. The S4 paper's key insight (Gu, Goel, Re, 2021) was that the choice of $A$ controls how the past is compressed into the state, and we can pick $A$ on principled grounds.
The principle comes from a 2020 paper, also from Gu, called HiPPO (High-order Polynomial Projection Operators). It asks: given a fixed-size state of dimension $N$, what is the optimal way to compress the entire history of a continuous signal into that state, in the sense of minimizing reconstruction error under some measure?
The answer is beautifully classical: project the history onto an orthogonal polynomial basis (typically the Legendre polynomials). The state vector $x$ becomes the coefficients of the projection. The matrix $A$ that maintains this projection as new input arrives is the HiPPO matrix — a specific lower-triangular matrix whose entries are simple functions of $i$ and $j$ (Legendre formula). The HiPPO matrix has the provable property that the state $x_k$ is, at every step, the optimal Legendre projection of the input history seen so far.
In other words, when you initialize $A$ to the HiPPO matrix, the SSM starts out with mathematically optimal long-range memory, then training can perturb $A$ as needed. This is what unlocked SSMs on the Long Range Arena benchmark, where S4 achieved the first decisive break past the 60% average bar that all prior architectures (Transformers included) had hit.
Quick check
What is the purpose of using a HiPPO matrix to initialize the SSM's A matrix?
The HiPPO matrix is dense and structured. Computing a length-$N$ kernel from it requires a clever Cauchy-kernel reformulation to stay numerically stable, plus some careful FFT manipulation. S4 worked but the implementation was fiddly.
The follow-up, S4D (Gu, Goel, Gupta, Re, 2022), made a striking observation: most of the empirical benefit of HiPPO survives if you replace the full matrix with a diagonal matrix whose entries are the eigenvalues of a related operator (specifically a HiPPO-LegS variant). With $A$ diagonal, $\bar A$ is also diagonal, and $\bar A^k$ is just elementwise exponentiation — the kernel can be built in $O(N)$ with no Cauchy machinery at all.
S4D matches S4 to within a small margin on Long Range Arena and is significantly faster and simpler to implement. Almost every later SSM (including Mamba) builds on the S4D structure rather than the full-matrix S4.
Loading visualization...
Run the cell. You will see that the convolutional and recurrent forms agree to machine precision, and the response to the second impulse is a translated copy of the response to the first. That last point is exactly the limitation we are about to attack.
Here is the structural weakness of every SSM we have built so far. The kernel $K$ is fixed — it depends on the weights but not on the input. Every token gets convolved with the same kernel. There is no mechanism for the model to say "this token is important, hold onto it" or "this token is filler, throw it away." The dynamics are input-invariant.
Contrast with attention. In a Transformer, the value at position $i$ is mixed into the output at position $j$ with a weight $\text{softmax}(q_j \cdot k_i)$, which depends on both tokens. Attention is content-based. SSMs of the S4 / S4D family are not — they are exactly as content-blind as a 1D convolutional layer.
This shows up empirically on associative recall tasks. Imagine the input is a sequence of (key, value) pairs followed by a query key, and the model must output the matching value. A Transformer solves this trivially (the query attends to the matching key). An S4 model fails: the kernel cannot route information from a specific past position to the present based on the current input.
This is the gap Mamba closes.
What Do You Think?
An associative recall task: the model sees a stream of (key, value) pairs, then a query key, and must produce the matching value. Which is better suited to solve this?
Mamba (Gu and Dao, December 2023) made one elegant change to the S4D recipe: it let the input parameterize the dynamics. Concretely, in every Mamba layer, three of the SSM's parameters are produced by small linear projections of the current input token:
Δk=softplus(WΔuk),Bk=WBuk,Ck=WCuk
With $\Delta$, $B$, and $C$ now varying per token, the model has a knob for every token to choose how much of the new input to integrate ($B_k$), how the state should be read out ($C_k$), and how quickly to advance time ($\Delta_k$). A small $\Delta_k$ means $\bar A_k \approx I$ — the state barely moves, the model effectively ignores this token. A large $\Delta_k$ means the state turns over fast — the model "writes" new information aggressively. The model has learned a soft, content-dependent attention-like mechanism without ever building an $N \times N$ matrix.
But this breakthrough comes with a serious computational price.
Quick check
Why does making B, C, and Delta input-dependent break the convolutional view of an SSM?
The standard convolutional speed trick (FFT of a precomputed kernel) is gone. But the recurrence is still associative: composing two linear updates $(\bar A_a, \bar B_a u_a)$ and $(\bar A_b, \bar B_b u_b)$ gives another linear update. Any associative operation can be evaluated by a parallel scan in $O(N)$ work and $O(\log N)$ depth — essentially the same algorithm as parallel prefix-sum.
The Mamba paper's engineering contribution is a fused selective scan kernel that keeps the intermediate states in SRAM on the GPU (never spilling them to HBM) and fuses the scan with the projections. The result: empirically, on a modern GPU, Mamba trains at speeds comparable to FlashAttention-2 Transformers up to roughly 16k context, and pulls ahead beyond that.
So we lose the FFT-convolutional trick but gain back parallelism through a different mechanism — the parallel scan — at the cost of needing custom CUDA. This is why every Mamba implementation you find ships a hand-tuned kernel: in pure PyTorch, the scan is dominated by memory traffic and the model is much slower.
Mamba-2 (Dao and Gu, 2024) is built on a deeper observation called State Space Duality (SSD). The authors show that a selective SSM with input-dependent $B$ and $C$ can be re-cast as a particular kind of structured matrix mixing — specifically, the output is the input multiplied by a semiseparable matrix. Semiseparable matrices are matrices whose blocks have low rank in a precise sense; multiplying by them can be done efficiently using a block-decomposed algorithm.
This re-casting gives Mamba-2 two practical benefits. First, it admits an efficient implementation that is friendly to tensor cores (the matrix-multiply units on modern GPUs), which the original Mamba's elementwise scan was not. Second, the SSD framing makes it natural to allow a much larger state dimension $N$ at roughly the same FLOPs — Mamba-2 routinely uses $N = 256$ where Mamba-1 used $N = 16$. A larger state means more memory for past tokens, which directly improves associative recall.
There is also an unexpected payoff: SSD reveals a structural connection between selective SSMs and a degenerate form of linear attention. Concretely, linear attention without softmax is a special case of an SSM with input-dependent $B$ and $C$ and an identity $A$. So in a real sense the SSM family and the linear-attention family are the same primitive viewed from different angles. The Transformer's quadratic softmax attention is the one that does not fit.
A pure Mamba is great at most things but has a real weakness: dense, content-addressed retrieval at very long context. The model has $O(N)$ state, but $N$ is a few hundred — not enough to losslessly remember tens of thousands of specific tokens. On "needle in a haystack" benchmarks where the answer requires copying a specific token from far back in the context, Transformers win.
The pragmatic response of the post-2023 community has been hybrids: stack Mamba layers (cheap, long-range) with a small number of attention layers (precise, content-addressed). The mixing ratios people have landed on are striking:
Jamba 1.5 (AI21, 52B active params): roughly 7 Mamba layers per 1 attention layer, in a Mixture-of-Experts setup. Handles 256k context with materially less KV cache than a pure Transformer.
Mamba-2-Hybrid (NVIDIA): only $\approx 10%$ of layers are attention; the rest are Mamba-2. Quality matches a same-size pure Transformer on most benchmarks at a fraction of the inference cost.
Falcon-Mamba (TII): a 7B pure Mamba model that competes with same-size Transformers on most benchmarks without any attention at all — but does worse than the hybrids on retrieval.
The empirical lesson is consistent: a little attention goes a long way. You can replace 90% of attention layers with Mamba and pay almost nothing in capability, while the 10% you keep buys you the specific content-addressed retrieval that pure Mamba lacks.
Now we can put SSMs on the same shelf as the other primitives and read off the trade-offs directly.
Primitive
Training cost
Inference per token
Persistent state
Long-range memory
Content-based routing
RNN (LSTM/GRU)
$O(N)$ sequential
$O(1)$
$O(d)$
Poor (vanishing grads)
Implicit, weak
CNN (1D)
$O(N \cdot k)$ parallel
$O(k)$
None (windowed)
Limited to receptive field
None
Transformer
$O(N^2 \cdot d)$ parallel
$O(N \cdot d)$ (KV-cache grows)
$O(N \cdot d)$ KV-cache
Excellent
Excellent (softmax attention)
S4 / S4D
$O(N \log N)$ parallel (FFT)
$O(1)$ with $O(N_{\text{state}})$ state
Read each row left-to-right and you can see why Mamba is exciting and also why no one has thrown out Transformers. Inference is constant-time per token regardless of context — that is a hard, qualitative advantage for streaming and long-document applications. Training is linear in $N$ rather than quadratic — that lets context length grow further than a Transformer at the same hardware budget. But content-based routing is soft and limited compared to attention, so for the tasks that actually depend on routing (retrieval, in-context learning, instruction-following at long context), Transformers and hybrids still dominate.
What Do You Think?
A Mamba model is asked to autoregressively generate one new token, given the existing context of length N. What is the asymptotic compute cost per generated token?
A useful decision rubric, distilled from the empirical results across Mamba-1, Mamba-2, Jamba, and the various hybrid papers:
Reach for Mamba / a Mamba-heavy hybrid when
The deployment cares about cost per generated token at long context. Mamba's $O(1)$ inference is decisive for streaming summarization, long-document QA agents, and any setting where the KV cache would dominate memory.
The task is modeling, not retrieval: language modeling, time-series forecasting, audio modeling, raw-waveform processing, DNA sequences.
Context is genuinely very long ($\geq 64{,}000$ tokens) and most of it is contextually relevant rather than acting as a literal store of facts to be queried.
Stay with a Transformer (or use only a little Mamba) when
The task is retrieval-heavy: needle-in-haystack, table lookup, code search by exact identifier. Pure Mamba loses here today.
The task relies heavily on in-context learning from few-shot examples. Transformers still outperform pure Mamba on standard ICL benchmarks at the same parameter count.
You need the strongest possible model at moderate context ($\leq 8{,}000$ tokens). Below the quadratic-attention pain point, the Transformer's content-based routing is free quality.
Use a hybrid when both pressures apply: long context and retrieval. The current empirical sweet spot is somewhere between 5% and 25% attention layers, with Mamba doing the heavy lifting on the rest. This is the architecture you should expect to see in most foundation models trained from 2025 onward.