-
背景
文章介绍了Transformers在自然语言处理任务中的成功，尤其是以GPT和Llama为代表的解码器模型，但指出了它们在效率上的问题，并指出在训练过程中核心关注层随序列长度二次方增长问题，及自回归生成时需要与序列长度线性增长的缓存问题。 -
已有的工作 有许多尝试通过近似核心关注层来解决Transformers效率问题的方法，而与此同时，结构化状态空间模型（SSMs）如Mamba表现出与Transformers相近甚至更好的性能，尤其在处理长范围任务时表现出色。但SSMs的开发似乎与社区改进Transformers的集体努力相脱节。
- 提出了一个新框架
-
挑战1：如何使 SSMs 与 Transformers 同样的算法和系统优化相结合 通过建立结构化SSMs和关注机制变体之间的理论联系，将原本为Transformers开发的算法和系统优化转移到SSMs上，从而实现比Transformers更好的表现，并在序列长度上有更高的效率。
-
挑战2：如何在保留SSMs线性复杂度和Transformers二次复杂度形式的优势的基础上将它们结合起来 提出了状态空间对偶性（SSD）框架，该框架通过结构化矩阵的抽象链接了结构化SSMs和注意力的变体，以便在序列模型表示、算法和实现上融合SSMs和注意力的优点。
-
SSD框架提出新的高效且易于实现的SSMs计算算法。这个新算法基于分块分解半可分离矩阵，充分利用线性 SSM 递归和二次方形式的双重特点，得到训练和推理的计算、内存使用、利用现代硬件上矩阵乘法单元等方面的最优权衡。SSD的专门实现比Mamba的优化选择扫描实现快2-8倍，允许更大的递归状态大小，并与优化的softmax注意力实施（FlashAttention-2）高度竞争，在序列长度为2K时相交，在序列长度为16K时速度快6倍。此外，Mamba-2架构通过稍微修改Mamba块，并与使用SSD作为内部SSM层的结合，展示了在与Mamba相同的设置下的Chinchilla缩放法则，发现其在困惑度和墙面时钟时间上均优于Mamba和Transformer++。实验结果表明Mamba-2在标准的下游评测中匹敌或超越了Mamba和开源Transformers。
本论文展示了一个全新的状态空间对偶性（SSD）框架，连接了结构化的状态空间模型（SSMs）和注意力机制变体。论文的主要贡献包括将原本针对Transformers的算法和系统优化应用到SSMs上，以及开发了一种新的SSD算法，有效提高了模型训练和推理的效率。Mamba-2架构作为最终产品，实现了理想的性能表现，为未来的深度学习模型设计和优化提供了新的方向。