Transformers are SSMs: 结构化状态空间对偶性框架下的通用模型与高效算法

在深度学习研究领域，Transformer模型于语言建模任务中取得了突破性进展。然而，近期兴起的状态空间模型（SSMs），例如Mamba，在处理小型至中型规模任务时已展现出与Transformer相当甚至更优的性能。本文揭示了这两类模型之间存在着紧密的关联，并借助结构化半可分矩阵的不同分解方式，构建了SSMs与注意力机制变体之间的理论联系体系。我们提出的状态空间对偶性（SSD）框架不仅拓展了对相关模型的理论认知，还为新型架构Mamba-2的设计提供了支持。Mamba-2的核心组件是对原有Mamba选择性SSM的优化，其运算效率提升了2至8倍，并且在语言建模任务中依旧维持着与Transformer相媲美的表现水平。

研究背景与问题提出

以GPT和Llama为代表的解码器模型在当前深度学习体系中占据着关键地位。为应对注意力机制在计算效率方面的诸多挑战，例如训练过程中序列长度带来的二次方复杂度问题，以及自回归生成阶段所需线性规模的缓存资源，研究者们提出了多种近似解决方案。与此同时，另一类序列模型——结构化状态空间模型（SSMs）也逐步获得关注。这类模型在训练时具备线性复杂度，在生成过程中仅需常数级别的状态存储，并且在处理长距离依赖任务时表现优异，在小到中等规模的语言建模任务中甚至可以与Transformer相媲美或超越其性能。

然而，SSMs的发展似乎与当前主流的Trans