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 的高效算法。










