def matrix_state_dual(Q, K, V, a):
"""Y = (L o Q K^T) V with L_ts = a_t ... a_{s+1} for t >= s, else 0."""
cum = jnp.cumsum(jnp.log(a))
logL = cum[:, None] - cum[None, :] # Sum of logs over (s, t]
causal = jnp.tril(jnp.ones((len(a), len(a)), bool))
L = jnp.exp(jnp.where(causal, logL, -jnp.inf))
return (L * (Q @ K.T)) @ V
with jax.default_matmul_precision('highest'):
err = (jnp.abs(matrix_state_dual(Q, K, V, a) - y_rec).max()
/ jnp.abs(y_rec).max())
print(f'relative deviation, dual vs recurrence: {float(err):.2e}')
assert err < 1e-3 # Two orders, one computation
ones = jnp.ones(T) # All a_t = 1
err = jnp.abs(matrix_state_dual(Q, K, V, ones)
- jnp.tril(Q @ K.T) @ V).max()
print(f'a_t = 1, deviation from masked linear attention: {float(err):.2e}')
assert err < 1e-6 # L is literally the causal maskThe Matrix State: From Linear Attention to Mamba-2
Dive into Deep Learning · §12.4
The matrix state
one recurrence, measured capacity · the decay ladder · state-space duality · chunked training · the family table
Two roads to one recurrence
Ch. 10 ended linear attention as a recurrence; ch. 12 built selective SSMs from continuous time. Same object:
\mathbf{S}_t = \mathbf{D}_t\, \mathbf{S}_{t-1} + \mathbf{k}_t \mathbf{v}_t^\top, \qquad \mathbf{o}_t = \mathbf{S}_t^\top \mathbf{q}_t
- Transition \mathbf{D}_t: what survives. Write \mathbf{k}_t\mathbf{v}_t^\top: what is filed.
- Linear attention: \mathbf{D}_t = \mathbf{I} — never forget.
What the modern family drops from ch. 10’s construction: the feature map \phi and the normalizer state \mathbf{z} (output normalization instead). One member keeps \mathbf{z} — wait for the mLSTM.
What the memory costs, measured
Store n unit-norm pairs, read key j: \mathbf{S}^\top \mathbf{k}_j = \mathbf{v}_j + \sum_{i \neq j} (\mathbf{k}_i^\top \mathbf{k}_j)\,\mathbf{v}_i
- Measured error on the (n-1)/d_k prediction — for i.i.d. isotropic unit keys; recall collapses past n \approx d_k.
- Correlated keys (dotted) revive the cancelled cross terms: width stops helping.
- Capacity is set by width, not time: the enemy of memory is other memories.
Forgetting: the decay ladder
Three one-line fixes for crowding, on the same recurrence:
\gamma\,\mathbf{S}_{t-1} \;\to\; a_t\,\mathbf{S}_{t-1} \;\to\; \mathrm{diag}(\boldsymbol{\alpha}_t)\,\mathbf{S}_{t-1}\;\; (+\;\mathbf{k}_t\mathbf{v}_t^\top)
- RetNet: fixed \gamma per head · Mamba-2: input-dependent scalar · GLA / RWKV-6: input-dependent diagonal — the selective SSM per channel.
Unroll it: the state-space duality
\mathbf{Y} = \big(\mathbf{L} \circ \mathbf{Q}\mathbf{K}^\top\big)\,\mathbf{V}, \qquad L_{ts} = a_t\, a_{t-1} \cdots a_{s+1}
The pivot: all a_t = 1 → \mathbf{L} is the causal mask. A gated linear recurrence is attention with a learned, decaying mask (SSD, Dao & Gu 2024).
- Quadratic mode: materialize \mathbf{L} \circ \mathbf{Q}\mathbf{K}^\top — attention.
- Linear mode: sweep the sum left to right — the recurrence, state \mathbf{S}_t.
- Two contraction orders, one computation — an identity for this scalar-gated (1-semiseparable) linear family, not softmax attention.
The duality, asserted
relative deviation, dual vs recurrence: 4.14e-06
a_t = 1, deviation from masked linear attention: 0.00e+00- Agreement to float rounding; at a_t = 1 the dual reduces exactly to ch. 10’s masked linear attention.
- Ch. 10’s promise is now a computation you have run.
Chunked: mostly matmul, a little scan
Partition the semiseparable matrix into chunks of C tokens:
- Diagonal blocks: C \times C masked-attention matmuls, all parallel.
- Off-diagonal blocks: rank \le d_k — a d_k \times d_v state passed chunk to chunk, T/C steps.
A dozen lines, verified
segsum builds the within-chunk mask in log space (-\infty above the diagonal, before any exponential):
def segsum(log_a):
"""Segment sums: (..., C) -> (..., C, C) with entries sum over (s, t]."""
C = log_a.shape[-1]
cum = jnp.cumsum(log_a, -1)
diff = cum[..., :, None] - cum[..., None, :]
keep = jnp.tril(jnp.ones((C, C), bool))
return jnp.where(keep, diff, -jnp.inf) # Mask in log space
def matrix_state_chunked(Q, K, V, a, chunk_size=64):
"""The same Y, chunk by chunk: attention inside, recurrence across."""
T, C = Q.shape[0], chunk_size
Qc, Kc = Q.reshape(-1, C, Q.shape[-1]), K.reshape(-1, C, K.shape[-1])
Vc, log_a = V.reshape(-1, C, V.shape[-1]), jnp.log(a).reshape(-1, C)
L = jnp.exp(segsum(log_a)) # (T/C, C, C)
b = jnp.exp(jnp.cumsum(log_a, -1)) # Decay since chunk entry
Y = (L * (Qc @ Kc.transpose(0, 2, 1))) @ Vc # Diagonal blocks
def boundary(S, chunk): # T/C boundary steps
Qb, Kb, Vb, bb, Yb = chunk
Yb = Yb + (Qb * bb[:, None]) @ S # Read the carried state
S = bb[-1] * S + (Kb * (bb[-1] / bb)[:, None]).T @ Vb
return S, Yb
S0 = jnp.zeros((Q.shape[-1], V.shape[-1]))
_, Y = jax.lax.scan(boundary, S0, (Qc, Kc, Vc, b, Y))
return Y.reshape(T, -1)
with jax.default_matmul_precision('highest'):
for chunk_size in [16, 64, 256]:
err = (jnp.abs(matrix_state_chunked(Q, K, V, a, chunk_size)
- y_rec).max() / jnp.abs(y_rec).max())
print(f'chunk size {chunk_size:4d}: relative deviation {float(err):.2e}')
assert err < 1e-3 # Chunked == recurrencechunk size 16: relative deviation 2.13e-07
chunk size 64: relative deviation 3.84e-07
chunk size 256: relative deviation 1.54e-06The trade, measured
C interpolates recurrence (C = 1) → quadratic dual (C = T):
form time (ms) temp (MiB)
recurrence 49.63 0.0
chunked (16) 2.80 1.1
chunked (64) 1.62 1.1
chunked (256) 0.46 5.0
chunked (1024) 0.37 21.0
quadratic dual 0.45 68.0
(NVIDIA GeForce RTX 4090, float32, jax 0.10.2; teaching implementations, all forms compiled)- Token loop: launch-bound, slowest by far. Dual: matmul-fast, T^2 memory.
- Chunked: matmul speed at an order of magnitude less memory.
What the hardware bought
- Tensor cores run matmuls; Mamba-1’s scan is elementwise (bandwidth-bound). Mamba-2’s scalar transition buys the chunked form.
- The paper’s empirical scaling observation: training several times faster and state size N: 16 → 64–256 at little extra wall clock — the capacity law says what that buys.
- Frontier: Mamba-3 upgrades the discretization itself (trapezoidal rule, complex states, MIMO).
per layer, per sequence, fp16:
matrix state, one 64x64 head: 0.01 MiB, constant
Mamba-2 2.7B, SSM 655,360 + conv 15,360 elements: 1.28 MiB, constant
GQA KV cache at n = 4,096: 16.00 MiB and growing
GQA KV cache at n = 131,072: 512.00 MiB and growing
GQA KV cache at n = 1,048,576: 4096.00 MiB and growingThe flat line vs. the growing cache — the hybrid-design problem of the chapter’s last section.
The family, so far
| model | transition \mathbf{D}_t | write |
|---|---|---|
| linear attention | \mathbf{I} | add \phi(\mathbf{k}_t)\mathbf{v}_t^\top |
| RetNet | \gamma\,\mathbf{I} fixed | add \mathbf{k}_t\mathbf{v}_t^\top |
| Mamba | input-dep. diagonal (per channel) | add \Delta_t u_t \mathbf{B}_t |
| GLA / RWKV-6 | input-dep. diagonal | add \mathbf{k}_t\mathbf{v}_t^\top |
| Mamba-2 | input-dep. scalar | add \mathbf{k}_t\mathbf{v}_t^\top |
| mLSTM | scalar exp-gates + normalizer | add i_t\,\mathbf{k}_t\mathbf{v}_t^\top |
Three lineages — attention, SSMs, RNNs — one recurrence. Read the write column: every row only adds.
mLSTM: the one that kept the normalizer
\mathbf{S}_t = f_t \mathbf{S}_{t-1} + i_t \mathbf{k}_t\mathbf{v}_t^\top, \quad \mathbf{z}_t = f_t \mathbf{z}_{t-1} + i_t \mathbf{k}_t, \quad \mathbf{o}_t = \frac{\mathbf{S}_t^\top \mathbf{q}_t}{\max(|\mathbf{z}_t^\top \mathbf{q}_t|, 1)}
Exponential gates i_t = \exp(\tilde{i}_t) overflow float32 in a few dozen steps; the fix is online softmax’s running max, carried as a state m_t.
d_k, d_v, T = 8, 8, 64
kk = jax.random.split(jax.random.key(2), 5)
q, k, v = (jax.random.normal(key, (T, d))
for key, d in zip(kk, (d_k, d_k, d_v)))
i_pre = jax.random.normal(kk[3], (T,)) # Gate pre-activations
f_pre = 2.0 + jax.random.normal(kk[4], (T,))
def mlstm_naive(q, k, v, i_pre, f_pre, dtype):
"""The unstabilized recurrence (NumPy), in a given precision."""
q, k, v = (np.asarray(x, dtype) for x in (q, k, v))
i_g, f_g = np.exp(np.asarray(i_pre, dtype)), np.exp(np.asarray(f_pre, dtype))
S, z = np.zeros((d_k, d_v), dtype), np.zeros(d_k, dtype)
outputs = []
for t in range(T):
S = f_g[t] * S + i_g[t] * np.outer(k[t], v[t])
z = f_g[t] * z + i_g[t] * k[t]
outputs.append(q[t] @ S / max(abs(q[t] @ z), 1.0))
return np.stack(outputs)
def mlstm_stabilized(q, k, v, i_pre, f_pre):
"""Carry m_t and rescale the past, as in online softmax."""
def step(carry, x):
S, z, m = carry
q_t, k_t, v_t, i_t, f_t = x
m_new = jnp.maximum(f_t + m, i_t)
f_g, i_g = jnp.exp(f_t + m - m_new), jnp.exp(i_t - m_new)
S = f_g * S + i_g * (k_t[:, None] * v_t[None, :])
z = f_g * z + i_g * k_t
o = q_t @ S / jnp.maximum(jnp.abs(q_t @ z), jnp.exp(-m_new))
return (S, z, m_new), o
init = (jnp.zeros((d_k, d_v)), jnp.zeros(d_k), -jnp.inf)
return jax.lax.scan(step, init, (q, k, v, i_pre, f_pre))[1]
exact = mlstm_naive(q, k, v, i_pre, f_pre, np.float64)
with np.errstate(over='ignore', invalid='ignore'): # The overflow is the point
naive = mlstm_naive(q, k, v, i_pre, f_pre, np.float32)
bad = ~np.isfinite(naive).all(-1) # argmax(bool) would report 0
first_bad = int(bad.argmax()) if bad.any() else 'none'
stab = mlstm_stabilized(q, k, v, i_pre, f_pre)
err = jnp.abs(stab - exact).max() / jnp.abs(exact).max()
print(f'float32 unstabilized: first non-finite output at step {first_bad}')
print(f'float32 stabilized vs float64: relative deviation {float(err):.2e}')
assert bool(jnp.isfinite(stab).all()) and err < 1e-3float32 unstabilized: first non-finite output at step 47
float32 stabilized vs float64: relative deviation 1.38e-06The column left open
- Decay manages crowding; it does not cure it — forgetting is indiscriminate, interference is specific.
- Not one model in the table can edit what it wrote about \mathbf{k} without decaying everything else.
The fast weight programmers’ repair (1990s → 2021): read what the memory holds at \mathbf{k}_t, write only the correction.
A memory that edits: DeltaNet, next section.