Technical Deep Dives

State Space Models and the Post-Transformer Era

How state space models went from S4 and Mamba to Mamba-3's complex-valued states and MIMO, and why the 2026 answer is hybrid: Nemotron 3, Granite 4.0 and Qwen3-Next all keep a thin minority of attention layers.

Gabriele Masetti · 2026-08-17

The Quadratic Wall

Every transformer decoder pays the same tax: computing self-attention over a context of length n costs O(n²) time and memory during prefill, because every token attends to every other token. Autoregressive decoding is usually reformulated with a key-value (KV) cache so each new token costs O(n) rather than recomputing the full attention matrix, but the trick trades compute for memory: the cache grows linearly with context and must be read from GPU HBM at every step, making long-context decoding memory-bandwidth bound rather than compute bound.

For a mid-size model running at a context of 128K tokens, the KV cache alone can run into tens of gigabytes. Doubling the context does not just double the work; it changes which resource you run out of first. That is the specific, mechanical problem — not "transformers are slow" in the abstract, but a quadratic prefill cost plus a linearly growing, bandwidth-hungry cache — that state space models were built to route around.

From Convolutions to State Spaces: S4

The state space model (SSM) lineage traces to HiPPO (Gu, Dao, Ermon, Rudra and Ré, NeurIPS 2020), a framework for compressing a continuous signal's history online into a fixed-size vector by projection onto a polynomial basis. Its contribution was a specific matrix, HiPPO-LegS, for updating that compressed memory so it captures long-range dependencies without the vanishing gradients that plague vanilla RNNs.

Albert Gu, Karan Goel, and Christopher Ré built on this with S4 — "Efficiently Modeling Long Sequences with Structured State Spaces" (arXiv:2111.00396, ICLR 2022). S4 treats a sequence as the discretization of a continuous-time linear dynamical system, x'(t) = Ax(t) + Bu(t), y(t) = Cx(t) + Du(t), with A initialized from HiPPO.

What makes it practical is that a discretized linear time-invariant (LTI) SSM is mathematically equivalent to a global convolution: because A, B and C are fixed across timesteps, the whole output sequence computes as one long convolution kernel applied via FFT, in O(n log n) rather than O(n²).

S4 reported state-of-the-art results across every task in the Long Range Arena benchmark, 91% accuracy on sequential CIFAR-10 with no data augmentation, and roughly 60x faster generation than comparable autoregressive models. SSMs were not a theoretical curiosity: they could beat transformers on tasks designed to test long-range reasoning, while costing less.

But LTI is also S4's ceiling. Because A, B and C don't change with the input, the model applies identical dynamics to every sequence regardless of content — it can weight recent against distant tokens, but cannot decide, token by token, "this one matters, forget the rest." That is a content-based capability transformers get for free from query-key dot products, and S4 lacks it.

The Selection Mechanism: Why Mamba Works

Mamba (Gu and Dao, "Mamba: Linear-Time Sequence Modeling with Selective State Spaces," arXiv:2312.00752, 2023) closes that gap by making the SSM parameters input-dependent. Instead of fixed A, B, C, Mamba computes the discretization step size Δ, the input matrix B, and the output matrix C as functions of the current token via learned linear projections.

That is the "selection mechanism": the model can now modulate, per timestep, how much of the incoming token gets written into the recurrent state and how much of the state gets read out, conditioned on content rather than position alone. Concretely, a large Δ lets the model reset and focus on the current input (behaving more like it's ignoring history), while a small Δ preserves prior state — and because Δ, B, C are now functions of x_t, the model learns when to do each.

The cost of selectivity is that the model is no longer time-invariant, so it loses the convolution trick that made S4 fast — a naive implementation runs the recurrence step by step, which is slow on GPUs and memory-hungry if the expanded (batch, length, hidden dimension, state dimension) tensor is materialized in HBM.

Mamba's second contribution is a hardware-aware algorithm that keeps it practical: a parallel associative scan restores parallelism across the sequence dimension despite the recurrence, and fused kernels compute the large intermediate state tensor in on-chip SRAM rather than writing it to HBM, recomputing activations on the backward pass rather than storing them.

The net result is a model with no attention and no MLP blocks that scales linearly in sequence length, reportedly delivers around 5x higher inference throughput than similarly sized transformers, and matches or exceeds transformer quality up to a few billion parameters in the original paper's evaluations.

Mamba-2 and Structured State Space Duality

Mamba-2 (Dao and Gu, "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality," arXiv:2405.21060, ICML 2024) doesn't just iterate on Mamba's engineering — it establishes a formal equivalence between a restricted class of SSMs and attention. The paper shows that when the state-transition matrix is constrained to be a scalar times the identity (rather than a general diagonal), the resulting SSM is mathematically the same computation as masked linear attention with a specific 1-semiseparable causal mask.

That equivalence — structured state space duality (SSD) — means the identical sequence transformation admits two algorithmic realizations: the linear-time O(n) recurrent scan, or an O(n²) matmul-based attention-like form.

The practical payoff is architectural and systems-level. Because the quadratic form reduces to block matrix multiplications, Mamba-2's core layer can lean on tensor cores the way transformer attention does rather than depending solely on the sequential scan — the paper reports the SSD-based layer running 2–8x faster than Mamba-1's selective scan while remaining competitive with transformers on language modeling. SSD also imports transformer training techniques — tensor parallelism, sequence parallelism, variable-length sequence handling — into the SSM world, narrowing the infrastructure gap that made transformers easier to train at scale.

Mamba-3: Complex States and MIMO

The lineage did not stop at SSD. Mamba-3 (Lahoti, Li, Chen, Wang, Bick, Kolter, Dao and Gu, arXiv:2603.15569, submitted March 16, 2026) is written from what its authors call an inference-first perspective, and targets the two things linear models had been trading away: state tracking, and hardware efficiency as opposed to the theoretical kind.

Three changes carry the paper. A more expressive recurrence derived from the SSM discretization itself replaces Mamba-2's; a complex-valued state update rule gives the recurrence richer state tracking than a real-valued one can express; and a multi-input, multi-output formulation widens the read and write paths into the state without lengthening decode latency.

At the 1.5B scale Mamba-3 improves average downstream accuracy by 0.6 percentage points over the next best linear model, Gated DeltaNet, and the MIMO variant adds another 1.2 points for a 1.8-point total. Across state-size experiments it matches Mamba-2's perplexity using half the state — the number that matters for deployment, because state size is what a linear model spends instead of a KV cache.

Hybrids: Mixing Attention Back In

If selective SSMs are so efficient, why not go all-in and drop attention entirely? In practice the strongest deployed systems mix the two. What looked in 2024 like two independent efforts converging on the same answer is now the default recipe across a dozen shipping models.

AI21's Jamba was the first of them to ship at scale: a hybrid Transformer-Mamba mixture-of-experts model that interleaves blocks of transformer attention with blocks of Mamba layers at roughly a 1:7 ratio, adding MoE feed-forward layers every couple of blocks to expand capacity without proportionally increasing active compute. Jamba 1.5 shipped in 2024 in two sizes, Large at 94B active parameters and Mini at 12B, both at an effective context length of 256K tokens.

The design logic is explicit: most sequence-processing work goes to cheap, linear-cost Mamba layers, while a sparse scattering of full-attention layers preserves the capabilities pure SSMs are weak at (more on that below), keeping the KV cache small relative to an all-attention model at the same context length.

Google DeepMind's Griffin (De et al., arXiv:2402.19427, 2024) arrives at a structurally similar answer from a different direction. Hawk, its pure-recurrent sibling, uses a gated linear recurrent layer called RG-LRU and was reported to exceed Mamba's downstream performance; Griffin alternates RG-LRU blocks with local, sliding-window attention rather than full global attention.

The paper reports that Griffin matches Llama-2's performance while being trained on more than 6x fewer tokens, extrapolates to sequence lengths well beyond what it saw during training, was scaled up to 14B parameters, and — during inference — achieves markedly lower latency and higher throughput than a comparable transformer while matching transformer-level hardware efficiency during training.

By 2026 the pattern is everywhere in open weights, with the linear module varying by house. NVIDIA's Nemotron 3 line uses Mamba-2 inside a mixture-of-experts hybrid: Nemotron 3 Super is a 120B-parameter model with about 12B active, pretrained on 25 trillion tokens at a 1M-token context and built for agentic reasoning. IBM's Granite 4.0, released October 2, 2025, stacks Mamba-2 and transformer blocks sequentially at a 9:1 ratio across four sizes from 3B to 32B. Qwen3-Next 80B-A3B and Qwen3.5 use a 3:1 pattern of Gated DeltaNet to gated attention rather than Mamba.

The ratios differ, the linear primitive differs, the conclusion does not: a thin minority of attention layers, a large majority of linear-cost ones.

RWKV: A Parallel Lineage

RWKV (Peng et al., "RWKV: Reinventing RNNs for the Transformer Era," arXiv:2305.13048, 2023) reaches a similar destination — linear-cost sequence modeling — from a lineage rooted in linear attention and RNNs rather than continuous-time state spaces. Its core blocks are time-mixing and channel-mixing sub-blocks built from four interacting components, among them a "receptance" gate (R) governing how much past information passes forward and a trainable, decaying positional weight (W).

The architecture is formulated so that it can be trained in a transformer-like parallel mode but executed at inference time as a plain RNN, meaning inference runs in constant memory and constant per-token compute rather than needing a growing KV cache. RWKV has continued to iterate through RWKV-5, -6 and -7, progressively adopting more expressive, matrix-valued state transitions — convergent evolution with the direction SSD points toward, from a different theoretical starting point.

Where SSMs Still Lose: Copying and Retrieval

The tradeoff is not free, and it has been characterized precisely rather than anecdotally. Jelassi et al., "Repeat After Me: Transformers are Better than State Space Models at Copying" (ICML 2024, arXiv:2402.01032), prove that a two-layer transformer can copy input strings of exponential length, while fixed-state recurrent models are constrained by their bounded latent state: unlike a KV cache, which grows with context and lets the model look up any prior token exactly, an SSM's state vector has a fixed capacity that must summarize arbitrarily long history through lossy compression.

Empirically, a Mamba model failed to learn a synthetic copying task that a similarly sized transformer learned quickly, and needed roughly 100x more training data to close the gap. On pretrained models, the authors found transformers "dramatically" outperform SSMs at copying and retrieving specific information from context.

The result explains the entire hybrid-architecture trend. Precise in-context retrieval — quoting a passage verbatim, looking up a specific fact planted earlier in a long prompt, following an exact reference — is exactly the kind of task where a fixed-size recurrent state is information-theoretically disadvantaged relative to an attention mechanism with an ever-growing, exactly addressable cache. Jamba's sparse attention layers and Griffin's local-attention blocks aren't cosmetic; they're targeted patches for the one class of task where pure SSMs have a proven, provable weakness.

The Practical Calculus

The resulting picture is more nuanced than "linear beats quadratic." Pure SSMs (S4, Mamba, RWKV) win decisively on raw long-context throughput, memory footprint, and streaming inference, and Mamba-2's SSD reformulation shows the efficiency gap over hand-tuned scan kernels can itself shrink further with better algorithms. But they carry a structural weakness at exact copying and retrieval that scales with context length rather than shrinking with better training.

Transformers retain the opposite profile: expensive and memory-hungry at long context, but able to address any prior token exactly via attention. Hybrid designs are no longer a hedge but the default, from Nemotron 3 and Granite 4.0 to Qwen3-Next: keep the bulk of computation linear-cost and reserve a thin layer of full or local attention specifically where exact in-context lookup matters.

Whether a single-architecture design can close the copying gap without reintroducing quadratic cost is still open, but less open than it was. Mamba-3's complex-valued state update was aimed squarely at state tracking, and matching Mamba-2's perplexity at half the state suggests the bounded-state penalty is softer than the 2024 results implied. None of that repeals the Jelassi bound, which is about capacity rather than engineering. The realistic reading is that the attention fraction in a hybrid keeps shrinking rather than reaching zero.