From HIPPO polynomial projection to Mamba's selective state-space mechanism, Mamba-2's SSD duality with linear attention, parallel scan training, RWKV, and an honest accounting of where SSMs beat and where they lose to transformers.
Transformers quadratic attention has a 12-word summary: every token attends to every other token — O(T² · D) memory, O(T² · D) FLOPs. For sequences of length 1M this is prohibitive. The question SSM researchers asked: can we replace attention entirely with a recurrence that has O(T) inference cost and still captures long-range dependencies?
O(T²) compute and memory during training. Recurrence-free → all tokens computed in parallel. Unlimited receptive field at each layer. KV cache grows linearly at inference: O(T · D) per layer.
O(T) compute at inference: one state update per step, fixed-size state. Training requires either sequential recurrence (slow) or parallel scan (fast). Receptive field theoretically infinite but empirically limited by state size.
The key failure of earlier RNNs (LSTM, GRU) was not recurrence itself but vanishing gradients over long sequences, and the difficulty of learning which past information to keep. The SSM line of work addresses this through structured state-space theory and, in Mamba, through input-dependent selectivity.
The SSM community frames the problem around three modes of computation for a sequence model: recurrent (O(1) state, sequential), convolutional (efficient parallel training via FFT), and attention (full quadratic, optimal recall). S4 and Mamba trade attention-class recall for O(1) inference state. Mamba-2 proves a formal duality between SSMs and linear attention, collapsing two seemingly different research lines.
The Structured State Space sequence model (S4) (Gu et al., NeurIPS 2021) begins with a continuous-time linear dynamical system and discretises it for use on sequences. The state-space equations are:
# Continuous form:
x'(t) = A x(t) + B u(t)
y(t) = C x(t) + D u(t)
# Discretised (zero-order hold, step size Δ):
A¯ = exp(Δ A)
B¯ = (Δ A)-1 (exp(Δ A) - I) · Δ B
xt = A¯ xt-1 + B¯ ut
yt = C xtThe key insight of HIPPO (High-order Polynomial Projection Operators, Gu et al. 2020) is a mathematical construction for the A matrix such that the state vector x maintains an optimal polynomial approximation of the input history. Concretely, HIPPO-LegS projects the input onto Legendre polynomials with a sliding-window measure. This gives S4 its long-range memory at initialisation — before any gradient descent.
Given fixed A, B, C, the SSM recurrence unrolls to a convolution with kernel K = (C B¯, C A¯ B¯, C A¯² B¯, ...). This can be computed in O(T log T) via FFT during training — much faster than sequential recurrence. At inference the recurrence is used directly in O(1) per step.
S4 parameterises A as a diagonal-plus-low-rank (DPLR) matrix. Later work showed that purely diagonal A (S4D, DSS) gives comparable performance with simpler implementation. Mamba builds on this foundation.
Gu & Dao released Mamba in December 2023 (arXiv 2312.00752). The defining innovation: make B, C, and Δ functions of the input token. In S4, these are fixed matrices. In Mamba they are linear projections of the input:
# u: input sequence (B, L, D_in)
# Project to SSM dimension D_state:
B = linear(u) # (B, L, N) — state mixing matrix, per token
C = linear(u) # (B, L, N) — output mixing matrix, per token
Δ = softplus(linear(u)) # (B, L, D) — time-step, per token and channel
# A is fixed (diagonal, negative-real for stability)
# Discretise A¯, B¯ using Δ (same formula as S4)
A¯ = exp(Δ * A) # (B, L, D, N)
B¯ = Δ * B # simplified ZOH for diagonal A
# Selective scan: state update
ht = A¯t * ht-1 + B¯t * ut
yt = Ct · htIn S4, the state update A¯ is a fixed decay regardless of input content. Mamba's Δ lets the model “reset” the state on semantically important tokens (large Δ ≈ full state replacement) and “coast” on irrelevant tokens (small Δ ≈ pure copy). This is the mechanism that gives Mamba selective copying / associative recall, which S4 cannot do.
Dao & Gu (2024, arXiv 2405.21060) introduced Mamba-2 with the Structured State-Space Duality (SSD) framework. The central theorem: a specific class of SSMs is mathematically equivalent to a form of linear attention with a semiseparable matrix.
| View | Formulation | Efficient mode |
|---|---|---|
| SSM (recurrent) | ht = At ht-1 + Bt xt | O(TND) sequential |
| Linear attention | y = (L ⊙ QKT) V, L lower-triangular semiseparable | O(T²N) quadratic chunk |
| Hybrid chunk scan | Process intra-chunk as matmul, inter-chunk as recurrence | O(TC² + T/C · N) optimal |
Mamba-2 adopts a multi-head formulation (analogous to multi-head attention) over the SSM state: state dim N is factored as N = P × H where H is number of heads. This enables tensor-core-friendly matmul kernels.
At is reduced to a scalar per head per token (not a full D×N matrix). This simplification, which is key to the SSD duality, reduces parameters but costs little in practice: the scalar gives per-head forgetting rates.
The duality means Mamba-2 can be trained in chunk mode using the same CUDA matmul kernels as flash attention, with chunks of size C=64 or C=128. Intra-chunk is a dense attention; cross-chunk is recurrent SSM. This gives near-linear scaling with a ~1.8× throughput improvement over Mamba-1 at the same parameter count.
The selective SSM recurrence is fundamentally sequential: ht depends on ht-1. Naively this means O(T) kernel launches on a GPU — catastrophic for training throughput. Two techniques unlock parallelism:
The SSM recurrence is a linear recurrence: h_t = a_t * h_{t-1} + b_t. Linear recurrences can be solved in O(log T) depth via a parallel prefix (scan) computation, with O(T) total work — like a parallel prefix sum. On GPU this occupies all streaming multiprocessors simultaneously.
# Up-sweep (reduce phase): combine pairs → T/2 → T/4 → ... → 1
for d in range(log2(T)):
for k in range(0, T, 2d+1):
x[k + 2d+1 - 1] = combine(x[k + 2d - 1], x[k + 2d+1 - 1])
# Down-sweep (distribute phase): reverse direction
for d in reversed(range(log2(T))):
for k in range(0, T, 2d+1):
x[k + 2d - 1], x[k + 2d+1 - 1] = \
x[k + 2d+1 - 1], combine(...)Mamba fuses the discretisation, the scan, and the output projection into a single CUDA kernel. This avoids writing O(TND) intermediate states to HBM. All intermediate values stay in SRAM (L2 cache). The fused kernel is the selective_scan_cuda function in the Mamba repo (mamba_ssm/ops/selective_scan_interface.py).
Rather than storing all intermediate hidden states for the backward pass, Mamba recomputes them during backprop (similar to gradient checkpointing). This trades FLOPs for memory — acceptable because the scan FLOPs are cheap compared to attention.
RWKV (Receptance Weighted Key Value, Peng et al. 2023, arXiv 2305.13048) arrived at a similar destination from a different direction: starting from linear attention, not state-space models. RWKV-4 and RWKV-6 are fully open-weights models up to 14B parameters, trained by the EleutherAI-adjacent RWKV foundation.
| Version | Params | Key change |
|---|---|---|
| RWKV-4 (2023) | up to 14B | Time-mixing + channel-mixing; additive recurrence |
| RWKV-5 Eagle (2024) | up to 7B | Multi-headed time-mix; improved long-context recall |
| RWKV-6 Finch (2024) | up to 14B | Data-dependent receptance/key; closes recall gap with Mamba |
| RWKV-7 (2025) | up to 7B | Dynamic state update; state transform generalises SSM and linear attn |
# w: channel-wise learned decay (negative, time-independent)
# u: per-channel “bonus” for current token
# k, v, r: key, value, receptance projections from input
wkvt = (exp(u + kt) · vt + sums<t exp((t-1-s)w + ks) · vs) /
(exp(u + kt) + sums<t exp((t-1-s)w + ks))
ot = sigmoid(rt) · layer_norm(wkvt)RWKV's time-mixing uses a fixed per-channel decay w (learned but input-independent), making it analogous to S4 rather than Mamba-1. RWKV-6 introduced data-dependent w, closing the gap. The two projects converged on the same core insight independently: selectivity is required for strong recall benchmarks.
The empirical case for SSMs is clearest in domains where O(1) inference state and linear compute matter more than per-token recall precision.
At sequence length 1M, a 3B Mamba model processes tokens at 5× the throughput of a 3B transformer (no KV cache growth). This makes Mamba attractive for continuous-stream applications: transcription, genome analysis, long document streaming.
Transformer KV cache grows as O(T · L · D · 2). Mamba state is O(L · D · N) — constant in T. For a 1M-token context, this is the difference between ~16 GB KV cache and ~200 MB recurrent state.
Mamba was applied to raw audio (8 kHz, 50k-token sequences) and DNA sequences (human genome ~3B bases). Both are long, with local structure. SSMs handle these without chunking tricks that attention requires.
Single-token generation latency for Mamba is O(1) (one state update). A transformer must compute attention over the full KV cache. At 4K+ context, Mamba is measurably faster per generated token on a single GPU.
Jamba (arXiv 2403.19887) is a 52B hybrid model alternating Mamba and attention layers in a 1:7 ratio — one attention layer per 7 Mamba layers. It fits on a single 80GB GPU at BF16 (impossible for a pure-transformer of equivalent quality) while matching Mixtral 8x7B on standard NLP benchmarks. Covered in Arch 05.
SSMs have known, measurable failure modes. Being honest about these is as important as the wins.
The standard synthetic benchmark for in-context copying. Input: a sequence of (key, value) pairs, then query keys. Task: retrieve the associated values. Mamba-1 underperforms transformers on MQAR at comparable model sizes. Mamba-2 and hybrid architectures close the gap substantially.
Finding a specific fact buried in a 100K-token document. Pure Mamba models degrade noticeably beyond 32K tokens on fact retrieval tasks. Transformers with proper positional encodings (YaRN/RoPE) remain flat to 128K+. This is the core argument for hybrid architectures.
At model sizes ≥7B and training tokens ≥1T, transformer perplexity on standard language modelling benchmarks (The Pile, C4) is still ~5% lower than Mamba at matched FLOPs. The gap narrows with Mamba-2 and hybrid models but has not been fully closed as of Q1 2025.
Most RLHF pipelines (DPO, PPO, GRPO) have been developed for transformers. Mamba's non-differentiable state resets make reward-model fine-tuning and PPO value estimation subtler. Very little published work on post-training MoE-free Mamba models at 7B+.
Pure-SSM models are not yet a drop-in replacement for transformers in general-purpose chat and coding. Hybrid SSM-attention (Jamba, Zamba, Zephyr-Mamba) looks more promising: attention layers handle the high-recall tasks, Mamba handles the long-range cheap computation. Arch 05 covers these hybrid designs in detail.
Deck 03 addresses the transformer's own long-context problem: how positional encodings (RoPE, YaRN, ALiBi) and attention sparsity (sliding window, ring attention) allow transformers to reach 1M+ tokens without the architectural change that SSMs represent.