An Empirical Study of Mamba-based Language Models
This study compares 8B Mamba, Mamba-2, and Transformer models trained on up to 3.5T tokens, showing hybrid models outperform in speed and long-sequence tasks.
Key Findings
Methodology
Unified training of 8B models (Mamba, Mamba-2, Transformer) on datasets up to 3.5T tokens, with identical hyperparameters. Designed pure and hybrid architectures (including attention and MLP layers). Used NVIDIA Megatron-LM for training and inference optimization. Evaluated performance across diverse NLP tasks, including short and long-context benchmarks, measuring accuracy, inference speed, and sequence length support.
Key Results
- Pure SSM models outperform Transformers on many short tasks (e.g., WinoGrande, PIQA), but lag on tasks requiring copying or long-context reasoning (e.g., MMLU, Phonebook). After training on 3.5T tokens, Mamba-2 narrows the MMLU gap to 1.37 points. Hybrid models (Mamba-2-Hybrid) surpass Transformers on all 12 short tasks (+2.65 points), with inference speeds up to 8Γ faster and support for sequences up to 128K tokens.
- Long-sequence evaluations show hybrid models excel in multi-document QA and retrieval tasks, maintaining performance beyond trained sequence lengths. The 128K sequence model performs Phonebook perfectly with over 150K tokens, demonstrating robust long-range reasoning.
- These results highlight that combining SSM and attention layers yields models that are both fast and capable of handling extremely long inputs, addressing key bottlenecks in current NLP systems.
Significance
This work advances the application of linear-time state-space models in large-scale NLP, demonstrating their potential to offer faster, more efficient inference without sacrificing accuracy. The hybrid architecture bridges the gap between speed and capability, paving the way for scalable, energy-efficient models suitable for real-world deployment. It challenges the dominance of pure Transformers, suggesting new directions for model design that balance performance, efficiency, and long-range reasoning. The open-sourcing of models and code fosters community engagement and further innovation.
Technical Contribution
Introduced multi-layer hybrid architectures combining SSM (Mamba, Mamba-2) with attention and MLP layers, optimized for hardware-aware training. Demonstrated linear complexity and reduced memory footprint, enabling support for sequences up to 128K tokens. Validated large-scale training on datasets up to 3.5T tokens, with comprehensive benchmarking showing competitive or superior performance. Provided insights into the trade-offs between pure SSM, hybrid, and Transformer models, establishing new benchmarks for long-context NLP.
Novelty
First large-scale (8B parameters, 3.5T tokens) systematic comparison of pure SSM and Transformer models, revealing the strengths of hybrid architectures. Innovatively integrated SSM with attention and MLP layers, enabling efficient long-sequence processing. Demonstrated that hybrid models can outperform Transformers in both speed and long-range reasoning, a significant step forward in scalable NLP model design.
Limitations
- Pure SSM models still underperform in copying and in-context learning tasks, indicating limitations in information retrieval and dynamic memory capabilities. Further architectural refinement is needed.
- While long-sequence support is improved, the models' accuracy on extremely long inputs (beyond 128K) and in complex reasoning scenarios remains limited, requiring ongoing research.
- Training large hybrid models demands substantial computational resources, which may hinder widespread adoption. Future work should focus on reducing costs and improving training efficiency.
Future Work
Future efforts will enhance pure SSM architectures to improve copying and in-context learning, explore more efficient training algorithms, and extend long-sequence capabilities further. Integrating multi-modal data and refining hybrid strategies could unlock broader applications. Additionally, optimizing hardware utilization and reducing training costs will be key to scaling these models for industrial deployment.
AI Executive Summary
This research presents a comprehensive comparison of 8B parameter Mamba, Mamba-2, and Transformer models trained on datasets up to 3.5 trillion tokens. By designing both pure and hybrid architectures, the study evaluates their performance across a wide range of NLP tasks, including short-text understanding and long-sequence reasoning. Results reveal that pure SSM models can match or surpass Transformers on many standard benchmarks, yet they face challenges in tasks requiring strong copying or in-context learning, such as MMLU and Phonebook tasks. Notably, training on larger datasets (3.5T tokens) narrows these gaps significantly.
The most striking finding is that hybrid models, combining Mamba-2 with attention and MLP layers, outperform pure models and Transformer baselines across all short tasks, with an average gain of 2.65 points. These hybrid models also support sequences up to 128K tokens, enabling perfect performance on long-context tasks like Phonebook with over 150K tokens. The inference speed of hybrid models can be up to 8 times faster, demonstrating their efficiency advantage.
This work underscores the potential of integrating state-space models with traditional attention mechanisms to build scalable, fast, and capable language models. Such architectures are promising for applications demanding long-range context understanding, such as multi-document QA, knowledge retrieval, and extended dialogue systems. Despite these advances, limitations remain in pure SSM's copying and reasoning abilities, and the high computational cost of training large hybrid models. Future research will focus on further improving long-sequence generalization, reducing training costs, and exploring multi-modal extensions.
Open-sourcing the models and training code aims to accelerate community adoption and innovation. Overall, this study marks a significant step toward more efficient and scalable NLP systems, challenging the dominance of conventional Transformer architectures and opening new avenues for research and industrial deployment.
Deep Dive
Abstract
Selective state-space models (SSMs) like Mamba overcome some of the shortcomings of Transformers, such as quadratic computational complexity with sequence length and large inference-time memory requirements from the key-value cache. Moreover, recent studies have shown that SSMs can match or exceed the language modeling capabilities of Transformers, making them an attractive alternative. In a controlled setting (e.g., same data), however, studies so far have only presented small scale experiments comparing SSMs to Transformers. To understand the strengths and weaknesses of these architectures at larger scales, we present a direct comparison between 8B-parameter Mamba, Mamba-2, and Transformer models trained on the same datasets of up to 3.5T tokens. We also compare these models to a hybrid architecture consisting of 43% Mamba-2, 7% attention, and 50% MLP layers (Mamba-2-Hybrid). Using a diverse set of tasks, we answer the question of whether Mamba models can match Transformers at larger training budgets. Our results show that while pure SSMs match or exceed Transformers on many tasks, they lag behind Transformers on tasks which require strong copying or in-context learning abilities (e.g., 5-shot MMLU, Phonebook) or long-context reasoning. In contrast, we find that the 8B Mamba-2-Hybrid exceeds the 8B Transformer on all 12 standard tasks we evaluated (+2.65 points on average) and is predicted to be up to 8x faster when generating tokens at inference time. To validate long-context capabilities, we provide additional experiments evaluating variants of the Mamba-2-Hybrid and Transformer extended to support 16K, 32K, and 128K sequences. On an additional 23 long-context tasks, the hybrid model continues to closely match or exceed the Transformer on average. To enable further study, we release the checkpoints as well as the code used to train our models as part of NVIDIA's Megatron-LM project.