[读论文]Transformers are SSMs_transformers are ssms: generalized models and effi

Notation

$T$: Sequence length/ time length
$$:

摘要

虽然transformer一直是深度学习在语言建模方面成功的主要架构，但状态空间模型(ssm)，如Mamba，最近被证明在中小规模上与transformer相匹配或优于transformer。这些模型族实际上是非常密切相关的，并在ssm和注意力变体之间发展了一个丰富的理论联系框架，通过对一类经过充分研究的结构化半可分矩阵的各种分解连接起来。状态空间对偶(SSD)框架使我们能够设计一个新的架构(Mamba-2)，其核心层是Mamba选择性SSM的改进，速度快了2-8倍，同时在语言建模方面继续与transformer竞争。

介绍

Transformer存在效率问题。例如在训练期间按序列长度进行二次缩放，以及在自回归生成期间需要按序列长度线性的缓存大小。而一类可供选择的序列模型，即结构化状态空间模型(SSMs)，在训练期间序列长度呈线性扩展，在生成期间状态大小恒定。同时后者在长程任务上表现出强大的性能，最近在小到中等规模的语言建模上与transformer不相上下或击败了transformer。然而，ssm的发展似乎与社区改进transformer的集体努力无关，例如从理论上理解它们，以及在现代硬件上优化它们。因此，与transformer相比，理解和实验ssm更加困难，从算法和系统的角度来看，训练ssm像transformer一样有效仍然具有挑战性。

本文的主要目标是在结构化ssm和注意力变体之间建立丰富的理论联系。这将使我们能够将最初为transformer开发的算法和系统优化转移到ssm，以建立性能比transformer更好的基础模型，同时在序列长度上更有效地扩展。该方向的一个里程碑式贡献是线性注意力(LA)框架，通过展示二次核注意力的"对偶形式"和特定线性递归之间的等价性，导出了自回归注意力和线性rnn之间的联系。这种对偶性允许新的功能，如具有高效的可并行化训练和高效的自回归推理的能力。同样，本文提供了将线性复杂度ssm与二次复杂度形式联系起来的多个观点，以结合ssm和attention的优势。

状态空间对偶性。所提出的连接结构化ssm和注意力变体的框架，称为结构化状态空间对偶(SSD)，是通过对结构化矩阵的抽象来建立的(结构化矩阵是具有次二次参数和乘法复杂度的矩阵)。本文开发了两个广泛的框架来表示序列模型，一个作为矩阵变换，一个作为张量收缩，它们都揭示了对偶性的不同角度。我们的技术贡献包括:

• 展示了状态空间模型和被广泛研究的结构化矩阵族(半可分离矩阵)之间的等价性(第3节)。这种联系是所提出框架的核心，揭示了ssm的新属性和算法。本文的中心思想是计算状态空间模型的不同方法可以被重新定义为结构化矩阵上的各种矩阵乘法算法。
• 显著改进了线性注意力理论。本文首先通过张量收缩的语言提供了其递归形式的简明证明，然后将其推广到一个新的结构化掩码注意力(SMA)族(第4节)。
• 连接了SSM和SMA，表明它们有一个彼此对偶的大交集，具有类似SSM的线性和类似注意力的二次形式(第5节)。还证明了任何具有快速递归形式的核注意力方法都必须是SSM。

除了其内在的理论价值外，该框架为理解和改进序列模型开辟了一套广泛的方向。

高效的算法。首先，该框架暴露了新的高效和易于实现的算法，用于计算SSM(第6节)。提出了一种新的SSD算法，基于半可分矩阵的分块分解，同时利用了线性SSM递归和二次对偶形式，在所有主要效率轴上获得了最佳折衷(例如训练和推理计算、内存使用，以及在现代硬件上利用矩阵乘法单元的能力)。SSD的专用实现比Mamba的优化选择性扫描实现快2 - 8倍，同时允许更大的循环状态大小(8倍或更高的Mamba大小，以最小的减慢)。SSD与softmax注意力(FlashAttention-2 (Dao 2024))的优化实现相比具有很强的竞争力，在序列长度为2K时交叉速度提高了6倍，在序列长度为16K时交叉速度提高了6倍。

架构设计。采用ssm等新架构的一个主要障碍是为transformer量身定制的生态系统，如用于大规模训练的硬件高效优化和并行技术。所提出框架允许使用既定的注意力惯例和技术为ssm建立架构设计选择的词汇表，并进一步改进(第7节)。例如，引入了从多头注意力(MHA)到ssm的头的模拟。我们表明这样的Mamba架构是一个多输入SSM (MIS)，被证明类似于多值注意力(MVA)，并比较了具有不同头部结构的Mamba的其他变体。

我们还使用这些想法对Mamba块进行轻微修改，它允许实现张量并行。主要思想包括引入组值注意力(GVA)头部结构，并将所有数据依赖的投影移动到区块开始时并行发生。

将改进后的并行Mamba块与使用SSD作为内层SSM层相结合，得到Mamba-2架构。本文研究了Mamba-2和Mamba在相同环境下的Chinchilla缩放定律，发现在困惑度和壁钟时间上，它帕累托优于Mamba和Transformer++。在Pile上训练了一系列不同大小的Mamba-2模型，表明它在标准的下游评估中匹配或优于Mamba和开源transformer。例如，在桩上300B token上训练的2.7B参数的Mamba-2优于在相同数据集上训练的Mamba-2.8B, Pythia-2.8B，甚至Pythia-6.9B。

系统优化。SSD框架将ssm和transformer连接起来，使我们能够利用为transformer开发的大量系统优化工作(第8节)。

• 例如，张量并行(TP)是一种重要的模型并行技术，通过在同一节点上的gpu上划分每一层来训练大型Transformer模型。我们将Mamba-2设计为TP友好型，将每个块的同步点数量减少了一半。
• 对于非常长的序列，其激活不适合在一个设备上，针对注意力块开发了序列并行性。本文描述了如何通过在设备之间传递循环状态，用序列并行来训练ssm，特别是Mamba-2。
• 为了对不同长度的示例进行微调，为了获得最佳效率，Transformer需要复杂的技术来删除填充标记并对可变长度序列进行关注。我们展示了如何用可变序列长度有效地训练Mamba-2，不需要填充标记。

第9节对Mamba-2在语言建模、训练效率和一个困难的多查询关联回忆任务上进行了实证验证(Arora, Eyuboglu, Zhang等人2024)。最后，在第10节中，提供了相关的扩展工作，并讨论了该框架开辟的潜在研究方向。
代码见https://github.com/state-spaces/mamba。

2背景及概述

2.1结构化状态空间模型(Structured SSM)

结构化状态空间序列模型(S4)是最近一类用于深度学习的序列模型，与rnn、cnn和经典状态空间模型有广泛的关系。它们的灵感来自于一个特定的连续系统(1)，该系统通过隐潜状态ℎ∈R(T,N)映射一个一维序列