Mamba and State Space Models: The Transformer Challenger

Self-attention made Transformers the default architecture for sequence modeling, but its cost grows quadratically with sequence length. State space models (SSMs) — and in particular Mamba — offer a linear-time alternative that keeps a compressed recurrent state instead of comparing every token to every other token. This article explains how SSMs work, what selectivity changes, why the hardware-aware parallel scan matters, how Mamba-2 reconnects SSMs to attention, and where these models still fall short.

Attention keeps every token and pays O(n²) to relate them; a selective SSM keeps a fixed-size state and pays O(n) to summarize them.

The quadratic cost of attention

A Transformer layer computes, for a sequence of length n, an attention matrix of shape n × n. Every token attends to every other token, so both compute and memory for that matrix scale as O(n²). For short prompts this is fine; for documents, genomes, high-resolution audio, or agent transcripts with hundreds of thousands of tokens, the quadratic term dominates and eventually becomes the bottleneck.

Two costs are worth separating:

• Training / prefill: processing a length-n sequence in parallel costs O(n²) time and, naively, O(n²) memory. FlashAttention removes the materialized matrix so memory becomes O(n), but the compute stays quadratic.
• Autoregressive decoding: generating one token requires attending back over all previous tokens. With a key/value (KV) cache this is O(n) work per step and an O(n) cache that grows with every token produced.

The growing KV cache is the practical pain point at long context: memory scales linearly with sequence length and with batch size, and it is read on every decode step. SSMs attack exactly this by replacing the ever-growing cache with a fixed-size state.

State space models, briefly

A continuous linear state space model maps an input signal x(t) to an output y(t) through a hidden state h(t):

h'(t) = A h(t) + B x(t)
y(t)  = C h(t) + D x(t)

Here A governs how the state evolves on its own, B writes the input into the state, and C reads the state out. To use this on discrete token sequences, we discretize with a step size Δ (delta), producing matrices Ā and B̄. The model then has two mathematically equivalent views:

• Recurrent view — compute step by step, ideal for inference:

  h_t = Ā h_{t-1} + B̄ x_t
  y_t = C h_t

• Convolutional view — when the parameters are fixed over time (linear time-invariant, LTI), the whole output is a convolution of the input with a single long kernel K, computable in parallel over the sequence.

This duality is the whole trick of early SSMs: train efficiently as a convolution, then run cheaply as a recurrence. The catch is that the convolutional view requires the parameters to be constant across time.

From S4 to structured recurrence

A naive SSM struggles with long-range dependencies: repeatedly multiplying by A either explodes or vanishes, much like a plain RNN. S4 (Structured State Spaces), from Gu, Goel, and Ré (2021), made SSMs actually work at long range by giving the state matrix a special structure.

• HiPPO initialization. The A matrix is initialized from HiPPO theory, which yields a state that optimally compresses the input history onto a basis of polynomials. Intuitively, the state becomes a running, decaying summary of everything seen so far.
• Diagonal / structured parameterization. S4 and its simpler successor S4D use a diagonal (plus low-rank) form that makes the long convolution kernel efficient to compute, avoiding an O(n²) blow-up in the kernel itself.

S4 set strong results on long-range benchmarks, but it is time-invariant: the same A, B, C are applied to every token regardless of content. That is what makes the convolution possible — and also what limits it. A time-invariant filter cannot decide to pay extra attention to one token and ignore another; it treats the sequence like a fixed signal-processing pipeline.

Selectivity: the core Mamba idea

Mamba (Gu and Dao, 2023) introduces the selective SSM, often called S6. The single conceptual change is this: make the SSM parameters functions of the input. Specifically, B, C, and the step size Δ are produced from the current token by small linear projections, rather than being fixed weights.

Why this matters:

• Content-based memory control. Because Δ is input-dependent, the model can effectively decide how much a given token updates the state. A large Δ lets a token strongly write into and reset the state; a small Δ lets the state coast, preserving older information. This is the mechanism by which Mamba "chooses what to remember and what to forget."
• Input-dependent read/write. Making B and C input-dependent means what gets written into the state and what gets read out both depend on the token itself — closer in spirit to attention's data-dependent routing than to a fixed convolution.

The cost of selectivity is exactly what S4 was avoiding: once the parameters vary per token, the model is no longer time-invariant, so the convolutional view disappears. There is no single fixed kernel anymore. Mamba is now a genuine, content-dependent recurrence — and recurrences are hard to parallelize on a GPU. That problem is solved not with math but with systems engineering.

Hardware-aware parallel scan

The linear recurrence h_t = Ā_t h_{t-1} + B̄_t x_t is an associative scan (a prefix-sum-like operation). Associative operations can be computed in parallel with a work-efficient scan (the Blelloch-style algorithm), turning what looks sequential into a logarithmic-depth parallel computation. Mamba uses a parallel scan so training and prefill remain parallel across the sequence, even without a convolution.

Just as important is the memory strategy. Mamba's implementation is hardware-aware: it fuses the discretization, the scan, and the output projection into a kernel that keeps the large intermediate states in fast on-chip SRAM and avoids writing the full expanded state to slow high-bandwidth memory (HBM). Combined with recomputation during the backward pass, this keeps the expanded state from ever fully materializing in HBM — the same "don't materialize the big matrix" philosophy that makes FlashAttention fast, applied to the scan.

Net effect: training is O(n) in sequence length (versus O(n²) for attention) and remains GPU-efficient despite being a recurrence.

Recurrent inference and constant memory

At generation time, Mamba switches back to the pure recurrent view. Each new token does a constant amount of work — one state update — and the state is a fixed-size tensor whose size does not depend on how many tokens have already been generated.

• O(1) time per generated token, versus O(n) per token for attention (which must consult the whole KV cache).
• O(1) memory in sequence length: a constant-size recurrent state replaces a KV cache that grows with every token.

For long-context autoregressive workloads this is the headline advantage: throughput does not degrade as the context grows, and memory stays flat. The trade-off is that all information about the past must be squeezed into that fixed-size state — a lossy compression, which is precisely where the recall limitations (below) come from.

Attention vs. recurrent SSM (diagram)

The contrast is fundamentally about what state is carried between steps. Attention carries the full set of past keys and values; a selective SSM carries one compressed state vector.

The KV cache branch keeps looping back and growing; the state branch loops back onto a fixed-size tensor. That single structural difference is the source of every complexity number in the comparison table below.

Mamba-2 and state-space duality

Mamba-2 (Dao and Gu, 2024) reframes selective SSMs through state-space duality (SSD). The key insight is that a class of SSMs can be written as matrix multiplications that are structurally equivalent to a form of attention with a specific structured (semiseparable) mask. In other words, SSMs and attention are two ends of a spectrum, not unrelated architectures.

• Two computation modes, one model. SSD lets the same layer be computed either as a linear-time recurrence/scan or as a quadratic-time matrix form. Training can use a block-decomposed matrix algorithm that leverages efficient matrix-multiply hardware (tensor cores), which the original Mamba scan did not exploit as directly.
• A larger, simpler state. Mamba-2 adopts a more restricted but hardware-friendlier parameterization (for example, a scalar-times-identity state transition), which allows much larger state dimensions and simpler, faster kernels.
• Conceptual payoff. SSD gives a shared vocabulary for reasoning about linear attention, selective SSMs, and softmax attention, and explains why techniques from one side (multi-head structure, tensor-parallel layouts) transfer to the other.

The practical takeaway: Mamba-2 keeps the linear-time inference story while making training simpler and better matched to modern accelerators.

Hybrids: Jamba and friends

In practice the strongest long-context models are frequently hybrids that interleave SSM layers with a smaller number of attention layers, capturing the strengths of both. The attention layers restore precise, content-addressed recall; the SSM layers carry most of the sequence cheaply.

• Jamba (AI21, 2024) interleaves Mamba and Transformer layers and adds mixture-of-experts (MoE) feed-forward blocks. The result targets very long context windows with a far smaller KV-cache footprint than a pure Transformer of comparable quality, because only the sparse attention layers maintain a cache.
• Other hybrids. Research lines such as Zamba and the broader "striped" or interleaved SSM-attention designs follow the same recipe: a majority of linear-time SSM blocks with a sparing amount of attention to preserve exact recall and in-context lookup.

The design heuristic that has emerged: use SSM layers as the cheap backbone for carrying context, and spend a small attention budget where exact retrieval genuinely matters.

Honest limits and when to reach for what

SSMs are not a strict upgrade over attention. The fixed-size state is a feature for efficiency and a liability for precision.

• Weaker exact recall and copying. Tasks that require verbatim retrieval — copying a specific earlier substring, associative recall of an arbitrary key-value pair seen once, needle-in-a-haystack lookups — are harder for pure SSMs. Attention can point directly at the relevant past token; an SSM must have preserved that detail in its compressed state, and it may not have. This is well documented empirically and follows from the information-theoretic limit of a bounded state.
• In-context learning nuances. Some in-context abilities that rely on attending to specific demonstrations can be weaker, which is a major reason hybrids retain a few attention layers.
• Ecosystem maturity. Attention has years of tooling, kernels, quantization recipes, and serving infrastructure; SSM support, while growing quickly, is less universal.

Rules of thumb:

• Very long streams where throughput and flat memory dominate (audio, DNA, long transcripts, streaming): SSM or hybrid.
• Workloads dominated by precise retrieval, tool/argument copying, or heavy in-context lookup: attention, or a hybrid with enough attention layers.
• General-purpose long-context LLMs today: hybrids are the pragmatic sweet spot.

Side-by-side comparison

Takeaways

• Attention's O(n²) cost and growing KV cache are the constraints SSMs are built to remove.
• S4 made SSMs work at long range via HiPPO-based structure, but stayed time-invariant.
• Mamba's selectivity makes SSM parameters input-dependent, trading the convolutional view for content-based memory control, recovered on hardware with a fused parallel scan.
• Inference is the standout win: O(1) time per token and constant memory, independent of context length.
• Mamba-2's state-space duality shows SSMs and attention are two views of one family, and makes training tensor-core-friendly.
• Pure SSMs trade away exact recall; hybrids like Jamba (Mamba + attention + MoE) are the practical way to get both efficiency and precise retrieval.