Mamba2：基于 SSD 框架的 SSM 新架构，速度提升 8 倍
本文深入解读 Mamba 系列的最新研究论文《Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality》。该论文提出了一个名为结构化状态空间对偶（Structured State Space Duality, SSD）的新框架，通过结构化矩阵的视角，统一了状态空间模型（State Space Model, SSM）和注意力机制。基于该框架设计的 Mamba-2，其核心是对 Mamba 中 SSM 进行了改进，在与 Transformer 表现相当的同时，速度提高了 2-8 倍。
1. 基本信息
论文题目: Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
作者: Tri Dao, Albert Gu (按照姓氏首字母排序)
作者研究单位:
- Department of Computer Science, Princeton University
- Machine Learning Department, Carnegie Mellon University
论文代码: https://github.com/state-spaces/mamba
2. 研究背景
在深度学习领域，Transformer 的注意力机制虽然取得了巨大成功，但也存在明显的性能瓶颈。具体而言，Transformer 的注意力机制计算复杂度随序列长度呈二次关系增长（O(N^2)），这导致在处理长序列时训练时间显著增加。此外，自回归生成过程中需要维护与序列长度线性增长的缓存（cache），限制了推理效率。
相比之下，结构化状态空间模型（SSMs）这一系列的模型展现出不同的特性。SSM 的训练时间与序列长度呈线性关系（O(N)），且在生成时状态大小为常数，这使得它们在理论上更适合处理长序列任务。如 Mamba 等早期工作，在小规模语言建模任务上已达到甚至超过 Transformer 的性能。然而，SSM 的发展与 Transformer 相关的工作割裂比较严重，在理论理解和优化实现上都比较困难，缺乏统一的理论框架来连接两者。
本文的核心贡献在于在 SSM 和 self-attention 变体之间建立了理论联系，填补了这一空白。
3. 主要内容
论文篇幅较长，共 52 页，以下为核心内容的梳理与解读。
3.1 SSD 框架与理论成果
本文提出 structured state space duality (SSD) 框架，在结构化 SSM 和注意力变体之间建立了联系。作者提出了两种广义序列模型表示方法：矩阵变换和张量收缩，分别揭示了对偶性的不同视角。以下是三个主要的理论结果：
- 证明了 SSM 与半分离矩阵（Semiseparable Matrices）的等价性。这是整个框架的核心，也表明了 SSM 的新性质和算法。本文的中心观点是，不同的 SSM 计算方法可以重新理解为在结构化矩阵上的矩阵乘法算法的不同。
深度解析：半分离矩阵是一类特殊的结构化矩阵。一个下三角矩阵称为 k-semiseparable 的，如果其严格下三角部分的任意子矩阵的秩不超过 k。这里的 k 称为半分离矩阵的阶或秩。半分离矩阵的一个重要性质是，虽然完整矩阵有 N^2 个元素，但其 SSS 表示只需 O(kN) 的参数，且在这个表示上可以实现矩阵乘法等基本操作的近似线性时间算法。本文的一个核心结果（定理 3.5）就是证明了：state space model 变换与 semiseparable 矩阵乘法是等价的。因此，计算 state space model 的不同算法，本质上对应着 semiseparable 矩阵乘法的不同算法。这个联系既揭示了 state space model 的新性质，也启发了新的计算 state space model 的高效算法。
