Mamba 2的发布是否可以撼动Transformer模型的AI大一统的江湖地位

Transformer模型

Transformer模型是深度学习领域的一种神经网络架构，特别适用于自然语言处理（NLP）任务。它由Vaswani等人在2017年的论文《Attention is All You Need》中提出。Transformer模型的关键创新在于其使用注意力机制，而不是传统的递归神经网络（RNN）或卷积神经网络（CNN）来处理序列数据。正是由于Transformer模型强大的注意力机制，让模型进行并行运算成功降低了训练时间，加速了模型的训练。

1.架构概述

Transformer模型由两个主要部分组成：编码器（Encoder）和解码器（Decoder）。

• 编码器：由多个相同的编码器层（通常是6层）堆叠而成。每个编码器层包含两个子层：多头自注意力机制（Multi-Head Self-Attention）和前馈神经网络（Feed Forward Neural Network），每个子层后都有一个Add & Norm层进行残差连接和规范化。
• 解码器：由多个相同的解码器层（通常也是6层）堆叠而成。每个解码器层包含三个子层：Masked多头自注意力机制、编码器-解码器多头注意力机制（Encoder-Decoder Attention）和前馈神经网络。和编码器一样，每个子层后都有一个Add & Norm层。

2.注意力机制

注意力机制是Transformer的核心组件，它允许模型在处理一个元素时关注输入序列中的其他元素。Transformer使用的是Scaled Dot-Product Attention：

• Q（Queries）：查询矩阵
• K（Keys）：键矩阵
• V（Values）：值矩阵
• d_k：键向量的维度，用于缩放点积的结果

以上过程展示了QKV三矩阵的来历，其中输入的 I 矩阵为已知矩阵，而Wq Wk Wv为未知矩阵，而注意力机制的过程便是需要优化其中的未知参数了。

注意力机制计算的过程也是极其简单，根据注意力机制的计算公式，其Q矩阵乘以K矩阵的转置矩阵，然后乘以一个缩放系数后，再乘以v矩阵，就得到了注意力机制的计算结果。

正是注意力机制的强大算法，不仅NLP领域各种任务中可以采用 transformer 模型，而在计算机视觉任务中也成功采用了 transformer 模型，且效果显著，正是这样的模型统一，让人工智能领域一度认为 transformer 模型将会承担模型大一统的任务。

Mamba模型的兴起

随着模型越来越大，训练数据越来越庞大，其输入序列也越来越长。而我们知道 transformer 模型随着输入序列长度，其模型复杂度成平方的增加。这就无意导致了大量的内存开销，虽然 transformer 模型的优点在于可以并行运算，但是较大的内存开销无法进行模型有效的训练，如何能够找到一个既能降低内存开销，又能并行运算的模型成为了大家研究的对象。

Mamba 是一种新的状态空间模型架构，在语言建模等信息密集数据上显示出良好的性能，而以前的二次模型在 Transformers 方面存在不足。Mamba基于结构化状态空间模型的，并使用FlashAttention进行高效的硬件设计和实现。

Mamba 享有快速推理（吞吐量比 Transformer 高 5 倍）和序列长度线性缩放，并且其性能在高达百万长度序列的实际数据上得到了很多的提高。 作为通用序列模型骨干，Mamba 在语言、音频和基因组学等多种任务上实现了最先进的性能。 在语言建模方面， Mamba-3B 模型在预训练和下游任务方面都优于相同大小的 Transformer模型，并且与两倍大小的 Transformer 相匹配。

而本次 2 位作者在Mamba 的基础上，更新了Mamba模型，发布了Mamba 2 。

而Mamba 2 模型的论文就叫transformer are SSMs，transformer 模型本身就是一个状态空间模型。作者用了大量的篇幅介绍了 transformer 模型与状态空间的关系，既然两者有着密切的联系，那么 transformer 模型就可以结合状态空间模型的优点，并行运算的同时降低内存开销。

上图是不是很熟悉，这个跟Retnet模型解决的大模型不可能三角很像，都是为了解决大模型的内存开销与计算复杂度问题。而根据其Mamba 2 的注意力机制的计算公式跟标准的注意力公式就是删除了 soft max 计算，取而代之的是乘以了一个 L 矩阵。而Retnet模型是一个衰减系数矩阵D.

而L 矩阵同样是一个下三角矩阵，一来屏蔽未来的输入信息，另外一个便是帮忙输入数据计算单词与单词之间的相似性。后期我们将在动画详解 transformer 专栏中详细介绍Retnet与Mamba模型，并用动画的方式来介绍相关的技术难点

Mamba 采用状态空间模型作为其核心架构，而非 Transformer 的自注意力机制。状态空间模型通过递归方式更新隐藏状态，能够更高效地处理长序列数据。Mamba 的架构允许在训练过程中进行高度并行化，从而显著缩短训练时间。与 Transformer 的平方复杂度不同，Mamba 在序列长度上的复杂度为线性，这意味着它在处理长序列时效率更高。虽然 Mamba 旨在提高效率，但初步研究表明，它在各种序列建模任务上都能保持与 Transformer 相当的性能。

Mamba的模型框架不仅可以进行高度并行化，而可以通过状态空间的递归方式更新模型状态，这就可以让模型可以进行递归方式的训练，降低内存开销。

Mamba2 的发布证明了 transformer 模型就是一个SSMs模型，完全可以采用递归的方式降低内存开销，这个跟 retnet模型的分块并行与块与块进行递归操作类似，都成功降低了内存开销。当然Mamba2模型是否在各个人工智能领域任务上都适合使用，需要各大模型的不断适配。而 transformer 模型已经被很多大模型成功使用，是否可以进行替代，需要模型与时间的验证。

https://arxiv.org/pdf/2405.21060

更多transformer，VIT，swin tranformer
参考头条号：人工智能研究所
v号：启示AI科技
 

 动画详解transformer  在线教程