State space model

In short

A state space model reads a sequence by keeping one hidden state and updating that state with each new input, instead of looking back at every earlier step.

Also called: SSM, State-space model
Topic: How AI models work
Level: Some background helps
First described: The continuous-time state equation is the classical state space model. Gu, Goel and Ré, October 2021, treat it as a control-theory model and build the S4 sequence layer on it¹. Gu and Dao cite Kalman (1960) as the classical source their structured models draw on².

In the continuous-time form used by Gu, Goel and Ré, a state space model maps an input u(t) to an output y(t) through a latent state x(t)¹. The state moves by x'(t) = A x(t) + B u(t), and the output is y(t) = C x(t) + D u(t)¹. They describe these models as a foundational tool in control theory, not a method invented for chatbots¹. Gu and Dao say the structured deep-learning versions take their inspiration from classical state space models, citing Kalman (1960)².

In discrete time the same model is a recurrence: the new state depends only on the previous state and the new input¹. When the rule does not change with time, it can also be a convolution and trained in parallel¹. Their 2021 model, S4, was a reparameterisation that made that convolution practical¹. On Long Range Arena, S4 solved Path-X, length 16,384, at 88% accuracy, where they say every earlier model was near chance, about 50%¹. On CIFAR-10 and WikiText-103 it generated about 60 times faster than standard autoregressive models¹.

A state space layer is not attention. Attention stores earlier tokens and looks back at them. The state here stays a fixed size, so cost grows linearly with length, but a small state can drop detail². Mamba, from Gu and Dao in 2023, lets the update depend on the current token, so the model can keep or forget information². Their Mamba block has no attention and no separate MLP block². They report about five times the inference throughput of Transformers. A 3 billion parameter Mamba beats Transformers of the same size and matches Transformers twice its size, in pre-training and downstream tests².

Products can mix the two. Jamba, from Lieber and colleagues, interleaves Transformer blocks with Mamba layers, which they call a recent state space model, plus mixture-of-experts in some blocks³. That paper's model has 12 billion active parameters and 52 billion in total, and was built to fit one 80GB GPU³. AI21's current docs do not say state space. They label the type "Joint Attention and Mamba" and list Jamba Large at 398 billion parameters, 94 billion active, with a 256,000 token context⁴.

Example

Reading a long contract, an attention layer can look back at earlier words. A state space layer folds each new word into the same hidden state and does not keep a growing cache, which is why S4 can also run as a plain recurrence¹.

Why it matters

It is the main current alternative to attention for long sequences: a fixed state and linear cost, with hybrids such as Jamba keeping attention in only some layers²³.

Questions people ask

Is a state space model the same as a Transformer?

No. A Transformer attention layer looks back across earlier tokens and stores them. A state space layer keeps one hidden state and updates it, so the cost grows linearly with length¹². Hybrids such as Jamba use both in different layers³.

What is Mamba?

Mamba is Gu and Dao's 2023 selective state space model. The update depends on the current token, so the model can pass information on or drop it. Their block uses no attention. They report five times the throughput of similar Transformers, and a 3 billion parameter model matching Transformers twice that size².

Did AI research invent state space models?

No. Gu, Goel and Ré describe them as a foundational model in control theory and other sciences¹. The Mamba paper points to Kalman's 1960 work as the classical source, and then uses the structured deep-learning form that starts with S4²¹.

What is Jamba?

Jamba is AI21's hybrid: Transformer layers interleaved with Mamba layers, plus mixture-of-experts³. The 2024 paper's model has 12 billion active parameters out of 52 billion. Current docs list a larger Jamba Large, 398 billion parameters with 94 billion active, and a 256,000 token context⁴.

Related terms

Sources

Back to the glossary