Mamba-2新架构出世一统江湖！普林斯顿CMU华人再出神作，性能狂飙8倍

【新智元导读】在开源社区引起「海啸」的Mamba架构，再次卷土重来!这次，Mamba-2顺利拿下ICML。通过统一SSM和注意力机制，Transformer和SSM直接成了「一家亲」，Mamba-2这是要一统江湖了?

年前，Mamba被顶会ICLR拒稿的消息曾引起轩然大波。

甚至有研究人员表示:如果这种工作都被拒了，那我们这些「小丑」要怎么办?

这次，新一代的Mamba-2卷土重来、再战顶会，顺利拿下了ICML2024!

仍是前作的两位大佬（换了个顺序），仍是熟悉的配方:

论文地址:https://arxiv.org/pdf/2405.21060

开源代码和模型权重:https://github.com/state-spaces/mamba

不同的是，作者在更高的视角上，统一了状态空间模型（SSM）和注意力机制(Attention)，也就是文章标题所说的「Transformers are SSMs」。

——这下咱们都是一家人了，不用动不动就「打生打死」了。

性能方面，Mamba-2采用了新的算法（SSD），比前代提速2-8倍，对比FlashAttention-2也不遑多让，在序列长度为2K时持平，之后便一路遥遥领先。

在Pile上使用300B token训练出的Mamba-2-2.7B，性能优于在同一数据集上训练的Mamba-2.8B、Pythia-2.8B，甚至是更大的Pythia-6.9B。

从理论上整合了SSM和Transformer，同等性能下，模型更小，消耗更低，速度更快。

更重要的是，能够利用GPU的硬件资源（矩阵乘法单元），以及针对Transformer的一系列优化。

——Mamba-2大有一统江湖之势。

1代Mamba，爆发式占领AI社区

事实上，关于1代Mamba的各种研究一直在爆发性地增长，arxiv已经被各种Mamba所占领，谷歌学术的引用量也达到了350多。

后续工作如雨后春笋一般冒出，包括视觉、基因组学、图表等的直接应用，以及回忆能力、上下文学习能力、形式语言表达能力等方面的研究。

作者兴奋地表示:「我们多年来一直在追求的高效序列模型研究路线，真正引起了机器学习社区的共鸣。」

唯一遗憾的是，Mamba遭到ICLR拒稿，所以关于Mamba到底有没有前途这个事也就被打上了问号。

现在，问题解决了，不但论文被接收了，而且还证明了Transformer和Mamba其实是一家人——

「你说我不行?那Transformer到底行不行?」

值得注意的是，之前很火的Vision Mamba以及另一篇关于Mamba的研究也杀入了ICML2024。

对于改进Mamba的初衷，作者表示，当前AI社区的大家都在努力解决Transformer的问题，尽管SSM的特性和效果都相当好，但却跟社区的努力方向不一致。

这次的Mamba-2可以把针对Transformer的优化都用上，不浪费大家的努力。

新架构一统江湖

在介绍新架构之前，小编先帮大家简单理一下背景。

状态空间模型SSM之所以如此令人着迷，是因为它们显得如此之「基础」。

比如，它们与序列模型的许多主要范式，都有着丰富的联系。

它们似乎抓住了连续、卷积和循环序列模型的本质，把所有这些元素都包含在了一个简单优雅的模型里。

不过，另一个主要的序列模型范式——注意力机制的变体，却更加无所不在。

然而SSM却总感觉和Attention是脱节的。

在这里，研究者们发出了「灵魂拷问」——SSM和注意力之间的概念联系是什么?有无可能将二者结合起来?

那就要从公式说起了。

状态空间模型SSM可以这么定义:

这是个微分方程，利用导数定义进行代换:

可以得到SSM的解:

这个东西就跟RNN一毛一样了:

所以可以认为SSM等价于RNN。

如果将RNN的递归结构展开，那么它又可以等价于卷积:

此时，便可以利用卷积的特性进行并行训练，而进行推理时又可以享受RNN带来的O（1）复杂度。

当然，好事不能让你全占了，这种结构仍然逃不过固有的梯度爆炸（或消失），以及难以胜任选择性复制和上下文学习等任务。

为此，Mamba在SSM的基础上加入了能够随输入变化的参数。

不过这样做的代价是失去了固定kernel带来的并行性，所以作者另辟蹊径，使用前缀和的方式来加速RNN的训练。

不过，从计算角度来看，Mamba在硬件效率上仍然远不如注意力机制。

原因在于，目前常用的GPU、TPU等加速器，是为矩阵乘法进行过专门优化的。

1代Mamba吃不到硬件矩阵运算单元的红利，尽管推理时有速度优势，但训练时问题就大了。

所以作者就想，我能不能把Mamba的计算重构成矩阵乘法呢?

于是，新一代的Mamba诞生了。

结构化状态空间对偶性:SSD

Mamba-2的核心，是结构化状态空间对偶性（State Space Duality，SSD）的概念:

1.SSD模型指的是一个特定的独立层，比如注意力层或状态空间模型（SSM），可以被整合到深度神经网络中;

2.SSD框架是一个用于推理该模型（以及更多理论连接）的通用框架;

3.SSD算法是一种比以前的SSM更高效地计算SSD层的算法。

SSD框架（红色，蓝色）:状态空间模型(即半分离矩阵)和结构化掩码注意力涵盖了大量高效的序列模型。它们的交集就是SSD模型(紫色)

原始的Mamba（或更准确地说，其核心「S6」层）实际上是一个具有对角结构的选择性状态空间模型(SSM)。

Mamba-2的SSD层只做了一个小改动:它进一步限制了对角矩阵