State Space Models Explained: Linear-Time Sequences with a Fixed-Size State

Memory often runs out first. A transformer keeps a KV cache, the saved keys and values of every token it has read, and the cache grows with each new token. State space models are a family of sequence models built to avoid that growth.

Picture a common setup. A team in Bengaluru runs a model that summarises long customer-support transcripts, about 32,768 tokens each, on one data-centre GPU, a graphics processor used for model work. The setup is made up, but every number below comes from our own runs or a cited paper.

Say 60 GiB of GPU memory is left for the cache after the model loads, a made-up round number. GiB means gibibyte, 1,073,741,824 bytes. With the made-up model shape used in this guide, each transcript's KV cache takes 4 GiB, so only 15 transcripts fit at once.

A state space model keeps one fixed-size state per transcript instead, 16 MiB in our made-up shape, whatever the length. MiB means mebibyte, 1,048,576 bytes. This guide shows how that works, why training stays fast, and the one job where it lets the team down: quoting an exact order number from deep inside a transcript.

Eight words carry the rest of this guide. Each one is tied to the transcript summariser.

• Token: a small piece of text, often a word or part of a word. The model reads each transcript as a sequence of tokens.
• Sequence model: a model that reads tokens in order and produces an output at each position, such as a guess for the next token.
• Attention: the step in a transformer where each new token compares itself with every earlier token's key and mixes in the values of the ones that match. A key is a token's label for matching; a value is the content it passes on.
• KV cache: the keys and values attention saves for every earlier token, in every layer, so they are not recomputed. It grows by the same number of bytes per token.
• State: a fixed-size list of numbers that summarises everything read so far. A state space model rewrites it once per token.
• Recurrence: a rule that computes the new state from the old state and the current token, one step at a time.
• Linear update: the new state is the old state times some factors plus the token times some factors. The factors may depend on the token, but the state is only scaled and added, never squashed.
• Channel: one of the numbers in the vector that represents a token inside the model. Many state space layers run a separate small state for each channel.

State space models carry a fixed-size state from token to token and update it with a linear rule. For token x_t, the new state is h_t = A_bar h_(t-1) + B_bar x_t, and the output is y_t = C h_t.

Read the rule as three jobs. A_bar decides how much of the old state survives, so it controls fading. B_bar decides how the new token is written in, and C decides how the state is read out.

The name comes from control theory, the engineering of systems such as circuits and motors, where a state space model describes a system by its internal state. Deep learning papers in 2021 turned it into a layer that mixes information along a sequence.

One 2021 paper (arXiv 2110.13985) showed that such a layer is closely related to three older families: recurrent networks, convolutions and continuous-time models. That three-way link is the key to everything below.

A real model stacks many of these layers. Inside a layer, each channel gets its own small state, and other parts of the layer mix the channels together.

An everyday picture helps. A cricket scorer does not record a video of every ball. After each ball they update a small scorecard of runs, wickets and overs, and anyone can read the score from that card.

Map it part by part. Each ball is a token, the scorecard is the state, the scorer's update habit is A_bar and B_bar, and reading the score aloud is C.

Where the picture breaks: a scorecard keeps exact totals that never fade. A model's state is a blend of learned numbers that fade over time, and nobody chose what each number means. And like a scorecard, it cannot tell you what happened on ball 37.

For long inputs, a lot. A KV cache grows by a fixed number of bytes for every token, while the state stays the same size at any length.

KV cache bytes per token = 2 x layers x KV heads x head size x bytes per number. The 2 counts one key and one value. A KV head is one of the parallel attention units that stores keys and values, and head size is the length of each key.

The paper behind vLLM's paged KV cache (arXiv 2309.06180) uses the same count for a model where every head stores keys and values. Its example needs about 800 KB per token and up to 1.6 GB for one 2,048-token request.

State bytes per sequence = layers x channels x state size x bytes per number. There is no token count anywhere in that formula.

Our made-up attention model has 32 layers, 8 KV heads, head size 128 and a 16-bit cache, 2 bytes per number, which gives 128 KiB per token. KiB means 1,024 bytes. Our made-up state space model has 32 layers, 8,192 channels, state size 16 and a 32-bit state, which gives 16 MiB.

The two are equal at only 128 tokens for these shapes. Past that, the KV cache is always bigger, and at 131,072 tokens it is 1,024 times bigger than the state.

Treat the 3,840 as a memory ceiling, not a real batch size. Weights, temporary working memory and compute limit the batch long before that, and we did not measure speed.

Work per new token follows the same pattern. A 2021 structured state space paper (arXiv 2111.00396) notes that in recurrent mode its model needs constant memory and computation per time step. Standard attention, by contrast, reads the whole cache for every new token.

The 2023 selective state space paper (arXiv 2312.00752) reports about 5 times higher generation throughput than transformers in its own tests. Throughput here means tokens generated per second.

They lose on exact recall: copying long text, or looking up one fact far back in the input. A fixed state must keep summarising, so a detail that was not kept when it was read cannot be fetched later.

A 2024 paper on copying (arXiv 2402.01032) proves a hard limit. If random strings carry more bits than the state can hold, a fixed-state model must copy more than half of them wrongly.

The same paper builds a two-layer transformer that can copy strings whose length grows exponentially with its number of heads. Attention can do this because it keeps every token and can look back at any of them.

Its experiments show the gap well before the hard limit. With models of about 160 million parameters (learned numbers), transformers needed about 100 times fewer training examples to learn copying than the best fixed-state models.

Trained on strings of up to 50 tokens, the fixed-state models dropped to near zero on longer strings. One transformer variant stayed almost perfect up to 1,000 tokens.

On a phone-book lookup with pretrained models (already trained on general text), the smallest transformer tested, 410 million parameters, beat the largest fixed-state model, 2.8 billion parameters. That held once the phone book had 70 or more entries.

Raw capacity is not the whole story. Our made-up 16 MiB state holds 134,217,728 bits, enough in theory for about 8.6 million tokens from a 50,000-token vocabulary, the set of tokens a model knows. The paper's point is that a model may not learn to pack its state that well.

Our own NumPy toy shows the same shape. A fixed state of 2,048 numbers recalled 99% or more of its key-value pairs up to 32 pairs, then fell to 92% at 64 pairs and 38% at 256. A KV cache stayed at 100%.

For the transcript summariser, a summary needs the gist, which a state can keep. Quoting the exact order number a customer typed at token 3,000 is the risky job.

Because the original model is written in continuous time, while a computer reads tokens in discrete steps. Discretisation turns the continuous rule into a per-token rule, using a step size called delta.

The continuous form is h'(t) = A h(t) + B x(t), with output y(t) = C h(t). Here h'(t) is how fast the state is changing at time t.

The selective state space paper uses zero-order hold, which assumes the input stays constant for the length of one step. That gives A_bar = exp(delta A) and B_bar = (delta A)^-1 (exp(delta A) - I) delta B.

Here exp is the exponential function applied to a matrix, and I is the identity matrix, the matrix version of the number 1. When A is diagonal, meaning each state slot fades on its own, both formulas become simple per-slot formulas.

For a diagonal A, A_bar is exp(delta a) for each slot a, and B_bar is (A_bar - 1) / a times B. Our code below uses exactly these lines. The 2021 structured paper used a different rule, the bilinear method, but the idea is the same.

Delta is the knob worth understanding. With A negative, a large delta makes A_bar close to 0, so the state drops the past and takes in the new token. A tiny delta makes A_bar close to 1 and B_bar close to 0, so the token is almost ignored.

When A_bar, B_bar and C are the same for every token, each output is a weighted sum of all earlier inputs with fixed weights. A fixed weighted sum slid along a sequence is a convolution, and its weights are called the kernel.

Unroll three steps from an empty state to see it. y_2 = C B_bar x_2 + C A_bar B_bar x_1 + C A_bar^2 B_bar x_0.

The weight on a token depends only on how far back it is, not on where it sits. So the kernel is K = (C B_bar, C A_bar B_bar, C A_bar^2 B_bar, and so on). Papers call this property linear time invariance, or LTI.

The 2021 structured paper put it simply: convolutions are efficient for training while recurrence is efficient for inference. In training, every token of the input is known at once, so one big convolution beats a slow loop.

Pause and predict: if the loop and the convolution are the same maths, how far apart will their outputs be on a random 12-token input? Guess before running the code.

import numpy as np

rng = np.random.default_rng(0)
N, L = 4, 12                          # state size, sequence length
A = -np.array([0.5, 1.0, 2.0, 4.0])   # diagonal A: how fast each state slot fades
B = rng.normal(size=N)                # how the input enters the state
C = rng.normal(size=N)                # how the state is read out
delta = 0.1                           # step size

# Zero-order hold discretisation, written for a diagonal A
A_bar = np.exp(delta * A)
B_bar = (A_bar - 1) / A * B

x = rng.normal(size=L)

# Recurrent view: one step at a time, fixed-size state h
h = np.zeros(N)
y_rec = []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = np.array(y_rec)

# Convolution view: build the kernel K once, then slide it over x
K = np.array([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = np.convolve(x, K)[:L]

print("kernel K[0:4]:", np.round(K[:4], 4))
print("max difference:", np.abs(y_rec - y_conv).max())

What just happened? It printed a maximum difference of 8.3e-17, which is floating-point rounding, so the two views agree. The first kernel weights were 0.0728, 0.0571, 0.0447 and 0.035, shrinking because the state fades.

For long inputs, the convolution runs with a fast Fourier transform, or FFT, a standard method that computes a convolution in about L log L operations for L tokens. Our 4,096-token FFT check matched the loop to 1.1e-16.

A selective model computes delta, B and C from each token, so it chooses what to write and what to ignore. The 2023 selective state space paper (arXiv 2312.00752) introduced this to handle content-based reasoning in language.

In that paper, each token passes through small linear layers that produce its own B, C and delta. Delta then goes through softplus, a smooth function that keeps it positive. A stays a learned constant: the paper argues that selecting delta is enough, because A only enters through exp(delta A).

The paper's selective copying task shows why this matters. The model must copy a few marked tokens scattered at random among filler, so it has to decide token by token. A time-invariant model treats every token the same way, whatever its content.

Here is that idea in miniature: one state slot, a 7 to remember, then 40 noisy tokens. The marker says which token matters, and in a real model a learned layer would play that role.

import numpy as np

rng = np.random.default_rng(1)
values = np.concatenate([[7.0], rng.normal(scale=5, size=40)])
marker = np.array([1.0] + [0.0] * 40)    # 1 = "remember this token"

A = -1.0                                 # one state slot, fades at rate 1
softplus = lambda z: np.log1p(np.exp(z))


def run(deltas):
    h = 0.0
    for x, d in zip(values, deltas):
        a_bar = np.exp(d * A)
        b_bar = (a_bar - 1) / A          # zero-order hold with B = 1
        h = a_bar * h + b_bar * x
    return h


fixed = run(np.full(len(values), 1.0))              # same step for every token
selective = run(softplus(12 * marker - 9))          # step depends on the input
print("time-invariant final state:", round(fixed, 3))
print("selective final state:     ", round(selective, 3))
print("step on the 7:", round(softplus(3.0), 3), " step on noise:", round(softplus(-9.0), 6))

What just happened? The time-invariant run ended at -0.18, because 40 noisy tokens washed out the 7. The selective run ended at 6.635: the 7 got a step of 3.049 and each noise token got 0.000123.

The price is the convolution trick. Once delta, B and C depend on the token, no single fixed kernel describes the layer, and the paper notes it loses the equivalence to convolutions.

Our check shows why. For a time-invariant update, the output for two inputs added together equals the sum of their outputs, to 2.2e-16. For our selective update it was off by 1.65.

With a parallel scan. Linear time means the work grows in step with the length: twice the tokens, twice the work. The update h_t = a_t h_(t-1) + b_t is still linear in the state, so a GPU can combine steps in any grouping and finish in about log2 L rounds.

The trick is a combine rule for two steps. Doing (a1, b1) and then (a2, b2) equals one step (a1 a2, a2 b1 + b2). That rule is associative, meaning the grouping does not change the answer, which is all a scan needs.

A scan, also called a prefix sum, produces every running total of a sequence. Guy Blelloch's classic text on prefix sums, first a 1990 Carnegie Mellon report, solves exactly this first-order recurrence with a scan over such pairs.

import numpy as np

rng = np.random.default_rng(2)
L = 1024
a = rng.uniform(0.8, 1.0, size=L)   # per-step decay (input-dependent in a selective SSM)
b = rng.normal(size=L)              # per-step input term

# Sequential: h_t = a_t * h_(t-1) + b_t, L dependent steps
h, seq = 0.0, []
for t in range(L):
    h = a[t] * h + b[t]
    seq.append(h)
seq = np.array(seq)

# Parallel scan: combine (a1, b1) then (a2, b2) into (a1*a2, a2*b1 + b2).
# That rule is associative, so we can combine in log2(L) rounds.
A_s, B_s = a.copy(), b.copy()
rounds = int(np.log2(L))            # L is a power of 2 here
for r in range(rounds):
    shift = 2**r
    A_prev = np.concatenate([np.ones(shift), A_s[:-shift]])
    B_prev = np.concatenate([np.zeros(shift), B_s[:-shift]])
    A_s, B_s = A_prev * A_s, A_s * B_prev + B_s   # each round is one vector operation

print("rounds:", rounds, "for", L, "steps")
print("max difference:", np.abs(B_s - seq).max())

What just happened? It took 10 rounds for 1,024 steps, and the largest difference from the plain loop was 3.6e-15. Each round is one vector operation over every position, which a GPU runs in parallel.

Our simple version touches every position in every round, so its total work is about L times log2 L. Blelloch's work-efficient scan keeps total work in the same order as the plain loop, which is where linear time comes from.

Why the memory layout matters on a GPU

The full set of A_bar and B_bar values is large: batch x length x channels x state size numbers. The selective paper avoids writing them to HBM, the GPU's large but slower main memory.

It loads the small inputs into SRAM, the fast on-chip memory, then discretises and scans there. Only the outputs go back to HBM.

For training, it does not save the intermediate states. It recomputes them in the backward pass, the step that computes gradients, and the paper reports memory use then matches an optimised attention layer.

The duality view: chunks and matrix multiplications

A 2024 paper (arXiv 2405.21060) shows that a selective state space layer is also one matrix multiplication, y = M x. M is lower-triangular, a square grid that is zero above its diagonal, so no token sees the future.

Softmax is the function attention uses to turn scores into weights that sum to 1. Masked attention is attention where each token sees only earlier tokens.

The paper restricts A to one number per step times the identity. Then M becomes a mask of decay factors, multiplied entry by entry with C times B transposed. That is the same form as masked attention without the softmax.

The paper notes that linear attention, which drops the softmax, can be expanded as a recurrence, and that is the bridge between the two views.

The paper splits the sequence into chunks, uses matrix multiplications inside each chunk and a scan across chunks. GPUs have dedicated units for matrix multiplication, which is why this helps.

It reports its layer runs 2 to 8 times faster than the earlier selective scan, and allows a state 8 times larger or more with little slowdown. Those are the paper's own measurements, on its own hardware.

Because a few attention layers bring back exact recall, while most layers keep the cheap fixed state. A hybrid stacks mostly state space layers with a small number of attention layers that keep a KV cache.

A 2024 empirical study (arXiv 2406.07887) trained 8-billion-parameter state space, transformer and hybrid models on the same data, up to 3.5 trillion tokens. Pure state space models matched or beat transformers on many tasks.

They trailed on copying, on in-context learning and on long-context reasoning. In-context learning means picking up a task from examples placed inside the prompt.

The study's hybrid used 43% state space layers, 7% attention layers and 50% MLP layers. An MLP layer is the per-token feed-forward block that transformers use too.

That hybrid beat the transformer on all 12 standard tasks the authors tested, by 2.65 points on average. It had 4 attention layers out of 56. They predict up to 8 times faster token generation, a prediction for their setup.

The phone-book lookup shows the recall fix most clearly. The pure state space models began to answer wrongly beyond about 500 tokens. The transformer stayed near 100% up to its 4,096-token training length, and the hybrid scored 100% up to 5,500 tokens.

Memory comes back only for the attention layers. Suppose 4 of our made-up 32 layers use attention. At 32,768 tokens the KV cache is then 512 MiB, plus 14 MiB of state, so 116 transcripts fit in 60 GiB instead of 15.

The SevyDevy guide on long-context attention costs covers the attention side of this trade-off in more depth. This guide stays on what replaces the KV cache and what that costs.

Store key-value pairs in a fixed-size state and read them back. The state here is a 32 x 64 grid built by adding one outer product per pair. An outer product multiplies every number of the key by every number of the value.

That is the duality paper's form with the fade set to 1, so it is a real state space update, not a strawman. The KV cache side simply keeps every key and value.

import numpy as np

rng = np.random.default_rng(3)
N, P, V = 32, 64, 64                     # state rows, value size, vocabulary
codebook = rng.choice([-1.0, 1.0], size=(V, P))


def trial(n_pairs, orthogonal):
    if orthogonal:                       # only called with n_pairs up to N
        keys = np.linalg.qr(rng.normal(size=(N, N)))[0][:n_pairs]
    else:
        keys = rng.normal(size=(n_pairs, N))
        keys /= np.linalg.norm(keys, axis=1, keepdims=True)
    ids = rng.integers(0, V, size=n_pairs)
    vals = codebook[ids]

    # Fixed-size state: S = sum of key-value outer products, always N x P numbers
    S = keys.T @ vals
    read = keys @ S                      # query with each key
    state_ok = (np.argmax(read @ codebook.T, axis=1) == ids).mean()

    # KV cache: keep every key and value, look up the best-matching key
    best = np.argmax(keys @ keys.T, axis=1)
    cache_ok = (best == np.arange(n_pairs)).mean()
    return state_ok, cache_ok


print("pairs | state, random keys | state, best keys | KV cache | cache numbers")
for n in [8, 16, 32, 64, 128, 256]:
    r = np.mean([trial(n, False)[0] for _ in range(50)])
    o = np.mean([trial(n, True)[0] for _ in range(50)]) if n in range(1, N + 1) else float("nan")
    c = np.mean([trial(n, False)[1] for _ in range(50)])
    print(f"{n:5d} | {r:18.2f} | {o:16.2f} | {c:8.2f} | {n * (N + P)}")
print("state numbers, always:", N * P)

What just happened? With perfectly separated keys, a state with 32 rows recalls all 32 pairs, but 32 dimensions have room for only 32 perfectly separated keys. Past that, every new pair blurs the old ones.

The KV cache never blurs because it keeps every pair, but its size grows with the pairs. That is the whole trade-off of this guide in one table.

This toy is linear and untrained, so do not read its numbers as a real model's accuracy. It shows the shape of the limit, and the copying paper shows the same shape in trained models.

How to check it worked

1. 1

   Install the pinned version

   Create a virtual environment and run pip install numpy==2.5.4. NumPy 2.5.4 is the latest release as of 28 October 2026.

2. 2

   Run the convolution check

   Save the first script and run it. Expect a maximum difference around 1e-16; anything above 1e-9 means a typo in the formulas.

3. 3

   Run the selective example

   Expect about -0.18 for the time-invariant run and 6.635 for the selective one. Change the 12 in softplus(12 * marker - 9) to 0: every token then gets the tiny step, the 7 is never written in, and the result is about 0.001.

4. 4

   Run the scan

   Expect 10 rounds for 1,024 steps and a difference near 1e-15. Set L to 4,096 and expect 12 rounds.

5. 5

   Run the recall test

   Expect the state column to fall below 100% after 32 pairs. Change N to 64 and the state stays at 99% or better up to 64 pairs, then falls to 92% at 128.

Start from the memory problem, then the fixed-size state, then how it trains, then where it fails. Interviewers want the trade-off and a fix, not a list of paper names.

1. 1

   Name the cost

   A KV cache grows by 2 x layers x KV heads x head size x bytes per token. At 128 KiB per token, 32,768 tokens need 4 GiB per sequence.

2. 2

   State the recurrence

   A state space layer keeps h_t = A_bar h_(t-1) + B_bar x_t and outputs C h_t. Memory per sequence is fixed, and each new token costs the same work.

3. 3

   Explain training

   Time-invariant layers train as one long convolution. Selective layers, whose delta, B and C depend on the token, train with a parallel scan in about log2 L rounds.

4. 4

   Admit the weakness

   A fixed state cannot keep every token, so exact copying and lookup degrade as inputs grow. In the copying paper, a 410-million-parameter transformer beat a 2.8-billion-parameter fixed-state model on phone-book lookup once the book had 70 or more entries.

5. 5

   Propose a hybrid and a test

   Mix a few attention layers into a mostly state space stack. Then measure exact-recall tasks from your own data before and after the change.