Mamba真比Transformer更优吗？Mamba原作者：两个都要！混合架构才是最优解
新智元报道
编辑：乔杨【新智元导读】Mamba模型由于匹敌Transformer的巨大潜力，在推出半年多的时间内引起了巨大关注。但在大规模预训练的场景下，这两个架构还未有「一较高低」的机会。最近，英伟达、CMU、普林斯顿等机构联合发表的实证研究论文填补了这个空白。
去年12月，CMU、普林斯顿的两位华人学者Albert Gu和Tri Dao一举推出了Mamba架构，向Transformer多年的霸主地位发起挑战。论文发表后的6个月中，两位作者发现，虽然Mamba很强大，但是大家依旧更关注各种Transformer的变体。
毕竟整个学术社区在注意力机制上深耕多年，从模型、标准库到算子、GPU，此时完全抛弃之前的研究、转向Mamba的SSM不太现实，也让Mamba架构显得非常格格不入。
于是，我们看到Mamba-2的论文在更高层面上将SSM和注意力机制统一了起来，同时相比Mamba-1实现了2～8倍的速度提升。
就在大家都期待着「王者归来」的Mamba-2与Transformer一决高下时，英伟达、威斯康星-麦迪逊大学、普林斯顿、CMU等多个机构的作者共同发表了一篇实证研究文章，发现基于Mamba架构的语言模型在长上下文任务上不敌Transformer。
其实不管出现哪种创新的方法或模型，有论文提出批评意见总是难免的。但细看这篇文章居然发现，Mamba的创造者Tri Dao和Albert Gu两人竟然也在作者列表中。
在此为两位科学家实事求是的精神点赞。
此外，作者列表中还能发掘到另一个华点——Albert Gu和Tri Dao都有了新title。
Albert Gu现任Cartesia AI的联合创始人兼首席科学家，他们最新的产品是实时语音交互API Cartesia Sonic。
Tri Dao是Together AI的创始科学家，该公司主要提供云服务，同时也贡献前沿的开源研究。
https://www.together.ai
接下来我们还是详细看看，这篇文章对Mamba和Transformer的能力具体做了哪些对比研究。
简介
在迄今为止的研究中（包括提出Mamba架构的论文），SSM与Transformer的对比都只进行了较小规模的实验（<3B参数，<1T token），这些结论在训练预算更大的情况下是否成立？
这篇技术报告就是要回答这个问题。作者分别训练出Mamba、Mamba-2、Mamba-2-Hybrid、Transformer等4种架构的8B参数模型，在35个NLP下游任务中对比性能。为了方便复现和进一步研究，用于训练Mamba、Mamba-2和Mamba-2-Hybrid的代码已经开源，而且研究团队还在HuggingFace上发布了Mamba-2 8B和Mamba-2-Hybrid 8B的模型权重（作为英伟达Megatron-LM框架和代码库的一部分）。
实验结果表明，虽然Mamba和Mamba-2更擅长建模语言，但在上下文学习方面，以及从上下文中回忆信息时，性能落后于Transformer模型。
尤其是在MMLU基准上，即使提高了训练数据的token数量，基于Mamba的模型依旧和Transformer有不小的差距。
Mamba vs. Transformer
用于评估的35个下游任务大致包含3个类别：
标准短上下文任务（12个）：HellaSwag、ARC-Easy、ARC-Challenge、MMLU、OpenBookQA、TruthfulQA等
自然长上下文任务（9个）：LongBench中的6个任务和LM Evaluation Harness框架中的3个任务
综合长上下文任务（14个）：RULER框架中的13个开源测试（包括「大海捞针」的8个变体）以及今年刚提出的「电话簿」（Phonebook）任务，旨在衡量模型在长输入文本中检索、跟踪、聚合信息的能力。
表2展示了经过1.1T数据训练后，纯SSM架构的Mamba和Mamba-2与Transformer模型的部分评估结果。
在常见任务上，Mamba和Mamba-2的性能都可以匹配甚至超过Transformer模型，但MMLU基准是一个例外。进行零样本或少样本学习时，Mamba-2相比Transformer分别有10分和17分的差距。
因为在1.1T数据集上Mamba模型的训练速度就已经比Mamba-2慢了将近3×（模型的状态维度较大），出于效率方面的考量，在3.5T数据集上只训练了Mamba-2模型和Transormer模型，部分结果如表3所示。
从表3可知，更多的训练数据有助于Mamba-2在MMLU任务上得到改进，5-shot分数的差距缩小到仅1.37分，其他任务上依旧全面领先Transformer。
Mamba折戟MMLU与电话簿任务
由于MMLU在一众下游任务的结果中显得如此反常，论文对此进行了更细致的拆解和讨论。
如上图所示，MMLU的任务类似于考试中的选择题，但在cloze格式中也可以不提供备选答案，以填空题的方式提供给模型。
表4中提供了MMLU按照格式细分后，3个模型各自的分数（用1.1T token训练）。在标准模式和选择题模式中，Mamba架构不敌Transformer，但在填空题模式中居然实现了分数反超。
结合表3中的结果，我们有理由推断，纯SSM模型和Transformer模型包含的知识内容应该是同等级别的，但前者需要更多的训练才能理解MMLU的前两种格式。
作者推断，这种差距可能源于Transformer强大的上下文学习能力，可以看到该模型从0-shot到5-shot的准确度提升非常明显。
此外，SSM模型可能无法直接将答案所需的知识路由到输出的单个答案token中（即ABCD选项的其中一个），而这正是自注意力层擅长的任务。
此外，Mamba系列模型在「电话簿」上的表现也并不理想，该任务旨在衡量模型通过少数示例进行上下文学习，以及从上下文中复制信息的能力。
下图展现了任务的两种变体，标准版是先提供整个电话簿，再给出目标查询；反转版则是先查询，再给电话簿。
Transformer在电话簿长度不超过预训练的上下文长度（4096）时，准确率接近100%，相比之下，Mamba和Mamba-2在输入序列达到500 token时就出现了显著的性能滑坡。
如果仔细观察Mamba系列的输出答案（图2b），可以发现SSM架构的模型并非完全无法记忆上下文信息，而是保留了一些模糊记忆，给出的电话号码通常有几位是正确的。
综合以上结果，我们可以将MMLU和「电话簿」任务确立为纯SSM架构模型的挑战性任务，并且推测出可能原因：这两个任务需要上下文学习、token间信息路由以及从上下文复制的能力，它们可能是Mamba系列模型的能力软肋。
SSM-Transformer混合架构
由于在MMLU和「电话簿」任务上看到了SSM架构的能力缺陷，作者想到——让SSM和Transformer强强联合，能够起到取长补短的效果？
于是他们将自注意力和MLP层添加到Mamba架构中，想看看模型能否克服上述问题。
论文首先报告了一系列消融实验的结果，通过对比在下游任务上的表现，探索出了能达到最佳性能的架构设计与参数（表6）。
56层的Mamba-2-Hybrid中包含4个（7.1%）自注意力层，24 个（42.9%）Mamba-2层和28个（50%）MLP 层，其中Mamba-2层使用与Mamba-2模型相同的参数。
自注意力、MLP层的数量以及MLP层扩展因子这些参数的选择并非随机，而是根据验证集上损失值结果（图4）进行的最优化设计。
消融实验的结果还显示，混合模型中不添加旋转位置编码（RoPE）能达到更好的下游任务性能（表5），而且Mamba层、自注意力层、MLP层的顺序也会影响模型能力。
首先，Mamba层必须出现在架构的开头，以确保模型自然地学习到位置信息。相比使用重复块模式，将自注意力和MLP均匀分散在整个模型是更好的配置。
而且通过计算验证集上的模型困惑度（perplexity）可以得知，相比多头注意力（MHA），使用组查询注意力层（GQA）能减少推理计算量和内存量，但几乎不会造成模型质量的下降。
效率方面，Mamba-2-Hybrid实现了29.9%的FLOP利用率（MFU），与Transfomer的30.7%基本相当。此外，前者有推理速度方面的巨大优势。
在长上下文情境中，受益于多个SSM层的存在，Mamba-2-Hybrid的token生成速度比Transformer加速了将近8×（图5）。
评估
测评发现，这种混合架构果然有了「取长补短」的效果，混合架构在5-shot MMLU测评中同时超过了单纯的Transformer和SSM架构，取得得了最高准确度（图6）。
在Mamba系列表现较差的「电话簿」任务上，Mamba-2-Hybrid可以在预训练上下文长度 (4K) 内以近乎完美的精度完成电话簿任务，还可以稍微超出该长度进行泛化，在最多5.5k token的电话簿上实现100%准确率。
甚至，Mamba-2-Hybrid的潜力还不止于此，当预训练长度扩展到128k并在4个自注意力层中使用全局注意力时，「电话簿」任务的100%准确率也延伸到了将近150k token。
结论
论文开头的评估结果表明，在更大训练预算的情况下，纯SSM模型依旧能在下游任务上超过Transformer，但上下文学习和信息检索能力有所局限。
基于此，作者提出的混合架构模型Mamba-2-Hybrid能够在提高效率的同时继续表现出比Transformer更强大的性能，并弥补了纯SSM架构的相关缺陷。
这项研究所展示的全面结果告诉我们，Mamba和Transformer这两种架构各有长短，也许并不需要其中一个取代另一个，将二者结合起来是一条值得探索的、有巨大潜力的路径。
参考资料：https://arxiv.org/pdf/2406.07887