State space models and Mamba

At 3 a.m., your model is processing the 128,001st token in a customer-support transcript. It needs to decide whether that token refers to an account number mentioned near the beginning.

A transformer can look back directly. That is its great trick. It is also the bill.

When all 128,000 positions are processed together, as during training or transformer prefill, the dense attention score shape for one head is 128,000 × 128,000 = 16,384,000,000 entries. In bfloat16, those scores alone occupy about 30.5 GiB per head. Across 32 heads in one layer, that is roughly 977 GiB, before storing activations, gradients, keys, values, or anything else. That 977 GiB figure is not the memory for a 32-layer model; it is the score storage for one such layer.

Autoregressive decoding is different. When the 128,001st token is generated, its new query normally produces a 1 × 128,000 score row against the keys in the key-value cache, not a new 128,000 × 128,000 matrix. That step does O(L) score and value work for context length L, while the cache itself keeps growing with L. Generate enough tokens and the accumulated decode work can still become expensive.

Efficient kernels such as FlashAttention avoid materialising the entire score matrix. They reduce memory traffic and can make attention much faster. They do not make every comparison disappear. Full-sequence training and prefill still have dense all-pairs interaction, and decode still compares each new query with a growing cache.

A different family of sequence models takes a more radical approach: do not keep the transcript. Keep a learned summary of it.

That family is called state space models, or SSMs. Mamba is a modern selective SSM designed to make that summary depend on the current input.

One system, two views

A state space model borrows a simple idea from dynamical systems: describe the past with a hidden state, a vector updated whenever a new input arrives.

A continuous-time linear system is commonly written as:

dx/dt = A x + B u

y = C x + D u

Here u is the input, x is the hidden state, and y is the output. A controls how old state persists or decays, B writes input into the state, C reads information out, and D provides a direct input-to-output path.

After discretisation, the system becomes:

x_t = A_bar x_(t-1) + B_bar u_t

y_t = C x_t + D u_t

The bars indicate per-step parameters derived from the continuous-time system. In a classical linear time-invariant SSM, these parameters do not change from token to token.

At inference, processing token t requires only the previous state and the new input. The carried state uses constant recurrent-state memory for streaming inference: its storage does not grow with sequence length. The work per step still depends on the state and channel dimensions.

Unrolling the recurrence, starting from a general initial state x_0, gives:

y_t = C A_bar^t x_0 + D u_t + Σ_(k=0)^t C A_bar^k B_bar u_(t-k)

The first term is information already present before the sequence began. The second, D u_t, bypasses the state. The sum contains the input history carried through the state.

For the zero-initial-state case, define the effective impulse-response kernel:

K_0 = D + C B_bar

K_1 = C A_bar B_bar

K_2 = C A_bar^2 B_bar

The same system can therefore be computed as a causal convolution, with each output a weighted sum of the current and earlier inputs.

This is the central insight: a fixed linear SSM has a recurrent inference view and a parallel convolution view for training.

The recurrence is ideal for streaming inference. The convolution removes the sequential dependency during training, so all positions can be processed together on a GPU. Depending on the parameterisation, structured algorithms or FFTs make long convolutions practical; selective SSMs instead use a linear-time scan. Avoiding an attention score matrix alone does not guarantee subquadratic training.

A worked example

Start with one scalar state and one scalar output:

x_t = 0.8 x_(t-1) + 0.5 u_t

y_t = x_t

The 0.8 means that 80 percent of the previous state survives each step. The 0.5 controls how strongly new input is written.

For input [2, 0, 3], with initial state zero:

• At step 1: x_1 = 0.8 × 0 + 0.5 × 2 = 1.0
• At step 2: x_2 = 0.8 × 1.0 + 0.5 × 0 = 0.8
• At step 3: x_3 = 0.8 × 0.8 + 0.5 × 3 = 2.14

The convolution kernel is:

[0.5, 0.4, 0.32, ...]

Here D = 0, so the zero-lag term is C B_bar = 0.5; later terms multiply by additional powers of 0.8.

The third convolution output is:

0.32 × 2 + 0.4 × 0 + 0.5 × 3 = 2.14

One calculation marched through time; the other multiplied the sequence by a causal kernel.

import torch

a = 0.8
b = 0.5
u = torch.tensor([2.0, 0.0, 3.0], dtype=torch.float64)

# Recurrent view
x = torch.tensor(0.0, dtype=torch.float64)
y_recurrent = []

for token in u:
    x = a * x + b * token
    y_recurrent.append(x)

y_recurrent = torch.stack(y_recurrent)

# Convolution view, written directly so the indexing is visible
kernel = torch.tensor([b, a * b, a * a * b], dtype=torch.float64)
y_convolution = torch.zeros_like(u)

for t in range(u.numel()):
    for lag in range(t + 1):
        y_convolution[t] += kernel[lag] * u[t - lag]

print(y_recurrent)
print(y_convolution)

Both printed tensors are:

tensor([1.0000, 0.8000, 2.1400], dtype=torch.float64)
tensor([1.0000, 0.8000, 2.1400], dtype=torch.float64)

Real SSM layers use vectors, many channels, and structured matrices. The mechanism is the same.

Why naive RNNs lost the long-memory race

A vanilla RNN also keeps a constant-size state. But a dependency from 1,000 steps ago is carried by something like A_bar^1000.

If relevant eigenvalues have magnitude below one, that signal shrinks toward zero. If they exceed one, it grows until the state or gradient becomes unstable. This causes vanishing and exploding gradients. LSTMs add controlled paths that preserve information more easily, but they still train through a sequential recurrence.

The fix was to give transition dynamics useful timescales from the start.

HiPPO, short for High-order Polynomial Projection Operators, provides the key intuition. Instead of storing an unstructured stream, state coordinates behave like coefficients of an approximation to recent history. New observations update that approximation, while different coordinates preserve different temporal scales. A few coefficients might track an audio signal’s level, slope, and finer structure.

HiPPO-inspired transition matrices give an SSM a principled spread of memory timescales. Structured SSMs such as S4 combine those dynamics with matrix structure that makes long sequences practical.

HiPPO does not make finite state lossless. It makes the loss useful and trainable.

Mamba makes memory selective

A fixed SSM has one kernel for every input sequence. Whether a token is a customer ID, punctuation, or boilerplate, the same transition dynamics process it. The model can learn average behaviour, but it cannot decide from the token itself to pause, write strongly, or ignore.

Mamba makes important SSM parameters input-dependent. In the original design, projections of the current input produce values related to:

• Δ_t, a step size controlling how much the state changes or persists
• B_t, how the current input enters the state
• C_t, how the state is read out

The learned transition structure A remains a shared backbone, while the effective update changes with the input:

x_t = A_bar_t x_(t-1) + B_bar_t u_t

y_t = C_t x_t

For example, in a long log, a selective model can strongly write an identifier into its state, suppress routine messages, and later read the relevant coordinates. A fixed SSM has no explicit input-dependent “this token is important” switch; it must express the behaviour through one input-independent kernel.

Once B_t, C_t, or the step size depends on the input, there is no single convolution kernel to precompute. The effect of one token depends on later input-dependent updates. Mamba therefore uses a selective scan, a hardware-aware way to evaluate the recurrence efficiently.

Each step applies an affine transformation to the state. Two affine transformations can be composed into one, and that composition is associative. Associativity lets a parallel scan combine chunks on a GPU instead of running a slow loop over every token. The scan preserves the input-dependent recurrence and constant recurrent-state memory for streaming inference; it does not turn the model back into a convolution.

The honest trade-off

An SSM compresses the past. Attention addresses the past.

If a task requires copying an arbitrary 200-character string exactly from 80,000 tokens ago, a bounded state must preserve enough information about that string in a small vector. Attention can point directly to the original characters. Selectivity helps Mamba keep the right region, but it does not turn compression into a searchable archive.

Architecture	Long-stream memory	Arbitrary old-token recall	Training over positions	Best fit
Full attention	Grows with context	Strong	Highly parallel, quadratic interaction	Moderate context and precise lookup
Sparse attention	Lower than full attention	Depends on the pattern	Parallel with restricted interactions	Known local or global structure
LSTM or GRU	Constant recurrent-state memory for streaming inference	Limited by hidden state	Sequential recurrence	Small streaming systems
Fixed SSM	Constant recurrent-state memory for streaming inference	Good for learned temporal patterns, weaker for exact lookup	Parallel convolution or structured scan	Long smooth sequences
Selective SSM such as Mamba	Constant recurrent-state memory for streaming inference	Better content filtering, still compressed	Parallel hardware-aware scan	Long streams and token-dependent memory

There is no universal winner. A language model that must retrieve names, code fragments, or exact quotations may benefit from attention. An audio model processing millions of samples may benefit more from bounded state and local temporal dynamics. Hybrid stacks combine both: attention handles sharp, content-addressed retrieval, while SSM layers provide cheap long-range propagation.

When to choose one

Choose a selective SSM when the sequence is genuinely long, arrives continuously, or has strong temporal structure: audio, sensor data, event logs, biological sequences, and long-running streams.

Choose full attention when exact retrieval from arbitrary positions is central and the context is manageable. Choose sparse attention when the access pattern is known well enough to encode. Choose an LSTM or GRU when the model is small or the deployment target is constrained.

Use a hybrid when the task needs both bounded long-range propagation and exact lookup.

What to remember

• Attention keeps many past tokens available for direct lookup; an SSM compresses them into a fixed-size state.
• A classical linear SSM is both a recurrence for constant recurrent-state memory during streaming inference and a convolution for parallel training.
• HiPPO-style structure gives the state useful memory timescales instead of asking a naive RNN to discover stable long memory from nothing.
• Mamba makes the state update selective by deriving transition, write, and read behaviour from the current input.
• Selective SSMs are strong for very long and streaming sequences, but they can lose against attention when exact arbitrary-token recall is the task.

Quick check

0/3

Q1

Its transition matrices are fixed, so unrolling the recurrence produces one input-independent causal kernel.Convolution removes the hidden state from the model.The model stores every previous token in a compressed key-value cache.The recurrence is run backward during training.

Check answer

Q2

It must store the full attention matrix.Its input-dependent updates remove the single fixed convolution kernel, so training needs an efficient selective scan.It can no longer maintain a recurrent state at inference.It requires a separate model for every sequence length.

Check answer

Q3

Full attention everywhere, because its cache stays constant during streaming.A fixed SSM only, because compression always preserves exact quotes.A selective SSM for the long stream plus some attention capacity for precise retrieval.An LSTM with no state carried between audio chunks.

Check answer