An Empirical Study of Mamba-based Language Models

选择性状态空间模型（SSMs），例如Mamba，克服了Transformers的一些缺点，例如随着序列长度的增加而呈二次计算复杂度和来自键值缓存的大量推理时间内存需求。此外，最近的研究表明，SSMs可以匹配或超越Transformers的语言建模能力，使它们成为一种有吸引力的替代方案。然而，在受控环境中（例如相同的数据），到目前为止的研究只展示了比较小规模的实验，比较了SSMs和Transformers。为了了解这些架构在更大规模上的优势和劣势，我们在相同的数据集上训练8B参数的Mamba、Mamba-2和Transformer模型，并进行直接比较，数据集的规模可达3.5T个标记。我们还将这些模型与混合架构进行比较，该架构由43％的Mamba-2、7％的注意力和50％的MLP层（Mamba-2-Hybrid）组成。通过使用多样化的任务，我们回答了Mamba模型是否能够在更大的训练预算下匹配Transformers的问题。我们的结果表明，纯SSMs在许多任务上可以匹配或超越Transformers，但在需要强大的复制或上下文学习能力（例如，5-shot MMLU，电话簿）或长上下文推理的任务上落后于Transformers。相反，我们发现8B的Mamba-2-Hybrid在我们评估的12个标准任务中均超过了8B的Transformer（平均+2.65分），并且在生成标记时预计速度可提高高达8倍。为了验证长上下文能力，我们提供了额外的实验，评估了Mamba-2-Hybrid和Transformer的变体，以支持16K、32K和128K序列。在额外的23个长上下文任务中，混合模型在平均上仍然能够与或超越Transformer。为了进一步研究，我们作为NVIDIA的Megatron-LM项目的一部分发布了检查点以及用于训练我们的模型的代码。

对比研究Selective state-space models (SSMs)和Transformers的性能差异，探究在大规模训练预算下，Mamba模型能否与Transformer匹敌？

论文通过对比8B-parameter Mamba、Mamba-2、Transformer和Mamba-2-Hybrid模型在多个任务上的表现，发现纯SSM模型在需要强大的复制或上下文学习能力的任务上表现不如Transformer，而Mamba-2-Hybrid模型在所有12个标准任务上都超越了Transformer，并且在生成标记时预计会快8倍。

论文使用了多个数据集和任务来评估模型的性能，同时开源了代码和检查点，方便进一步研究。实验结果表明，Mamba-2-Hybrid模型在长上下文任务上的表现也与Transformer相当甚至更好。

最近的相关研究包括《Attention Is All You Need》、《BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding》等。