MAMBA介绍：一种新的可能超过Transformer的AI架构

03-06 1211阅读

有人说，“理解了人类的语言，就理解了世界”。一直以来，人工智能领域的学者和工程师们都试图让机器学习人类的语言和说话方式，但进展始终不大。因为人类的语言太复杂，太多样，而组成它背后的机制，往往又充满着不可名状的规律。

过去人们在自然语言处理中多采用 RNN 循环神经网络，它十分类似于人类逻辑上对语言的理解，即：强调上下文顺序、前后文逻辑关系。但是这种顺序方式让 RNN 无法实现并行计算，也就是说，它的速度十分缓慢，而规模也很难扩大。

直到 2017 年 6 月 12 日，一篇名为“Attention is All You Need”的论文被提交到预印论文平台 arXiv 上。一切从此改变。Transformer 的提出直接导致了现在的生成式 AI 风暴。机器好像在一瞬间就学会了如何与人类自如交流。Transformer点石成金的魔力，主要在于它彻底抛弃了前面提到的 RNN 循环神经网络这套逻辑，它完全由自注意力机制组成。大家都有过这样的经验，打乱一个句子中字词序顺，很多时候不并响影你对句子的解理。这是因为人脑在处理信息时会区分权重，也就是说，我们的注意力总是被最重要的东西吸引走，次要的细节则被忽略。Transformer 正是模仿了这一点，使它能够自动学习输入的序列中不同位置之间的依赖关系，并计算其相关性（而不是对整个输入进行编码）。这让针对序列的建模变得更加容易和精准。

尽管如此，随着模型规模的扩展和需要处理的序列不断变长，Transformer 的局限性也逐渐凸显。一个很明显的缺陷是：Transformer 模型中自注意力机制的计算量会随着上下文长度的增加呈平方级增长，比如上下文增加 32 倍时，计算量可能会增长 1000 倍，计算效率非常低。现在如日中天的ChatGPT大模型就有一大痛点：处理长文本算力消耗巨大。背后原因其实就是Transformer架构中注意力机制的二次复杂度。为了克服这些缺陷，研究者们开发出了很多注意力机制的高效变体，但这往往以牺牲其有效性特为代价。到目前为止，这些变体都还没有被证明能在不同领域发挥有效作用。

最近，卡内基梅隆大学机器学习系助理教授 Albert Gu和普林斯顿大学计算机科学系即将上任的助理教授Tri Dao，联合提出一项名为「MAMBA」的研究似乎打破了这一局面。这篇论文的预印本本月初分布在arXiv网站上：

图一: MAMBA预印本论文截图

论文的第一作者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 也有过合作。

另外，这项研究的模型代码和预训练的检查点是开源的，参见以下链接：https://github.com/state-spaces/mamba.

图二：MAMBA下载页面截图

方法创新

MAMBA论文的第 3.1 节介绍了如何利用合成任务的直觉来启发选择机制，第 3.2 节解释了如何将这一机制纳入状态空间模型。由此产生的时变 SSM (State Space Mode) 不能使用卷积，导致了高效计算的技术难题。研究者采用了一种硬件感知算法，利用当前硬件的内存层次结构来克服这一难题（第 3.3 节）。第 3.4 节描述了一个简单的 SSM 架构，不需要注意力，甚至不需要 MLP 块。第 3.5 节讨论了选择机制的一些其他特性。

(1) 选择机制

本文的研究者首先发现了此前模型的一个关键局限：以依赖输入的方式高效选择数据的能力（即关注或忽略特定输入）。

序列建模的一个基本方法是将上下文压缩到更小的状态，我们可以从这个角度来看待当下流行的序列模型。例如，注意力既高效又低效，因为它根本没有明确压缩上下文。这一点可以从自回归推理需要明确存储整个上下文（即 KV 缓存）这一事实中看出，这直接导致了 Transformer 缓慢的线性时间推理和二次时间训练。

递归模型的效率很高，因为它们的状态是有限的，这意味着恒定时间推理和线性时间训练。然而，它们的高效性受限于这种状态对上下文的压缩程度。

为了理解这一原理，图三展示了两个合成任务的运行示例：

图三: MAMBA中两个合成任务的运行示例。(左) 复制任务的标准版本涉及输入和输出元素之间的恒定间距，很容易通过线性递归和全局卷积等时不变模型来解决。(右上) 选择性复制任务在输入之间具有随机间隔，并且在需要时改变模型，该模型可以根据输入的内容选择性地记住或忽略输入。(右下) 联想回忆的一个例子，需要根据上下文检索答案，这是大语言模型的一项关键能力。

本文作者设计了一种简单的选择机制，根据输入对 SSM 参数进行参数化。这样，模型就能过滤掉无关信息，并无限期地记住相关信息。

例如将选择机制纳入模型的一种方法就是让影响序列交互的参数（如 RNN 的递归动力学或 CNN 的卷积核）与输入相关。下图中算法 1 和 2 展示了本文使用的主要选择机制。其主要区别在于，该方法只需将几个参数 ∆，B，C 设置为输入函数，并在整个过程中改变张量形状。这些参数现在都有一个长度维度 L ，意味着模型已经从时间不变变为时间可变。

图四：SSM和SSM+selection两种算法的流程图。

(2) 硬件感知算法

上述变化对模型的计算提出了技术挑战。所有先前的 SSM 模型都必须是时间和输入不变的，这样才能提高计算效率。为此，本文作者采用了一种硬件感知算法，通过扫描而不是卷积来计算模型，但不会将扩展状态具体化，以避免在 GPU 存储器层次结构的不同级别之间进行 IO 访问。由此产生的实现方法在理论上（与所有基于卷积的 SSM 的伪线性相比，在序列长度上呈线性缩放）和现有硬件上都比以前的方法更快（在 A100 GPU 上可快达 3 倍）。

图五: 结构化的SSM通过更高维度的潜在状态h（例如，二进制操作= 4）独立地映射输入变量的每个通道（例如𝐷= 5）到输出变量的每个通道（例如，二进制操作= 4）。先前的SSM通过需要时不变性的替代计算路径来避免实现这个大的有效状态（𝐷二进制操作，倍批大小的变量和序列长度的变量𝐿）:（∆，a, B, C）参数在时间上是恒定的。我们的选择机制增加了依赖输入的动态，这也需要一个谨慎的硬件感知算法，只在更有效的GPU内存层次结构中实现扩展状态。

(3) 算法架构

MAMBA将先前的 SSM 架构设计与 Transformer 的 MLP 块合并为一个块，从而简化了深度序列模型架构，形成了一种包含选择性状态空间的简单、同质的架构设计（MAMBA）。

与结构化 SSM 一样，选择性 SSM 也是一种独立的序列变换，可以灵活地融入神经网络。H3 架构是著名的同质化架构设计的基础，通常由线性注意力启发的块和 MLP（多层感知器）块交错组成。

如图六所示，本文作者简化了这一架构，将这两个部分合二为一，均匀堆叠。他们受到门控注意力单元（GAU）的启发，该单元也对注意力做了类似的处理。

总而言之，选择性 SSM 以及 Mamba 架构的扩展是完全递归模型，几个关键特性使其适合作为在序列上运行的通用基础模型的骨干：

高质量：选择性为语言和基因组学等密集模型带来了强大的性能。
快速训练和推理：在训练过程中，计算量和内存与序列长度成线性关系，而在推理过程中，由于不需要缓存以前的元素，自回归展开模型每一步只需要恒定的时间。
长上下文：质量和效率共同提高了实际数据的性能，序列长度可达 100 万。

图六: 简化块设计结合了H3块与MLP块。与H3块相比，MAMBA用激活函数取代了第一个乘法门。与MLP块相比，MAMBA在主分支中添加了一个SSM。

实验评估

MAMBA论文对该架构进行了多方位测试。这些实证验证了 MAMBA 作为通用序列基础模型骨干的潜力。无论是在预训练质量还是特定领域的任务性能方面，MAMBA 都能在多种类型的模态和环境中发挥作用。

(1) 合成任务

在复制和感应头等重要的语言模型合成任务上，MAMBA 不仅能轻松解决，而且能推断出无限长的解决方案（>100 万 token）。

(2) 音频和基因组学

在音频波形和 DNA 序列建模方面，Mamba 在预训练质量和下游指标方面都优于 SaShiMi、Hyena、Transformer 等先前的 SOTA 模型（例如，在具有挑战性的语音生成数据集上将 FID 降低了一半以上）。在这两种情况下，它的性能随着上下文长度的增加而提高，最高可达百万长度的序列。

(3) 语言建模

Mamba 是首个线性时间序列模型，在预训练复杂度和下游评估方面都真正达到了 Transformer 质量的性能。通过多达 1B 参数的缩放规律，研究者发现 Mamba 的性能超过了大量基线模型，包括 LLaMa 这种非常强大的现代 Transformer 训练配方。

（4）速度和显存基准测试

下图展示了scan操作（状态扩展N = 16）速度，以及Mamba端到端推理吞吐量的基准测试。

测试结果显示，当序列长度超过2k时，高效的SSM scan比目前最优秀的注意力机制——FlashAttention-2还要快。而且，比起PyTorch标准的scan实现，速度提升更是高达20到40倍。由于没有键值（KV）缓存，因此Mamba可以支持更大的批处理大小，从而使推理吞吐量比同等规模Transformer高了4到5倍。

举个例子，一个未经训练的69亿参数的Mamba（Mamba-6.9B），在推理处理能力上可以超过仅有13亿参数、规模小5倍的Transformer模型。

与大多数深度序列模型一样，显存使用量与激活张量的大小成正比。表15显示，Mamba的显存需求与经过优化的Transformer相当。

表一: Mamba的内存占用可与最优化的Transformer相媲美（125M模型的结果）。