再战Transformer！原作者带队的Mamba 2来了，新架构训练效率大幅提升

自 2017 年被提出以来，Transformer 已经成为 AI 大模型的主流架构，一直稳居语言建模方面 C 位。
但随着模型规模的扩展和需要处理的序列不断变长，Transformer 的局限性也逐渐凸显。

---

一个很明显的缺陷是：Transformer 模型中自注意力机制的计算量会随着上下文长度的增加呈平方级增长。
几个月前，Mamba 的出现打破了这一局面，它可以随上下文长度的增加实现线性扩展。随着 Mamba 的发布，这些状态空间模型 (SSM) 在中小型规模上已经实现了与 Transformers 匹敌，甚至超越 Transformers。
Mamba 的作者只有两位，一位是卡内基梅隆大学机器学习系助理教授 Albert Gu，另一位是 Together.AI 首席科学家、普林斯顿大学计算机科学助理教授 Tri Dao。
Mamba 面世之后的这段时间里，社区反应热烈。可惜的是，Mamba 的论文却惨遭 ICLR 拒稿，让一众研究者颇感意外。
仅仅六个月后，原作者带队，更强大的 Mamba 2 正式发布了。

• 论文地址：https://arxiv.org/pdf/2405.21060
• GitHub 地址：https://github.com/state-spaces/mamba
• 论文标题：Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

总体而言，本文提出了 SSD（state space duality）框架，基于此，研究者设计了一个新的体系架构 Mamba-2，其核心层是对 Mamba 的选择性 SSM 的改进，速度提高了 2-8 倍，同时在语言建模方面继续与 Transformers 竞争。
Tri Dao 表示，他们构建了一个丰富的 SSD 理论框架，许多线性注意力变体和 SSM 是等效的，由此产生的模型 Mamba-2 比 Mamba-1 更好、更快。

Mamba-2 的新算法使其能够利用更大的状态维度 (16 → 256)，同时训练速度更快。在需要更大状态容量的任务上，例如 MQAR 任务，它比 Mamba-1 有了显著的改进。

此外研究者还发现，最近新出的混合模型（Jamba、Zamba）增加了一些注意力层来提高模型质量。基于这些发现，研究者将 4-6 个注意力层与 Mamba-2 层混合，其表现优于 Transformer++ 和纯 Mamba-2，因而得出注意力和 SSM 是互补的。

这项研究的贡献概括为：

本文展示了状态空间模型与一类称为半可分矩阵的结构化矩阵族之间的等价性。这一联系是 Mamba-2 框架的核心，揭示了状态空间模型的新属性和算法。
本文显著改进了线性注意力理论，首先通过张量收缩的语言对其循环形式提供了一个明确的证明，然后将其推广到一种新的结构化掩码注意力（SMA）家族。
本文将 SSM（状态空间模型）和 SMA（结构化掩码注意力）联系起来，显示它们有一个很大的交集，彼此是对偶的，同时具有 SSM 式的线性形式和类似注意力的二次方形式。本文还证明了任何具有快速循环形式的核注意方法都是 SSM。
除了内在的理论价值外，研究者所提出的框架为理解和改进序列模型开辟了广阔的方向。
在算法层面。所提框架为计算 SSM 提供了新的高效且易于实现的算法。本文提出了一种基于半可分离矩阵块分解的 SSD 算法，该算法利用了 SSM 线性递推和二次对偶形式，在所有主要效率轴上获得了最优的权衡。基于 SSD 的实现比 Mamba 的优化选择性扫描实现快 2 到 8 倍，同时允许使用更大的循环状态大小（是 Mamba 的 8 倍甚至更高，且几乎不影响速度）。SSD 与优化过的 softmax 注意力实现（FlashAttention-2）具有高度竞争力，在序列长度 2k 时性能相当，在序列长度 16K 时速度快 6 倍。

架构设计。采用 SSM 等新架构的一个主要障碍是针对 Transformers 量身定制的生态系统，例如用于大规模训练的硬件高效优化和并行技术。本文框架允许使用已建立的惯例和技术来构建 SSM 的架构设计选择词汇表，并进一步改进它们。
本文还对 Mamba 块做了一些修改，这些修改允许实现张量并行，其主要思想包括引入分组值注意力 (GVA，grouped-value attention) 头结构。
将修改后的并行 Mamba 块与作为内部 SSM 层的 SSD 结合使用，形成了 Mamba-2 架构。研究者在与 Mamba 相同的设置中研究了 Mamba-2 的 Chinchilla 扩展法则，发现它在困惑度和实际运行时间方面均优于 Mamba 和 Transformer++。研究者还在 Pile 数据集上训练了一系列 Mamba-2 模型，结果显示 Mamba-2 在标准下游评估中匹配或超过 Mamba 和开源的 Transformers。例如，在 Pile 上训练了 3000 亿 token 的 2.7B 参数的 Mamba-2 在性能上超过了在同一数据集上训练的 2.8B 参数的 Mamba 和 Pythia 以及 6.9B 参数的 Pythia。

系统优化：SSD 框架连接 SSM 和 transformer，允许利用为 transformer 开发的丰富的系统优化工作。

SSD 层

Mamba-2 的核心贡献是新的 SSD（state space dual）层。SSD 层可以被定义为选择性 SSM 的特例。与 Mamba 相比，Mamba-2 的改动会略微降低表达能力，但却显著提高了训练效率，特别是允许在现代加速器上使用矩阵乘法单元。

SSD 层的对偶注意力：

除了最新的 SSD 层，研究者也对 Mamba 的神经网络架构做了一些小的改变，Mamba-2 架构如下所示。

Mamba-2 在网络架构上的主要变化是从顺序生成变为并行生成 SSM 参数，并且 Mamba-2 更适合张量并行等扩展方法。
通过提供状态空间模型的显式矩阵变换形式，研究团队揭示了理解和使用它们的新方法。从计算的角度来看，任何计算状态空间模型前向传播的方法都可以看作是半可分离矩阵上的矩阵乘法算法。半可分离矩阵视角为 SSD 提供了一个视角，其中双重模式分别指的是线性时间半可分离矩阵乘法算法和二次时间朴素矩阵乘法。

研究团队定义了结构化状态空间模型和结构化注意力，讨论了它们的属性，并表明它们都有二次算法和线性算法。

自最初的 Mamba 论文研究了合成任务 —— 如：合成复制和归纳 Head 以来，许多后续工作开始研究更难的关联回忆任务。由 Zoology 和 Based 系列工作引入的 MQAR（multi-query associative recall）任务已成为事实上的标准。

通过运行一个比文献中通常报告的版本要难得多的任务，该团队发现 Mamba-2 明显优于 Mamba-1，而改善性能的一个原因是状态大小（比 Mamba-1 大约 16 倍）。

在这篇文章中，作者深入探讨了模型背后的理论。

从两个完全不同的角度推导出 SSD 的「对偶性」：

• 一个从 SSM 的角度出发；
• 另一个从注意力机制的角度出发。

SSD 框架提供了状态空间模型、注意力机制和结构化矩阵之间丰富的联系。
虽然 SSD 模型可以被视为框架内每个分支的具体实例，但 SSD 框架本身更加通用，为未来的工作开辟了许多方向。

SSD 算法

通常，矩阵乘法（matmul）的 FLOPs 速度要比非矩阵乘法 FLOPs 快得多（高达 16 倍）：A100 GPU 具有 312 TFLOPS 的 BF16 矩阵乘法性能，但只有 19 TFLOPS 的 FP32 算术性能，而 H100 具有 989 TFLOPS 的 BF16 矩阵乘法性能，但只有 67 TFLOPS 的 FP32 算术性能。
Mamba-2 的主要目标之一是「利用张量核心加速 SSM」。

在绑定参数并引入 Head 结构后，Mamba-1 中的 SSM 变成了 SSD，这是一种更具限制性的形式，具有类似注意力的公式。并且由于 SSD 连接 SSM 和结

感谢阅读！如果您对AI的更多资讯感兴趣，可以查看更多AI文章：GPTNB。