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.

选择性状态空间模型（SSMs）如 Mamba 克服了 Transformer 的一些缺点，例如与序列长度呈二次增长的计算复杂度和从键值缓存中获取大量的推理时间内存需求。此外，最近的研究显示，SSMs 可以达到或超越 Transformer 的语言建模能力，使其成为一种有吸引力的替代选择。然而，迄今为止的研究只在相同数据的受控环境中进行了小规模实验，比较了 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 倍。为了验证长期上下文能力，我们进行了其他实验，评估了支持 16K、32K 和 128K 序列的 Mamba-2-Hybrid 和 Transformer 的变体。在 23 个额外的长期上下文任务中，混合模型在平均水平上继续紧密匹配或超越了 Transformer。为了进一步的研究，我们将检查点以及用于训练我们的模型的代码作为 NVIDIA 的 Megatron-LM 项目的一部分发布。