WEBVTT

NOTE AI Engineering Visualised: https://understandingdata.com/ai-engineering-visualised/state-space-models/

1
00:00:00.245 --> 00:00:18.334
Attention pays for everything it has read: each new token scores itself against every cached key, so eight tokens take 36 scores, and by 4,096 tokens one layer has computed 8.4 million scores and holds a 32 MiB cache.

2
00:00:18.460 --> 00:00:40.118
A state-space model reads the same tokens through one running state instead: each step keeps 0.8 of the state and adds 0.2 of the token, so the inputs 5, 0, 0 and 10 leave it at 1, then 0.8, then 0.64 and then 2.51, which fades again through the zeros that follow.

3
00:00:40.120 --> 00:01:03.346
That state never grows: a Mamba layer of the same width holds 128 KiB at any length, attention’s cache matches it at 16 tokens and is 256 times larger by 4,096, and each new token costs one fixed update instead of reading every earlier key back from memory.

4
00:01:03.346 --> 00:01:23.152
Those two numbers come from a step size, Δ, because the model is defined in continuous time and Δ sets how far each token moves it: a Δ of 0.1 keeps 90% of the state and barely writes, while a Δ of 2 keeps 14% and nearly overwrites it.

5
00:01:23.180 --> 00:01:46.325
What makes Mamba different is selectivity: S4 uses one Δ for every token, but Mamba computes Δ, B and C from each token, so in this toy the filler gets a tiny Δ and slips past, the key gets a large one, and 91% of the 7 survives nine fillers, against 0.4% with one fixed Δ.

6
00:01:46.325 --> 00:02:04.655
Training cannot wait for one token at a time, but each step is linear, so two steps compose into one of the same form, and a parallel scan finds all eight states in three rounds instead of seven, while inference simply runs the recurrence, one update per token.

7
00:02:04.680 --> 00:02:25.904
The price is exact recall: a fixed state is a summary, so when this toy writes 64 key and value pairs into a 16 by 16 state, it returns the right value only 28% of the time, while attention, which keeps every key, finds all 64, and copying fails the same way.

8
00:02:25.904 --> 00:02:43.531
So production models mix the two: Jamba puts one attention layer among every seven Mamba layers, which keeps an exact lookup and cuts the cache to an eighth, and before you pick a long-context model, test recall and copying at your real lengths, not just perplexity.