Research

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

← All research
Preprint Third-party arXiv·December 1, 2023

This is third-party work — not authored by or affiliated with DarcStar Technologies.

By Albert Gu, Tri Dao

Research

Abstract

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× 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.

DarcStar commentary

Our perspective on this work and how it informs our research — not part of the original paper.

Mamba is the paper we point to when someone assumes attention is the only substrate. It isn't a faster attention — it's a bet that you don't need attention at all: a selective state-space model, linear in sequence length, with no attention and no MLP blocks, that still matches or beats Transformers at real scale. We don't treat attention as the destination, so we read this as a serious data point about what else works.

The insight that makes it go is selective. Earlier state-space models were fast but content-blind — the same recurrence regardless of what the token was — and that's exactly why they lost to attention on language. Mamba lets the state-space parameters depend on the input, so the model can choose what to keep and what to forget as it moves along the sequence. That's the content-based reasoning attention gets for free, recovered inside a recurrent model — and paid for with a hardware-aware parallel scan so the loss of convolutional shortcuts doesn't sink throughput.

Where we'd stay measured is on the word “matches.” Parity with attention is empirical and moves with scale and modality; a linear-time recurrent state is a strict information bottleneck, and whether that bottleneck costs you depends entirely on the task — a thing to measure, not to assume from a throughput number. In practice the strongest systems increasingly combine attention and SSM layers rather than crown one, and that's the reading we take: Mamba doesn't prove attention is obsolete, it proves the recurrent substrate is a real option — and once it's a real option, which substrate you reach for becomes a choice you make per job, not a default you inherit.