再战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 正式颁布了。

总体而言，本文提出了 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 的「对偶性」：

SSD 框架提供了形态空间模型、留意力机制和结构化矩阵之间丰盛的咨询。

只管 SSD 模型可以被视为框架内每个分支的详细实例，但 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 和结构化矩阵，计算 SSM 的高效算法间接对应于「token-mixing」或「sequence-mixing」矩阵 M 的不同分解。

因此，可以经过寻觅代替的矩阵乘法方式，例如经过各种方式对其启动分解，从而创立计算 SSM 的新算法。

经过精心选用块大小，对这个矩阵启动便捷块分解，就可以集 SSD 线性递归和二次留意力对偶方式的两种长处于一身。

而这也就是 SSD 算法的来源，它有 4 个步骤，并且关于这个算法有两种齐全不同的诠释。

SSD 算法：分块矩阵分解

首先将半可分 SSM 矩阵划分为大小为 Q×Q 的块，而后，应用半分矩阵的性质来分解每个低秩的非对角块：

SSD 算法：分块和形态传递

该算法的另一种诠释触及「推理 SSM 如何在实践序列上启动操作」。

首先将输入序列宰割成大小为 Q 的块，步骤可以分为：

可以看到，大部分算法（步骤 1、2 和 4）应用了矩阵乘法（因此应用了张量外围），而且可以并行计算。

只要步骤 3 须要扫描，但它只操作一个十分短的序列，通常只要要很少期间。

系统及裁减优化

张量并行

经常使用张量并行对 Mamba-1 启动大规模训练的一项艰巨是，每层都须要 2 次 all-reduce，而在 Transformer 中，每个留意力或 MLP 层只要 1 次 all-reduce。这是由于 SSM 的一些参数是外部激活的函数，而不是层的输入函数。在 Mamba-2 中，由于驳回了「并行投影」结构，一切 SSM 参数都是层输入的函数，因此可以轻松地将张量并行运行于输入投影：将输入投影和输入投影矩阵宰割成 2、4、8 个碎片，详细取决于张量并行度。钻研者经常使用 grouped norm，分组数除以张量并行度，这样每个 GPU 都能独自成功归一化。这些变动造成每层只要 1 次 all-reduce，而不是 2 次。

序列并行

在对超长序列启动训练时，或者须要沿着序列长度启动宰割，并将不同部分调配给不同的设施。序列并行重要有两种方式：关于残差和归一化操作，用 reduce-scatter、残差 + 归一化、而后 all-gather，取代张量并行中的 all-reduce。由于 Mamba-2 经常使用与 Transformer 相反的残差和归一化结构，因此这种方式的序列并行无需修正即可间接运行。关于留意力或 SSM 操作，又称高低文并行（CP）。关于留意力，可以经常使用环形留意力沿序列维度启动宰割。关于 Mamba-2，SSD 框架再次提供了协助：经常使用相反的蒯分解，可以让每个 GPU 计算其本地输入和最终形态，而后在降级每个 GPU 的最终输入之前，在 GPU 之间传递形态（经常使用发送 / 接纳通讯原语）。

试验结果

该钻研在 MQAR 的一种具备应战性的版本上，经常使用更难的义务、更长的序列和更小的模型启动了对比试验。基线包括规范的多头 softmax 留意力以及 Based 架构，试验结果如图 8 所示。

下表显示了 Mamba-2 在一系列下游零样本评价义务上的功能：

感兴味的读者可以阅读论文原文，了解更多钻研内容。

原文链接:​ ​​ ​