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.
- Assemble a full Mamba block (expand → causal conv → SiLU → S6 → gate → contract) and prove it is a causal, drop-in replacement for an attention block.
- Show why SSMs challenge attention: linear vs quadratic compute, constant vs growing decode memory.
- Build a hybrid decoder that interleaves attention and Mamba layers by a schedule (the Jamba and Griffin designs), and prove it stays causal.
- Read the KV-cache win as exact arithmetic — why a 1:7 stack holds ⅛ the cache — and see why a hybrid needs both mixers (local attention forgets past its window; the recurrent state remembers).
Prerequisites
This module requires familiarity with:
- Module 05: Attention — the quadratic-cost mechanism SSMs are the alternative to.
- Module 06: Transformer — the block an SSM layer replaces.
- Module 08: Generation — the KV cache whose growth SSMs sidestep.
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:
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: TrueBridge 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)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.
- Select nothing. With every token skipped (\Delta \approx 0), the state barely moves — the model ignores the whole stream. Selectivity off = no memory writes.
- 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.
- 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: TrueThe Full Mamba Block
The selective SSM is the engine, but you cannot yet stack it where an attention block used to sit. SelectiveSSM maps (batch, len, d_model) to the same shape, but a real Mamba block wraps that engine in five more pieces so it behaves like a self-contained sequence mixer — the same role a TransformerBlock plays. The block is a gated MLP with a scan in the middle:
- Expand. An input projection lifts the stream from d to 2d_\text{inner} (d_\text{inner}=E\,d, Mamba’s E=2) and splits it into a main branch x and a gate branch z.
- Convolve. A short causal depthwise convolution lets each channel mix a few adjacent tokens — a cheap local pre-filter before the global recurrence.
- Activate. A
SiLUnonlinearity on the conv output. - Scan. The selective SSM you just built — \Delta, B, C read from the conv’d x, A=-\exp(A_{\log}).
- Gate. Multiply the SSM output by \text{SiLU}(z) — the same gated-MLP idea as m06’s SwiGLU, letting the network suppress or pass each channel.
- Contract. An output projection maps d_\text{inner} back to d.
The headline is what all this preserves: the whole block is causal. The conv is left-padded so it never reads ahead, and the scan runs strictly left-to-right, so output position t depends only on inputs \le t. That is exactly the property a causal-attention block has — which is why a Mamba block is a drop-in replacement for one. MambaBlock in ssm.py builds all six stages and reuses the very selective_scan from above; it adds no new state-space idea, only the wrapper that turns the idea into a layer.
The Math: Six Stages
Writing u_t \in \mathbb{R}^{d} for the input at position t and dropping the batch axis:
\begin{aligned} [\,x, z\,] &= u\,W_\text{in}^\top, & W_\text{in} &\in \mathbb{R}^{2d_\text{inner}\times d} \\ \tilde{x} &= \text{SiLU}\!\big(\text{CausalConv1d}(x)\big), & &\text{(depthwise, width } d_\text{conv}) \\ [\,\Delta_\text{lr}, B, C\,] &= \tilde{x}\,W_x^\top, \quad \Delta = \text{softplus}(\Delta_\text{lr} W_{\Delta}^\top) \\ y &= \text{SelectiveSSM}(\tilde{x};\,\Delta, A, B, C, D) \\ \text{out} &= \big(y \odot \text{SiLU}(z)\big)\,W_\text{out}^\top, & W_\text{out} &\in \mathbb{R}^{d\times d_\text{inner}} \end{aligned}
\Delta is projected in two steps — a low-rank map to \text{dt\_rank}=\lceil d/16\rceil dimensions, then up to d_\text{inner} — Mamba’s parameter-thrifty way to make the per-channel step size input-dependent. Walk the tensor through each stage and watch it widen at the expand, ride through the middle, and contract back:
The bar heights encode tensor width. Step to stage 2 and the stream doubles (the expand); it holds that width through conv, activation, scan, and gate; then stage 7 halves it back to d_\text{model}. All the state-space work happens in the wide interior; the two projections are just the on- and off-ramps.
The Causal Convolution
Every stage but one you have already built. The new mechanism is the causal depthwise convolution — a deliberately tiny convolution (d_\text{conv}=4 in Mamba) that gives each channel a short-range view before the SSM’s long-range scan. Two design choices define it:
- Depthwise. Each channel is convolved with its own length-k kernel; channels never mix here (that is the projections’ job). So the weight is a small (d_\text{inner}, k) matrix, not a full (d_\text{inner}, d_\text{inner}, k) one.
- Causal. The signal is left-padded by k-1 and cropped back, so output t reads only x_{t-k+1},\dots,x_t and can never see the future. This is what keeps the whole block autoregressive-safe.
Drag t and watch the length-k window slide along the input — it always ends at “now” and never crosses the dashed line into the future:
- Drag t to 0. The window has only one valid tap (the rest are the left pad, treated as zero) — the first token can look back at nothing.
- Drag t to the end. The window is full: four taps, all in the past. The dashed “now” line always sits at the right edge of the highlighted block.
- There is never a highlighted cell to the right of “now.” That is the whole point — the conv, like the scan after it, cannot peek ahead.
Code: A Mamba Block from Scratch
MambaBlock in ssm.py assembles the six stages and reuses selective_scan for the middle. It preserves shape, so it drops into a residual stack exactly where an attention block would:
from ssm import MambaBlock
import torch
torch.manual_seed(0)
block = MambaBlock(d_model=16) # expand=2 → d_inner=32, d_conv=4, N=16
u = torch.randn(1, 20, 16)
y = block(u)
print("input shape:", tuple(u.shape))
print("output shape:", tuple(y.shape), "— same width, a drop-in mixer")
print("d_inner:", block.d_inner, " dt_rank:", block.dt_rank)input shape: (1, 20, 16)
output shape: (1, 20, 16) — same width, a drop-in mixer
d_inner: 32 dt_rank: 1The property that makes it a causal mixer — safe to train with next-token prediction — is exact. Edit a token in the future and every earlier output is untouched, bit for bit:
# Causality firewall: perturb a future token, watch earlier outputs stay identical.
t = 12
u_future = u.clone()
u_future[:, t, :] += 4.0 # change only token t
y_future = block(u_future)
print("outputs before t identical:", torch.equal(y[:, :t], y_future[:, :t]))
print("token t itself did change: ", not torch.equal(y[:, t], y_future[:, t]))outputs before t identical: True
token t itself did change: TrueBecause both the depthwise conv (left-padded) and the selective scan (left-to-right) respect the arrow of time, the composed block does too — no attention mask required. A full model stacks many of these, each wrapped in an RMSNorm and a residual (just like m06 wraps attention), which is the recipe hybrid SSM/attention models use.
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)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.
Hybrid Stacks: Interleaving the Two Mixers
The cost curves make SSMs look like a strict upgrade — so why does every strong Mamba-era model still keep some attention? Because the two mixers have opposite weaknesses, and the best architectures use each where it wins.
Intuition: Recall vs. Cost
Think about what each layer can remember:
- Attention reads any earlier token exactly. Query i can put all its weight on token 3 a thousand steps back and copy it verbatim. That perfect recall is what makes in-context learning and long-range lookup work — but it costs O(L^2) compute and a KV cache that grows with every token.
- A Mamba layer compresses the whole past into a fixed-size state. That is why it is cheap and constant-memory — but the state is a lossy summary, so exact recall of one specific far-away token is harder.
A hybrid stack interleaves a few attention layers among many Mamba layers. The recurrent majority carries the sequence cheaply; the attention minority supplies precise recall — and, crucially, only those few layers pay the KV-cache cost. Two real 2024 designs anchor the idea:
| Model | Attention | Interleave | KV cache |
|---|---|---|---|
| Jamba (AI21) | global | a:m = 1:7 — 1 attention layer per 8 |
8× smaller than all-attention (paper: 4 GB vs 128 GB at 256K) |
| Griffin (DeepMind) | local, window 1024 | 2 recurrent : 1 attention | constant in L — even attention caches only a window |
A hybrid is not a compromise — it is a division of labour. Jamba keeps global attention but on only ⅛ of layers, so its cache is 8× smaller. Griffin goes further: by making its attention local (a sliding window), even those layers cache a fixed number of tokens, so the total KV cache stops growing with context entirely.
The Schedule
A hybrid is defined by a schedule — a list saying whether each layer is "attn" or "mamba". make_schedule in hybrid.py builds the archetypes: one attention layer per period of P (P=8 for Jamba, P=3 for Griffin), placed mid-block so the stack starts and ends recurrent.
from hybrid import make_schedule, attention_fraction, kv_cache_reduction
for name in ["jamba", "griffin", "all_mamba"]:
sched = make_schedule(8, name)
print(f"{name:9s} {sched}")
print(f" attention fraction = {attention_fraction(sched):.3f}, "
f"KV cache {kv_cache_reduction(sched):.1f}× smaller than dense\n")jamba ['mamba', 'mamba', 'mamba', 'mamba', 'attn', 'mamba', 'mamba', 'mamba']
attention fraction = 0.125, KV cache 8.0× smaller than dense
griffin ['mamba', 'attn', 'mamba', 'mamba', 'attn', 'mamba', 'mamba', 'attn']
attention fraction = 0.375, KV cache 2.7× smaller than dense
all_mamba ['mamba', 'mamba', 'mamba', 'mamba', 'mamba', 'mamba', 'mamba', 'mamba']
attention fraction = 0.000, KV cache inf× smaller than dense
Drive the design space below: pick an archetype and watch the layer strip, the attention fraction, and the KV-cache bar (versus an all-attention stack of the same depth) update. Notice how few attention layers a hybrid actually needs.
from hybrid import (
make_schedule,
schedule_grid,
kv_growth_curve,
demonstrate_locality,
GRIFFIN_WINDOW,
)
_REF_L, _D, _NL = 8192, 512, 32
_grids = {
"jamba": schedule_grid(make_schedule(_NL, "jamba"), _REF_L, _D),
"griffin": schedule_grid(make_schedule(_NL, "griffin"), _REF_L, _D, window=GRIFFIN_WINDOW),
"all_attention": schedule_grid(make_schedule(_NL, "all_attention"), _REF_L, _D),
"all_mamba": schedule_grid(make_schedule(_NL, "all_mamba"), _REF_L, _D),
}
_mini = schedule_grid(make_schedule(8, "jamba"), 4096, _D)
_curve = kv_growth_curve([2 ** k for k in range(9, 19)], d_model=_D, n_layers=_NL)
ojs_define(
hybrid_grids=_grids,
hybrid_mini=_mini,
hybrid_curve=_curve,
hybrid_loc=demonstrate_locality(),
)- Jamba → Griffin: attention drops from ⅛ of layers to ⅓, but Griffin’s cache is still smaller at long context — because its window caps each layer’s cache. The next plot shows why.
- all_mamba: the KV cache vanishes entirely — but a pure-SSM model has no exact recall. The attention layers are what buy that back.
Anchor 1: A Hybrid Is Still Causal
Whatever the schedule, the stack must never let an output peek at a future token — otherwise next-token training is a lie. Both mixers are causal (attention via its mask, Mamba via its left-padded conv and left-to-right scan), so any interleaving of them is too. HybridLM in hybrid.py assembles the stack pre-norm (x ← x + mixer(RMSNorm(x)), exactly as m06 wraps attention) with a tied LM head:
import torch
from hybrid import HybridLM, make_schedule
torch.manual_seed(0)
schedule = make_schedule(8, "jamba") # 7 Mamba + 1 global-attention layer
lm = HybridLM(vocab_size=32, d_model=16, schedule=schedule).eval()
ids = torch.randint(0, 32, (1, 12))
with torch.no_grad():
logits = lm(ids)
ids_future = ids.clone()
ids_future[0, 6] = (int(ids[0, 6]) + 1) % 32 # edit a FUTURE token
logits_future = lm(ids_future)
print("logits before pos 6 identical:", torch.equal(logits[:, :6], logits_future[:, :6]))
print("pos 6 itself changed: ", not torch.equal(logits[:, 6], logits_future[:, 6]))logits before pos 6 identical: True
pos 6 itself changed: TrueAnchor 2: The KV-Cache Win, as Arithmetic
Only attention layers cache K and V; a Mamba layer carries a fixed d\times N state instead. So a hybrid’s cache is exactly the dense cache scaled by the attention fraction — a 1:7 Jamba stack holds 1/8 of it. And with Griffin’s sliding window, each attention layer caches at most window tokens, so the total flattens once the context passes the window. Drive the context length:
from hybrid import demonstrate_kv_savings
s = demonstrate_kv_savings(length=4096, d_model=512, n_layers=32)
print(f"dense KV : {s['dense'] / 1e6:8.1f} MB")
print(f"Jamba KV : {s['jamba'] / 1e6:8.1f} MB ({s['jamba_reduction']:.0f}× smaller)")
print(f"Griffin : {s['griffin'] / 1e6:8.1f} MB (flat in L: {s['griffin_flat']})")dense KV : 268.4 MB
Jamba KV : 33.6 MB (8× smaller)
Griffin : 23.1 MB (flat in L: True)The 8× is not a tuned result — it is the schedule. One attention layer per eight means one-eighth of the caches, full stop. The GQA and exact widths in Jamba’s paper turn that ratio into the headline 4 GB-vs-128 GB figure; the ratio itself is architecture, and it is what hybrid_kv_cache_bytes reproduces.
Anchor 3: Why a Hybrid Needs Both Mixers
If Mamba is cheap and attention is expensive, why not drop attention entirely? Run the layers side by side. Griffin’s attention is local: query i only sees the last window tokens. Perturb a token just outside that window and the attention output does not flinch — it never saw it. But the same edit does move a Mamba layer’s output, because its recurrent state carried that token forward. The recurrent layers are the hybrid’s long-range memory; the local-attention layers are its precise short-range recall.
from hybrid import demonstrate_locality
loc = demonstrate_locality(window=4, length=16)
print(f"perturb token {loc['p']} (just outside the {loc['window']}-token window of query {loc['i']}):")
print(f" local attention output at {loc['i']} unchanged : {loc['local_forgets']} ← it forgot")
print(f" GLOBAL attention would use it : {loc['global_uses_it']} ← locality, not position")
print(f" Mamba output at {loc['i']} changed : {loc['recurrent_remembers']} ← the state remembered")perturb token 11 (just outside the 4-token window of query 15):
local attention output at 15 unchanged : True ← it forgot
GLOBAL attention would use it : True ← locality, not position
Mamba output at 15 changed : True ← the state rememberedThis is the whole argument for hybrids in one experiment: neither mixer alone is enough at long context — local attention forgets, and a pure recurrence can’t recall exactly — so the strongest models keep both. Step through a Jamba mini-stack to see what each layer carries:
Common Pitfalls
When building SSMs, watch out for:
- 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.
- 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.
- 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.
- Δ 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).
- 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.
- A non-causal conv in the block. Use a centered or right-padded convolution and output t suddenly sees future tokens — the block leaks the answer and can no longer be trained with next-token prediction. The conv must be left-padded by k-1 and cropped, exactly like
causal_depthwise_conv1d. - Forgetting the gate branch. The block’s second projection output z is not decoration: without the \odot\,\text{SiLU}(z) gate the block is just an SSM with projections and loses the multiplicative, content-dependent control that makes gated architectures expressive.
- Thinking a hybrid halves the cache. The KV cache scales with the attention fraction, not with “some attention vs none.” A 1:7 Jamba stack holds 1/8 of the dense cache, not 1/2; the reduction is n_\text{layers}/n_\text{attn}. Count the attention layers, not the presence of attention.
- Dropping attention entirely to save more. A pure-SSM stack has the smallest cache (zero) but weaker exact recall — its state is a lossy summary. Likewise a pure local-attention stack forgets everything past its window. Hybrids exist because each mixer alone leaves a gap the other fills; removing one reopens it.
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:Exercise 4: Prove the block is causal
from ssm import MambaBlock
import torch
# Build a MambaBlock(d_model=8) and a random input (1, 16, 8). Run it once. Then
# perturb a single FUTURE token (say position 10) and run it again. Show every
# output position < 10 is bit-for-bit identical (torch.equal), and that position 10
# changed. Then break causality on purpose: swap causal_depthwise_conv1d's left pad
# for a symmetric pad and watch an earlier output move — the property is fragile.
# Your implementation here:Exercise 5: Design your own schedule
from hybrid import make_schedule, attention_fraction, hybrid_kv_cache_bytes
import torch
# Build a HybridLM with a schedule you design (try make_schedule(12, 4) — one
# attention layer every 4). Compare its KV-cache bytes at L=8192, d_model=512 to
# both an all-attention and an all-mamba stack of the same depth, and report its
# attention fraction. Which schedule would you pick to fit a 200K-token context in
# 16 GB, and why? (Hint: kv_cache_reduction is n_layers / n_attn.)
# Your implementation here:Exercise 6: Local attention forgets, the state remembers
from hybrid import CausalSelfAttention
from ssm import MambaBlock
import torch
# Reproduce the locality anchor by hand. Build a local CausalSelfAttention(16, 4,
# window=4) and a MambaBlock(16). Feed both a random (1, 20, 16) input, then edit a
# token far outside the window (e.g. position 2, query position 19). Show the local
# attention output at 19 is unchanged (torch.equal) while the Mamba output moves —
# then explain in one sentence why a Griffin stack still handles long-range recall.
# Your implementation here:Summary
Key takeaways:
- 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.
- 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).
- 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.
- 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.
- 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.
- The Mamba block turns the S6 core into a layer. Expand → causal depthwise conv → SiLU → selective scan → \odot\,\text{SiLU}(z) gate → contract. Because the conv is left-padded and the scan runs left-to-right, the whole block is causal — a drop-in replacement for a causal-attention block, no mask needed.
- 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.
- Hybrid stacks interleave the two mixers. A schedule of mostly-Mamba layers with a few attention layers (Jamba’s 1:7 global, Griffin’s 2:1 local) keeps the attention only where precise recall matters. The stack stays causal for any schedule, and the KV cache scales with the attention fraction — a 1:7 stack holds ⅛ of the dense cache; local windows make it flat in L.
- Each mixer covers the other’s weakness. Local attention forgets past its window; a fixed recurrent state can’t recall one far token exactly. Provably, an out-of-window edit leaves local attention unchanged but moves the recurrent state — which is exactly why the strongest efficient models keep both.
What’s Next
You now have a from-scratch alternative to attention and a from-scratch hybrid that interleaves it with attention — the way real Mamba-era models (Jamba, Griffin) are actually built. The frontier from here is training such a stack end to end, folding MoE into the hybrid (Jamba puts experts on the FFN of every other layer, tying this to m11), and the other sub-quadratic mixers next door in Module 22: Linear & Recurrent Attention — whose gated variants slot into the very same schedule. 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:
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces — Gu & Dao (2023), the selective SSM (S6), hardware-aware scan, and ~5× throughput built toward here.
- Efficiently Modeling Long Sequences with Structured State Spaces (S4) — Gu, Goel & Ré (2021), the structured-SSM predecessor that first made SSMs competitive on long-range tasks.
- HiPPO: Recurrent Memory with Optimal Polynomial Projections — Gu et al. (2020), the theory of how to initialize A for long memory.
- Transformers are SSMs: Generalized Models and Efficient Algorithms (Mamba-2) — Dao & Gu (2024), the duality tying state-space models and attention together.
- Jamba: A Hybrid Transformer-Mamba Language Model — Lieber et al. (2024), the 1:7 attention:Mamba interleave with MoE, and the 8× KV-cache reduction at 256K context reproduced here.
- Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models — De et al. (2024), the recurrent + sliding-window-attention hybrid whose fixed window keeps the cache flat in context length.
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.