新架构Mamba更新二代！作者：别争了，数学上Transformer和SSM是一回事

1,022 0 0

文章摘要

【关键词】 Mamba-2模型、SSM与Transformer、状态空间维度、训练速度提升、注意力机制

研究人员推出了Mamba-2模型，该模型在状态空间模型（SSM）与Transformer架构之间建立了新的联系。通过对两种模型的数学基础进行深入分析，研究团队发现，尽管SSM的核心是一个线性时变系统，而Transformer的核心是注意力机制，它们实际上都可以被表示为可半分离矩阵的变换。这一发现揭示了SSM与带选择性的注意力机制之间的紧密联系，并由此提出了“结构化状态空间二元性”（SSD）的理论框架。

基于SSD，Mamba-2模型的状态空间维度扩展至原来的8倍，训练速度提升了50%。这种扩展使得模型能够学习更强的表示能力，并在诸如硬关联召回任务等需要更大状态容量的任务上表现显著改进。此外，通过利用GPU的存储层次结构和块分解矩阵乘法，Mamba-2在保持高效训练的同时，也提高了模型的质量。

研究还发现，通过SSD框架，可以将Transformer架构多年来积累的优化方法引入SSM，如张量并行和序列并行，从而扩展模型规模和序列长度。实验结果显示，混合了4-6个注意力层和Mamba-2层的模型，在性能上甚至超越了原版Transformer++和纯Mamba-2模型，这表明注意力机制和SSM可以互为补充。

研究者Tri Dao强调了博客文章对理解论文的重要性，并提供了相关资源和代码，以促进对Mamba-2模型和状态空间二元性理论的进一步探索。此项研究的论文已发表，并且代码和模型权重也已开放获取，为相关领域的研究提供了新的工具和方法。