State space models and Mamba

Replace attention with structured state space models (S4, Mamba) for linear-time sequence mixing.

⏱ 26 minPrerequisites:The Transformer

Self-attention is $O(N^2)$ in sequence length and quickly becomes a bottleneck. State space models (SSMs) offer a different sequence-mixing operation that runs in $O(N)$ time and $O(N)$ memory while matching or beating Transformer quality on long-context tasks. The line runs from continuous-time linear systems through S4 to Mamba, which is the modern hardware-aware variant powering long-context language models.

The starting point: linear state space systems

A continuous-time linear SSM maps input $u(t) \in \mathbb{R}$ to output $y(t) \in \mathbb{R}$ via a hidden state $h(t) \in \mathbb{R}^d$ :

\dot h(t) = A\, h(t) + B\, u(t), \qquad y(t) = C\, h(t) + D\, u(t).

Discretize with step size $\Delta$ :

h_k = \bar A\, h_{k-1} + \bar B\, u_k, \qquad y_k = C\, h_k.

This is a linear recurrent network: hidden state evolves linearly, output is a linear projection. For input sequence $u_1, \dots, u_N$ , the output can be computed:

Recurrently: $O(N)$ time, $O(d)$ memory per step.
As a convolution: $y = K * u$ for an appropriately defined convolution kernel $K$ . Used for parallel training.

So SSMs are both linear RNNs (for sequential inference) and convolutions (for parallel training).

S4 (Structured State Space)

Gu et al. (2022)[1]. The state matrix $A$ is parameterized as a structured matrix (HiPPO matrices) that gives the SSM a useful prior over long-range dependencies.

Key contributions:

HiPPO initialization: $A$ is initialized to project the input history onto Legendre polynomials, giving the model a built-in inductive bias for compressing the past.
Diagonal parameterization: with the right basis, $A$ is diagonal, and the convolution can be computed efficiently in the frequency domain.
Long Range Arena benchmark wins: S4 dominated on synthetic long-context tasks where Transformers struggled.

S4 was the proof of concept that an alternative to attention could compete with Transformers on long-context modeling.

Mamba (2024)

Gu and Dao[2]. The breakthrough that made SSMs competitive with Transformers for language modeling.

Two innovations:

Selective SSM: the state matrices $A, B, C$ become functions of the input, not fixed parameters. The model selects what to remember based on the current token. This is the SSM analog of attention's content-based addressing.
Hardware-aware kernel: a custom CUDA kernel that fuses the SSM computation, exploits the GPU memory hierarchy, and avoids materializing the hidden state for the whole sequence.

The resulting Mamba block is a drop-in replacement for the attention sublayer of a Transformer:

x = x + Mamba(LayerNorm(x))
x = x + FFN(LayerNorm(x))

Mamba achieves Transformer-quality language modeling at sequence lengths where attention becomes infeasible (100K+ tokens), with constant-time per-token inference (no growing KV cache).

Mamba-2 (2024)

Refinement of Mamba with:

State space duality (SSD): a unified theoretical framework showing SSMs and attention are special cases of a more general family.
Larger state dimensions: efficient via the SSD reformulation.
Improvements in modeling quality at scale.

The Mamba-2 paper showed that careful SSM design can match the parameter count vs perplexity curve of well-tuned Transformers.

Why SSMs might (or might not) replace attention

Strengths:

Linear scaling in sequence length (training and inference).
Constant inference memory: no growing KV cache.
Strong long-range modeling: HiPPO and selective mechanisms.
Hardware-friendly: dense matrix-vector products, cache-friendly.

Weaknesses or open questions:

Local copy/lookup tasks that attention nails: SSMs sometimes underperform. Hybrid architectures (next-to-last D3 lesson) address this.
In-context learning scaling laws: less well-studied than Transformers.
Less mature tooling: the Transformer ecosystem (HuggingFace, PyTorch primitives, training recipes) is decades ahead.

Mamba and its descendants are viable; whether they replace attention or coexist alongside it is an open question as of 2024.

The relationship to RNNs

SSMs are linear RNNs with a specific parameterization. Compared to LSTM/GRU:

Linear vs gated nonlinear updates.
Structured matrices with mathematical properties (HiPPO) vs general gates.
Parallel training via convolution view vs sequential.

The selective version (Mamba) adds back input-dependence, recovering some of the gating flexibility of LSTM while preserving the parallel training story.

An SSM's convolution kernel is a sum of decaying modes λᵏ: eigenvalues near 1 decay slowly and carry long-range memory, while small ones forget almost immediately.

Exercises

Exercise 1. Unroll the discrete SSM recurrence $h_k = \bar A h_{k-1} + \bar B u_k$ , $y_k = C h_k$ (with $h_0 = 0$ ) and show that the output is a convolution $y_k = \sum_{j=0}^{k-1} K_j\, u_{k-j}$ . Give the convolution kernel $K_j$ in terms of $\bar A, \bar B, C$ .

Solution

Step 1. Apply the recurrence repeatedly from $h_0 = 0$ :

h_k = \sum_{i=1}^{k} \bar A^{\,k-i}\, \bar B\, u_i.

Each input $u_i$ enters at step $i$ and is propagated forward by $k - i$ applications of $\bar A$ .

Step 2. Project to the output with $C$ :

y_k = C h_k = \sum_{i=1}^{k} C \bar A^{\,k-i} \bar B\, u_i.

Step 3. Re-index with $j = k - i$ (so $i = k - j$ , $j$ runs from $0$ to $k-1$ ):

y_k = \sum_{j=0}^{k-1} \left(C \bar A^{\,j} \bar B\right) u_{k-j}.

This is a convolution with kernel $K_j = C \bar A^{\,j} \bar B$ . The kernel is fixed (input-independent) for a non-selective SSM, which is exactly why training can run as a single convolution in parallel rather than stepping the recurrence.

Exercise 2. Take the SSM diagonal so $\bar A = \operatorname{diag}(\lambda_1, \dots, \lambda_d)$ with each $|\lambda_i| < 1$ . Show that the kernel $K_j = C \bar A^{\,j} \bar B$ is a sum of decaying geometric modes $\lambda_i^{\,j}$ , and explain what an eigenvalue near 1 versus near 0 does to the model's memory. (This is the behavior the lesson's figure plots.)

Solution

Step 1. With $\bar A$ diagonal, $\bar A^{\,j} = \operatorname{diag}(\lambda_1^{\,j}, \dots, \lambda_d^{\,j})$ . Writing $C = (c_1, \dots, c_d)$ as a row and $\bar B = (b_1, \dots, b_d)^\top$ as a column,

K_j = C \bar A^{\,j} \bar B = \sum_{i=1}^{d} c_i\, \lambda_i^{\,j}\, b_i.

The kernel at lag $j$ is a weighted sum of the powers $\lambda_i^{\,j}$ , one geometric sequence per state coordinate.

Step 2. Each mode $\lambda_i^{\,j}$ decays geometrically in the lag $j$ because $|\lambda_i| < 1$ . The decay rate is set by $|\lambda_i|$ : the half-life of mode $i$ is the lag $j$ at which $|\lambda_i|^{\,j} = \tfrac12$ , namely $j = \log(1/2) / \log|\lambda_i|$ .

Step 3. An eigenvalue with $|\lambda_i|$ close to 1 decays slowly, so its mode keeps a large weight at large lag $j$ : that coordinate carries long-range memory, mixing inputs from far in the past into the current output. An eigenvalue close to 0 decays almost immediately ( $\lambda_i^1 \approx 0$ ), so its mode forgets after one or two steps and contributes only local context. The spectrum of $\bar A$ is the model's memory spectrum, which is what HiPPO initialization is chosen to shape.

Exercise 3. Mamba makes $B$ and $C$ (and the step $\Delta$ ) functions of the input token, calling the result a selective SSM. Explain why this input-dependence breaks the convolution view of Exercise 1, and what the model gives up and gains by accepting that. Contrast the asymptotic inference cost of a selective SSM with softmax attention for generating a sequence of length $N$ .

Solution

Step 1. The convolution view of Exercise 1 relied on a fixed kernel $K_j = C \bar A^{\,j} \bar B$ , the same at every position. When $\bar B, C, \Delta$ depend on the input token $u_k$ , the effective kernel changes from step to step, so there is no single $K$ to convolve with: the output is no longer a convolution of the input with a position-independent filter.

Step 2. What it gives up: the simple parallel-convolution training path. Mamba recovers parallelism with a hardware-aware parallel-scan kernel instead of an FFT convolution, which is more engineering but still $O(N)$ work. What it gains: content-based selection. Like attention's content addressing, the model can choose what to write into and read from its state based on the current token, which fixes copy/lookup tasks that a fixed-kernel SSM handles poorly.

Step 3. Inference cost for length $N$ . A selective SSM keeps a fixed-size hidden state $h_k \in \mathbb{R}^d$ and updates it in $O(d)$ per token, so generating $N$ tokens is $O(Nd)$ time with $O(d)$ memory: constant per-token cost, no growing cache. Softmax attention stores a KV cache that grows to length $N$ and scores each new query against the whole cache, an $O(Nd)$ per-token cost and $O(N^2 d)$ total, with $O(Nd)$ cache memory. The SSM's constant per-token cost is the long-context advantage.

The big picture

The next lesson covers another path to sub-quadratic sequence mixing: linear attention, which keeps the QKV machinery but replaces the softmax with a kernel feature map that lets the computation be reordered to linear cost.

References

[1]Gu, A., Goel, K., & Ré, C. (2022). "Efficiently modeling long sequences with structured state spaces (S4)". ICLR. arXiv:2111.00396. HiPPO-structured SSMs with the convolution/recurrence duality.↩
[2]Gu, A., & Dao, T. (2023). "Mamba: linear-time sequence modeling with selective state spaces". arXiv:2312.00752. Input-dependent (selective) SSM with a hardware-aware kernel.↩