ElemedexA pocket encyclopedia of AI architectures
No. 043

Mamba

Selective state-space model (S6)

Recurrent Serpent

A snake that never looks back: everything it has read it carries digested in its belly, in a fixed-size state. With every bite it decides what to swallow and what to let pass. It is very fast on very long texts, but ask it to repeat word for word something from a thousand pages back and things get muddled.

Height
64 blocks (Mamba-2.8B)
Weight
≈ 5.6 GB (2.8 B in fp16)
Parameters
130 M – 2.8 B (released models)
Context
Trained on 2,048 tokens; fixed state, no KV cache; up to 1 M on synthetic tasks
Training data
The Pile, 300 billion tokens
Year
2023
Generation
6 · Assistants, reasoning, and agents
Creators
Albert Gu (Carnegie Mellon) & Tri Dao (Princeton)
Kind
Architecture

Base stats

Total
319

Approximate, ordinal scales meant for comparing species with each other; the hard numbers are in the data sheet.

Weaknesses

  • Copying or retrieving information verbatim from a long context: a fixed state cannot keep everything.
  • Little tested at scale: the original paper stops at 2.8 B parameters.
  • A poorer ecosystem than transformers' in fine-tuning recipes, tricks and tools.
  • On continuous signals, such as raw audio, selection can hurt what time-invariant SSMs did well.

Habitat

It was raised on The Pile (300 billion tokens), a human genome and audio recordings. Today it lives in labs and in hybrid models such as Jamba, where it shares the body with attention layers.

Moves

Abilities

Selection
The parameters Δ, B and C depend on the input: the model decides, token by token, what to keep in the state and what to take in.
Parallel scan
The recurrence is computed with a parallel scan in the GPU's fast memory, without materializing the expanded state.
Fixed-size state
Generating each token costs the same, whether it is the first or the millionth.

Anatomy

Mamba block×64in every channelTokensEmbeddingsRMSNormLinear projection (×2)Short convolution (k = 4)Selective SSM× SiLU gateOutput projectionNext tokenInput xₜΔₜ, Bₜ, Cₜ = f(xₜ)hₜ = Āₜ·hₜ₋₁ + B̄ₜ·xₜyₜ = Cₜ·hₜMamba architectureInside the SSM, at step t

Left: the Mamba block, which replaces both attention and the feedforward network; the residual wraps the whole block. Right: the selective heart. Because Δ, B and C depend on the current token, the model decides how much of the previous state to keep and how much of the token to take in; the state h has a fixed size.

Evolution chain

  • a selective state space: recurrence that trains in parallel015 LSTM
043 Mamba

Mamba is a sequence architecture that does without attention: it processes text with a state-space model, a form of recurrence that carries the history compressed into a fixed-size state. Gu and Dao introduced it in December 2023 as the first linear-time alternative to match transformers at language modeling: their 2.8-billion-parameter model beat transformers of the same size and matched ones twice as large (Gu & Dao, 2023). Alammar & Grootendorst (2024) list it among the new architectures that aim for transformer-level performance with larger context windows or faster inference (2024, p. 24).

How it works

A state-space model, in its discrete form, is a linear recurrence: hₜ = Ā·hₜ₋₁ + B̄·xₜ, yₜ = C·hₜ. The vector h is the hidden state; the idea comes from control theory, from the Kalman filter (Gu & Dao, 2023, sec. 2). Its predecessors (S4 and kin) had parameters that were fixed over time, which allowed computing them as one big convolution and training them in parallel, but prevented them from deciding what to remember based on content: they solved the task of copying a sequence at a fixed distance, but not “selective copying,” where only certain tokens scattered among noise must be retained (Gu & Dao, 2023, sec. 3.1).

Mamba's novelty is selection: Δ, B and C become functions of the input, so that at every token the model decides how much of the state to keep and how much of the token to take in. The authors show that the gates of classic recurrent networks are a special case of this mechanism (Gu & Dao, 2023, sec. 3.5). The price is losing the convolution; they pay it with a parallel scan that runs in the GPU's fast memory without materializing the expanded state (Gu & Dao, 2023, sec. 3.3). The resulting block (linear projection, short convolution, selective SSM and multiplicative gate) is stacked homogeneously, with no separate attention or feedforward layers (Gu & Dao, 2023, sec. 3.4).

What changed

Compared with the LSTM, the recurrence is linear in the state, and that is what makes it trainable in parallel. Compared with the Transformer, the authors frame the problem as compression: attention is effective and inefficient because it does not compress anything (it stores the whole context in the KV cache, which grows with every token), whereas a recurrent network is efficient because it compresses everything into a finite state, but is only as good as that compression (Gu & Dao, 2023, sec. 3.1). Mamba bets on compressing with judgment. The result: compute and memory linear in length during training, constant cost per token during generation, about five times the generation throughput of a similarly sized transformer, and improvements with contexts of up to a million elements on synthetic tasks and DNA (Gu & Dao, 2023, sec. 1).

Limits

The original paper stops at 2.8 billion parameters and 300 billion tokens; its authors warn that it remains to be seen whether the advantage holds at larger scales (Gu & Dao, 2023, sec. 5). Besides, a fixed state cannot keep everything. Jelassi et al. (2024) proved that a two-layer transformer can copy strings of exponential length while these models are limited by their state, and found that in practice they do much worse at copying and retrieving information from context. Industry's answer was hybrids: Jamba (Lieber et al., 2024) interleaves one attention layer for every seven Mamba layers and shrinks the KV cache eightfold relative to a plain transformer. Since then hybrids have grown from an experiment into a family of their own: SSM hybrid (Jamba, Qwen3.5).

Sources

  • Alammar, J., & Grootendorst, M. (2024). Hands-on large language models: Language understanding and generation. O’Reilly Media. §
  • Buckner, C. J. (2024). From deep learning to rational machines: What the history of philosophy can teach us about the future of artificial intelligence. Oxford University Press. https://doi.org/10.1093/oso/9780197653302.001.0001 §
  • Gu, A., & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces (arXiv:2312.00752). arXiv. https://doi.org/10.48550/arXiv.2312.00752 §
  • Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. In Proceedings of the 41st International Conference on Machine Learning (Vol. 235, pp. 21502–21521). PMLR. https://proceedings.mlr.press/v235/jelassi24a.html §
  • Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev-Shwartz, S., Abend, O., Alon, R., Asida, T., Bergman, A., Glozman, R., Gokhman, M., Manevich, A., Ratner, N., Rozen, N., . . . Shoham, Y. (2024). Jamba: A hybrid Transformer-Mamba language model (arXiv:2403.19887). arXiv. https://doi.org/10.48550/arXiv.2403.19887 §