Model architectures
What problem do state-space models like Mamba solve?
State-space models, with Mamba the best-known example, model a sequence by carrying a compact state forward token by token. That makes their cost grow linearly with length, rather than quadratically as attention does, which matters most on very long inputs. The trade is weaker exact recall, which is why many recent models pair the two.
Attention has a cost that grows badly with length. To let every token look at every other token, the work rises with the square of the input: double the context and you roughly quadruple the attention compute — a tenfold increase in context costs roughly a hundredfold increase in attention compute, all else equal. For short prompts this is fine. For very long documents, audio, or genomic data it becomes the bottleneck. State-space models (SSMs) take a different route. Instead of comparing all pairs of tokens, they carry a compact state forward through the sequence, updating it one token at a time, the way an older recurrent network did. The cost of that scales linearly with length, not quadratically.
Mamba is the model that made this competitive with transformers on language. Earlier state-space designs scaled well but lagged on text, and the paper’s authors traced the weakness to the state being fixed: it could not choose what to keep and what to forget based on the actual input. Mamba’s fix is a selective mechanism that lets the model decide, per token, what passes into the state. With that change it reported up to 5x higher throughput than a transformer, with the million-token results coming from its audio and genomics benchmarks rather than language modeling itself.[1] The appeal is plain. On long inputs an SSM can be much faster and lighter on memory than attention, because its running state stays the same size no matter how long the sequence gets.
The catch is also about that fixed-size state. Attention can reach back and pull an exact earlier token whenever it needs one; a state that must compress everything it has seen into a constant-size summary will sometimes lose that precision. This is why the most practical recent designs do not pick a side. Jamba, for instance, interleaves Mamba layers with a smaller number of attention layers, keeping attention’s exact recall while letting the state-space layers carry most of the long-range load cheaply.[2] The hybrid is covered on the next page.
References
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arXiv:2312.00752) — Gu & Dao
- Jamba: A Hybrid Transformer-Mamba Language Model (arXiv:2403.19887) — Lieber et al., AI21 Labs
Common questions
- Is Mamba a transformer?
- No. Mamba is a state-space model, a different design that carries a running state forward rather than using attention to compare all tokens. It was built specifically to avoid attention's quadratic cost on long sequences.
- If Mamba is faster, why are most models still transformers?
- Attention is very good at exact recall: pulling a specific earlier token back when needed. A fixed-size running state can lose that detail. Many recent models keep some attention layers and add state-space layers to get both, rather than dropping attention entirely.