L9.4 · LESSON 4 OF 6

State-space models

Explain expert routing, long-context costs, state-space models and multimodal token streams under explicit assumptions.

THE IDEA

Compressing the past into a state

A state-space model processes a sequence with a recurrence: a compressed state is updated at each step, h_t = a·h_(t−1) + b·x_t, and the output is read from it. It never builds an n × n attention matrix, and its state has a fixed size no matter how long the input is. The price is that the past must fit in that state.

decay a0.5

Click an input to cycle it through 0, 1, 2.

step	1	2	3	4	5	6
input x
state / output	2	2	1	0.5	0.25	0.125

h_t = a·h_(t−1) + x_t The whole history is compressed into one number (a vector in real models) whose size does not depend on how long the input is. attention would store 2 × n × d numbers; a state-space layer stores d.

Selective variants let the decay and input weights depend on the current input, so the model can choose to keep or overwrite its state. Hybrids interleave such layers with attention.

Set the decay and the inputs. Compare with the attention cache, which grows with the input.

Predict first. Must every sequence model keep a full attention matrix?

Chapter context

New architectures alter where computation and memory are spent. Compare them by tracing one request through their actual operations rather than assuming every improvement reduces every cost.

TRY IT YOURSELF

Part 1. Do it hands-on

Trace a state-space layer by hand, then write the scan, a selective variant and the unrolled form.

decay a0.5

Click an input to cycle it through 0, 1, 2.

step	1	2	3	4	5	6
input x
state / output	2	2	1	0.5	0.25	0.125

Selective variants let the decay and input weights depend on the current input, so the model can choose to keep or overwrite its state. Hybrids interleave such layers with attention.

outputs for inputs [2, 1, 0], a = 0.5:

output at step 5, a = 0.9, input [1,0,0,0,0] (4 decimals) attention cache numbers (n = 1000, d = 16)

To do: A state-space layer: h_t = 0.5·h_(t−1) + x_t, output y_t = h_t, h_0 = 0. Inputs [2, 1, 0]. Type the three outputs
To do: With the decay 0.9 instead, a single input 1 followed by zeros: the output at step 5 is 0.9⁴. Type it (4 decimals)
To do: An attention layer must remember a key and a value of width 16 for each of 1000 tokens: 2 × 1000 × 16 numbers. A state-space layer keeps a state of fixed size 16. How many numbers does the attention cache hold?

Part 2. Do it in code

ssm_scan(a, b, c, x): h_t = a·h_(t−1) + b·x_t with h_0 = 0 and output y_t = c·h_t. Return the array of outputs.
selective_scan(x, a_fn, b, c): like ssm_scan, but the decay at step t depends on the input: a_t = a_fn(x_t).
ssm_unrolled(a, b, c, x): compute the same outputs without a loop over state: y_t = Σ_{k≤t} c·b·a^(t−k)·x_k.
state_vs_cache(n, d): return (state size, attention cache numbers) = (d, 2·n·d).

import numpy as np

def ssm_scan(a, b, c, x):
    pass

def selective_scan(x, a_fn, b, c):
    pass

def ssm_unrolled(a, b, c, x):
    pass

def state_vs_cache(n, d):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Must every sequence model keep a full attention matrix?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a toy MoE with shared weights S, E equal experts of P weights and k experts active per token:

N_{total}=S+EP,\quad N_{active}\approx S+kP.

This is a parameter-accounting model, not a full FLOP or communication model. Naively materialized attention scores use $T^2$ entries per head.

Worked example

With 2 billion shared parameters and eight experts of 1 billion each, total parameters are 10 billion. If two experts are selected per token, approximately 4 billion parameters participate along that token's path. The other experts still require storage or loading. A 4096×4096 naive score matrix contains 16,777,216 entries per head.

Reference code

The following example uses the Python standard library and does not call a remote service.

shared, experts, per_expert, active=2,8,1,2
assert shared+experts*per_expert == 10
assert shared+active*per_expert == 4
assert 4096**2 == 16777216

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Active parameters and resident parameters differ. Linear sequence complexity does not make every model property linear. Context-window capacity and demonstrated context use are separate claims.