Transformers 是 SSM：通过结构化状态空间对偶的通用模型与高效算法

Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

特里·道 Tri Dao · · 2024-05-31 · arXiv:2405.21060 ↗

打开互动全文版（逐段中英对照 + 图/公式 + 论文问答）→

摘要 · Abstract

虽然 Transformer 一直是深度学习在语言建模上取得成功的主要架构，但最近状态空间模型（如 Mamba）在小到中等规模上已被证明可以匹敌或超越 Transformer。我们表明这两类模型实际上密切相关，并通过一系列结构化的半可分矩阵的分解，发展了一个丰富的理论框架来连接 SSM 和多种注意力变体。我们的状态空间对偶（SSD）框架使我们能够设计一种新架构（Mamba-2），其核心层是对 Mamba 选择性 SSM 的改进，速度提高 2-8 倍，同时在语言建模上继续与 Transformer 竞争。

While Transformers have been the main architecture behind deep learning's success in language modeling, state-space models (SSMs) such as Mamba have recently been shown to match or outperform Transformers at small to medium scale. We show that these families of models are actually quite closely related, and develop a rich framework of theoretical connections between SSMs and variants of attention, connected through various decompositions of a well-studied class of structured semiseparable matrices. Our state space duality (SSD) framework allows us to design a new architecture (Mamba-2) whose core layer is an a refinement of Mamba's selective SSM that is 2-8X faster, while continuing to be competitive with Transformers on language modeling.

核心贡献 · Key contributions

• 建立 SSM 与半可分矩阵变换之间的等价性，揭示新的理论联系。
  Establish equivalence between SSMs and semiseparable matrix transformations, revealing new theoretical connections.
• 提出状态空间对偶（SSD）框架，通过矩阵分解统一 SSM、注意力与结构化矩阵。
  Develop state space duality framework unifying SSMs, attention, and structured matrices through matrix decompositions.
• 推出 Mamba-2，其核心 SSD 层比 Mamba 快 2-8 倍，同时与 Transformer 保持竞争力。
  Introduce Mamba-2 whose core SSD layer is 2-8× faster than Mamba while competitive with Transformers.
• 提出结构化掩码注意力（SMA），将线性注意力推广到任意结构化掩码。
  Propose structured masked attention generalizing linear attention to arbitrary structured masks.
• 证明高效自回归注意力需要半可分掩码，刻画了有界阶自回归注意力的特征。
  Prove efficient autoregressive attention requires semiseparable masks, characterizing autoregressive attention with bounded order.
• 开发硬件高效的分块分解算法，训练 FLOPs 为 O(TN²)，以矩阵乘法为主导。
  Develop hardware-efficient block-decomposition algorithm with O(TN²) training FLOPs dominated by matrix multiplications.

局限 · Limitations

• A 矩阵的标量乘单位阵简化相比通用选择性 SSM 降低了表达能力。
  Scalar-times-identity simplification of A matrix reduces expressive power versus general selective SSMs.
• 对偶形式去掉了 softmax，偏离标准注意力；表达性影响尚未充分分析。
  Softmax is dropped in dual form, deviating from standard attention; expressivity implications remain unanalyzed.
• 对偶性仅在标量 SSM 与 1-半可分 SMA 情形建立，不覆盖通用 SSM 族。
  Duality established for scalar-identity SSM and 1-semiseparable SMA, not for general SSM families.
• 实证验证限于中小规模，大规模下的 Scaling 行为尚未证实。
  Empirical validation limited to small-to-medium scale; scaling at large scale not demonstrated.
• 稠密无结构 A 矩阵仍需昂贵预处理，限制了通用结果的实际适用性。
  Dense unstructured A matrices still incur expensive preprocessing, limiting practical applicability of general results.

论文章节 · Sections（共 36）

• 摘要 Abstract
• 2.1 Structured State Space Models 2.1 Structured State Space Models
• 2.2 注意力机制 2.2 Attention
• 2.3 结构化矩阵 2.3 Structured Matrices
• 2.4 概述：结构化状态空间对偶性 2.4 Overview: Structured State Space Duality
• 2.5 Notation 2.5 Notation
• 3 状态空间模型是结构化矩阵 3 State Space Models are Structured Matrices
• 3.2 Semiseparable Matrices 3.2 Semiseparable Matrices
• 3.3 State Space Models are Semiseparable Matrices 3.3 State Space Models are Semiseparable Matrices
• 4.1 注意力机制框架 4.1 The Attention Framework
• 2 Normalizing the 𝐺 matrix on the S axis. 2 Normalizing the 𝐺 matrix on the S axis.
• 4.2 线性注意力 4.2 Linear Attention
• 4.3 结构化掩码注意力 4.3 Structured Masked Attention
• 5 状态空间对偶性 5 State Space Duality
• 5.1 Scalar-Identity Structured State Space Models 5.1 Scalar-Identity Structured State Space Models
• 5.3 Structured State-Space Duality (SSD) 5.3 Structured State-Space Duality (SSD)
• 6.1 对角块 6.1 Diagonal Blocks
• 6.2 Low-Rank Blocks 6.2 Low-Rank Blocks
• 6.3 计算成本 6.3 Computational Cost
• 7 The Mamba-2 Architecture 7 The Mamba-2 Architecture
• 7.1 Block Design 7.1 Block Design
• 7.3 线性注意力启发的其他 SSD 扩展 7.3 Other SSD Extensions from Linear Attention
• 8 SSM 的系统优化 8 Systems Optimization for SSMs
• 8.1 Tensor Parallel 8.1 Tensor Parallel
• 8.3 可变长度 8.3 Variable Length
• 9 实证验证 9 Empirical Validation
• 9.1 Synthetics: Associative Recall 9.1 Synthetics: Associative Recall
• 9.2 语言建模 9.2 Language Modeling
• 2 Mamba-2：64 个 SSD 层 2 Mamba-2: 64 SSD layers.
• 9.3 速度基准 9.3 Speed Benchmarks
• 9.4 架构消融 9.4 Architecture Ablations
• 相关工作与讨论 10 Related Work and Discussion
• 10.1 状态空间模型 10.1 State Space Models
• 10.4 相关模型 10.4 Related Models
• 11 Conclusion 11 Conclusion
• 注 Note

阅读逐段中英对照全文 →