A Survey of Mamba

TL;DR

The paper surveys Mamba, a selective State Space Model-based backbone designed to achieve near-linear scaling with sequence length while matching Transformer-level modeling power. It introduces Mamba-1 and Mamba-2, including HiPPO-based memory initialization, input-dependent selection, and State Space Duality (SSD) to connect SSMs with attention, plus hardware-aware computation strategies. The review covers advancements in block design, scanning modes, memory management, and extensive data adaptation across sequential, non-sequential, and multimodal domains, with numerous example models and applications spanning NLP, vision, speech, drug discovery, and robotics. It also discusses challenges such as memory, generalization, and trust, and proposes directions like foundation-model development, cross-pollination of Transformer techniques, PEFT, continual learning, and retrieval-augmented generation to enhance Mamba's practical impact.

Abstract

As one of the most representative DL techniques, Transformer architecture has empowered numerous advanced models, especially the large language models (LLMs) that comprise billions of parameters, becoming a cornerstone in deep learning. Despite the impressive achievements, Transformers still face inherent limitations, particularly the time-consuming inference resulting from the quadratic computation complexity of attention calculation. Recently, a novel architecture named Mamba, drawing inspiration from classical state space models (SSMs), has emerged as a promising alternative for building foundation models, delivering comparable modeling abilities to Transformers while preserving near-linear scalability concerning sequence length. This has sparked an increasing number of studies actively exploring Mamba's potential to achieve impressive performance across diverse domains. Given such rapid evolution, there is a critical need for a systematic review that consolidates existing Mamba-empowered models, offering a comprehensive understanding of this emerging model architecture. In this survey, we therefore conduct an in-depth investigation of recent Mamba-associated studies, covering three main aspects: the advancements of Mamba-based models, the techniques of adapting Mamba to diverse data, and the applications where Mamba can excel. Specifically, we first review the foundational knowledge of various representative deep learning models and the details of Mamba-1&2 as preliminaries. Then, to showcase the significance of Mamba for AI, we comprehensively review the related studies focusing on Mamba models' architecture design, data adaptability, and applications. Finally, we present a discussion of current limitations and explore various promising research directions to provide deeper insights for future investigations.

Figures (7)

• Figure 1: Examples of the applications of Mamba-based models for different downstream tasks.
• Figure 2: An illustration of representative model architectures, namely Recurrent Neural Network (RNN), Transformer, and State Space Model (SSM). (a) RNNs function within a nonlinear recurrent framework, facilitating rapid outputs during auto-regressive inference. (b) Transformers execute matrix multiplications concurrently across numerous query-key pairs, facilitating parallel training. (c) SSMs exhibit versatility by accommodating both recurrent and convolutional computations due to their linear nature. This fusion harnesses the strengths of RNNs and Transformers, allowing SSMs for recurrent inference and parallel training. Despite this, traditional time-invariant SSMs fall short in context-aware modeling, resulting in diminished performance in specific tasks.
• Figure 3: Overview of the Selective State Space Model with hardware-aware state expansions. The selective mechanism introduces input-dependent dynamics, while the hardware computation algorithm optimizes memory usage within the more efficient levels of GPU memory hierarchy.
• Figure 4: The block architectures of Mamba-1 and Mamba-2.
• Figure 5: Representative examples of improved Mamba models based on the perspective of block design: (a) Integration methods combine orthogonal architectural designs (e.g., Transformer) with Mamba, leading to enhanced model performance and increased throughput, all while upholding a manageable memory footprint xu2024integratingpilault2024block; (b) Substitution methods enhance the learning capabilities of standard learning frameworks (e.g., UNet) by integrating the Mamba block as a substitute for their primary layers yue2024medmambaliao2024lightm; (c) Modification methods employ cutting-edge techniques, such as Mix-of-Expert (MoE), to refine the Mamba block lieber2024jamba.