五倍吞吐量，性能全面包围Transformer：新架构Mamba引爆AI圈

文章来源：机器之心

屹立不倒的 Transformer 迎来了一个强劲竞争者。

图片来源：由无界 AI 生成

在别的领域，如果你想形容一个东西非常重要，你可能将其形容为「撑起了某领域的半壁江山」。但在 AI 大模型领域，Transformer 架构不能这么形容，因为它几乎撑起了「整个江山」。

自 2017 年被提出以来，Transformer 已经成为 AI 大模型的主流架构，但随着模型规模的扩展和需要处理的序列不断变长，Transformer 的局限性也逐渐凸显。一个很明显的缺陷是：Transformer 模型中自注意力机制的计算量会随着上下文长度的增加呈平方级增长，比如上下文增加 32 倍时，计算量可能会增长 1000 倍，计算效率非常低。

为了克服这些缺陷，研究者们开发出了很多注意力机制的高效变体，但这往往以牺牲其有效性特为代价。到目前为止，这些变体都还没有被证明能在不同领域发挥有效作用。

最近，一项名为「Mamba」的研究似乎打破了这一局面。

在这篇论文中，研究者提出了一种新的架构 ——「选择性状态空间模型（ selective state space model）」。它在多个方面改进了先前的工作。

作者表示，「Mamba」在语言建模方面可以媲美甚至击败 Transformer。而且，它可以随上下文长度的增加实现线性扩展，其性能在实际数据中可提高到百万 token 长度序列，并实现 5 倍的推理吞吐量提升。

消息一出，人们纷纷点赞，有人表示已经迫不及待想要把它用在大模型上了。

作为通用序列模型的骨干，Mamba 在语言、音频和基因组学等多种模态中都达到了 SOTA 性能。在语言建模方面，无论是预训练还是下游评估，他们的 Mamba-3B 模型都优于同等规模的 Transformer 模型，并能与两倍于其规模的 Transformer 模型相媲美。

这篇论文的作者只有两位，一位是卡内基梅隆大学机器学习系助理教授 Albert Gu，另一位是 Together.AI 首席科学家、普林斯顿大学计算机科学助理教授（即将上任）Tri Dao。

左：Albert Gu；右：Tri Dao。

Albert Gu 表示，这项研究的一个重要创新是引入了一个名为「选择性 SSM」的架构，该架构是 Albert Gu 此前主导研发的 S4 架构（Structured State Spaces for Sequence Modeling ，用于序列建模的结构化状态空间）的一个简单泛化，可以有选择地决定关注还是忽略传入的输入。一个「小小的改变」—— 让某些参数成为输入的函数，结果却非常有效。

值得一提的是，S4 是一个非常成功的架构。此前，它成功地对  Long Range Arena (LRA) 中的长程依赖进行了建模，并成为首个在 Path-X 上获得高于平均性能的模型。更具体地说，S4 是一类用于深度学习的序列模型，与 RNN、CNN 和经典的状态空间模型（State Space Model，SSM）广泛相关。SSM 是独立的序列转换，可被整合到端到端神经网络架构中（ SSM 架构有时也称 SSNN，它与 SSM 层的关系就像 CNN 与线性卷积层的关系一样）。Mamba 论文也讨论了一些著名的 SSM 架构，比如 Linear attention、H3、Hyena、RetNet、RWKV，其中许多也将作为论文研究的基线。Mamba 的成功让 Albert Gu 对 SSM 的未来充满了信心。

Tri Dao 则是 FlashAttention、Flash Attention v2、Flash-Decoding的作者。FlashAttention 是一种对注意力计算进行重新排序并利用经典技术（平铺、重新计算）加快速度并将内存使用从序列长度的二次减少到线性的算法。Flash Attention v2、Flash-Decoding 都是建立在 Flash Attention 基础上的后续工作，把大模型的长文本推理效率不断推向极限。在 Mamba 之前，Tri Dao 和 Albert Gu 也有过合作。