State-space models: Mamba and others
In short: State-space models (SSM) such as Mamba read sequences with a fixed, running state instead of attention over all earlier tokens. The cost grows only linearly with length.
In more detail: Transformers are expensive on long texts because attention grows quadratically and the KV cache grows along. SSMs keep a fixed-size summary.
In detail
- Idea:
hₜ = A·hₜ₋₁ + B·xₜ,yₜ = C·hₜ. Like an RNN but built so that training can be computed in parallel (convolution or scan). - S4: early model for very long sequences.
- Mamba (2023): makes A, B, C depend on the input (“selective”), so the model decides what to remember. Implemented close to the hardware.
- Mamba-2, RWKV, RetNet, linear attention: related approaches with linear cost.
- Strengths: constant memory when generating, fast on very long texts, audio, genome data.
- Weaknesses: reproducing details exactly from the context (copying, lookup) is harder than with attention.
- Hybrids: models mix SSM layers with a few attention layers (Jamba, Nemotron-H), combining speed and accuracy.
- Comparison: the KV cache grows with length in transformers, not in SSMs.
See also: Recurrent networks and LSTM, Context window