[论文解读] An Empirical Study of Mamba-based Language Models

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

Language, Linguistics, Cultural Analysis被引用 9

一句话总结

本研究比较了8B参数的Mamba、Mamba-2和Transformer模型（以及一个8B的Mamba-2-Hybrid），在多达3.5T个token上训练，在35个NLP任务和长上下文基准上进行评估，以评估扩展性、复制、上下文内学习和混合架构。

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.

研究动机与目标

• 评估在大规模（8B参数，最高3.5T tokens）下，基于Mamba的语言模型相对于Transformer基线的表现。
• 研究纯SSM（Mamba/Mamba-2）在标准和长上下文任务上的优势与弱点。
• 探讨混合的Mamba-Transformer架构在保持推理效率的同时是否能够缩小纯SSM模型的性能差距。
• 提供开放的基准、检查点和代码，以实现可重复性和进一步研究。

提出的方法

• 通过在相同数据、超参数和评估设置下，直接进行苹果-to-苹果的比较，训练Mamba、Mamba-2、Mamba-2-Hybrid和Transformer模型。
• 在12个标准短上下文任务和23个长上下文任务上，使用开放基准套件（LM Evaluation Harness、LongBench、RULER）进行评估。
• 分析MMLU在三种格式（标准、目标文本选择、填空）中的表现，以探测上下文学习格式。
• 通过消融研究设计混合架构，分配Mamba-2、自注意力和MLP层以实现最佳性能。
• 研究纯模型和混合模型的长上下文扩展（最高128K token）对性能的影响。
• 通过NVIDIA Megatron-LM和Hugging Face发布训练代码和模型权重。

实验结果

研究问题

• RQ1 8B参数的Mamba和Mamba-2在使用高达3.5T token的受控条件下，对比Transformer在标准NLP任务上的性能是否可以匹配？
• RQ2 在需要上下文学习、复制或长上下文推理的任务中，纯SSM模型的具体弱点是什么？
• RQ3 混合的Mamba-Transformer架构是否能够在保持推理效率优势的同时缩小纯SSM模型的差距？
• RQ4 长上下文扩展（16K、32K、128K）如何影响纯SSM和混合模型在标准和长上下文基准上的表现？
• RQ5 Mamba-2-Hybrid架构是否在推理速度和可扩展性方面相较于纯Transformer具有实际的优势？

主要发现

• Pure SSM模型（Mamba/Mamba-2）在许多标准任务上能够匹配或超过Transformer，但在MMLU上落后（尤其是短期情景）以及像Phonebook这样的复制任务上表现较弱。
• 使用3.5T token训练Mamba-2显著缩小MMLU与Transformer的差距，并且在3.5T时，短上下文基准上平均可超过Transformer。
• 一个8B参数的Mamba-2-Hybrid（24个Mamba-2、4个自注意力、28个MLP）在所有12个短上下文任务上超越8B参数的Transformer（平均+2.65点），在长上下文推理下速度可达8倍以上提升。
• Mamba-2-Hybrid在23个长上下文任务上的平均表现，与Transformer基准接近或超越，长上下文扩展与纯模型的表现接近或超越。
• Phonebook风格的复制任务显示纯SSM模型在上下文长度超过约500个标记时难以完成上下文内复制，而Transformer在预训练上下文长度高达4096时仍表现良好。
• 混合模型采用分布式注意力/MLP层，表现强劲，消融研究提示大约8%的自注意力层和30-50%的MLP层是有效配置；对于大规模混合模型，RoPE位置嵌入并非必需，且在长上下文中可省略。
• 推理加速：Mamba-2-Hybrid在长上下文下生成标记的速度远超Transformer，在强基线的推理时延上具有实用MFU的可比性。