【他山之石】Mamba 2模型携SSD算法，让长序列处理更高效，一举解决AI大模型的效率痛点！

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注！

自2017年首次亮相，Transformer架构便迅速崛起，成为人工智能大模型的中流砥柱，特别是在语言建模领域占据着核心地位。

然而，随着技术的演进和对处理更长序列需求的增长，Transformer的计算效率问题逐渐成为研究焦点，尤其是其自注意力机制计算成本随序列长度平方增长的局限性。

2023年，Mamba模型的出现标志着这一领域的重要突破，它通过实现计算复杂度的线性扩展，成功缓解了Transformer面对长序列处理的挑战。这一创新之作由两位杰出学者共同打造：卡内基梅隆大学的Albert Gu教授和Together.AI的首席科学家Tri Dao教授，后者同时也担任普林斯顿大学的助理教授。

尽管Mamba的初始成果遭遇ICLR的拒绝，但其引发的学术和社区兴趣并未减退。紧随其后，Mamba的进化版本——Mamba 2，在短短半年内发布，该版本通过结构上的优化，实现了2到8倍的性能提升，同时在语言处理能力上与Transformer保持并驾齐驱，甚至在某些指标上超越。

在他们的论文《Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality》中，研究人员揭示了SSD（State Space Duality）框架，为Transformer和状态空间模型之间的关系提供了深刻的理论基础。

Mamba 2的推出，不仅证实了模型结构的灵活性，还通过增加状态维度至256，显著增强了处理复杂任务的能力，特别是在MQAR等需要高容量状态的任务上，与初始版本相比，进步显著。

此外，研究团队的实验显示，通过在Mamba 2架构中融入4到6层注意力机制，创造出混合模型，这种结合不仅超越了纯Mamba 2和Transformer++，还强调了注意力机制与SSM结构的互补性，为未来模型设计提供了新的思路。

研究贡献

在研究成果中，研究团队揭示了SSM与半可分矩阵家族之间的深刻等价性，这一发现构成了Mamba-2的框架基础。这一框架不仅深化了对SSM属性和算法的理解，还通过线性注意力理论的创新性应用，开创了结构化掩码注意力（SMA）的新领域。SMA的引入不仅证明了与SSM的交集和对偶性质，还揭示了任何快速循环核注意方法的SSM本质。

理论贡献之外，该研究在实践层面也取得了突破。通过提出基于半可分离矩阵块分解的SSD算法，研究者为SSM的计算提供了前所未有的效率，实现了算法的优化。

SSD算法的引入，相较于Mamba的选择性扫描，速度提升达到2到8倍，同时支持更大循环状态的使用，而不会牺牲计算效率。与当前高效的softmax注意力实现FlashAttention-2相比，SSD在处理长序列时展现出了显著优势，特别是在序列长度达到16K时，速度提升达到6倍。

在架构设计的挑战上，此研究通过Mamba-2框架，为SSM的集成扫清了道路，允许利用现有的Transformer优化技术，并通过引入分组值注意力（GVA）头结构，实现了架构的并行化改进。

结合SSD作为内部层，Mamba-2展示了在大规模模型训练中的优越性能，尤其是在Chinchilla扩展法则下，不仅在理论困惑度上，更在实际运行效率上超越了Mamba和Transformer++。

实验结果显示，经过Mamba-2架构训练的模型，在Pile数据集上的表现令人瞩目。特别是2.7B参数的Mamba-2，尽管参数量略低于某些竞争模型，但在3000亿token的训练后，其性能超越了包括2.8B参数的Mamba和6.9B参数的Pythia在内的多个开源Transformer模型，证明了Mamba-2在序列模型设计与优化中的强大潜力。

核心创新——SSD

在研究中，Mamba-2模型引入的SSD层有效地桥接了SSM和Transformer架构。这一整合不仅增强了模型的表达能力，还显著提高了训练效率，尤其是在利用现代加速器进行矩阵乘法运算方面。

Mamba-2的核心创新在于其新的SSD层，这是选择性SSM的一种特例。与Mamba相比，Mamba-2在表达能力上略有下降，但在训练效率上却有了显著提升。这一改进使得Mamba-2能够更好地利用现代加速器上的矩阵乘法单元。

此外，研究团队还对Mamba的神经网络架构进行了一些调整。Mamba-2的网络架构从顺序生成SSM参数转变为并行生成，这使得Mamba-2更加适合张量并行等扩展方法。

研究团队还揭示了理解和使用状态空间模型的新方法，通过提供状态空间模型的显式矩阵变换形式。从计算的角度来看，任何计算状态空间模型前向传播的方法都可以看作是半可分离矩阵上的矩阵乘法算法。这一视角为SSD提供了一个新的理解方式，其中双重模式分别指的是线性时间半可分离矩阵乘法算法和二次时间朴素矩阵乘法。

此外，研究团队还定义了结构化状态空间模型和结构化注意力，并讨论了它们的属性。他们表明，这些模型和注意力机制都具有二次算法和线性算法。

自从Mamba的初始论文研究了合成任务，如合成复制和归纳Head以来，许多后续工作开始研究更难的关联回忆任务。由Zoology和Based系列工作引入的MQAR（多查询关联回忆）任务已成为事实上的标准。

通过运行一个比文献中通常报告的版本要难得多的任务，该团队发现Mamba-2明显优于Mamba-1。其中一个原因是状态大小（比Mamba-1大约16倍）。

在这篇文章中，作者深入探讨了模型背后的理论，并从两个完全不同的角度推导出SSD的“对偶性”：一个从SSM的角度出发，另一个从注意力机制的角度出发。SSD框架提供了状态空间模型、注意力机制和结构化矩阵之间丰富的联系。

虽然SSD模型可以被视为框架内每个分支的具体实例，但SSD框架本身更加通用，为未来的工作开辟了许多方向。

SSD算法

在探讨SSD算法时，研究团队首先注意到矩阵乘法在计算效率上的显著差异。以A100 GPU和H100 GPU为例，它们的BF16矩阵乘法性能分别达到312 TFLOPS和989 TFLOPS，而FP32算术性能则分别为19 TFLOPS和67 TFLOPS。这表明，矩阵乘法的FLOPs速度远超过非矩阵乘法FLOPs，速度比高达16倍。

Mamba-2模型的一个主要目标是利用张量核心来加速SSM。在Mamba-1模型中，通过绑定参数并引入Head结构，SSM转变为了SSD，这是一种更具限制性的形式，其具有类似注意力的公式。SSD通过连接SSM和结构化矩阵，使得计算SSM的高效算法直接对应于“token-mixing”或“sequence-mixing”矩阵M的不同分解。

为了创建计算SSM的新算法，我们可以探索替代的矩阵乘法方式，例如通过不同的分解方法。通过精心选择块大小，对这个矩阵进行简单的块分解，我们可以结合SSD线性递归和二次注意力对偶形式的优势。

这便是SSD算法的起源，对于这个算法有两种完全不同的解释。

SSD算法的关键在于分块矩阵分解。

首先，将半可分SSM矩阵划分为大小为Q×Q的块。然后，利用半分矩阵的性质来分解每个低秩的非对角块：

- 橙色部分：每个对角块是一个更小的半可分矩阵，可以以任何喜欢的方式计算这个乘法，特别是使用SSD的二次（类似注意力机制）形式。

- 绿色部分：总共有T/Q个不同的绿色块，通过批处理矩阵乘法来计算。

- 黄色部分：黄色项本身是一个1-半可分矩阵，这一步等价于对某些修改后的A因子的SSM扫描。

- 蓝色部分：与绿色部分类似，通过批处理矩阵乘法来计算。

通过这种方式，SSD算法既提高了计算效率，又保持了模型的结构优势。

深入解析分块与状态传递机制

SSD算法是一种高效的序列处理算法，其核心在于分块和状态传递。通过这一机制，算法能够实现对输入序列的高效处理，特别是在处理大规模数据时展现出卓越的性能。

在SSD算法中，首先需要将输入序列分割成大小为Q的块。这一步骤不仅有助于简化数据处理过程，还能提高算法的并行处理能力。在每个块内部，算法会计算局部输出，并假设初始状态为0。通过矩阵乘法等计算方式，算法能够快速地得出每个块的输出。

除了局部输出外，SSD算法还需要计算每个块的最终状态。这一步骤同样基于初始状态为0的假设，通过一系列计算得出每个块的最终状态。四、状态传递与全局优化

在SSD算法中，状态传递是一个关键步骤。通过递归计算所有块的最终状态，算法能够考虑到所有先前输入的影响，从而得出每个块的实际最终状态。这一步骤可以采用并行或顺序扫描等方式进行，以确保算法的高效性和准确性。

最后，SSD算法会根据每个块的真实初始状态（在状态传递步骤中计算得出），计算其对于最终输出的贡献。这一步骤确保了算法能够准确地反映出输入序列的整体特征，从而得出准确的输出结果。

SSD算法中的大部分步骤（包括分块内部输出、块状态计算和输出状态计算）都可以利用矩阵乘法等计算方式进行并行处理。这种并行计算能力使得算法在处理大规模数据时能够显著提高性能。同时，由于状态传递步骤只操作一个非常短的序列，因此通常只需要很少的时间即可完成。

张量并行优化在Mamba-2模型中的应用

研究者们对Mamba-2模型中的张量并行进行了优化，显著提升了训练效率。这一优化主要解决了在Mamba-1模型中，每层需要执行两次all-reduce操作的难题。

张量并行的改进点

在Mamba-2中，研究者们采用了“并行投影”结构，使得所有SSM参数都成为层输入的函数。通过这一设计，研究者们能够轻松地将张量并行应用于输入投影，将输入投影和输出投影矩阵分割成多个碎片，具体数量取决于张量并行度。

同时，为了维持归一化的正确性，研究者们还引入了grouped norm，其中分组数除以张量并行度，确保每个GPU都能独立完成归一化操作。这些改进使得Mamba-2模型在训练过程中，每层仅需执行一次all-reduce操作，极大地提高了训练效率。

序列并行的两种形式

在处理超长序列时，序列并行成为了一种有效的解决方案。对于Mamba-2模型而言，序列并行主要有两种形式：一种是对于残差和归一化操作的优化，通过reduce-scatter、残差+归一化、然后all-gather的方式取代了传统的all-reduce操作；

另一种是针对注意力或SSM操作的优化，称为上下文并行（CP）。在注意力机制中，研究者们利用环形注意力沿序列维度进行分割，提高了计算效率。而对于Mamba-2模型中的SSD框架，它再次发挥了关键作用，使得每个GPU都能独立计算其本地输出和最终状态，并通过发送/接收通信原语在GPU之间传递状态信息。

在本次研究中，研究者们采用了MQAR的一种更具挑战性的版本进行实验，该版本任务难度更高、序列更长，且使用的模型规模更小。研究者们将标准的多头softmax注意力机制和Based架构作为基线模型，来进行对比实验。

论文地址：https://arxiv.org/pdf/2405.21060

GitHub 地址：https://github.com/state-spaces/mamba

论文标题：Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。