Mamba
In plain wordsA 2023 alternative to the transformer design that processes text with a continually updated summary which decides what to keep or forget, letting it handle long sequences quickly.
A neural sequence architecture whose blocks contain no attention and no separate MLP, only a selective state space modelstate space model, n.: A type of sequence model that reads text like a running tally, updating a fixed-size internal summary at each step, so its memory and cost do not balloon as texts get longer. combined with gating and a short convolution. 'Selective' means that the step size and the input and output projections are computed from the current token, so the layer chooses, token by token, what to write into its fixed-size state and what to let fade. Input-dependent parameters forfeit the trick by which earlier SSMs were trained as one long convolution; Mamba instead evaluates the recurrence with a parallel scan engineered to keep the large expanded state in fast on-chip GPU memory. Its authors reported generation throughput five times that of transformers, cost linear in sequence length, and a 3-billion-parameter model on a level with transformers twice as large. Mamba-2 (2024) restricted the state matrix to a scaled identity, revealing a 'state space duality' with linear attentionlinear attention, n.: A family of faster substitutes for attention whose cost grows in step with the length of the text rather than with its square, by keeping a running summary instead of comparing every pair of words. that permits faster matrix-multiplication-based training. Mamba layers now appear mostly in hybridhybrid architecture, n.: An AI model design that mixes two kinds of building blocks — usually a few standard attention layers among cheaper, memory-saving layers — to keep most of the quality at lower cost on long inputs. models.
We replaced three of every four attention layers with Mamba-2 layers and kept quality flat.
Origin
Introduced by Albert Gu (Carnegie Mellon University) and Tri Dao (Princeton University) in 'Mamba: Linear-Time Sequence Modeling with Selective State Spaces' (arXiv, December 2023; COLM 2024); Tri Dao and Albert Gu followed with Mamba-2 in May 2024.
Related
- state space modelstate space model, n.: A type of sequence model that reads text like a running tally, updating a fixed-size internal summary at each step, so its memory and cost do not balloon as texts get longer.
- hybrid architecturehybrid architecture, n.: An AI model design that mixes two kinds of building blocks — usually a few standard attention layers among cheaper, memory-saving layers — to keep most of the quality at lower cost on long inputs.
- linear attentionlinear attention, n.: A family of faster substitutes for attention whose cost grows in step with the length of the text rather than with its square, by keeping a running summary instead of comparing every pair of words.
- RWKVRWKV, prop. n.: An open-source language model design that works like an old-style recurrent network when writing text, using constant memory, yet can be trained quickly in parallel like a transformer.
- transformertransformer, n.: The design behind today's chatbots and much of modern AI. It handles text by letting every word look at every other word and decide which of them matter for understanding it.
- recurrent neural networkrecurrent neural network, n.: A neural network that reads a sequence, such as a sentence or an audio clip, one piece at a time, carrying a running summary forward so that what came earlier can shape how later pieces are understood.
- transformertransformer, n.: The design behind today's chatbots and much of modern AI. It handles text by letting every word look at every other word and decide which of them matter for understanding it.
Sources
- Gu, A. & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752. arxiv.org/abs/2312.00752
- Dao, T. & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. ICML 2024.
Mentioned in
recurrent neural networkrecurrent neural network, n.: A neural network that reads a sequence, such as a sentence or an audio clip, one piece at a time, carrying a running summary forward so that what came earlier can shape how later pieces are understood., hybrid architecturehybrid architecture, n.: An AI model design that mixes two kinds of building blocks — usually a few standard attention layers among cheaper, memory-saving layers — to keep most of the quality at lower cost on long inputs., linear attentionlinear attention, n.: A family of faster substitutes for attention whose cost grows in step with the length of the text rather than with its square, by keeping a running summary instead of comparing every pair of words., RWKVRWKV, prop. n.: An open-source language model design that works like an old-style recurrent network when writing text, using constant memory, yet can be trained quickly in parallel like a transformer., state space modelstate space model, n.: A type of sequence model that reads text like a running tally, updating a fixed-size internal summary at each step, so its memory and cost do not balloon as texts get longer.