arXiv:2312.00752cs.LGcs.AI2023-12被引 8.8k

Mamba用选择性状态空间模型实现线性长序列建模,性能超越同规模Transformer。

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

  • 引入输入相关参数的可选状态空间,实现内容感知的信息传递与遗忘。
  • 在百万级序列上保持线性扩展,推理速度比Transformer快5倍。
  • 无需注意力或MLP,适用于文本、音频、基因组等多模态任务。

基础模型如今推动深度学习诸多前沿应用,几乎全部基于Transformer架构及其核心注意力机制。尽管已有多种亚二次时间复杂度的模型(如线性注意力、门控卷积、递归模型和结构化状态空间模型)试图缓解Transformer在长序列上的计算效率问题,但在语言等重要模态上表现仍不及注意力模型。我们识别出这些模型的关键弱点在于缺乏内容感知推理能力,并提出多项改进:首先,将状态空间模型参数设为输入的函数,使模型能根据当前标记选择性地沿序列长度维度传播或遗忘信息;其次,尽管此改动阻碍了高效卷积的应用,我们设计了一种硬件感知的并行递归算法。我们将这些可选状态空间模型整合进一个无注意力、甚至无MLP块的简化端到端神经网络架构中(Mamba)。Mamba具备快速推理能力(比Transformer高5倍吞吐量),且序列长度呈线性扩展,在真实数据上性能可提升至百万级长度序列。作为通用序列模型骨干,Mamba在语言、音频和基因组等多个模态上达到顶尖水平。在语言建模任务中,我们的Mamba-3B模型在预训练和下游评估中均优于同规模Transformer,且性能相当于两倍大小的Transformer。

原文摘要 · Abstract (English)

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.

序列建模状态空间线性复杂度Mamba

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。