Albert Gu, Tri Dao
9 min
Modern foundation models rely heavily on the Transformer architecture and its core attention mechanism. While Transformers are powerful, their attention mechanism results in quadratic computational complexity with respect to sequence length, making them inefficient for very long sequences. This paper investigates whether structured state space models (SSMs) can be improved to match Transformer quality while retaining the linear-time efficiency of recurrent models.
Prior SSMs often struggled with discrete data like language because they were linear time-invariant (LTI) systems, meaning they processed all tokens the same way. The authors introduce two primary innovations:
The Mamba architecture—a simplified, homogeneous design without standard attention or MLP blocks—yields several key results:
By successfully combining the modeling power of Transformers with the linear-time efficiency of recurrent models, Mamba offers a viable alternative to the Transformer architecture for long-sequence tasks. This is particularly promising for domains like genomics, long-form audio, and high-resolution video where the quadratic cost of standard attention is a significant barrier to progress.
Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5$\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.
Alex: It makes several parameters depend on the current input. Those parameters control how information enters memory, how memory persists, and how it contributes to the output. The authors call this selectivity: content-dependent control over information moving along the sequence.
Sam: Neural networks already multiply signals by learned gates. Why is this more than another gate in the architecture?
Alex: The paper distinguishes multiplying features at one position from changing what passes between positions. Selection does the latter. Its connection is closer to traditional recurrent-network gates, which regulate whether new information replaces or preserves the hidden state.
Sam: So the relevant distinction is not whether a model contains multiplication. It is whether the input can change what the model carries forward.
Alex: The controlled experiments support that distinction. Keeping the surrounding architecture but replacing the fixed state space layer with a selective one sharply improves selective copying. In the language ablations, that replacement also brings the main improvement.
Sam: But the fixed rules helped make earlier models fast. If those rules now change at every position, where does the efficiency come from?
Alex: That is the second contribution. Selection removes the efficient fixed-convolution route. The authors instead use a parallel scan, which combines successive memory updates through parallelizable operations.
Sam: Parallel computation addresses the order of updates. What prevents all those intermediate memories from filling the graphics processor?
Alex: They avoid writing the expanded memory states into the processor’s large, slower memory. Instead, they combine operations and produce those states in faster local memory. During training, they recompute intermediate states for the backward pass rather than store them all.
Sam: Then the contribution has two linked pieces: a more flexible memory rule, and an implementation that makes that flexibility affordable.
Alex: Yes. They put those selective layers into a repeated, simplified block to make Mamba. It has no attention blocks or separate multilayer perceptron blocks, though it still uses learned projections, nonlinearities, and a short convolution.
Sam: Let’s separate demonstrations of memory from evidence about useful language models. What does the controlled task establish?
Alex: One task tests associative recall: after seeing “Harry Potter,” predict “Potter” when “Harry” reappears. Mamba trained on sequences of 256 tokens. It retained perfect accuracy when tested at roughly a million tokens, far beyond its training length.
Sam: That is a clear test of length extrapolation. But remembering a constructed association does not establish general language understanding.
Alex: The language experiments address a different claim. On text prediction, Mamba matches a strong Transformer recipe incorporating architectural and training improvements. The authors also compare pretrained models on standard zero-shot tasks, meaning evaluation without task-specific training.
Sam: Were the comparisons close enough to isolate architecture, or did the models receive different training data?
Alex: The strongest matched-data comparisons include Pythia and RWKV, another recurrent model. They used the same dataset, tokenizer, and training-token count as Mamba, though RWKV used shorter training context. Mamba’s roughly three-billion-parameter model generally matched Transformer baselines around twice its size.
Sam: And the speed claim matters only if we know what was measured. Is it faster for an individual request, or processing more requests together?
Alex: The headline is throughput: total generated tokens per second. Mamba achieved roughly four to five times higher throughput than a similarly sized standard Transformer in their benchmark. Avoiding the growing cache allowed much larger batches on the same graphics processor.
Sam: So that is not a universal multiplier for every deployment, especially against any possible optimized Transformer.
Alex: Correct. The end-to-end comparison used a specified standard Transformer implementation and a fixed generation setup. Separately, their core-operation benchmark compared the scan against optimized attention. Those are different comparisons, not interchangeable speed claims.
Sam: What about the long-context applications? Does processing a longer sequence actually improve predictions?
Alex: In human-genome pretraining, Mamba’s prediction quality improved as context increased to million-length sequences. The tested HyenaDNA baseline instead worsened with longer context. That supports the proposed benefit of filtering irrelevant history, within this dataset and training setup.
Sam: Audio seems like a useful counterexample. Smooth waveforms might benefit from stable dynamics rather than constant selection.
Alex: The paper finds that tradeoff. Mamba-based audio architectures performed well, but detailed waveform experiments showed that adding selection could hurt performance. Fixed-dynamics layers worked better near the raw signal. Selection is not uniformly beneficial across modalities.
Sam: And the biggest unresolved issue for language is scale?
Alex: The authors say larger-scale quality remains untested. Their trained language models do not reach the seven-billion-parameter scale discussed in the limitations. A larger Mamba used for throughput measurements was untrained, so its speed says nothing about its language quality.
Sam: Given that boundary, who should read the full document, and where should they begin?
Alex: Researchers designing language or long-context sequence models should read it. Start with “Selection as a Means of Compression,” then the language comparisons and selective-layer ablations. Implementation researchers should read the selective-scan section; audio researchers should check the waveform parameterization ablations before adopting selection.
Sam: And what should everyone else carry away?
Alex: Efficient memory is not just a smaller store. It needs a way to decide what deserves to stay.
Sam: Thanks for listening.