[论文解读] Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao|arXiv (Cornell University)|Dec 1, 2023

Neural Networks and Applications被引用 967

一句话总结

Mamba 引入了对输入依赖的选择性状态空间模型（SSMs），实现线性时间序列建模，推理吞吐量比 Transformer 高出 5×，在语言、音频和基因组学等领域表现出色。在语言任务中，Mamba-3B 超越同等规模的 Transformer，在规模为其两倍时与之相当。

ABSTRACT

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$ imes$ 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.

研究动机与目标

• 通过引入具内容感知的选择性状态空间（SSMs）来解决 Transformer 在处理长序列时的低效问题。
• 在不使用注意力或 MLP 的情况下实现线性时间序列建模，同时在多模态下保持准确性。
• 开发面向硬件的选择性 SSM 并行算法，在循环模式下仍保持高效。

提出的方法

• 使 SSM 参数成为输入地址的函数，以实现沿序列的信息流离散、基于内容的调制。
• 将选择性 SSMs 集成到一个简单的一体化端到端骨干网络中，不使用注意力或 MLP（即 Mamba 模型）。
• 设计一个面向硬件的并行循环算法，在缺乏标准卷积优化的情况下仍能维持快速推理。
• 在推理阶段展示随着序列长度线性扩展，并实现相较于 Transformer 的 5× 吞吐提升。

实验结果

研究问题

• RQ1 基于内容的（地址相关的）选择性 SSM 是否能在处理长序列的推理方面优于传统的亚二次模型？
• RQ2 选择性 SSM 是否能实现线性时间序列建模，并在语言、音频和基因组学等领域达到具有竞争力的性能？
• RQ3 在无需注意力/MLP 的情况下，在循环式架构中实现硬件友好型并行并保持吞吐量是否可行？
• RQ4 在语言建模方面，Mamba 与同等规模及更大规模的 Transformers 相比表现如何？

主要发现

• Mamba 以 5× 高于 Transformer 的吞吐量实现快速推理。
• 该模型随序列长度线性扩展。
• Mamba 在语言、音频和基因组学模态上展现了最前沿的结果。
• 在语言建模方面，Mamba-3B 在同等规模上优于 Transformer，并在其两倍规模上与 Transformer 相当，同时在预训练和下游评估中均如此。