Mamba and State Space Models

What It Does

A language model has to keep something from every token it reads in order to go on writing. Transformer attention keeps the keys and values of every earlier token, and they pile up (see the plate “The KV Cache and PagedAttention”). A state space model does it differently: it folds what it has read into a state of fixed size and updates it once per token.

The 2023 Mamba paper (Gu and Dao) makes that update depend on each token. It reports 4 to 5 times the generation throughput of a Transformer of similar size, and, trained on the same data, a 2.8-⁠billion-parameter Mamba averaging 63.3 on six zero-shot tests, against 59.1 for the same-size Pythia and 61.7 for Pythia-⁠6.9B, more than twice its size.

On the stand, a reading machine. A paper tape printed with the tokens to read runs from the reel on the right, under the reading head, onto the reel on the left. Behind the head is the state: 8 rows of 16 rods, each rod’s height a number in the state, never more rods and never fewer. At the back, a ghosted rack holds what attention would keep: one more pair of pins for every token read.

Generation throughput: 4⁠–⁠5 times a Transformer of similar size · average score: Mamba-⁠2.8B 63.3, Pythia-⁠2.8B 59.1, Pythia-⁠6.9B 61.7

A Cache That Keeps Growing

To write each new token, attention uses the keys and values of every token before it, so all of them are kept: the KV cache. It is as long as the text read so far.

Take Pythia-⁠2.8B, the Transformer the Mamba paper compares against: 32 layers, keys and values of 2,560 numbers each per layer, 16 bits a number. One token takes 2 × 32 × 2,560 × 2 = 327,680 bytes; at 2,048 tokens, the longest it was trained on, one text’s cache is 671 MB.

The rack at the back is that cache: each token read drops in a pair of pins, the round head its key, the square head its value. When a text ends the rack empties and the next text starts from nothing. The state in front stays at its 128 rods throughout.

Pythia-⁠2.8B: 327,680 bytes a token · 2,048 tokens: 671 MB · it keeps everything it reads

A State of Fixed Size

State space models come from control and signal processing. A system holds a set of internal numbers, its state h; each input x updates it once, h ← Ā·h + B̄·x, and the output is read as y = C·h. Ā sets how much of the old state stays, B̄ how the input is written in, C what is read out.

However many tokens it reads, the state stays the same size. Mamba-⁠2.8B’s state is 6,553,600 numbers per text, 13.1 MB at 16 bits, as large after the first token as after the millionth. It has 64 layers of 5,120 channels, each channel with 16 numbers, plus a short convolution that holds the last 4 inputs. Pythia-⁠2.8B’s cache passes it at the 40th token and is 51 times larger at 2,048.

The state on the stand has 8 rows of 16, 128 rods. A row is a channel and a column one number of its state; 16 is Mamba’s default state size, and a real layer has thousands of channels. A rod above the midline is positive, below it negative.

Mamba-⁠2.8B: 6,553,600 numbers, 13.1 MB, whatever the length · Pythia-⁠2.8B’s cache passes it at token 40 and is 51 times larger at 2,048

Every Step the Same

The structured state space models before Mamba, such as S4 (Gu, Goel and Ré, ICLR 2022), used the same Δ, A, B and C for every token. The Mamba paper calls this linear time invariance: the model’s dynamics do not change over time.

Time invariance has a payoff: the whole output can be written as the input convolved with one fixed kernel, computed in one go for training, then run step by step for generation at the same cost per token. The price is that it cannot choose by content: a token to remember and a token that does not matter go into the state the same way, and older content fades at the same rate whatever comes next.

The stand switches to that rule: for every token, the valves on the right open to the same height. The springs under the rods are A, pulling every rod back to the midline, harder towards the right: A runs from −1 to −16, Mamba’s initial values.

The same Δ, A, B and C for every token · trainable as one convolution · no choosing by content

Remembering Selectively

The Mamba paper uses two synthetic tasks to show what time-invariant models lack. Selective copying: in a sequence 4,096 long, 16 tokens to remember sit among noise tokens at random places, and must be repeated in order at the end. Induction heads: having seen “Harry Potter”, the next “Harry” should be followed by “Potter”.

On selective copying, two-layer models score 18.3 % with S4 alone, 97.0 % with the selective version (the paper calls it S6), and 99.8 % inside Mamba’s block. On induction heads, trained on sequences of 256, Mamba is perfect at every length tested up to 1,048,576 tokens, 4,096 times longer; the attention models tested had fallen to 26.6⁠–⁠67.6 % by 1,024 tokens and were not tested beyond 16,384 for lack of memory.

The tape now carries such a sequence: one “moon” to remember, then a row of dots. Under Mamba’s rule, the rods “moon” wrote into stay red.

Selective copying: S4 18.3 % · S6 97.0 % · Mamba 99.8 % · induction heads: trained at 256, perfect up to 1,048,576

Chosen by the Token

Mamba’s change is that Δ, B and C are no longer fixed parameters but computed from the current token, each by a small linear projection. A stays fixed, but how much of the old state stays is exp(Δ·A), so when Δ changes, so does that.

The paper reads Δ this way: a large Δ resets the state and focuses on the current input; a small Δ keeps the state and ignores the current input. B decides whether and how a token is written into the state, C how the state is read into the output. With a state of a single number the update is exactly a recurrent network’s gate, h ← (1 − g)·h + g·x (the paper’s Theorem 1).

The row of red handwheels on the right is each row’s Δ: the higher it rises, the longer this step. The comb in front is B, the one behind is C; both change with every token. In the paper’s ablation (about 350 million parameters), perplexity is 8.71 with Δ, B and C all selective, 9.81 with only Δ, and 10.93 with none.

Large Δ: reset and take in this token · small Δ: keep the state and let it pass · with one number of state, a gate

A Scan on the Chip

Once Δ, B and C change with each token, the sequence can no longer be one convolution; it has to run as a recurrence. Done directly, the whole state of every channel at every token would be written to the GPU’s main memory, N times larger than the input and output (N is the state size).

Mamba uses three classic techniques. Kernel fusion: Δ, A, B and C are read from main memory into the chip’s fast SRAM, discretised and run through the recurrence there, and only the output is written back. A parallel scan for the recurrence. Recomputation: the intermediate states that training’s backward pass needs are not stored but computed again. (On main memory and SRAM, see the plate “FlashAttention and GPU Memory”.)

The paper measured this scan at 20 to 40 times the speed of an unfused PyTorch scan and, comparing that one operation alone, faster than FlashAttention-⁠2 beyond sequences of 2K (A100 80GB, batch 1). Training memory is about that of a same-size Transformer with FlashAttention.

The full state is never written out · discretised and scanned in SRAM · 20⁠–⁠40 times a PyTorch scan

How Much Faster, How Much Better

Generation, on one A100 80GB, after a 2,048-⁠token prompt, 128 tokens out: Mamba 1.4B peaks at 1,814 tokens a second (batch 128); the Transformer 1.3B peaks at 364 (batch 32, beyond which it runs out of memory). At batch 1 it is 140 against 79, 1.8 times. The paper writes 4 to 5 times: with no KV cache, Mamba can run much larger batches.

Quality: Mamba and Pythia share the tokenizer and the Pile data, 300 billion tokens each. At 2.8 billion parameters the six zero-shot tests average 63.3 against 59.1, and perplexity on the Pile is 6.22 against 6.73.

The paper states its limits too: the models it evaluated are small, the largest trained one has 2.8 billion parameters, and whether Mamba still compares favourably at larger sizes remains to be assessed.

Generation: up to 1,814 against 364 tokens a second · batch 1: 140 against 79 · average 63.3 against 59.1 on the same data

Mamba-⁠2: Two Algorithms

Mamba-⁠2 (Dao and Gu, ICML 2024) restricts A to a number times the identity. Then what a layer does to a whole sequence can be written as one lower-triangular matrix: entry i, j is C at step i times B at step j times all the decay in between. The paper calls this state space duality.

The same matrix can be computed two ways. Row by row as a recurrence, it is a state space model, with work growing linearly in length; all at once, it looks like masked attention, C as the queries, B the keys, the input the values, with no softmax and quadratic work. Mamba-⁠2 cuts the sequence into chunks: matrix multiplications inside each chunk, which the GPU’s tensor cores do best, and only the state passed from chunk to chunk.

The paper measured this layer at 2 to 8 times the speed of Mamba’s scan, faster than FlashAttention-⁠2 beyond 2K, with room for a much larger state: the released models keep 128 numbers per channel where Mamba keeps 16. It also notes that the duality does not cover standard softmax attention.

One layer = one lower-triangular matrix · run row by row, or multiplied in chunks · state size 16 → 128

Mixed Models

A state of fixed size has a cost. Jelassi and colleagues (ICML 2024) proved that a model whose state has too few bits must err when copying a long random string; and looking up a phone book, models trained on the Pile showed a 410-⁠million-parameter Pythia beating a 2.8-⁠billion-parameter Mamba once the book had 70 entries or more. In NVIDIA’s 2024 study of 8B models trained on the same data, pure Mamba and Mamba-⁠2 also lagged on five-shot MMLU and phone-book lookup.

Hence hybrids: mostly Mamba layers, with a few attention layers. AI21’s Jamba (2024) has one attention layer in every 8, with 52 billion parameters of which 12 billion are used per token; its Table 1, at 256K tokens: a 4 GB KV cache, against 32 GB for Mixtral and 128 GB for Llama-⁠2 7B. NVIDIA’s Nemotron-⁠H-⁠8B has 52 layers, 24 of them Mamba-⁠2 and 4 attention.

The Mamba-⁠2 paper found the same: replacing 6 of a 2.7B model’s 64 layers with attention brought perplexity on the Pile from 6.09 to 5.95. Its guess is that attention layers look earlier tokens up directly, so not everything has to be squeezed into the state.

Jamba: 1 attention layer to 7 Mamba layers · KV cache at 256K tokens: 4 GB against Mixtral’s 32 GB

What It Does Not Say

A state of fixed size is not a memory of unlimited length. It keeps only what its update rule picks, compressed, and a token once read cannot be looked at again. The paper’s language models were trained on sequences of 2,048 tokens; its million-token results come from a two-layer synthetic task and from DNA and audio.

Linear time is not faster everywhere. Both kinds of model read a prompt all at once; what Mamba saves is the work and memory per generated token. At batch 1, Mamba 1.4B is 1.8 times as fast as the Transformer 1.3B, not 5.

Nor is selection better everywhere. The paper writes that it fixes time-invariant models’ weakness on discrete data such as text and DNA, but can hold them back where they excel, on continuous signals such as audio.

A fixed state is a lossy one · no looking back at what was read · the saving is in every generated token

Taken Apart

The machine comes apart: the tape and its two reels, the reading head, the 128 rods and the springs under them, the frame, the Δ valves on the right, the B comb in front and the C comb behind, the outputs on the left; at the back, the rack of what attention would keep.

The head reads one token and computes Δ, B and C; the valves, the two combs and the rods move once. That is one step.

Tape · reading head · Δ valves · B comb · rods · springs A · frame · C comb · outputs y · key–value rack