Mamba: Linear-Time Sequence Modeling with Selective State Spaces

基础模型，现在主要用于深度学习中的许多令人兴奋的应用，几乎普遍基于Transformer架构及其核心注意力模块。许多次线性时间的架构（如线性注意力，门控卷积和循环模型以及结构状态空间模型（SSM））已经被开发出来，以解决Transformer在长序列上的计算效率问题，但它们在重要的模态（如语言）上的表现不如注意力。我们确定这些模型的一个关键弱点是它们无法进行基于内容的推理，并进行了几项改进。首先，简单地让SSM参数成为输入的函数，就可以解决它们在离散模态下的弱点，允许模型根据当前标记有选择地沿序列长度维度传播或遗忘信息。其次，即使这种改变防止了使用高效的卷积，我们也设计了一种硬件感知的并行算法，在循环模式下进行。我们将这些选择性的SSM集成到一个简化的端到端神经网络架构中，甚至没有注意力或MLP块（Mamba）。Mamba具有快速推理（比Transformer高5倍的吞吐量）和序列长度的线性扩展，并且在真实数据上的性能提高了，可以处理长度达百万的序列。作为一种通用的序列模型骨干，Mamba在多个模态（如语言，音频和基因组学）上实现了最先进的性能。在语言建模方面，我们的Mamba-3B模型在预训练和下游评估中均优于相同大小的Transformer，并且与两倍大小的Transformer相当。

本文试图通过提出一种新的序列模型Mamba来解决长序列上Transformer架构的计算效率问题，并在语言、音频和基因组等多个领域中实现最先进的性能。

Mamba模型通过将SSM参数作为输入的函数来实现内容为基础的推理，从而解决了传统模型在离散模态下的弱点，并设计了一种硬件感知并行算法，在递归模式下实现了快速推理。

Mamba模型在长序列上具有线性扩展性，并且在多个领域中实现了最先进的性能，包括语言、音频和基因组。在语言建模方面，Mamba-3B模型在预训练和下游评估中均优于相同大小的Transformer，并且与两倍大小的Transformer相当。研究者还开源了代码。

与本文相关的研究包括使用线性注意力、门控卷积和循环模型的子二次时间架构，以及结构化状态空间模型（SSMs）。