[论文解读] 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 Biology被引用 3

一句话总结

Falcon Mamba 7B 是首个基于 Mamba 架构、完全无需注意力机制的 70 亿参数语言模型，使用 5.8 万亿 token 进行训练。它在性能上超越了领先的开源模型如 Mistral 7B、Llama3.1 8B 和 Falcon2 11B，同时保持线性推理成本和低内存占用，适用于长上下文场景，证明纯 Mamba 模型在规模上可与或超越基于 Transformer 的最先进模型。

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.

研究动机与目标

• 证明纯 Mamba 架构的语言模型可在大规模下实现与最先进 Transformer 模型相当或更优的性能。
• 探究无需注意力机制的架构是否能在通用语言理解与生成任务中与混合 Mamba-Transformer 模型相媲美或超越。
• 验证纯 Mamba 模型在长上下文生成中实现可行性与优势，包括恒定内存使用和高推理效率。
• 提供一个完全开源、采用宽松许可证的 70 亿参数模型，以推动注意力之外高效、可扩展序列建模的研究。

提出的方法

• 该模型采用纯 Mamba 架构，包含 64 层、4096 隐状态维度，以及一个具有 16 维状态表示的状态空间模型（SSM）。
• 通过解耦输入和输出嵌入，提升 70 亿参数规模下的模型灵活性与性能。
• 在经过精心筛选的数据混合体上使用 5.8 万亿 token 进行训练，全程不使用注意力机制。
• 在因果卷积前后的填充标记处应用左对齐填充，并对这些填充标记的隐藏状态置零，以防止生成过程中产生干扰。
• 通过仔细的权重初始化和降低学习率敏感性，增强训练稳定性，解决 Mamba 训练中常见的损失尖峰问题。
• 通过 Hugging Face 和 llama.cpp 部署，支持 GGUF 量化和 CPU 推理，实现广泛可及性。

Figure 1: Data mixtures across training stages

实验结果

研究问题

• RQ1 纯 Mamba 架构的语言模型是否能在 70 亿参数规模下实现与最先进开源 Transformer 模型相当的性能？
• RQ2 完全无需注意力机制的架构是否能在长上下文长度下保持线性推理成本和内存效率，即使没有混合注意力组件？
• RQ3 与混合 Mamba-Transformer 模型及其他 SSM 架构（如 RWKV 或 Griffin）相比，纯 Mamba 模型的性能表现如何？
• RQ4 与 Transformer 相比，Mamba 模型在上下文学习能力上可能存在局限，数据质量和规模在多大程度上可弥补这一缺陷？
• RQ5 哪些关键的训练与架构实践能够稳定训练过程，并在纯 Mamba 模型中实现高性能？

主要发现

• Falcon Mamba 7B 在 Open LLM Leaderboard 上超越了 Mistral 7B、Llama3.1 8B 和 Falcon2 11B，证明其在无注意力机制下仍具备竞争力。
• 其性能与 RecurrentGemma 9B 及 RWKV-v6 Finch 7B/14B 相当或更优，表明在多种架构设计下均表现强劲。
• 无论序列长度如何，模型均保持恒定的内存占用和吞吐量，支持高效的长上下文生成。
• 尽管训练时上下文长度为中等（8k token），模型仍展现出强大的泛化与推理能力，表明其具备长上下文优化潜力。
• 在推理速度上显著快于同类 Transformer 模型，尤其在长序列场景下内存占用大幅降低。
• 该模型以宽松的 Apache 2.0 基础许可证公开发布，并附带负责任使用政策，可通过 Hugging Face 和 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.