Source-linked AI summary
An Empirical Study of Mamba-based Language Models
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi, Bryan Catanzaro
TL;DR
The paper addresses the limited evidence from large, controlled comparisons between SSMs and Transformers. It directly evaluates matched 8B Mamba, Mamba-2, hybrid, and Transformer models across standard and long-context tasks. Pure SSMs struggle with retrieval and in-context learning, whereas Mamba-2-Hybrid broadly matches or exceeds Transformers.
Problem
Large-scale apples-to-apples evidence is limited on whether SSM language models can match Transformers when trained on the same data and larger token budgets.
Method
The authors compare matched 8B Mamba, Mamba-2, Mamba-2-Hybrid, and Transformer models trained on up to 3.5T tokens across 35 downstream tasks and extended long-context evaluations.
Results
Pure SSMs match or exceed Transformers on many tasks but lag on context retrieval and in-context learning, while Mamba-2-Hybrid exceeds Transformers across the evaluated common benchmarks and remains competitive at long context.
Takeaways & Limitations
Hybrid SSM-Transformer models show potential to provide inference-time speedups without accuracy degradation relative to Transformers.
Takeaways & Limitations
Long-context hybrid results may depend on prompt formatting and on a continued-pretraining recipe that packs unrelated documents together.
Abstract
from arXiv · showhide
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.
1 Introduction
The report conducts a controlled large-scale comparison of Mamba variants and Transformers, finding that pure SSMs perform well broadly but struggle with context retrieval and in-context learning. A hybrid Mamba-2 architecture improves these weaknesses while retaining the promise of efficient inference.
- Motivation: Self-attention scales quadratically with sequence length during training and requires substantial inference-time memory for long-sequence generation.These costs motivate studying SSM alternatives.
- Study scope: 8B Mamba, Mamba-2, hybrid, and Transformer models are compared under matched hyperparameters, datasets, and evaluation pipelines using up to 3.5T training tokens.The study evaluates 35 downstream language-modeling tasks with the same evaluation setup.
- Pure SSM results: Mamba and Mamba-2 match or exceed Transformers on many downstream tasks but lag on in-context learning, context retrieval, five-shot MMLU, Phonebook, and standard long-context tasks.After 1.1T tokens, both pure SSMs are nearly 15 points below the Transformer on five-shot MMLU; additional training only partially reduces this gap.
- Hybrid models: The 8B Mamba-2-Hybrid uses 24 Mamba-2, 4 self-attention, and 28 MLP layers, and exceeds its Transformer counterpart on all 12 evaluated short-context benchmarks after 3.5T tokens.On MMLU, it achieves five-shot accuracy 3.5 points higher than the Transformer.
- Long-context extensions: On 23 long-context evaluations, 16K- and 32K-context hybrid models closely match or exceed Transformer baselines on average, especially for retrieving, tracking, and aggregating information.Three multi-document question-answering tasks remained challenging.
- Efficiency: Mamba-2 requires one all-reduce per layer under tensor model parallelism, compared with two for Mamba.This is an implementation-level communication difference highlighted in Figure 1.
2 Preliminaries
The report describes its Megatron-LM implementation, training data, model architectures, and evaluation suites. Experiments use multilingual, code-containing datasets and standard, natural long-context, and synthetic long-context tasks.
- Implementation: The implementation supports tensor, sequence, and pipeline parallelism, with pipeline parallelism available for Mamba-2.Hybrid models combine SSM, self-attention, and MLP layers supported by Megatron-LM.
- Implementation: Mamba tensor-parallel layers require two all-reduces per block, whereas Mamba-2 and Transformer layers require one all-reduce per layer.The difference affects communication overhead during large-scale training.
- Training Data: Models are trained on 1.1T- and 3.5T-token datasets, each containing 70% English, 15% non-English, and 15% code.The datasets are predecessors of the Nemotron-4 training dataset.
- Evaluation Tasks and Setup: Standard short-context evaluation covers 12 tasks using LM Evaluation Harness, including MMLU, WinoGrande, PIQA, HellaSwag, and generation-based NQ and SquadV2.Reported metrics include accuracy, normalized accuracy, exact match, and F1.
- Evaluation Tasks and Setup: Natural long-context evaluation uses NarrativeQA, Qasper, and QuALITY, each requiring answers based on a long input document.The tasks use F1 or normalized accuracy depending on the benchmark.
- Evaluation Tasks and Setup: Synthetic long-context evaluation uses Phonebook and 13 RULER generation tasks covering retrieval, tracking, aggregation, and question answering.Accuracy is reported on 400 synthetic RULER samples.
3 Mamba and Mamba-2 Compared to Transformers
At 8B parameters, pure Mamba and Mamba-2 match or exceed Transformers on many standard tasks but struggle with MMLU formatting, copying, and context recall. Longer training narrows MMLU differences, while Phonebook failures persist beyond short contexts.
- 3.3.1 Downstream Language Modeling Tasks: Pure SSM models match or exceed Transformers on many standard natural-language tasks, including higher average accuracy across five evaluations.The five tasks are WinoGrande, PIQA, HellaSwag, ARC-Easy, and ARC-Challenge.
- 3.3.1 Downstream Language Modeling Tasks: 17 points: Mamba-2 trails the Transformer on five-shot MMLU after 1.1T training tokens.With 3.5T tokens, the MMLU gap narrows to 1.37 points.
- 3.3.1 Downstream Language Modeling Tasks: 3.5T-token training lets Mamba-2 approach Transformer-level MMLU accuracy and exceed the Transformer in average task accuracy.Only pure Mamba-2 and Transformer models were trained at this scale because larger-scale Mamba training was inefficient.
- 3.3.2 A Closer Look at MMLU: Pure SSM models struggle with standard and choice-text MMLU but match or exceed the Transformer in the cloze formulation.The results suggest that SSM models contain comparable knowledge but need more training to understand some multiple-choice formats.
- 3.3.3 Copying Tasks: Beyond approximately 500-token Phonebook inputs, Mamba and Mamba-2 accuracy declines, while the 8B Transformer remains near 100% through its 4096-token pretraining context.The Mamba-2 degradation persists even after training on 3.5T tokens.
- 3.3.3 Copying Tasks: Pure SSMs exhibit fuzzy Phonebook memory: their predicted numbers often share correctly positioned digits with the answer without exact recall.Changing to a reversed prompt improves SSM accuracy but does not remove degradation for shorter-than-4096-token phone books.
- 3 Mamba and Mamba-2 Compared to Transformers: Pure SSM models remain challenging to train for tasks requiring in-context learning, information routing, and copying from context.These weaknesses are identified particularly for smaller-token-horizon MMLU and Phonebook.
4 Hybrid Mamba-Transformer Models
The study designs and evaluates hybrid Mamba-Transformer models that combine Mamba-2, self-attention, and MLP layers. The 8B hybrid exceeds the Transformer on standard tasks and remains competitive on long-context evaluations, while some prompt-formatting and multi-document tasks remain challenging.
- Architecture design: Validation loss is minimized with roughly 8% attention layers, while 50% MLP layers make training 20% faster than 5%.
- Efficiency: 29.9% MFU for the hybrid is close to the corresponding Transformer’s 30.7% MFU on 1,024 H100 GPUs.
- Short-context evaluation: The hybrid achieves higher accuracy than the Transformer on all 12 common tasks, improving by 2.65 points on average.
- Long-context evaluation: On long-context evaluations, hybrid models closely match or exceed Transformers on average, especially for retrieval, tracking, aggregation, Phonebook, and NIAH tasks.
- Limitations: Prompt formatting may affect hybrids more than Transformers, and continued pretraining with packed unrelated documents may confuse SSM layers on some tasks.
5 Related Work
Related work has explored Mamba-attention and other hybrid architectures for improving accuracy, efficiency, and long-sequence generalization. This study differs by examining larger models and finding that component proportions, rather than a fixed layer pattern, are key.
- Mamba-attention hybrids: Prior Mamba-attention work studied in-context learning but used models up to 77M parameters, while Jamba and Zamba reached 7B scale.
- Other hybrid architectures: Other hybrids combine attention with gated linear recurrences or convolutions to improve latency, perplexity, or longer-sequence generalization.
- Architectural distinction: Unlike manually patterned prior hybrids, this study finds that the relative proportions of Mamba, attention, and MLP components are the key architectural factor.
6 Conclusion
The report compares 8B Mamba, Mamba-2, hybrid, and Transformer models under matched large-scale training conditions. Pure SSMs remain competitive on most tasks but struggle with retrieval and in-context learning, whereas hybrids achieve higher benchmark accuracy and retain strong long-context performance.
- The study directly compares 8B Mamba, Mamba-2, Mamba-2-Hybrid, and Transformer models trained on up to 3.5T tokens.
- Pure SSMs match or exceed Transformers on most downstream tasks but are challenged by context-based retrieval and in-context learning.
- Mamba-2-Hybrid models reach higher accuracy than Transformers on all evaluated common benchmarks and remain strong at 16K and 32K contexts.
- The results support the potential for SSM-based hybrids to provide inference-time speedups without accuracy degradation relative to Transformers.
A Hybrid Layer Allocation Algorithm
The hybrid layer allocation algorithm converts target attention and MLP ratios into an ordered layer-type pattern. It distributes attention and MLP layers across Mamba layers to keep intervening Mamba sections approximately balanced.
- The algorithm takes total layer count, target attention ratio, and target MLP ratio as inputs and outputs a layer-type list.
- Attention layers are placed so contiguous Mamba runs are as equal as possible, with Mamba runs at both sequence ends.
- It rounds the target attention ratio, assigns the remaining layers to Mamba, and creates one more Mamba section than attention layers.
- The algorithm then rounds the target MLP count and distributes MLP layers among positions initially assigned to Mamba.
- Table 11 gives examples of layer patterns generated by the allocation algorithm.