再战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。
…
…
…

---

---

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