摘要
基础模型如今正驱动着深度学习领域绝大多数激动人心的应用，几乎全部基于Transformer架构及其核心的注意力机制。为解决Transformer在处理长序列时的计算效率问题，研究者提出了多种亚二次时间复杂度的替代架构，如线性注意力、门控卷积、递归模型以及结构化状态空间模型（Structured State Space Models, SSMs）。然而，这些模型在语言等关键模态上的表现仍不及注意力机制。我们识别出这类模型的一个核心缺陷：难以进行基于内容的推理。为此，我们提出多项改进。首先，通过将SSM的参数设计为输入位置（即序列索引）的函数，有效缓解了其在离散模态上的表现不足问题，使模型能够根据当前词元（token）的内容，有选择性地在序列长度维度上传播或遗忘信息。其次，尽管这一改进限制了高效卷积的使用，我们仍设计了一种面向硬件优化的并行算法，使其在递归模式下仍能高效运行。我们将这些具备选择性机制的SSM集成到一种简化的端到端神经网络架构中，该架构完全摒弃了注意力机制甚至多层感知机（MLP）模块，由此构建出名为Mamba的新型模型。Mamba在推理速度上显著优于Transformer（吞吐量提升5倍），且在序列长度上实现线性扩展。其性能在真实数据上可稳定扩展至百万级序列长度。作为通用序列建模骨干网络，Mamba在语言、音频和基因组等多个模态上均达到当前最先进水平。在语言建模任务中，我们的Mamba-3B模型在预训练和下游评估中均超越同等规模的Transformer模型，并达到两倍规模Transformer的性能水平。