[Paper Review] An Empirical Study of Mamba-based Language Models

Roger Waleffe, Wonmin Byeon+14 more|arXiv (Cornell University)|Jun 12, 2024

Language, Linguistics, Cultural Analysis9 citations

TL;DR

This work compares 8B-parameter Mamba, Mamba-2, and Transformer models (plus an 8B Mamba-2-Hybrid) trained on up to 3.5T tokens, evaluating them on 35 NLP tasks and long-context benchmarks to assess scaling, copying, in-context learning, and hybrid architectures.

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.

Motivation & Objective

• Assess how Mamba-based language models perform relative to Transformer baselines at large scale (8B parameters, up to 3.5T tokens).
• Investigate strengths and weaknesses of pure SSM (Mamba/Mamba-2) on standard and long-context tasks.
• Explore whether a hybrid Mamba-Transformer architecture can close performance gaps while retaining efficiency benefits.
• Provide open benchmarks, checkpoints, and code to enable reproducibility and further study.

Proposed method

• Direct apples-to-apples comparison by training Mamba, Mamba-2, Mamba-2-Hybrid, and Transformer models with the same data, hyperparameters, and evaluation setup.
• Evaluation on 12 standard short-context tasks and 23 long-context tasks using open benchmark suites (LM Evaluation Harness, LongBench, RULER).
• Analysis of MMLU across three formats (standard, choice-text-in-targets, and cloze) to probe in-context learning formats.
• Ablation studies to design hybrid architectures distributing Mamba-2, self-attention, and MLP layers for optimal performance.
• Investigation of long-context extensions up to 128K tokens for both pure and hybrid models.
• Release of training code and model weights via NVIDIA Megatron-LM and Hugging Face.

Experimental results

Research questions

• RQ1 Can 8B-parameter Mamba and Mamba-2 match Transformer performance on standard NLP tasks when trained on large token budgets (up to 3.5T tokens) under controlled conditions?
• RQ2 What are the specific weaknesses of pure SSM models in tasks requiring in-context learning, copying, or long-context reasoning?
• RQ3 Can a hybrid Mamba-Transformer architecture close the gaps observed for pure SSM models while preserving efficiency benefits during inference?
• RQ4 How do long-context extensions (16K, 32K, 128K) affect the performance of pure SSM and hybrid models on standard and long-context benchmarks?
• RQ5 Do Mamba-2-Hybrid architectures demonstrate practical inference speedups and scalability advantages over pure Transformers?

Key findings

• Pure SSM models (Mamba/Mamba-2) can match or exceed Transformers on many standard tasks, but lag on MMLU (especially with shorter horizons) and copying tasks like Phonebook.
• Training Mamba-2 with 3.5T tokens substantially closes the MMLU gap to Transformer and can surpass Transformer on average in short-context benchmarks at 3.5T.
• An 8B-parameter Mamba-2-Hybrid (24 Mamba-2, 4 self-attention, 28 MLP) exceeds the 8B-parameter Transformer on all 12 short-context tasks evaluated (+2.65 average points) and can be up to 8x faster at inference for long contexts.
• Long-context extensions (16K and 32K) for the Mamba-2-Hybrid closely match or exceed Transformer baselines on average across 23 long-context tasks.
• Phonebook-style copying tasks reveal pure SSM models struggle with in-context copying beyond ~500 tokens, whereas the Transformer handles pretraining context lengths up to 4096.
• Hybrid models with distributed attention/MLP layers show strong performance, with ablations suggesting around 8% self-attention layers and 30-50% MLP layers as effective configurations; RoPE position embeddings are not essential for large hybrids and may be omitted for long contexts.
• Inference speedup: Mamba-2-Hybrid can generate tokens significantly faster than Transformers on long contexts, with practical MFU comparable to strong Transformer baselines.