Module 19: State-Space Models (Mamba)

Introduction

Every module so far is built on attention. Attention is powerful, but it pays a tax: comparing every token to every other token costs \(O(L^2)\) compute, and generating with a growing KV cache (m08) costs memory that climbs with the context. Double the context and attention quadruples its work.

A state-space model (SSM) is the leading alternative — the architecture, made practical by Mamba, that many believe can rival the Transformer. An SSM reads a sequence the way an RNN does: it carries a fixed-size hidden state forward, one token at a time, folding each new token into that state. Because the state never grows, an SSM runs in linear time (\(O(L)\)) and constant memory per generated token — no cache that swells with the context.

The catch that sank classic RNNs was training speed and long-range memory. SSMs fix both: a time-invariant SSM can be run as a parallel convolution for training, and a carefully structured state matrix gives it long memory. Mamba’s final ingredient — selectivity — lets the state decide what to remember based on content, closing the last quality gap with attention.

Why it matters for LLMs:

  • Long context, cheaply. Linear scaling makes million-token sequences tractable where attention’s \(L^2\) wall does not.
  • Fast generation. A constant-size state means each new token costs the same, no matter how long the context — Mamba reports ~5× the throughput of a Transformer.
  • A different lens. Understanding SSMs shows you which parts of “the Transformer” are essential and which are just one choice among several.

What You’ll Learn

After this module, you can:

  • Explain a state-space model as a linear recurrence over a fixed-size state.
  • Discretize a continuous SSM with zero-order hold and run the discrete recurrence.
  • Run the same SSM as a recurrence and as a convolution, and see they are identical — the trick that trains in parallel and infers cheaply.
  • Explain why that recurrence is an associative scan that parallelizes in \(O(\log L)\) depth.
  • Build a selective SSM (Mamba’s S6) whose \(\Delta\), \(B\), \(C\) depend on the input.
  • Show why SSMs challenge attention: linear vs quadratic compute, constant vs growing decode memory.

Prerequisites

This module requires familiarity with:

Intuition: A Sequence as a Running State

Forget attention for a moment. Imagine reading a sentence with a small notebook: after each word you update a fixed set of numbers — your running summary — and you never allow the notebook to get bigger. That notebook is the hidden state \(h\), and the rule for updating it is a state-space model:

\[ h_t = \bar{A}\, h_{t-1} + \bar{B}\, x_t, \qquad y_t = C\, h_t. \]

\(\bar{A}\) says how much of the old state to keep (a decay), \(\bar{B}\) says how much of the new token to write in, and \(C\) reads an output off the state. That is the whole model. Feed it a single spike at \(t=0\) and watch the state light up, then decay — step through its impulse response:

NoteKey Insight

The state is a fixed-size summary of everything seen so far. Attention keeps every past token around (the KV cache) and re-reads them; an SSM keeps one running state and updates it. That single design choice is where the linear cost comes from.

The Math: The State-Space Recurrence

SSMs start in continuous time, borrowed from control theory. A state vector \(h(t) \in \mathbb{R}^N\) evolves under a linear system driven by the scalar input \(x(t)\):

\[ h'(t) = A\, h(t) + B\, x(t), \qquad y(t) = C\, h(t). \]

To use it on a discrete token stream we discretize it with a step size \(\Delta\), using the standard zero-order hold (ZOH) rule (it assumes the input is held constant across each step). For a diagonal \(A\) — a vector of scalars \(a\), which is exactly what Mamba uses — this is element-wise:

\[ \bar{A} = \exp(\Delta A), \qquad \bar{B} = A^{-1}\!\left(\exp(\Delta A) - I\right) B . \]

Two things to notice. First, if every \(a < 0\) then \(\bar{A} = \exp(\Delta a) \in (0, 1)\) — a genuine decay factor, so the state forgets old inputs gracefully and never blows up. Second, as \(a \to 0\) the \(\bar{B}\) factor tends to \(\Delta\), so \(\bar{B} \approx \Delta B\) — the simple forward-Euler limit that Mamba’s kernel actually uses. Substituting gives the discrete recurrence from the intuition:

\[ h_t = \bar{A}\, h_{t-1} + \bar{B}\, x_t, \qquad y_t = C\, h_t . \]

discretize in ssm.py implements the exact ZOH formula (with a stable small-\(a\) fallback so the Euler limit is numerically safe):

import torch
from ssm import discretize

A = torch.tensor([-0.5, -1.0, -3.0])   # diagonal state matrix (must be < 0 to decay)
B = torch.tensor([1.0, 1.0, 1.0])
A_bar, B_bar = discretize(A, B, delta=0.5)
print("Ā (decay, in (0,1)):", [round(v, 3) for v in A_bar.tolist()])
print("B̄ (input gain):     ", [round(v, 3) for v in B_bar.tolist()])
Ā (decay, in (0,1)): [0.779, 0.607, 0.223]
B̄ (input gain):      [0.442, 0.393, 0.259]

Code: The Recurrent Scan

Given the discrete coefficients, running the SSM is a short loop — the inference face, \(O(L)\) time and \(O(N)\) memory. ssm_recurrent returns both the outputs and the hidden state after each step:

from ssm import ssm_recurrent

A_bar = torch.tensor([0.5])            # a single-state SSM, decay 0.5 per step
B_bar = torch.tensor([1.0])
C = torch.tensor([1.0])

x = torch.tensor([1.0, 0.0, 0.0, 0.0, 0.0])   # a unit impulse at t=0
y, states = ssm_recurrent(A_bar, B_bar, C, x)
print("impulse response:", [round(v, 4) for v in y.tolist()])   # 1, .5, .25, ...
impulse response: [1.0, 0.5, 0.25, 0.125, 0.0625]

The output halves every step: the state remembers the spike, fading by \(\bar{A}=0.5\) each time. That decaying sequence is the SSM’s impulse response — hold onto it, because it is about to reappear as a convolution kernel.

Two Faces: Recurrent = Convolutional

Here is the trick that makes SSMs trainable. Unroll the recurrence from \(h_0 = 0\):

\[ y_t = \sum_{i=0}^{t} \big(C\, \bar{A}^{\,i}\, \bar{B}\big)\, x_{t-i} . \]

That is a causal convolution of the input with a fixed kernel

\[ \bar{K} = \big(C\bar{B},\; C\bar{A}\bar{B},\; C\bar{A}^2\bar{B},\; \dots\big) , \]

which is precisely the impulse response you just saw. So one SSM has two faces:

  • Recurrent — one step at a time, \(O(L)\), cheap at inference (no cache).
  • Convolutional — the whole sequence at once with a single kernel, parallel on the GPU, ideal for training.

Crucially, for time-invariant parameters they compute exactly the same function. ssm_convolution builds the kernel and convolves; it matches ssm_recurrent to floating-point tolerance:

from ssm import ssm_convolution, ssm_kernel

torch.manual_seed(0)
A_bar = torch.rand(4) * 0.9 + 0.05     # 4-dim state, each channel decays in (0.05, 0.95)
B_bar = torch.randn(4)
C = torch.randn(4)
x = torch.randn(16)

y_rec, _ = ssm_recurrent(A_bar, B_bar, C, x)
y_conv = ssm_convolution(A_bar, B_bar, C, x)
print("kernel (first 4 taps):", [round(v, 3) for v in ssm_kernel(A_bar, B_bar, C, 4).tolist()])
print("recurrent == convolutional:", torch.allclose(y_rec, y_conv, atol=1e-5))
kernel (first 4 taps): [0.942, 0.113, 0.074, 0.091]
recurrent == convolutional: True

Bridge the two outputs to a plot: the convolutional dots land exactly on the recurrent line, because they are the same function computed two ways.

pts = [{"t": t, "rec": float(y_rec[t]), "conv": float(y_conv[t])} for t in range(len(x))]
ojs_define(ssm_faces = pts)
NoteKey Insight

Train like a CNN, infer like an RNN. The convolutional face gives fast, parallel training on whole sequences; the recurrent face gives cheap, cache-free generation. Both are the same SSM — this duality is why SSMs are practical where plain RNNs were not.

The Parallel Scan

The convolution needs a fixed kernel, which only exists when the parameters do not change with position. There is a second, more general route to parallelism that does not need that assumption. Look at the recurrence as

\[ h_t = a_t\, h_{t-1} + b_t , \]

with \(a_t = \bar{A}\) and \(b_t = \bar{B}\, x_t\). Composing two consecutive steps is associative:

\[ (a_2, b_2) \circ (a_1, b_1) = (a_2 a_1,\; a_2 b_1 + b_2) . \]

Associativity is exactly the property a parallel prefix scan needs, so all of \(h_1, \dots, h_L\) can be computed in \(O(\log L)\) sequential depth instead of \(L\) serial steps. parallel_scan implements this and reproduces the sequential loop exactly:

from ssm import parallel_scan

a = torch.tensor([0.5, 0.9, 0.3, 0.7, 0.2])     # per-step decays
b = torch.tensor([1.0, 0.2, -0.5, 0.3, 1.0])    # per-step inputs
print("parallel scan h:", [round(v, 4) for v in parallel_scan(a, b).tolist()])

# Same numbers a plain left-to-right loop would give:
h, seq = 0.0, []
for i in range(5):
    h = a[i].item() * h + b[i].item()
    seq.append(round(h, 4))
print("sequential loop h:", seq)
parallel scan h: [1.0, 1.1, -0.17, 0.181, 1.0362]
sequential loop h: [1.0, 1.1, -0.17, 0.181, 1.0362]

This matters because it survives the next section. When Mamba makes the parameters input-dependent, the fixed-kernel convolution vanishes — but the scan does not (it never assumed constant \(a_t\)), so training stays parallel.

Selectivity: The Mamba Idea

Everything so far is time-invariant: \(\bar{A}\), \(\bar{B}\), \(C\) are the same at every step. That makes an SSM a fixed linear filter — and a fixed filter cannot do content-based reasoning. It cannot say “this token matters, remember it; that one is filler, skip it,” because its response to a token does not depend on the token.

Mamba’s insight (the S6 layer) is to make the parameters functions of the input: at each position, \(\Delta_t\), \(B_t\), and \(C_t\) are projected from \(x_t\). Now \(\Delta_t\) acts as an input-driven gate — a large \(\Delta_t\) writes the token firmly into the state and resets old memory; a small \(\Delta_t\) (via \(\bar{A} = \exp(\Delta_t a) \approx 1\), \(\bar{B} \approx 0\)) lets the token pass while the state holds its contents. The model selects what to keep.

Drive it: below, a stream of tokens flows into a one-dimensional state. Choose which tokens are “important” (high \(\Delta\)); the rest are near-skipped. Watch the state integrate only what you selected.

TipTry It!
  1. Select nothing. With every token skipped (\(\Delta \approx 0\)), the state barely moves — the model ignores the whole stream. Selectivity off = no memory writes.
  2. Select one token. The state jumps at that token and then holds (the skips have \(\bar{A}\approx 1\)), carrying that value forward — the SSM is copying the selected token across time.
  3. Select all. Now every token writes and the state churns like the time-invariant SSM. Selectivity is the dial between “ignore” and “integrate everything.”

Code: A Selective SSM (S6)

SelectiveSSM in ssm.py is a minimal, faithful S6 layer. The state matrix is diagonal and parameterized as \(A = -\exp(A_{\log})\) so it stays negative (stable) for any weights; \(\Delta\), \(B\), and \(C\) are linear projections of the input, with \(\Delta = \text{softplus}(\cdot)\) to keep it positive. The scan is the readable sequential selective_scan — Mamba’s engineering contribution is running that same scan fast on the GPU without ever materializing the state.

from ssm import SelectiveSSM

layer = SelectiveSSM(d_model=16, d_state=8)
x = torch.randn(2, 32, 16)          # (batch, length, channels)
y = layer(x)
print("input  shape:", tuple(x.shape))
print("output shape:", tuple(y.shape))

# Selectivity check: A is guaranteed negative (stable) whatever the parameters are.
A = -torch.exp(layer.A_log)
print("A all negative (stable):", bool((A < 0).all()))

# Because B, C, Δ are read from the input, two different inputs get different
# effective kernels — a fixed convolution could never do this.
y2 = layer(torch.randn(2, 32, 16))
print("input-dependent response:", not torch.allclose(y, y2))
input  shape: (2, 32, 16)
output shape: (2, 32, 16)
A all negative (stable): True
input-dependent response: True

A full Mamba block wraps this SSM with an input projection, a short causal depthwise convolution, and a SiLU gate (the same gated-MLP idea as m06’s SwiGLU) — but the state-space content is exactly what you built here.

Why They Challenge Attention

Now the payoff, made quantitative. A self-attention layer does two matmuls that scale with the sequence: \(QK^\top\) and the weighted sum over \(V\), each \(O(L^2 d)\) — quadratic. An SSM does an \(O(d\,N)\) update per token over \(L\) tokens — \(O(L\,d\,N)\), linear. At inference the difference is just as stark: attention’s KV cache holds \(O(L)\) vectors and grows with every token; an SSM carries a \(d\times N\) state that never grows. Drive the context length and watch the curves diverge:

from ssm import demonstrate_complexity
lengths = [512, 1024, 2048, 4096, 8192, 16384, 32768, 65536]
rows = demonstrate_complexity(lengths, d_model=1024, d_state=16)
ojs_define(ssm_cost = rows)
NoteKey Insight

The gap is not a constant factor — it widens with context. That is why SSMs are so attractive for long sequences: at 64K tokens the SSM does a tiny fraction of the compute and keeps a state thousands of times smaller than the KV cache. Mamba reports roughly 5× the generation throughput of a same-size Transformer.

Common Pitfalls

When building SSMs, watch out for:

  1. Forgetting \(A\) must be negative. Stability lives in \(\bar{A} = \exp(\Delta a) \in (0, 1)\), which needs \(a < 0\). Parameterize \(A = -\exp(A_{\log})\) so it can never drift non-negative and make the state explode.
  2. Discretization sign errors. ZOH is \(\bar{A}=\exp(\Delta A)\) (not \(\exp(-\Delta A)\)) and \(\bar{B}=A^{-1}(\exp(\Delta A)-I)B\). The near-\(a=0\) case needs the Euler limit \(\bar{B}\approx\Delta B\) or you divide by zero.
  3. Expecting a convolution after adding selectivity. The moment \(B\), \(C\), or \(\Delta\) depend on the input, the kernel is no longer fixed and the convolutional face is gone. Use the scan — it does not assume constant coefficients.
  4. Δ without softplus. \(\Delta\) is a step size and must be positive; feeding a raw linear projection lets it go negative and flips decay into growth. Mamba uses \(\Delta=\text{softplus}(\cdot)\).
  5. Confusing \(N\) with \(d\). The state dimension \(N\) (Mamba’s default is 16) is small and per-channel; the model width \(d\) is large. SSM cost is \(O(L\,d\,N)\) — linear in \(L\) because \(N\) is a small constant, not a second sequence axis.

Exercises

Exercise 1: The impulse response is the kernel

from ssm import ssm_recurrent, ssm_kernel
import torch

# Pick any diagonal SSM. Feed it a unit impulse (x = [1, 0, 0, ...]) through
# ssm_recurrent, and separately build ssm_kernel of the same length. Show the two
# are identical — the impulse response IS the convolution kernel. Then explain, in
# one sentence, why that means recurrent and convolutional forms must agree.

# Your implementation here:

Exercise 2: ZOH vs forward Euler

from ssm import discretize
import torch

# discretize uses exact ZOH. Write a forward-Euler discretization by hand
# (A_bar = 1 + delta*A, B_bar = delta*B) and compare to ZOH for a = -1 across
# delta in {0.1, 0.5, 1.0, 2.0}. Where do they agree, and where does Euler break
# (hint: what happens to 1 + delta*A when delta*A < -1)?

# Your implementation here:

Exercise 3: Selective copy

from ssm import selective_scan
import torch

# Build a batch of 1, length 8, d_model 1 input. Set delta HIGH (e.g. 1.5) at one
# chosen position and LOW (e.g. 0.02) everywhere else, with B = C = 1 and A = -1.
# Show the output "copies" the chosen token's value and holds it afterwards — the
# selective-copy task Mamba is designed to solve and a time-invariant SSM cannot.

# Your implementation here:

Summary

Key takeaways:

  1. An SSM is a linear recurrence over a fixed-size state. \(h_t = \bar{A}h_{t-1} + \bar{B}x_t\), \(y_t = Ch_t\) — a running summary that never grows, giving \(O(L)\) time and \(O(1)\) decode memory.
  2. Discretization connects continuous to discrete. Zero-order hold turns \((A, B)\) into \(\bar{A}=\exp(\Delta A)\) and \(\bar{B}=A^{-1}(\exp(\Delta A)-I)B\); \(A<0\) makes \(\bar{A}\) a stable decay in \((0,1)\).
  3. One SSM has two faces. Time-invariant, it is both a cheap recurrence (inference) and a parallel convolution (training) with kernel \(\bar{K} = (C\bar{B}, C\bar{A}\bar{B}, \dots)\) — provably the same function.
  4. The recurrence is an associative scan. \(h_t = a_t h_{t-1} + b_t\) composes associatively, so training parallelizes in \(O(\log L)\) depth even without a fixed kernel.
  5. Selectivity (Mamba’s S6) makes it content-aware. Letting \(\Delta, B, C\) depend on the input turns a fixed filter into a model that gates what to remember — at the cost of the convolution, recovered by the selective scan.
  6. SSMs challenge attention on cost. Linear vs quadratic compute and constant vs growing decode memory; the advantage widens with context, which is why SSMs are a serious contender for long-sequence modeling.

What’s Next

You now have a from-scratch alternative to attention itself — the last major architectural axis this book opens. The frontier from here is hybrid stacks that interleave SSM and attention layers (each covering the other’s weakness), and other sub-quadratic sequence mixers (linear and recurrent attention, RWKV, RetNet). The throughline is the one this module makes concrete: the Transformer is one point in a design space, and knowing the space is how you read — and build — what comes next.

Going Deeper

Core Papers:

Practical Resources:

  • The Annotated S4 — Rush & Karamcheti, a from-scratch, runnable walkthrough of structured state spaces.
  • state-spaces/mamba — the official reference implementation of the selective scan.