8.7 — State Space Models (Mamba) — brief

1. Why this session exists

Every architecture in Chapter 8 so far is a Transformer with one sub-module modified. State space models (SSMs), and their leading instantiation Mamba (Gu & Dao, 2023[3], arXiv:2312.00752), are different: they replace attention entirely with a recurrent sequence operator. Per token they need no KV cache, and per generated token they run in $O(d^2)$ time instead of attention's $O(T d)$ (with the cached-K, V matmul). For long-context inference this is a substantial wall-clock advantage.

The question is whether SSMs can match attention on quality at scale. As of 2026 the answer is: pure SSMs are competitive but not strictly better than Transformers at the frontier. Hybrid architectures (§8.8) that combine SSM and attention layers are an active research direction and have produced models that match dense Transformers at iso-compute. So SSMs matter both as a standalone variant and, more importantly, as the recurrent half of every hybrid model.

This session is brief on purpose. We define the continuous SSM, discretise it to a recurrence, motivate the selective (input-dependent) version that became Mamba, count its compute and memory, and place it in the 2026 architectural landscape. The Mamba paper is dense; we will skim its hardware-aware "parallel scan" details and refer the reader to the paper for the kernel-level implementation.

2. The continuous-time SSM

Fix state dimension $n$ and signal dimension (here we treat single-channel input; multi-channel SSMs are a per-channel collection of single-channel models). A linear time-invariant state space model is the continuous-time system

with state $h(t) \in \mathbb{R}^n$, input $x(t) \in \mathbb{R}$, output $y(t) \in \mathbb{R}$, and matrices $A \in \mathbb{R}^{n \times n}, B \in \mathbb{R}^{n \times 1}, C \in \mathbb{R}^{1 \times n}$. This is the canonical control-theory SSM (Kalman 1960; not specific to deep learning)[1].

The matrix $A$ defines the intrinsic dynamics of the hidden state — how it evolves in the absence of input. $B$ defines how new input enters the state. $C$ defines what part of the state is observed as output.

2.1 The convolutional view

Solving the ODE (assuming $h(0) = 0$):

So $y$ is the convolution of $x$ with the kernel $K(\tau) = C e^{A \tau} B$. The kernel is determined entirely by $(A, B, C)$. This is the operative observation that links SSMs to convolutions and, by Fourier, to fast algorithms.

3. Discretisation

To use SSMs on discrete token sequences, we discretise the ODE with step size $\Delta$. The Zero-Order Hold (ZOH) discretisation, which Mamba uses:

so the recurrence becomes

This is a linear recurrent network. Given the state $h_{t-1}$ and input $x_t$, the new state $h_t$ requires one $n \times n$ matrix-vector multiply (cost $n^2$) plus one $n \times 1$ matrix-vector multiply (cost $n$); the output requires one $1 \times n$ multiply (cost $n$). Per-step cost: $O(n^2)$.

The crucial structural difference from a standard RNN (5.5) is that the recurrence is linear — no per-step non-linearity inside the state update. Non-linearity is added externally (a SwiGLU-like block applied to $y_t$). Linearity is what enables the parallel-scan formulation of §5.

3.1 Structured $A$: from S4 to Mamba

The naive $n \times n$ matrix $A$ has $n^2$ free parameters per channel, which is impractical at the channel counts ($d_\text{model} = 4096$, etc.) of LLM-scale models. The S4 paper (Gu et al., 2022, arXiv:2111.00396)[2] introduced structured $A$ matrices — diagonal-plus-low-rank — that have $O(n)$ parameters and admit fast convolutional evaluation via FFTs.

Mamba uses a simpler structured form: $A$ is diagonal with negative real entries, so the recurrence is a per-state-coordinate exponential decay:

This has only $n$ parameters per channel, and each coordinate of the state evolves independently: $h_t^{(i)} = e^{\Delta a_i} h_{t-1}^{(i)} + (\text{input term})_i$. Different coordinates capture different timescales of state retention — small $|a_i|$ means slow decay (long memory), large $|a_i|$ means fast decay (short memory). Empirically, training discovers a spread of timescales without needing explicit constraints.

4. Selective SSM: Mamba

The S4 line gave good signal-processing models but lost to Transformers on language tasks. The Mamba paper's diagnosis: a time-invariant SSM (with fixed $A, B, C$) cannot do content-dependent state updates. For language, the model often needs to decide "keep this token's information in state" vs "forget it" based on the token itself — what attention's content-based addressing does, and what a content-blind SSM cannot.

Selective SSM (Mamba's contribution): make $\Delta, B, C$ functions of the input.

The recurrence becomes

with $\bar B_t$ the input-dependent discretised $B_t$. The matrix $A$ stays time-invariant (so the eigenstructure is preserved); only the coupling into and out of state, and the time step, vary per token.

The interpretation: $\Delta_t$ controls "how much time has passed" between tokens — large $\Delta_t$ means strong decay of past state (fast forgetting), small $\Delta_t$ means weak decay (long memory). $B_t$ controls how the current token's content enters the state — the model can choose to gate the token in (large $B_t$) or skip it (small $B_t$). $C_t$ controls what part of the state is read out at this step.

This is the SSM analogue of attention's content-addressed gathering: per token, the model dynamically chooses what to remember, how strongly, and what to output.

4.1 Why this destroys the convolutional structure

Time-invariant SSMs are convolutions (§2.1), so they can be computed in $O(T \log T)$ via FFT. Selective SSMs are time-varying convolutions — the kernel depends on the input — and admit no direct FFT trick.

Mamba's solution is the parallel scan: the linear recurrence $h_t = \bar A_t h_{t-1} + \bar B_t x_t$ is an "associative scan" over the sequence (the operator $(h, c) \mapsto \bar A c + h$ is associative when you collect the $\bar A_t, \bar B_t, x_t$ correctly into pairs). Associative scans can be computed in $O(\log T)$ depth with $O(T)$ work using the standard parallel-prefix-sum algorithm. The Mamba paper implements this with a hardware-aware CUDA kernel that minimises HBM↔SRAM transfers, similar in spirit to FlashAttention but for the scan operation.

Per-token training compute: $O(d \cdot n)$ for state update, $O(d \cdot n)$ for output projection — same per-token cost as a Transformer FFN of width $d$ (no quadratic-in-$T$ term, because there is no attention matrix to compute).

5. Compute and memory: the headline comparison

Per-layer cost comparison for a single token at position $T$:

| Operator | Train: FLOPs per token | Inference: FLOPs per generated token | KV cache | |---|---|---|---| | MHA (with FlashAttention) | $\Theta(T d)$ | $\Theta(T d)$ | $4 d T$ bytes | | GQA-$8$ | $\Theta(T d)$ | $\Theta(T d)$ | $4 (d/8) T$ bytes | | MLA | $\Theta(T d)$ | $\Theta(T d + r d)$ | $2 (r + d_h^R) T$ bytes | | Selective SSM | $\Theta(d n)$ | $\Theta(d n)$ | $0$ (state is $d n$ scalars, fixed not growing) |

The headline: SSM compute per generated token is constant in $T$. There is no cache that grows with sequence length. The hidden state has fixed dimension $d \cdot n$, independent of $T$.

This is the structural reason for SSMs' interest in long-context generation:

• At $T = 1{,}000$, a Transformer reads $\sim 10$ MB of KV cache per generated token (LLaMA-7B-class). An SSM reads $\sim 100$ KB of state. The SSM is $\sim 100\times$ faster at the memory-bandwidth level.
• At $T = 100{,}000$, the Transformer reads $\sim 1$ GB per generated token. The SSM still reads $\sim 100$ KB. The SSM is $\sim 10{,}000\times$ faster.
• The crossover point at which an SSM is "free" is essentially at $T = 0$; the constant cost is small enough that SSMs are competitive with attention even at short sequences.

5.1 But quality?

The Mamba paper's main empirical claim: at iso-parameters, Mamba matches Transformer perplexity on Pile-style pretraining and outperforms on long-context benchmarks (where the Transformer's effective context degrades). Subsequent work (Mamba-2, Dao & Gu 2024[4]; Hymba; etc.) has confirmed and extended this at the 1B–3B scale.

But at the frontier ($\sim 100$B+ activated parameters), pure SSMs have not yet matched dense Transformer LLMs on standard benchmarks. The leading hypotheses for the residual gap:

1. In-context retrieval. Attention's content-addressed lookup is exact (or nearly so, modulo softmax); an SSM's state is a lossy compression of the past. Tasks that require pulling a specific past token verbatim (e.g., needle-in-haystack) are easier for attention.
2. Heuristic inductive biases of attention. Decades of attention-related architectural tuning have given the field a host of small improvements (RoPE, GQA, sliding-window) that SSMs have not yet seen full attention to.
3. Training stability. SSMs at very large scale have known numerical-stability issues that the field has not fully worked through.

The current state of the art is to combine attention and SSM, which is the subject of §8.8 (hybrids).

6. The block

A Mamba block, in shape, is

x ── RMSNorm ── linear up-proj (to 2*d) ── split ──┐
                                                    │
       ┌── Selective SSM (with selective Δ, B, C) ──┤
       │                                            │
       └── SiLU ──────────── elementwise * ────────┘──── linear down-proj (back to d) ── + (residual) ── x_out

In words: pre-norm, up-project to $2d$, split into two halves of width $d$. Apply the selective SSM to one half. Apply SiLU to the other half. Elementwise-multiply (a gating like SwiGLU). Down-project back to $d$. Residual add. There is no separate FFN block: the SSM block already includes its own non-linearity and gating, so a Mamba layer is one block (not the two of attention + FFN in a Transformer).

This means a Mamba "layer" at fixed parameter count corresponds to about 2× the number of layers in a Transformer with the same parameters. Mamba-3B has 64 layers; a Transformer-3B has $\sim 32$. The depth doubles but each layer is cheaper.

7. What we keep going forward

1. An SSM is a linear (or near-linear) recurrent operator $h_t = \bar A h_{t-1} + \bar B x_t$, $y_t = C h_t$. Discretised from a continuous-time linear ODE. With diagonal $A$, the state coordinates evolve at independent timescales.

2. Selective SSM (Mamba) makes $\Delta, B, C$ input-dependent. This is the SSM analogue of attention's content-addressed lookup. Time-invariant SSMs (S4) lost to Transformers on language; selective SSMs are competitive.

3. SSMs have $O(1)$ KV-like state, not $O(T)$. Decode per token is constant in $T$, vs Transformer's $O(T)$. The wall-clock advantage at long $T$ is enormous.

4. At iso-parameters, Mamba matches Transformer at $\lesssim 3$B and underperforms slightly at frontier scale. The leading hypothesis is that exact in-context retrieval needs attention; a lossy state compression does not suffice.

5. The 2026 frontier uses hybrids. Interleave Mamba and attention layers (Jamba, Samba) to combine the per-token-throughput advantage of SSMs with the recall accuracy of attention. This is the subject of §8.8.