[Paper Review] Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao|arXiv (Cornell University)|Dec 1, 2023

Neural Networks and ApplicationsComputer Science967 citations

TL;DR

Mamba introduces selective state-space models that are input-dependent, enabling linear-time sequence modeling with 5× higher inference throughput than Transformers and strong performance across language, audio, and genomics. In language, Mamba-3B outperforms similarly sized Transformers and matches them at twice the size.

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$ imes$ 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.

Motivation & Objective

• Address Transformer inefficiency on long sequences by introducing content-aware selective state spaces (SSMs).
• Enable linear-time sequence modeling without attention or MLPs while preserving accuracy across modalities.
• Develop a hardware-aware parallel algorithm for selective SSMs that preserves efficiency in recurrent mode.

Proposed method

• Make SSM parameters functions of input addresses to enable discrete, content-based modulation of information flow along the sequence.
• Integrate selective SSMs into a simple end-to-end backbone without attention or MLPs (the Mamba model).
• Design a hardware-aware parallel recurrent algorithm to maintain fast inference despite the lack of standard convolutional optimization.
• Demonstrate linear scaling with sequence length and 5× throughput improvement over Transformers during inference.

Experimental results

Research questions

• RQ1 Can content-based (address-dependent) selective SSMs improve long-sequence reasoning compared to traditional subquadratic models?
• RQ2 Do selective SSMs enable linear-time sequence modeling with competitive performance across language, audio, and genomics?
• RQ3 Is it possible to achieve hardware-friendly parallelism in a recurrent-style architecture without attention/MLPs while maintaining throughput?
• RQ4 How does Mamba perform in language modeling relative to Transformers of the same size and at larger sizes?

Key findings

• Mamba achieves fast inference with 5× higher throughput than Transformers.
• The model scales linearly with sequence length.
• Mamba demonstrates state-of-the-art results across language, audio, and genomics modalities.
• In language modeling, Mamba-3B outperforms Transformers of the same size and matches Transformers twice its size in both pretraining and downstream evaluation.