[论文解读] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

Tri Dao, Albert G. Gu|arXiv (Cornell University)|May 31, 2024

Neural Networks and Applications被引用 66

一句话总结

该论文在结构化状态空间模型（SSMs）与注意力之间建立正式联系，引入结构化状态空间对偶性（SSD），以推导出高效的基于SSM的体系结构（Mamba-2），在语言建模方面能与 Transformers 相媲美，同时具备可扩展性优势。

ABSTRACT

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.

研究动机与目标

• 激发将 SSMs 与 Transformers 桥接的动机，以改善长序列的可扩展性和效率。
• 开发一个理论框架（SSD），将 SSMs 与结构化矩阵和注意力变体联系起来。
• 提出新的算法（SSD），将线性递推和二次对偶形式统一以实现高效计算。
• 提出利用 SSD 和结构化注意力实现更高吞吐量的架构设计（Mamba-2）。

提出的方法

• 将 SSMs 形式化为矩阵变换，并将其与半可分矩阵（SSS 表示）联系起来。
• 引入 Structured State Space Duality (SSD) 以统一线性（递推）和二次（类似注意力）形式。
• 使用结构化矩阵将线性注意力推广为 Structured Masked Attention (SMA)。
• 基于半可分矩阵的块分解来开发 SSD 算法，以实现硬件效率。
• 提出具备分组值注意力和并行数据投影的 Mamba-2 架构。
• 分析计算权衡并提供面向硬件的实现指南。

实验结果

研究问题

• RQ1 如何将 SSMs 表示为半可分矩阵，以揭示与注意力相关的算法对偶性？
• RQ2 一个统一的 SSD 框架是否能够为 SSMs 提供与 Transformers 相当的高效训练和推理？
• RQ3 哪些修改（如 Mamba-2）能够在不影响语言模型性能的前提下，使 SSMs 针对现代硬件得到优化？
• RQ4 结构化掩蔽注意力变体如何与线性注意力框架相关并实现推广？

主要发现

• SSMs 等价于半可分矩阵，在合适的结构下实现 O(TN) 计算。
• SSD 为线性递推与类似注意力的二次形式之间架起桥梁，使高效混合成为可能。
• 一种新的 SSD 算法在 Mamba 中的表现优于 selective scan，同时能够实现更大的递推状态。
• 具备并行数据投影和 SMA 变体的 Mamba-2 架构在标准下游任务上达到或超过 Mamba 和开源 Transformers。
• 在语言建模、效率指标和一个困难的回忆任务上的实证验证支持具有竞争力的性能。