State Space Models and Mamba

A state space layer is a recurrent layer whose update is linear, which is the property that lets it be trained in parallel like a transformer and run with a fixed-size state like an RNN. Mamba’s addition is to make the recurrence depend on the input token, which is what turned a signal-processing construction into something that works on language.

It is a linear recurrence

Strip away the control-theory vocabulary and a state space layer is two lines:

h_t = A * h_{t-1} + B * x_t      carry the state forward
y_t = C * h_t                    read something out of it

h is a state vector, A decides how much of the old state survives, B decides how much of the new input enters, and C decides what gets read out. The difference from an LSTM is the missing nonlinearity: there is no tanh wrapped around the state update. That absence is not a simplification for the reader’s benefit. It is the entire trick.

In practice A is diagonal, so “multiply by A” is one scalar per channel rather than a matrix multiply, and the diagonal entries are parameterised as negative reals passed through an exponential so that each channel decays rather than blows up. A channel with a slow decay remembers for a long time; a channel with a fast one is effectively local. A layer holds many channels with different decay rates, which is how one layer covers several timescales at once.

Why a linear recurrence can be trained in parallel

Unroll it. Because there is no nonlinearity in the way, the state at step t can be written as a plain weighted sum of every input so far:

h_t = B*x_t + A*B*x_{t-1} + A^2*B*x_{t-2} + ... + A^(t-1)*B*x_1

y_t = sum over k of  (C * A^k * B) * x_{t-k}

That is a convolution. The kernel is the sequence C*B, C*A*B, C*A^2*B, ..., and if A, B and C do not change with the input, the kernel can be computed once and applied to the whole sequence in one shot — with an FFT if it is long enough. Training touches every position at the same time, which is exactly the property recurrent networks did not have. That construction is S4, published by Gu, Goel and Ré in 2021.

The same layer has two faces. At training time it is a convolution over the whole sequence, parallel. At inference time it is the two-line recurrence, one step per token, with a fixed state. Same weights, two execution modes, chosen for whichever costs less.

Selective state: the change that made it work on text

Fixed A, B and C means the layer treats every token identically. It cannot decide that this token is a name worth keeping and that one is filler, because nothing in the update depends on what the token is. That limitation shows up sharply on tasks like “copy the identifier that appeared earlier”, which attention finds trivial and a time-invariant recurrence cannot do at all.

Mamba (Gu and Dao, 2023) makes B, C and the step size delta functions of the current input:

delta_t = softplus(W_delta @ x_t)     per-token step size
B_t     = W_B @ x_t                   per-token input gate
C_t     = W_C @ x_t                   per-token output gate

A_t = exp(delta_t * A)                decay, now input-dependent
h_t = A_t * h_{t-1} + delta_t * B_t * x_t

A large delta means “this token matters: reset toward it and write it into the state”. A small one means “skip this, keep what you had”. That is content-based memory, and it is the thing the fixed version was missing.

The price is immediate: with A now varying per token, the unrolled kernel is different at every position and the convolution trick is gone. Mamba recovers parallel training with a work-efficient associative scan — a tree-shaped reduction that computes all T states in O(log T) sequential depth — implemented as a fused kernel that keeps the state in on-chip SRAM instead of writing it to HBM. The paper calls this hardware-aware, and it is the same style of argument as FlashAttention: the algorithm was chosen around the memory hierarchy.

What a Mamba block actually contains

The recurrence is not the whole layer. A block looks like this:

x  -> linear up-projection (expansion 2) -> two branches

  branch A:  depthwise conv1d (kernel 4) -> SiLU -> selective SSM
  branch B:  SiLU                                    (a gate)

  y = branch_A * branch_B  ->  linear down-projection

Three details in there do real work and none is obvious from the recurrence alone.

• The short convolution. A four-wide depthwise conv1d sits in front of the SSM because a decaying recurrence is surprisingly bad at “the token immediately before this one” — the state blends the recent past rather than distinguishing it. A four-position convolution is the cheapest possible fix and it costs almost nothing.
• The multiplicative gate. Branch B is the same construction as the gated MLP in a modern transformer block. It means a Mamba block performs the roles that attention and the MLP perform in a transformer, in one unit.
• The parameter budget per block is smaller. With expansion 2, the projections cost roughly 6 * D^2 per block, against something in the region of 12 * D^2 for a transformer block with attention plus a gated MLP. Mamba stacks therefore use about twice as many blocks for the same parameter budget, which is why a layer count from one family is not comparable with a layer count from the other.

The state, in bytes, against a KV cache

This is the comparison that decides whether the architecture is worth anything, and it is pure arithmetic. Take a mid-size model: 48 layers, model dimension 2,048, and serve it in 16-bit.

A Mamba layer’s state

expansion factor 2   -> inner dimension = 4,096
state size N         = 16          (Mamba's default)

state per layer  = 4,096 * 16 = 65,536 values
                 = 131,072 bytes at fp16   (128 KB)
state per model  = 128 KB * 48 = 6.1 MB per sequence

...and this number does not change with sequence length.

A transformer’s KV cache

48 layers, 8 KV heads (grouped-query), head_dim 128, fp16

bytes per token = 2 (K and V) * 48 * 8 * 128 * 2
                = 196,608 bytes   (192 KB per token)

  4,096 tokens ->   0.77 GB
 32,768 tokens ->   6.1  GB
131,072 tokens ->  24.6  GB

Six megabytes, flat, against six gigabytes at 32k tokens. Three orders of magnitude at a context length people actually use, and the gap widens linearly from there. That is the argument, and no benchmark is needed to make it.

The consequence for serving is concurrency. A machine with 80 GB of memory and a 14 GB model has about 66 GB for cache; at 192 KB per token that is roughly ten concurrent 32k-token sessions. With a constant 6 MB state it is thousands.

What it costs

• The state is lossy and finite. 65,536 numbers per layer have to hold everything the model still needs from a sequence of any length. Attention does not compress: every earlier token is still individually addressable in the cache. This is why the reported weakness of pure SSMs is exact recall — retrieving a verbatim string from far back — and it is a structural weakness, not a tuning problem.
• No prefix caching in the usual sense. A shared system prompt can be replayed to produce a state, but the whole ecosystem of cached input tokens and per-token cache reuse is built around KV caches.
• Kernels are the moat. The parallel scan needs a custom fused kernel to be fast. Every inference server, every quantisation format and every serving optimisation of the last five years was written for attention, and an SSM model gets none of it for free.
• Two execution paths to keep correct. A model that trains as a scan and serves as a recurrence has two implementations that must agree numerically, and 16-bit accumulation over thousands of steps is where they stop agreeing.

Where adoption actually stands

Honestly: state space layers are shipping, mostly inside hybrids, and they have not displaced attention.

Pure SSM language models exist at small and medium scale — Mamba and Mamba-2 from the original authors, Codestral Mamba from Mistral in 2024, Falcon Mamba from TII in 2024. The models that reach production scale mix layer types instead: Jamba from AI21 (2024) interleaves Mamba and attention blocks, and NVIDIA’s Nemotron-H (2025) is a Mamba-transformer hybrid. That pattern is consistent enough to be the finding: a few attention layers are kept because a small amount of exact recall turns out to be worth its memory.