Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Albert Gu Alphabetical by first name. Tri Dao††
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 $\ti
中文速览
传统的结构化状态空间模型（SSM）虽然能以线性复杂度处理长序列，却因参数固定、无法根据输入内容动态筛选信息，在语言等离散模态上表现远不如Transformer。Mamba的核心思路是把SSM的关键参数（B、C、Δ）改为依赖当前输入动态生成，让模型能选择性地保留或遗忘序列中的信息——这一机制被称为"选择性状态空间"（Selective SSM）。由于这种输入相关的参数化打破了原有卷积并行计算的前提，作者专门设计了一种感知硬件内存层级的并行扫描算法，在不显式展开隐状态的情况下高效完成训练，同时推理时以纯递归方式运行，速度是同规模Transformer的5倍。实验表明，基于此构建的Mamba-3B语言模型在预训练困惑度和下游常识推理任务上均优于同等参数量的Transformer，甚至能与两倍参数量的模型持平，并在音频、基因组等多个领域刷新了最优结果——这是首个在语言建模上真正追平Transformer质量的线性时间序列模型，对构建高效长上下文基础模型具有重要意义。
原文 arXiv:2312.00752;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2312.00752v2