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。