Jun, 2024
基于蟒蛇的语言模型的实证研究
An Empirical Study of Mamba-based Language Models
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti...
TL;DR选择性状态空间模型（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 项目的一部分发布。