EMZETT.
Login

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