Mamba:具有选择性状态空间的线性时间序列建模

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

特里·道 Tri Dao · · 2023-12-01 · arXiv:2312.00752 ↗ · 被引 7852

打开互动全文版(逐段中英对照 + 图/公式 + 论文问答)→

摘要 · Abstract

基础模型如今支撑着深度学习中大多数令人兴奋的应用,它们几乎普遍基于 Transformer 架构及其核心的注意力模块。为了解决 Transformer 在长序列上的计算效率问题,研究者开发了许多次二次时间架构,如线性注意力、门控卷积和循环模型以及结构化状态空间模型(SSM),但这些模型在语言等重要模态上的表现不如注意力。我们发现这些模型的一个关键弱点在于它们无法进行基于内容的推理,并对此做出了若干改进。首先,让 SSM 参数成为输入的函数,解决了它们在离散模态上的弱点,使得模型能够根据当前 token 沿着序列长度维度选择性地传播或遗忘信息。其次,尽管这一改变阻止了高效卷积的使用,我们设计了一种硬件感知的循环模式并行算法。我们将这些选择性 SSM 集成到一个简化的端到端神经网络架构中,该架构不包含注意力甚至 MLP 块(Mamba)。Mamba 具有快速推理(吞吐量比 Transformer 高 5 倍)和序列长度线性扩展的特点,在长达百万长度的真实数据上性能提升。作为一个通用序列模型骨干,Mamba 在语言、音频和基因组学等多个模态上取得了最先进的性能。在语言建模方面,我们的 Mamba-3B 模型在预训练和下游评估中均优于相同大小的 Transformer,并与两倍大小的 Transformer 表现相当。

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.

核心贡献 · Key contributions

局限 · Limitations

论文章节 · Sections(共 33)

阅读逐段中英对照全文 →