Post

Mamba: linear-time language models without attention

Mamba replaces the quadratic attention of Transformers with a selective state-space model that costs one unit per token. A close reading of the December 2023 paper by Albert Gu and Tri Dao, and what its linear-time backbone changes for long context.

The Transformer’s attention is both its strength and its tax. Every token must compare itself with every other token, so the cost of a sequence grows quadratically with its length. Feed a model a million tokens and the attention computation dominates everything else, even though most of those token-to-token comparisons are uninteresting most of the time. A paper posted to arXiv on 1 December 2023 argues that the fix is not a better attention approximation but a different class of model altogether. ‘Mamba: Linear-Time Sequence Modeling with Selective State Spaces’, by Albert Gu and Tri Dao, rebuilds the backbone from state-space models that have lived in deep learning’s background for a decade.

State-space models, or SSMs, descend from classical control theory. They maintain a small hidden state and update it as each new input arrives, which makes inference cheap and memory usage bounded by the state rather than by the sequence length. Earlier SSMs, including Gu’s own S4, matched Transformers on long-range tasks such as speech and protein sequences but lagged badly on language. Gu and Dao diagnose the weakness precisely: those models applied the same parameters to every input, so they could not decide, token by token, what to keep and what to discard. Language requires content-based reasoning. A model needs to remember a name until it matters and then forget it.

Mamba’s answer is to make the SSM parameters functions of the input. The model now selectively propagates or suppresses information along the sequence dimension depending on the current token, a selection mechanism that behaves like a content-aware recall gate. That change breaks the fast convolutional implementation that older SSMs relied on, because a convolution cannot vary its behaviour with the signal. In response the authors design a hardware-aware parallel algorithm that trains in parallel on a GPU and runs as a recurrent model at inference time, exploiting the fact that the selective recurrence is cheap once the parameters are input-dependent.

The resulting network has no attention at all, and not even the usual feed-forward blocks. It is a stack of selective SSM layers plus convolutions and gating, a simplified end-to-end architecture. On language modelling, the reported Mamba-3B outperforms a Transformer of the same size and matches one twice as large, in both pretraining perplexity and downstream evaluation. Inference throughput is reported as five times higher than an equivalent Transformer, and the cost scales linearly with sequence length. The model keeps working on real data up to million-token contexts, where quadratic attention has already become impractical.

What the paper makes elegant is the unification rather than the gimmick. A companion observation, structured state-space duality, shows that the best of these models are mathematically close to attention in disguise. Mamba is the engineering consequence: if SSMs can be made selective, they can do language, and the long-context problem stops being a tax on every token. The honest qualifications are that the headline numbers come from the authors’ own open comparisons, that closed-source frontier systems move faster than any published baseline, and that the architecture has yet to prove itself on the full spectrum of agentic tool use. In mid 2024, the companion work on hybrid models confirmed that the selective idea scales, but a pure Mamba backbone is still an open design problem rather than a finished replacement.

The long arc of the argument is nonetheless hard to dismiss. A sequence model that costs one unit per token is a different object from one that costs its square, and on 1 December 2023 that object finally learned language well enough to matter. The field’s response over the following year, a rush of state-space and linear-attention successors, is the usual sign that a cleaner statement of the same problem has been found. Whether Mamba displaces the Transformer or merely becomes one option among many, the quadratic floor has been shown to be optional.

This post is licensed under CC BY 4.0 by the author.