Mamba (deep learning architecture)

Mamba^[a] is a deep learning architecture focused on sequence modeling. It was developed by researchers from Carnegie Mellon University and Princeton University to address some limitations of transformer models, especially in processing long sequences. It is based on the structured state space sequence (S4) model.^[2][3][4]

To enable handling long data sequences, Mamba incorporates S4.^[2] S4 can effectively and efficiently model long dependencies by combining continuous time, and recurrent, and convolutional models. These enable it to handle irregularly sampled data, unbounded context, and remain computationally efficient during training and inferencing.^[5]

Mamba introduces significant enhancements to S4, particularly in its treatment of time-variant operations. It adopts a unique selection mechanism that adapts structured state space model (SSM) parameters based on the input.^[6][2] This enables Mamba to selectively focus on relevant information within sequences, effectively filtering out less pertinent data. The model transitions from a time-invariant to a time-varying framework, which impacts both computation and efficiency.^[2][7]

Mamba employs a hardware-aware algorithm that exploits GPUs, by using kernel fusion, parallel scan, and recomputation.^[2] The implementation avoids materializing expanded states in memory-intensive layers, thereby improving performance and memory usage. The result is significantly more efficient in processing long sequences compared to transformers.^[2][7]

Additionally, Mamba simplifies its architecture by integrating the SSM design with MLP blocks, resulting in a homogeneous and streamlined structure, furthering the model's capability for general sequence modeling across data types that include language, audio, and genomics, while maintaining efficiency in both training and inference.^[2]

• Selective state spaces (SSM): The core of Mamba, SSMs are recurrent models that selectively process information based on the current input. This allows them to focus on relevant information and discard irrelevant data.^[2]
• Simplified architecture: Mamba replaces the complex attention and MLP blocks of transformers with a single, unified SSM block. This aims to reduce computational complexity and improve inference speed.^[2]
• Hardware-aware parallelism: Mamba uses a recurrent mode with a parallel algorithm specifically designed for hardware efficiency, potentially further enhancing its performance.^[2]

Mamba-2 serves as a successor to Mamba by introducing a new theoretical and computational framework called Structured State Space Duality (SSD). This contribution acts as a mathematical bridge between SSMs and Transformers. Specifically, a connection to the attention mechanism. This is in response to concerns over the difficulty in training SSMs compared to Transformers. SSD gives Mamba-2 the ability to inherit many system-level optimizations for Transformers, while maintaining linear-time scalability. ^[8]

Mamba-2 is designed to leverage many of the system and algorithmic optimizations that have been developed for Transformers. Mamba-2 does this through mathematical properties gained from Structured State Space Duality (SSD). Additionally, Mamba-2 introduces a parallel block to further connect the architecture to the attention mechanism and improve scalability.

• SSD Layer: The main contribution of structured state space duality in Mamba-2 is through the SSD layer. In Mamba-1, the state space A matrix is restricted to a diagonal matrix to improve computation. SSD further restricts the A matrix to be a scalar times the identity. This scalar restriction allows for the mathematical dual forms to arise.
• SSD Framework: The researchers further propose the SSD framework to better reason about the model. Firstly, SSD can be viewed through a structured matrix transformation framework. Many sequence models can be seen as a sequence of matrix transformations. State space models such as Mamba are among these models. This leads to a linear formulation of their output as sequential matrix multiplications. The addition of the SSD to Mamba further allows this formulation to be transformed into a quadratic form directly related to masked attention. This is the "duality" in the matrix transformation framework. SSD can also be viewed through a structured attention framework. This utilizes causal linear attention to show that structured masked attention has similar dual linear and quadratic modes. Linear attention is a method to compute the attention algorithm in linear time. It also has a recurrent form similar to SSMs. Mamba-2 bridges these two frameworks by showing that structured masked attention is equivalent to the scalar-identity formulation.
• SSD Algorithm: The SSD algorithm allows for a more hardware-efficient computation of the SSD model. It does this by rewriting computations to leverage matrix multiplication operations. This is due to matrix multiplications FLOPs being more efficient than non-matrix multiplication FLOPs when leveraging tensor cores. It does this by writing the matrix transformations of SSMs as semiseparable matrices as shown in the SSD framework.^[8][9]

Operating on byte-sized tokens, transformers scale poorly as every token must "attend" to every other token leading to O(n2) scaling laws, as a result, Transformers opt to use subword tokenization to reduce the number of tokens in text, however, this leads to very large vocabulary tables and word embeddings.

This research investigates a novel approach to language modeling, MambaByte, which departs from the standard token-based methods. Unlike traditional models that rely on breaking text into discrete units, MambaByte directly processes raw byte sequences. This eliminates the need for tokenization, potentially offering several advantages:^[10]

• Language independence: Tokenization often relies on language-specific rules and vocabulary, limiting applicability across diverse languages. MambaByte's byte-level representation allows it to handle different languages without language-specific adaptations.
• Removes the bias of subword tokenisation: where common subwords are overrepresented and rare or new words are underrepresented or split into less meaningful units. This can affect the model's understanding and generation capabilities, particularly for languages with rich morphology or tokens not well-represented in the training data.
• Simplicity in preprocessing: It simplifies the preprocessing pipeline by eliminating the need for complex tokenization and vocabulary management, reducing the preprocessing steps and potential errors.

Subword tokenisation introduces a number of quirks in LLMs, such as failure modes where LLMs can't spell words, reverse certain words, handle rare tokens, which are not present in byte-level tokenisation.^[11]

MoE Mamba represents a pioneering integration of the mixture of experts (MoE) technique with the Mamba architecture, enhancing the efficiency and scalability of SSMs in language modeling. This model leverages the strengths of both MoE and SSMs, achieving significant gains in training efficiency—requiring 2.2 times fewer training steps than its predecessor, Mamba, while maintaining competitive performance. MoE Mamba showcases improved efficiency and effectiveness by combining selective state space modeling with expert-based processing, offering a promising avenue for future research in scaling SSMs to handle tens of billions of parameters. The model's design involves alternating Mamba and MoE layers, allowing it to efficiently integrate the entire sequence context and apply the most relevant expert for each token.^[12][13]

Vision Mamba (Vim) integrates SSMs with visual data processing, employing bidirectional Mamba blocks for visual sequence encoding. This method reduces the computational demands typically associated with self-attention in visual tasks. Tested on ImageNet classification, COCO object detection, and ADE20k semantic segmentation, Vim showcases enhanced performance and efficiency and is capable of handling high-resolution images with lower computational resources. This positions Vim as a scalable model for future advancements in visual representation learning.^[14]

Jamba is a novel architecture built on a hybrid transformer and mamba SSM architecture developed by AI21 Labs with 52 billion parameters, making it the largest Mamba variant created so far. It has a context window of 256k tokens.^[15]

• Language modeling
• Transformer (machine learning model)
• State-space model
• Recurrent neural network