作者: Albert Gu, Tri Dao
提交日期: 2023年12月1日
最新修订日期: 2024年5月31日
摘要:
目前驱动深度学习中最激动人心应用的基础模型，几乎都基于Transformer架构及其核心注意力模块。许多亚二次时间复杂度的架构，如线性注意力、门控卷积和循环模型，以及结构化状态空间模型（SSMs），已被开发出来以解决Transformer在长序列上的计算效率低下问题，但它们在语言等重要模态上的表现尚未达到注意力的水平。我们指出，此类模型的一个关键弱点是其无法执行基于内容的推理，并提出了若干改进。首先，简单地让SSM参数成为输入的函数，解决了它们在离散模态上的弱点，使模型能够根据当前标记，在序列长度维度上选择性地传播或遗忘信息。其次，尽管这一改变阻碍了高效卷积的使用，但我们设计了一种硬件感知的循环模式并行算法。我们将这些选择性SSM集成到一个简化的端到端神经网络架构中，该架构没有注意力甚至MLP块（Mamba）。Mamba具有快速的推理速度（吞吐量比Transformer高5倍）和序列长度的线性缩放特性，其性能在真实数据上可提升至百万长度序列。作为一种通用的序列模型主干，Mamba在语言、音频和基因组学等多种模态上均达到了最先进的性能。在语言建模方面，我们的Mamba-3B模型在预训练和下游评估中，均优于同等规模的Transformer，并与两倍于其规模的Transformer性能相当。
主题分类:
- Machine Learning (cs.LG)
- Artificial Intelligence (cs.AI)
arXiv标识符: arXiv:2312.00752