梦晨 发自 凹非寺
量子位 | 公众号 QbitAI
Transformer 挑战者、新架构 Mamba，刚刚更新了第二代：
Mamba-2
，状态空间扩大 8 倍，训练速度提高 50%！
更重要的是，团队研究发现原来
Transformer
和状态空间模型
（SSM）
竟然是近亲？？
两大主流序列建模架构，在此统一了。
没错，这篇论文的提出的重磅发现：
Transformer 中的注意力机制与 SSM 存在着非常紧密的数学联系。
团队通过提出一个叫
结构化状态空间二元性
（Structured State Space Duality，SSD）
的理论框架，把这两大模型家族统一了起来。
Mamba 一代论文年初被 ICLR 拒稿，当时还
让许多学者集体破防，引起一阵热议
。
这次二代论文在理论和实验上都更丰富了，成功入选 ICML 2024。
作者依然是
Albert Gu
和
Tri Dao
两位。
他们透露，论文题目中“Transformers are SSMs”是致敬了 4 年前的线性注意力经典论文“Transformers are RNNs”。
那么，SSM 和注意力机制究竟是怎么联系起来的，Mamba-2 模型层面又做出哪些改进？
Transformer 的核心组件是
注意力机制
，SSM 模型的核心则是一个
线性时变系统
。
两者看似不相关，但论文指出：它们都可以表示成
可半分离矩阵
（Semiseparable Matrices）
的变换。
先从 SSM 的视角来看。
SSM 本身就定义了一个线性映射，恰好对应了一个半可分离矩阵。
半可分离矩阵有着特殊的低秩结构，这种结构又恰好对应了 SSM 模型中的状态变量。
于是，矩阵乘法就相当于 SSM 的线性时变系统了。
带选择性的 SSM 本质上就是一种广义线性注意力机制。
从注意力的视角看又如何？
团队试图以更抽象方式来刻画注意力机制的本质，毕竟“Softmax 自注意力”只是众多可能形式中的一种。
更一般地，
任意带掩码的注意力机制，都可以表示为 4 个张量的缩并
(Contraction)
。
其中 QKV 对应注意力中的 query，key，value，L 对应掩码矩阵。
借助这一联系，它们在线性注意力的基础上提出了
结构化掩码注意力 SMA
（Structured Masked Attention）。
当注意力的掩码矩阵是半可分离的，就与 SSM 等价了。
基于这个发现，作者进一步推导出两种等价的计算形式，这就是本文核心思想”状态空间二元性” SSD 的由来。
基于 SSD 思想的新算法，Mamba-2 支持
更大的状态维度
（从 16 扩大到 256）
，从而学习更强的表示能力。
新方法基于块分解矩阵乘法，利用了 GPU 的存储层次结构，提高训练速度。
架构设计上，Mamba-2
简化了块的设计
，同时受注意力启发做出一些改动，借鉴多头注意力创建了多输入 SSM 。
有了与注意力之间的联系，SSD 还可以轻松
将 Transformer 架构多年来积累起来的优化方法引入 SSM
。
比如引入
张量并行和序列并行
，扩展到更大的模型和更长的序列。
又比如引入
可变序列长度
，以实现更快的微调和推理。
Mamba-2 的 SSD 层比 Mamba-1 中的关联扫描快很多，使团队能够增加状态维度并提高模型质量。
