Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5$\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.
翻译:基础模型如今驱动着深度学习领域大多数激动人心的应用,它们几乎普遍基于Transformer架构及其核心注意力模块。为应对Transformer在处理长序列时的计算效率低下问题,已发展出许多亚二次时间复杂度的架构,如线性注意力、门控卷积与循环模型,以及结构化状态空间模型(SSMs),但这些模型在语言等重要模态上的表现始终不及注意力机制。我们发现此类模型的关键缺陷在于无法执行基于内容的推理,并据此提出若干改进。首先,通过使SSM参数成为输入的函数,有效解决了模型在处理离散模态时的弱点,使其能够根据当前标记选择性地沿序列长度维度传播或遗忘信息。其次,尽管这一改变阻碍了高效卷积运算的使用,我们设计了一种硬件感知的循环模式并行算法。我们将这些选择性SSM集成至一个简化的端到端神经网络架构中,该架构无需注意力机制甚至MLP模块(Mamba)。Mamba具备快速推理能力(吞吐量比Transformer高5倍)和序列长度的线性缩放特性,其在实际数据上的性能可提升至百万长度序列。作为通用序列模型主干,Mamba在语言、音频和基因组学等多种模态上均实现了最先进的性能。在语言建模任务中,我们的Mamba-3B模型在预训练和下游评估中均优于同等规模的Transformer,并达到两倍规模Transformer的性能水平。