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

本文主要是介绍五倍吞吐量，性能全面包围 Transformer：新架构 Mamba 引爆AI圈，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

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

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

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

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

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

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

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

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

Github 链接：https://github.com/state-spaces/mamba

论文链接：https://arxiv.org/ftp/arxiv/papers/2312/2312.00752.pdf

方法创新

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

选择机制

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

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

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

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

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

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