Mamba vs. Transformer：最新研究揭示混合架构才是最优解

Mamba vs. Transformer：最新研究揭示混合架构才是最优解

引用

澎湃

1.

https://www.thepaper.cn/newsDetail_forward_28061205

Mamba模型自推出以来就因其独特架构和潜在优势引起了广泛关注。最近，由英伟达、CMU等机构联合进行的一项实证研究，对Mamba和Transformer进行了全面对比。研究发现，虽然Mamba在某些任务上表现出色，但在上下文学习和信息检索方面存在局限。基于此，研究者提出了一个混合架构Mamba-2-Hybrid，结合了两者的优势，展现了更好的性能和效率。

去年12月，CMU和普林斯顿的两位华人学者提出了Mamba模型，该模型完全抛弃了传统的注意力机制和MLP模块，具有上下文长度线性缩放、推理速度快等优点。然而，在实际应用中，研究者发现学术界对各种Transformer变体的关注度更高，这促使他们对Mamba架构进行了进一步优化。

实验设计与数据集

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

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

实验结果

Mamba vs. Transformer

在常见任务上，Mamba和Mamba-2的性能可以匹敌甚至超过Transformer，但在MMLU基准测试中存在显著差距。即使增加训练数据量，Mamba-2在MMLU上的表现也仅能接近Transformer，但无法超越。

Mamba在MMLU和"电话簿"任务上的表现

研究发现，Mamba在MMLU和"电话簿"任务上的表现较差。MMLU任务中，Mamba在选择题模式下表现不佳，但在填空题模式下却能反超Transformer。这表明Mamba可能在上下文学习和信息路由方面存在局限。

在"电话簿"任务中，Mamba在输入序列达到500 token时就开始出现性能下降，而Transformer在预训练上下文长度内（4096 token）仍能保持接近100%的准确率。

混合架构Mamba-2-Hybrid

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

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

实验结果显示，这种混合架构在多个基准测试中都表现出色：

• 在5-shot MMLU测评中取得最高准确度
• 在效率更高的同时保持了比Transformer更强大的性能
• 在"电话簿"任务上实现了100%的准确率，甚至在预训练长度扩展到128k时仍能保持高精度

结论

研究证实，Mamba和Transformer各有优劣。单纯追求一个架构取代另一个可能不是最佳选择，将两者优势结合的混合架构Mamba-2-Hybrid展现出了巨大潜力。

参考资料：
[1] https://arxiv.org/pdf/2406.07887

本文原文来自澎湃新闻