Mamba真比Transformer更优吗？最新研究揭示混合架构才是最优解

Mamba真比Transformer更优吗？最新研究揭示混合架构才是最优解

引用

澎湃

1.

https://m.thepaper.cn/newsDetail_forward_28061205

Mamba模型自推出以来就因其独特架构和潜在优势引起了广泛关注。最近，来自英伟达、CMU、普林斯顿等机构的研究团队对Mamba和Transformer进行了全面对比研究，发现两者各有优劣。研究还提出了一种混合架构Mamba-2-Hybrid，结合了两者的优点，展现了更好的性能。

去年12月，CMU和普林斯顿的两位华人学者提出了Mamba模型，该模型完全抛弃了传统的注意力机制和MLP模块，采用状态空间模型（SSM）架构。Mamba模型具有上下文长度线性缩放、推理速度快等优点，一度被视为Transformer的潜在替代者。

然而，研究团队在后续研究中发现，虽然Mamba模型在某些任务上表现出色，但在长上下文任务和信息检索方面存在局限性。因此，他们提出了Mamba-2-Hybrid混合架构，将SSM和注意力机制相结合，以期发挥两者的优势。

实验设置

研究团队训练了四种不同架构的8B参数模型：Mamba、Mamba-2、Mamba-2-Hybrid和Transformer。训练数据集包括1.1T和3.5T两个规模，均由70%英语、15%非英语和15%代码组成。实验涵盖了35个NLP下游任务，分为三类：

• 标准短上下文任务（12个）
• 自然长上下文任务（9个）
• 综合长上下文任务（14个）

实验结果

Mamba vs. Transformer

在大多数标准任务上，Mamba和Mamba-2的性能可以与Transformer媲美，但在MMLU基准测试中存在显著差距。具体来说：

• 在1.1T数据集上，Mamba-2在零样本和少样本学习中分别落后Transformer 10分和17分。
• 增加训练数据到3.5T后，Mamba-2在MMLU任务上的5-shot分数差距缩小到1.37分。

Mamba在特定任务上的表现

研究团队对MMLU和"电话簿"任务进行了深入分析：

• MMLU：Mamba在选择题模式下表现较差，但在填空题模式中反而优于Transformer。这表明Mamba可能在上下文学习和信息路由方面存在局限性。
• 电话簿任务：Mamba在处理长上下文信息时表现不佳，尤其是在需要精确复制信息的场景中。

混合架构Mamba-2-Hybrid

为了解决Mamba的局限性，研究团队提出了Mamba-2-Hybrid架构，该架构包含：

• 24个Mamba-2层
• 4个自注意力层
• 28个MLP层

实验结果显示：

• 在5-shot MMLU测评中，Mamba-2-Hybrid超越了纯Transformer和纯SSM模型。
• 在长上下文任务中，Mamba-2-Hybrid的token生成速度比Transformer快8倍。
• 在"电话簿"任务中，Mamba-2-Hybrid在4K预训练长度内可达到100%准确率，并能泛化到5.5K token长度。

结论

研究结果表明，Mamba和Transformer各有优劣，单纯追求替代关系可能不是最佳选择。将两者优点结合的混合架构Mamba-2-Hybrid展现了更好的性能和效率，为未来模型设计提供了新的思路。

参考资料

• Mamba vs. Transformer实证研究论文
• Mamba-2论文
• 原始Mamba论文