arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2312.00752cs.LGcs.AI

Mamba：基于选择性状态空间的线性时间序列建模

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Carnegie Mellon University（卡内基梅隆大学）；Princeton University（普林斯顿大学）

Albert Gu, Tri Dao

阅读 PDF
arXiv 页面 ↗HTML 原文 ↗

AI总结：

针对 Transformer 在长序列上的计算低效问题，提出基于选择性状态空间的 Mamba 架构，通过输入依赖的参数实现内容推理，在多模态任务上取得最优，且推理更快、序列扩展呈线性。

摘要：

基础模型现在驱动着深度学习中大多数令人兴奋的应用，它们几乎普遍基于 Transformer 架构及其核心注意力模块。为了解决 Transformer 在长序列上的计算低效问题，人们开发了许多亚二次时间架构，如线性注意力、门控卷积和循环模型，以及结构化状态空间模型（SSMs），但它们在语言等重要模态上的表现不如注意力。我们发现此类模型的一个关键弱点是无法执行基于内容的推理，并进行了几项改进。首先，简单地让 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.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。

↑