《Transformer 就是 SSM：通过结构化状态空间对偶得到的广义模型与高效算法》

要解决的问题

建立在注意力（attention）之上的模型，在效率上有两个随序列长度增长的问题。训练时，注意力对序列长度呈二次增长，因为每个位置都要和其他所有位置比较；自回归生成时，模型要保留一份过去键值（key 与 value）的缓存，缓存随序列长度线性增长 ^[1]。结构化状态空间模型（structured state space model，SSM）走的是另一条路：它一步一步读入序列，把读过的内容压缩进一个固定大小的状态里，于是训练开销随序列长度线性增长，生成时也只需要一个大小恒定的状态 ^[1]。一段时间里，这两类模型被当作各有长短的两条路线来介绍，这篇论文要问的是：这种对立到底有多真实。

对偶框架

论文提出结构化状态空间对偶（state space duality，SSD）框架，通过一类被充分研究过的结构化半可分矩阵（structured semiseparable matrix）的分解，把结构化 SSM 与各种注意力变体联系起来 ^[1]。所谓半可分矩阵，粗略地说，是那种非对角块秩很低的矩阵：它内部结构很强，存储和乘一个向量的代价都远低于一般稠密矩阵。框架的出发点在于，两族模型的计算都可以写成乘上这类矩阵——SSM 用一步步的递推去乘，注意力式的层则用矩阵形式或分块形式去乘。这样看，两者不是互相竞争的两种机制，而是同一个对象的两种组织方式：一种在生成时便宜，另一种便于在现代硬件上并行训练。论文的核心主张正是：这两个模型族其实关系密切，可以在同一个框架里描述 ^[1]。

Mamba-2

这个框架不只是描述性的，它还给出一个新架构 Mamba-2，其核心层是对 Mamba 的选择性 SSM（selective SSM）的改进 ^[1]。这一核心层比 Mamba 的核心层快 2 到 8 倍，同时在语言建模上与 Transformer 相比仍然有竞争力 ^[1]。核心层是对既有设计的改进而不是替换，因此 Mamba-2 处在从 Mamba 延伸出来的那条线上，而不是站在注意力的对立面。

标题的来历

论文标题在向 Katharopoulos 等人 2020 年的论文《Transformers are RNNs》致敬，那篇论文最早指出了自回归注意力与一种特定线性递推之间的等价关系 ^[1]。两个标题说的都是「等同」，不是「相似」：重点不是一类模型看起来像另一类，而是在合适的描述下，它们做的是同一件事。把「RNN」换成「SSM」，就是把这一结论从循环网络推到了状态空间模型。

作者与发表

论文作者是普林斯顿大学计算机科学系的 Tri Dao 和卡内基梅隆大学机器学习系的 Albert Gu，2024 年 5 月 31 日提交到 arXiv，编号 2405.21060 ^[1]。论文发表在 ICML 2024 ^[1]。

为什么重要

从一个角度看，这篇论文讲的是速度：让状态空间层快到可以和 Transformer 用同样的硬件、同样的数据去训练。从另一个角度看，它改的是领域的地图。注意力和状态空间层此前常被描述成两个分开的家族，关于孰优孰劣的讨论，往往变成关于自己属于哪个家族的讨论。当注意力变体在这个框架里成为某类结构化矩阵计算的特例、SSM 层是同一套计算的另一种排布时，比较的含义就变了：人们可以问，计算的哪种排布更适合训练的哪个阶段、生成的哪个阶段，而不必先问哪个家族会赢。Mamba-2 核心层快 2 到 8 倍，就是这种描述方式变化带来的具体结果 ^[1]。