去年12月，CMU、普林斯頓的兩位華人學(xué)者Albert Gu和Tri Dao一舉推出了Mamba架構(gòu)，向Transformer多年的霸主地位發(fā)起挑戰(zhàn)。

論文地址：https://arxiv.org/abs/2312.00752

完全拋棄注意力機(jī)制和MLP模塊、上下文長度線性縮放、推理速度比Transformer快5倍…這些特點(diǎn)讓所有人都為之一振，Jim Fan大佬也發(fā)推贊嘆「為推翻Transformer的研究感到興奮」。

論文發(fā)表后的6個(gè)月中，兩位作者發(fā)現(xiàn)，雖然Mamba很強(qiáng)大，但是大家依舊更關(guān)注各種Transformer的變體。

畢竟整個(gè)學(xué)術(shù)社區(qū)在注意力機(jī)制上深耕多年，從模型、標(biāo)準(zhǔn)庫到算子、GPU，此時(shí)完全拋棄之前的研究、轉(zhuǎn)向Mamba的SSM不太現(xiàn)實(shí)，也讓Mamba架構(gòu)顯得非常格格不入。

于是，我們看到Mamba-2的論文在更高層面上將SSM和注意力機(jī)制統(tǒng)一了起來，同時(shí)相比Mamba-1實(shí)現(xiàn)了2～8倍的速度提升。

論文地址：https://arxiv.org/abs/2405.21060

就在大家都期待著「王者歸來」的Mamba-2與Transformer一決高下時(shí)，英偉達(dá)、威斯康星-麥迪遜大學(xué)、普林斯頓、CMU等多個(gè)機(jī)構(gòu)的作者共同發(fā)表了一篇實(shí)證研究文章，發(fā)現(xiàn)基于Mamba架構(gòu)的語言模型在長上下文任務(wù)上不敵Transformer。

其實(shí)不管出現(xiàn)哪種創(chuàng)新的方法或模型，有論文提出批評(píng)意見總是難免的。但細(xì)看這篇文章居然發(fā)現(xiàn)，Mamba的創(chuàng)造者Tri Dao和Albert Gu兩人竟然也在作者列表中。

論文地址：https://arxiv.org/abs/2406.07887

在此為兩位科學(xué)家實(shí)事求是的精神點(diǎn)贊。

此外，作者列表中還能發(fā)掘到另一個(gè)華點(diǎn)——Albert Gu和Tri Dao都有了新title。

Albert Gu現(xiàn)任Cartesia AI的聯(lián)合創(chuàng)始人兼首席科學(xué)家，他們最新的產(chǎn)品是實(shí)時(shí)語音交互API Cartesia Sonic。

https://cartesia.ai

Tri Dao是Together AI的創(chuàng)始科學(xué)家，該公司主要提供云服務(wù)，同時(shí)也貢獻(xiàn)前沿的開源研究。

https://www.together.ai

接下來我們還是詳細(xì)看看，這篇文章對(duì)Mamba和Transformer的能力具體做了哪些對(duì)比研究。

簡介

在迄今為止的研究中（包括提出Mamba架構(gòu)的論文），SSM與Transformer的對(duì)比都只進(jìn)行了較小規(guī)模的實(shí)驗(yàn)（<3B參數(shù)，<1T token），這些結(jié)論在訓(xùn)練預(yù)算更大的情況下是否成立？

這篇技術(shù)報(bào)告就是要回答這個(gè)問題。作者分別訓(xùn)練出Mamba、Mamba-2、Mamba-2-Hybrid、Transformer等4種架構(gòu)的8B參數(shù)模型，在35個(gè)NLP下游任務(wù)中對(duì)比性能。

訓(xùn)練數(shù)據(jù)包括1.1T和3.5T兩個(gè)數(shù)據(jù)集，都是英偉達(dá)用于訓(xùn)練Nemotron-4的數(shù)據(jù)集的前身，由70%英語、15%非英語和15%代碼組成

其中，Mamba-2-Hybrid是一個(gè)SSM-Transformer的混合架構(gòu)模型，包含24個(gè)Mamba-2層，以及均勻分布在整個(gè)模型中的4個(gè)自注意力層和28個(gè)MLP層。

總體而言，這項(xiàng)對(duì)比實(shí)驗(yàn)消除了比較不同LLM的常見困難，包括訓(xùn)練數(shù)據(jù)、分詞器、評(píng)估管道等方面，確保評(píng)估流程的標(biāo)準(zhǔn)和可重復(fù)性。

為了方便復(fù)現(xiàn)和進(jìn)一步研究，用于訓(xùn)練Mamba、Mamba-2和Mamba-2-Hybrid的代碼已經(jīng)開源，而且研究團(tuán)隊(duì)還在HuggingFace上發(fā)布了Mamba-2 8B和Mamba-2-Hybrid 8B的模型權(quán)重（作為英偉達(dá)Megatron-LM框架和代碼庫的一部分）。

https://huggingface.co/nvidia

實(shí)驗(yàn)結(jié)果表明，雖然Mamba和Mamba-2更擅長建模語言，但在上下文學(xué)習(xí)方面，以及從上下文中回憶信息時(shí)，性能落后于Transformer模型。

尤其是在MMLU基準(zhǔn)上，即使提高了訓(xùn)練數(shù)據(jù)的token數(shù)量，基于Mamba的模型依舊和Transformer有不小的差距。

Mamba vs. Transformer

用于評(píng)估的35個(gè)下游任務(wù)大致包含3個(gè)類別：

• 標(biāo)準(zhǔn)短上下文任務(wù)（12個(gè)）：HellaSwag、ARC-Easy、ARC-Challenge、MMLU、OpenBookQA、TruthfulQA等
• 自然長上下文任務(wù)（9個(gè)）：LongBench中的6個(gè)任務(wù)和LM Evaluation Harness框架中的3個(gè)任務(wù)
• 綜合長上下文任務(wù)（14個(gè)）：RULER框架中的13個(gè)開源測試（包括「大海撈針」的8個(gè)變體）以及今年剛提出的「電話簿」（Phonebook）任務(wù)，旨在衡量模型在長輸入文本中檢索、跟蹤、聚合信息的能力。

表2展示了經(jīng)過1.1T數(shù)據(jù)訓(xùn)練后，純SSM架構(gòu)的Mamba和Mamba-2與Transformer模型的部分評(píng)估結(jié)果。

在常見任務(wù)上，Mamba和Mamba-2的性能都可以匹配甚至超過Transformer模型，但MMLU基準(zhǔn)是一個(gè)例外。進(jìn)行零樣本或少樣本學(xué)習(xí)時(shí)，Mamba-2相比Transformer分別有10分和17分的差距。

因?yàn)樵?.1T數(shù)據(jù)集上Mamba模型的訓(xùn)練速度就已經(jīng)比Mamba-2慢了將近3×（模型的狀態(tài)維度較大），出于效率方面的考量，在3.5T數(shù)據(jù)集上只訓(xùn)練了Mamba-2模型和Transormer模型，部分結(jié)果如表3所示。

從表3可知，更多的訓(xùn)練數(shù)據(jù)有助于Mamba-2在MMLU任務(wù)上得到改進(jìn)，5-shot分?jǐn)?shù)的差距縮小到僅1.37分，其他任務(wù)上依舊全面領(lǐng)先Transformer。

Mamba折戟MMLU與電話簿任務(wù)

由于MMLU在一眾下游任務(wù)的結(jié)果中顯得如此反常，論文對(duì)此進(jìn)行了更細(xì)致的拆解和討論。

如上圖所示，MMLU的任務(wù)類似于考試中的選擇題，但在cloze格式中也可以不提供備選答案，以填空題的方式提供給模型。

表4中提供了MMLU按照格式細(xì)分后，3個(gè)模型各自的分?jǐn)?shù)（用1.1T token訓(xùn)練）。在標(biāo)準(zhǔn)模式和選擇題模式中，Mamba架構(gòu)不敵Transformer，但在填空題模式中居然實(shí)現(xiàn)了分?jǐn)?shù)反超。

結(jié)合表3中的結(jié)果，我們有理由推斷，純SSM模型和Transformer模型包含的知識(shí)內(nèi)容應(yīng)該是同等級(jí)別的，但前者需要更多的訓(xùn)練才能理解MMLU的前兩種格式。

作者推斷，這種差距可能源于Transformer強(qiáng)大的上下文學(xué)習(xí)能力，可以看到該模型從0-shot到5-shot的準(zhǔn)確度提升非常明顯。

此外，SSM模型可能無法直接將答案所需的知識(shí)路由到輸出的單個(gè)答案token中（即ABCD選項(xiàng)的其中一個(gè)），而這正是自注意力層擅長的任務(wù)。

此外，Mamba系列模型在「電話簿」上的表現(xiàn)也并不理想，該任務(wù)旨在衡量模型通過少數(shù)示例進(jìn)行上下文學(xué)習(xí)，以及從上下文中復(fù)制信息的能力。

下圖展現(xiàn)了任務(wù)的兩種變體，標(biāo)準(zhǔn)版是先提供整個(gè)電話簿，再給出目標(biāo)查詢；反轉(zhuǎn)版則是先查詢，再給電話簿。

圖3a、c分別展示了3個(gè)模型在這兩個(gè)任務(wù)變體上的準(zhǔn)確率。

Transformer在電話簿長度不超過預(yù)訓(xùn)練的上下文長度（4096）時(shí)，準(zhǔn)確率接近100%，相比之下，Mamba和Mamba-2在輸入序列達(dá)到500 token時(shí)就出現(xiàn)了顯著的性能滑坡。

如果仔細(xì)觀察Mamba系列的輸出答案（圖2b），可以發(fā)現(xiàn)SSM架構(gòu)的模型并非完全無法記憶上下文信息，而是保留了一些模糊記憶，給出的電話號(hào)碼通常有幾位是正確的。

綜合以上結(jié)果，我們可以將MMLU和「電話簿」任務(wù)確立為純SSM架構(gòu)模型的挑戰(zhàn)性任務(wù)，并且推測出可能原因：這兩個(gè)任務(wù)需要上下文學(xué)習(xí)、token間信息路由以及從上下文復(fù)制的能力，它們可能是Mamba系列模型的能力軟肋。

SSM-Transformer混合架構(gòu)

由于在MMLU和「電話簿」任務(wù)上看到了SSM架構(gòu)的能力缺陷，作者想到——讓SSM和Transformer強(qiáng)強(qiáng)聯(lián)合，能夠起到取長補(bǔ)短的效果？

于是他們將自注意力和MLP層添加到Mamba架構(gòu)中，想看看模型能否克服上述問題。

論文首先報(bào)告了一系列消融實(shí)驗(yàn)的結(jié)果，通過對(duì)比在下游任務(wù)上的表現(xiàn)，探索出了能達(dá)到最佳性能的架構(gòu)設(shè)計(jì)與參數(shù)（表6）。

56層的Mamba-2-Hybrid中包含4個(gè)（7.1%）自注意力層，24 個(gè)（42.9%）Mamba-2層和28個(gè)（50%）MLP 層，其中Mamba-2層使用與Mamba-2模型相同的參數(shù)。

自注意力、MLP層的數(shù)量以及MLP層擴(kuò)展因子這些參數(shù)的選擇并非隨機(jī)，而是根據(jù)驗(yàn)證集上損失值結(jié)果（圖4）進(jìn)行的最優(yōu)化設(shè)計(jì)。

消融實(shí)驗(yàn)的結(jié)果還顯示，混合模型中不添加旋轉(zhuǎn)位置編碼（RoPE）能達(dá)到更好的下游任務(wù)性能（表5），而且Mamba層、自注意力層、MLP層的順序也會(huì)影響模型能力。

首先，Mamba層必須出現(xiàn)在架構(gòu)的開頭，以確保模型自然地學(xué)習(xí)到位置信息。相比使用重復(fù)塊模式，將自注意力和MLP均勻分散在整個(gè)模型是更好的配置。

而且通過計(jì)算驗(yàn)證集上的模型困惑度（perplexity）可以得知，相比多頭注意力（MHA），使用組查詢注意力層（GQA）能減少推理計(jì)算量和內(nèi)存量，但幾乎不會(huì)造成模型質(zhì)量的下降。

效率方面，Mamba-2-Hybrid實(shí)現(xiàn)了29.9%的FLOP利用率（MFU），與Transfomer的30.7%基本相當(dāng)。此外，前者有推理速度方面的巨大優(yōu)勢。

在長上下文情境中，受益于多個(gè)SSM層的存在，Mamba-2-Hybrid的token生成速度比Transformer加速了將近8×（圖5）。

評(píng)估

測評(píng)發(fā)現(xiàn)，這種混合架構(gòu)果然有了「取長補(bǔ)短」的效果，混合架構(gòu)在5-shot MMLU測評(píng)中同時(shí)超過了單純的Transformer和SSM架構(gòu)，取得得了最高準(zhǔn)確度（圖6）。

從表7中的多個(gè)基準(zhǔn)總體來看，Mamba-2-Hybrid在效率更高的同時(shí)，性能也超過了Transformer模型。

相比Mamba-2，混合架構(gòu)的長上下文能力也得到了顯著提高（表10），在RULER基準(zhǔn)上的綜合任務(wù)、「大海撈針」任務(wù)的平均成績也都超過了Transformer。

在Mamba系列表現(xiàn)較差的「電話簿」任務(wù)上，Mamba-2-Hybrid可以在預(yù)訓(xùn)練上下文長度 (4K) 內(nèi)以近乎完美的精度完成電話簿任務(wù)，還可以稍微超出該長度進(jìn)行泛化，在最多5.5k token的電話簿上實(shí)現(xiàn)100%準(zhǔn)確率。

甚至，Mamba-2-Hybrid的潛力還不止于此，當(dāng)預(yù)訓(xùn)練長度擴(kuò)展到128k并在4個(gè)自注意力層中使用全局注意力時(shí)，「電話簿」任務(wù)的100%準(zhǔn)確率也延伸到了將近150k token。

結(jié)論

論文開頭的評(píng)估結(jié)果表明，在更大訓(xùn)練預(yù)算的情況下，純SSM模型依舊能在下游任務(wù)上超過Transformer，但上下文學(xué)習(xí)和信息檢索能力有所局限。

基于此，作者提出的混合架構(gòu)模型Mamba-2-Hybrid能夠在提高效率的同時(shí)繼續(xù)表現(xiàn)出比Transformer更強(qiáng)大的性能，并彌補(bǔ)了純SSM架構(gòu)的相關(guān)缺陷。

這項(xiàng)研究所展示的全面結(jié)果告訴我們，Mamba和Transformer這兩種架構(gòu)各有長短，也許并不需要其中一個(gè)取代另一個(gè)，將二者結(jié)合起來是一條值得探索的、有巨大潛力的路徑。