在AI大模型领域，若要评选一个“撑起半壁江山”的核心技术，Transformer架构无疑是最佳答案。自2017年问世以来，它已成为几乎所有主流大模型的基础架构。然而，随着模型规模持续扩大、处理序列长度不断增长，Transformer的短板也逐渐暴露——最突出的问题在于其自注意力机制的计算量随上下文长度呈平方级增长。简单估算：当上下文长度增加32倍时，计算量可能飙升近1000倍，效率瓶颈十分显著。 为突破这一困局，研究者们尝试了多种注意力机制的高效变体，但往往以牺牲有效性为代价。长期以来，这些变体未能真正在不同领域证明自身价值。直到近期，一项名为“Mamba”的研究似乎带来了转机。 该论文提出了一种全新架构——选择性状态空间模型（Selective State Space Model），在多个关键维度上实现了对先前工作的显著改进。作者直言：Mamba在语言建模方面不仅可与Transformer媲美，在某些场景下甚至能超越它。更关键的是，随着上下文长度增加，Mamba实现了线性扩展——在实际数据中，性能可延伸至百万token级别的序列，同时推理吞吐量提升5倍。 消息传出后，业界反响热烈，有研究者直言“已经等不及要将Mamba用于大模型了”。 作为通用序列模型的骨干，Mamba在语言、音频和基因组学等多种模态中均达到了SOTA性能。具体到语言建模，无论是在预训练还是下游评估中，Mamba-3B模型都显著优于同等规模的Transformer模型，甚至可与两倍于其规模的Transformer相媲美。 这篇论文仅有两位作者，却都是重量级人物：一位是卡内基梅隆大学机器学习系的助理教授Albert Gu，另一位是Together.AI首席科学家、即将就任普林斯顿大学计算机科学助理教授的Tri Dao。Albert Gu表示，研究的核心创新在于引入了“选择性SSM”架构——这是此前他主导研发的S4架构的一个简单泛化，使模型能有选择性地关注或忽略传入的输入。这一看似“小小的改变”——让某些参数成为输入的函数——却带来了惊人效果。 值得一提的是，S4本身已是一个非常成功的架构。它曾在Long Range Arena（LRA）的长程依赖建模任务中表现出色，并且是首个在Path-X上获得高于平均性能的模型。S4是一类用于深度学习的序列模型，与RNN、CNN和经典状态空间模型（SSM）广泛关联。Mamba论文还讨论了一些著名的SSM架构，如Linear attention、H3、Hyena、RetNet、RWKV——这些也被用作论文研究的基线。Mamba的成功让Albert Gu对SSM的未来充满信心。 另一位作者Tri Dao，是FlashAttention、Flash Attention v2、Flash-Decoding的创造者。FlashAttention通过重新排序注意力计算并利用平铺、重新计算等经典技术，将内存使用从序列长度的二次方降至线性。在Mamba之前，Tri Dao与Albert Gu已有合作，此次可谓强强联手。 这项研究的模型代码和预训练检查点已开源，相关链接可在论文和项目主页中获取。 论文链接：https://arxiv.org/ftp/arxiv/papers/2312/2312.00752.pdf 项目主页：https://github.com/state-spaces/mamba ### 方法创新 论文第3.1节介绍了如何利用合成任务的直觉启发选择机制，第3.2节解释了如何将这一机制融入状态空间模型。由此产生的时变SSM无法直接使用卷积，从而引发了高效计算的技术难题——研究者采用了一种硬件感知算法，利用当前硬件的内存层次结构来克服这一挑战（第3.3节）。第3.4节描述了一个简洁的SSM架构，无需注意力，甚至无需MLP块。第3.5节则探讨了选择机制的若干其他特性。 ### 选择机制 研究者发现此前模型的一个关键局限：缺乏以依赖输入的方式高效选择数据的能力——即能有选择地关注或忽略特定输入。序列建模的基本方法是将上下文压缩到更小的状态，我们可以从这个角度审视当下流行的序列模型。注意力机制既高效又低效，因为它根本没有明确压缩上下文。从自回归推理需要显式存储整个上下文（即KV缓存）就能看出这一点，这直接导致了Transformer缓慢的线性时间推理和二次时间训练。 递归模型的效率很高，因为它们的状态有限，意味着恒定时间推理和线性时间训练。但它们的高效性受限于这种状态对上下文的压缩程度。为了直观理解这一原理，下图展示了两个合成任务的运行示例： 研究者设计了一种简单的选择机制：根据输入对SSM参数进行参数化。这样一来，模型就能过滤掉无关信息，并无限期记住相关信息。实现选择机制的核心方法，是让影响序列交互的参数（如RNN的递归动力学或CNN的卷积核）变得与输入相关。算法1和2展示了本文使用的主要选择机制。关键差异在于：只需将几个参数（∆，B，C）设置为输入的函数，并在过程中改变张量形状。这些参数现在都具有长度维度L，意味着模型已从时间不变变为时间可变。 ### 硬件感知算法 上述变化对模型计算提出了新的技术挑战。所有先前的SSM模型都必须保持时间和输入不变，才能保证计算效率。为此，研究者采用了一种硬件感知算法——通过扫描而非卷积来计算模型，但又不将扩展状态具体化，从而避免在GPU存储器层次结构的不同级别之间频繁进行IO访问。最终实现的方法在理论上（与所有基于卷积的SSM的伪线性相比，在序列长度上呈线性缩放）和实际硬件上（在A100 GPU上可快达3倍）都比之前的方法更快。 ### 架构 研究者将先前的SSM架构设计与Transformer的MLP块合并为一个块，从而简化了深度序列模型架构，形成了一种包含选择性状态空间的简单、同质的架构设计（Mamba）。与结构化SSM一样，选择性SSM也是一种独立的序列变换，可灵活融入神经网络。H3架构是著名同质化架构设计的基础，通常由线性注意力启发的块和MLP（多层感知器）块交错组成。研究者的简化思路是：将这两个部分合二为一，均匀堆叠，如图3所示，这受到了门控注意力单元（GAU）的启发。 选择性SSM以及Mamba架构的扩展是完全递归模型，几个关键特性使其非常适合作为在序列上运行的通用基础模型的骨干： - **高质量**：选择性为语言和基因组学等密集模型带来了强大的性能。 - **快速训练与推理**：训练过程中，计算量和内存与序列长度成线性关系；推理过程中，由于无需缓存以前的元素，自回归展开模型每一步只需恒定时间。 - **长上下文**：质量与效率共同提升了实际数据的性能，序列长度可达100万。 ### 实验评估 实证结果验证了Mamba作为通用序列基础模型骨干的潜力——无论是在预训练质量还是特定领域的任务性能方面，Mamba都能在多种模态和环境中展现出色表现。 **合成任务**：在复制和感应头等重要的语言模型合成任务上，Mamba不仅能轻松解决，还能推断出超过100万token的无限长解决方案。 **音频和基因组学**：在音频波形和DNA序列建模方面，Mamba在预训练质量和下游指标上均优于SaShiMi、Hyena、Transformer等先前的SOTA模型。例如，在具有挑战性的语音生成数据集上，它成功将FID降低了一半以上。在这两种场景下，其性能随着上下文长度的增加而提升，最高可达百万长度的序列。 **语言建模**：Mamba是首个线性时间序列模型，在预训练复杂度和下游评估方面均真正达到了Transformer级别的性能。通过多达1B参数的缩放规律研究，研究者发现Mamba的性能超过了大量基线模型，包括LLaMa这种非常强大的现代Transformer训练配方。 与类似规模的Transformer相比，Mamba实现了5倍的生成吞吐量。此外，Mamba-3B的质量与两倍于其规模的Transformer相当——例如，与Pythia-3B相比，在常识推理上的平均值高出4分，甚至超过了Pythia-7B。