Mamba: Linear-Time Sequence Modeling with Selective State Spaces

The problem

Sequence models face a trade-off between how much of the earlier input a model can use and how expensive it is to process a long sequence. Attention compares every position with every other position, so its cost grows quickly as the sequence gets longer. A family of cheaper designs — linear attention, gated convolutions, structured state space models — was built to bring that cost down. The Mamba paper, submitted to arXiv on December 1, 2023 by Albert Gu and Tri Dao, argues that these cheaper architectures share a weakness: they cannot perform content-based reasoning^[1].

Content-based reasoning means deciding what to keep and what to discard on the basis of what the current input actually is, not merely on where it sits in the sequence. A model that always compresses its history with the same fixed rule will treat a filler word and a crucial name the same way. That is a real limitation, because much of what a language model must do — recall a fact stated earlier, ignore a distractor, copy a token — depends on the content of the tokens involved.

Selective state spaces

The starting point is the state space model (SSM), a classical way of describing a sequence with a hidden state that is updated step by step as new tokens arrive. Reading the sequence this way scales linearly with its length, since each new token costs about the same amount of work. The catch is that in earlier SSMs the update rule is fixed: the same matrices govern how the state moves and how much of the new input enters it, no matter what the input is.

Mamba's central move is to let the SSM parameters be functions of the input. Because the update depends on the token in front of the model, it can selectively propagate or forget information along the sequence length dimension depending on the current token^[1]. In effect, the model gains a data-dependent gate on its own memory: an uninformative token can be let through with little effect, while an important one can be written into the state and carried forward.

Architecture

These selective SSMs are integrated into a simplified end-to-end neural network architecture that contains no attention and not even MLP blocks; the authors call it Mamba^[1]. Keeping the design simple matters, because the point of the work is that the sequence-mixing mechanism, not extra depth or extra components, should carry the modelling burden.

A recurrent formulation is normally awkward for modern hardware, since it must be computed one step at a time. Mamba addresses this with a hardware-aware parallel algorithm in recurrent mode^[1], which allows the recurrence to be evaluated efficiently during training rather than only sequentially. The same structure also makes generation cheap, because producing a token requires only one update of the state instead of a fresh comparison against the whole context.

Results

The paper reports that Mamba enjoys fast inference, with 5x higher throughput than Transformers, and linear scaling in sequence length, with performance improving on real data up to million-length sequences^[1]. The second point is notable in practice: on many long-sequence models, adding more context stops helping or starts to hurt, whereas here the results on real data kept improving as the sequences grew towards a million tokens.

On language modeling, the authors report that Mamba-3B outperforms Transformers of the same size and matches Transformers twice its size, in both pretraining and downstream evaluation^[1].

Why it matters

The contribution is a design argument as much as a set of numbers. The paper identifies content-based reasoning as the specific capability that earlier subquadratic architectures lacked, shows that making the state space parameters input-dependent supplies it, and then reports that a network built entirely from such blocks is competitive with attention-based models at the same scale^[1]. That places Mamba among the architectures proposing an alternative to attention for long sequences, where the cost of attention is hardest to pay.