[Paper Review] Falcon Mamba: The First Competitive Attention-free 7B Language Model

Jingwei Zuo, Maksim Velikanov+5 more|arXiv (Cornell University)|Oct 7, 2024

DNA and Biological ComputingBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

Falcon Mamba 7B is the first competitive, fully attention-free 7B parameter language model based on the Mamba architecture, trained on 5.8 trillion tokens. It outperforms leading open-weight models like Mistral 7B, Llama3.1 8B, and Falcon2 11B, while maintaining linear inference cost and low memory usage for long contexts, proving pure Mamba models can match or exceed Transformer-based SoTA models at scale.

ABSTRACT

In this technical report, we present Falcon Mamba 7B, a new base large language model based on the novel Mamba architecture. Falcon Mamba 7B is trained on 5.8 trillion tokens with carefully selected data mixtures. As a pure Mamba-based model, Falcon Mamba 7B surpasses leading open-weight models based on Transformers, such as Mistral 7B, Llama3.1 8B, and Falcon2 11B. It is on par with Gemma 7B and outperforms models with different architecture designs, such as RecurrentGemma 9B and RWKV-v6 Finch 7B/14B. Currently, Falcon Mamba 7B is the best-performing Mamba model in the literature at this scale, surpassing both existing Mamba and hybrid Mamba-Transformer models, according to the Open LLM Leaderboard. Due to its architecture, Falcon Mamba 7B is significantly faster at inference and requires substantially less memory for long sequence generation. Despite recent studies suggesting that hybrid Mamba-Transformer models outperform pure architecture designs, we demonstrate that even the pure Mamba design can achieve similar, or even superior results compared to the Transformer and hybrid designs. We make the weights of our implementation of Falcon Mamba 7B publicly available on https://huggingface.co/tiiuae/falcon-mamba-7b, under a permissive license.

Motivation & Objective

• To demonstrate that a pure Mamba-based language model can achieve state-of-the-art performance comparable to or exceeding optimized Transformer models at scale.
• To investigate whether attention-free architectures can match or surpass hybrid Mamba-Transformer models in general language understanding and generation.
• To validate the feasibility and advantages of pure Mamba models in long-context generation with constant memory and high inference efficiency.
• To provide a fully open-weight, permissively licensed 7B model to advance research in efficient, scalable sequence modeling beyond attention.

Proposed method

• The model uses a pure Mamba architecture with 64 layers, 4096 hidden dimension, and a state space model (SSM) with a 16-dimensional state representation.
• It employs untied input and output embeddings to improve model flexibility and performance at the 7B scale.
• The model is trained on 5.8 trillion tokens using a carefully curated data mixture, with no attention mechanisms.
• Left-side padding is applied and hidden states are zeroed for padding tokens before and after the causal convolution to prevent interference during generation.
• Training stability is enhanced through careful weight initialization and reduced learning rate sensitivity, addressing common loss spikes in Mamba training.
• The model is deployed via Hugging Face and llama.cpp, supporting GGUF quantization and CPU inference for broad accessibility.

Figure 1: Data mixtures across training stages

Experimental results

Research questions

• RQ1 Can a pure Mamba-based language model achieve competitive performance with state-of-the-art open-weight Transformer models at 7B parameters?
• RQ2 Does a fully attention-free architecture maintain linear inference cost and memory efficiency across long context lengths, even without hybrid attention components?
• RQ3 How does the performance of a pure Mamba model compare to hybrid Mamba-Transformer models and other SSM-based architectures like RWKV or Griffin?
• RQ4 To what extent can data quality and scaling compensate for potential in-context learning limitations in Mamba models compared to Transformers?
• RQ5 What are the key training and architectural practices that stabilize training and enable high performance in pure Mamba models?

Key findings

• Falcon Mamba 7B surpasses Mistral 7B, Llama3.1 8B, and Falcon2 11B on the Open LLM Leaderboard, demonstrating competitive performance without attention mechanisms.
• It matches or exceeds the performance of RecurrentGemma 9B and RWKV-v6 Finch 7B/14B, confirming strong results across diverse architectural designs.
• The model maintains constant memory and throughput usage regardless of sequence length, enabling efficient long-context generation.
• Despite being trained on a medium context length (8k tokens), the model shows strong generalization and reasoning capabilities, indicating potential for long-context optimization.
• The model is significantly faster at inference and uses substantially less memory than comparable Transformers, especially for long sequences.
• The model is publicly released under a permissive Apache 2.0-based license with responsible use policy, enabling broad research and deployment access via Hugging Face and llama.cpp.

Figure 2: We vary the context length of the prompt to determine the maximum sequence length that could be processed without encountering an out-of-memory (OOM) error. To ensure a fair comparison, all models were configured with a rescaled vocabulary size.