Interactive Explainer

Mamba & State-Space Models, Visualised

A linear recurrence over a hidden state. Watch the state evolve as it processes a 256-step signal. Toggle the input-dependent Δ that gives Mamba its selectivity edge. See why O(N) is the right complexity for million-token contexts.

~18 min SSM · Mamba · Long Context

Prelude

The long-context bottleneck

Self-attention is $O(N^2)$ in sequence length. At 1k tokens that's fine; at 100k it's the dominant cost; at 1M it's impossible. State-space models (S4, S5, H3, Mamba) replace attention with a linear recurrence whose per-token cost is $O(1)$ regardless of sequence length.

The slogan. A discrete-time SSM is a fancy RNN with carefully chosen dynamics that lets you compute many tokens in parallel via a "scan". It keeps the linear-in-length cost of an RNN and regains the parallelism of attention.

Step 1

The recurrence

$h_t \in \mathbb{R}^d$ is the hidden state. $\bar A, \bar B$ are the discretised dynamics; $C$ is the output projection. For S4 / S5 these are fixed (or learned but input-independent). The big Mamba idea: make $\bar A, \bar B$ depend on the current input $x_t$. That single change gives "selective" SSMs that can decide which inputs to carry forward and which to forget.

Δ is the per-token step size; tiny Δ means "barely update the state" (good for skipping uninformative inputs); large Δ means "treat this input as important". Both are learned.

Step 2

Watch the state evolve

A 1-channel discrete SSM with $A = -0.2$ (slow decay), $B = C = 1$, and a Δ that you can either fix or modulate from the input. The bottom panel shows $|h_t|$ across the 256-step sequence; the top shows the input.

signal fixed Δ 0.30 selective Δ (Mamba)

Top: input x_t. Bottom: |h_t| (state magnitude). With selective Δ, the state ignores low-amplitude inputs and latches onto sparse important events.

Step 3

Cost vs sequence length

Mamba is linear in length; attention is quadratic. Below: synthetic operation count for both vs sequence length. The break-even point depends on hidden width, but for any sufficiently long sequence, attention loses.

Operation counts (log-log). Attention's $O(N^2)$ vs Mamba scan's $O(N \cdot d_{\text{state}})$.

Step 3 ½

The parallel scan — the trick that makes SSMs trainable at scale

A linear recurrence $h_t = A h_{t-1} + B x_t$ looks inherently sequential — each step depends on the previous. RNNs sit in this trap: their training is $O(N)$ wall-clock because each layer waits for the previous token's hidden state. SSMs escape via the parallel scan (Blelloch, 1990; Smith et al., 2023): an associative tree-reduction that computes all $h_1, \dots, h_N$ in $O(\log N)$ depth.

The trick: the operator $(c_1, M_1) \oplus (c_2, M_2) = (M_2 c_1 + c_2,\; M_2 M_1)$ is associative. Pair-wise reduce neighbours, then pair-wise reduce neighbours-of-neighbours, and so on. Each level halves the count; $\log_2 N$ levels finish the scan. On a GPU this matches matmul throughput.

Parallel-scan tree for N=8. Step 1 combines (1,2), (3,4), (5,6), (7,8). Step 2 combines (1:2, 3:4) and (5:6, 7:8). Step 3 combines (1:4, 5:8). After log₂N levels, every prefix is computed.

sequence length N 8 animate sequential RNN comparison

Wall-clock depth: sequential RNN is $O(N)$; parallel scan is $O(\log N)$. On a 1M-token context that's 1M sequential steps vs 20.

Why this is the SSM superpower. The parallel scan is the only reason SSMs (S4, S5, Mamba) can train as fast as Transformers despite being recurrent. Without it, "linear in length" stays a paper claim.

Step 4

Where each wins

Attention wins on: short sequences, tasks needing exact long-range copy / lookup, in-context learning of arbitrary patterns. Most LLM tasks at sub-32k context.
Mamba wins on: very long contexts (DNA, audio, long documents, time series). Linear scaling stops being a research detail and starts being the difference between feasible and not.
Hybrids win: Mamba-2, Jamba, Nemotron-H, Zamba interleave SSM and attention layers and beat the pure variants of either. The 2026 default.
Time-series forecasting. Mamba and S5 currently trade blows with PatchTST on long horizons. For data with hour / minute resolution stretching across years, this is the natural backbone.

Final takeaway. The right way to think of an SSM is "a parallel-friendly RNN with hand-tuned eigenvalues". The right way to think of Mamba on top is "the SSM gets to look at its input before deciding how much to update". Together they remove the $N^2$ ceiling on context length without sacrificing training speed.