深度學(xué)習(xí)進(jìn)入新紀(jì)元，Transformer的霸主地位，要被掀翻了？

2017年6月12日橫空出世，讓NLP直接變天，制霸自然語(yǔ)言領(lǐng)域多年的Transformer，終于要被新的架構(gòu)打破壟斷了。

Transformer雖強(qiáng)大，卻有一個(gè)致命的bug：核心注意力層無(wú)法擴(kuò)展到長(zhǎng)期上下文。

剛剛，CMU和普林斯頓的研究者發(fā)布了Mamba。這種SSM架構(gòu)在語(yǔ)言建模上與Transformers不相上下，而且還能線性擴(kuò)展，同時(shí)具有5倍的推理吞吐量！

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

論文一作Albert Gu表示，二次注意力對(duì)于信息密集型模型是必不可少的，但現(xiàn)在，再也不需要了！

論文一出，直接炸翻了AI社區(qū)。

英偉達(dá)首席科學(xué)家Jim Fan表示，自己一直期待能有人來(lái)推翻Transformer，并且對(duì)Albert Gu和Tri Dao多年以來(lái)做出替代Transformer序列架構(gòu)的嘗試表示感謝。

「你們做的研究太酷了，一會(huì)兒蹦出一個(gè)來(lái)，不能稍微停一下嗎！」

「湖人粉表示，對(duì)Mamba這個(gè)名字很滿意！」

對(duì)于這個(gè)架構(gòu)為何取名曼巴，作者也給出了解釋——

- 速度快：原因在于（1）序列長(zhǎng)度線性縮放的簡(jiǎn)單遞歸，（2）硬件感知設(shè)計(jì)和實(shí)現(xiàn)

- 致命性：它對(duì)序列建模問題具有致命的吸引力

- 就連發(fā)出的「聲音」都很像：其核心機(jī)制是結(jié)構(gòu)化狀態(tài)空間序列模型（S4）的最新演進(jìn)……SSSS

性能碾壓Transformer？

Mamba源自Albert Gu之前「結(jié)構(gòu)化狀態(tài)空間模型」的相關(guān)工作，可以看作是強(qiáng)大的循環(huán)運(yùn)算符。這就得以實(shí)現(xiàn)序列長(zhǎng)度的線性縮放和快速自回歸解碼。

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

然而，以前的遞歸模型的缺點(diǎn)是，它們的固定大小狀態(tài)難以壓縮上下文。

而Mamba的主要貢獻(xiàn)，就是引入了「選擇性SSM」，這是S4的簡(jiǎn)單泛化，可以選擇性地關(guān)注或忽略輸入。

這一小小的改變——只需讓某些參數(shù)成為輸入的函數(shù)——就能讓它立即解決對(duì)以往模型來(lái)說艱巨無(wú)比的任務(wù)。

例如，它可以無(wú)限長(zhǎng)地推斷出重要的「聯(lián)想回憶」任務(wù)的解決方案！(訓(xùn)練長(zhǎng)度256，測(cè)試長(zhǎng)度1M）

關(guān)鍵就在于：這一變化涉及到非同小可的效率權(quán)衡，S4的原始設(shè)計(jì)有著特定的原因。

在DNA和音頻等其他模態(tài)的真實(shí)數(shù)據(jù)上，Mamba的預(yù)訓(xùn)練性能超過了之前的專業(yè)基線（如HyenaDNA和SaShiMi）。

值得注意的是，無(wú)論在合成、DNA還是音頻數(shù)據(jù)中，隨著序列長(zhǎng)度達(dá)到1M+，Mamba的性能也在不斷提高！

而另一位一作Tri Dao介紹了如何利用硬件感知設(shè)計(jì)應(yīng)對(duì)這一挑戰(zhàn)，以及Mamba在語(yǔ)言方面的強(qiáng)大性能。

他表示，正如Albert所說，狀態(tài)空間模型（SSM）的特征，就是其固定大小的遞歸狀態(tài)。如果想實(shí)現(xiàn)更好的性能，就要求這種狀態(tài)更大，并且更具表現(xiàn)力。

不幸的是，因?yàn)檩^大的狀態(tài)太慢，會(huì)導(dǎo)致無(wú)法在實(shí)踐中使用遞歸進(jìn)行計(jì)算。

過去，曾有基于S4的SSM通過做出結(jié)構(gòu)假設(shè)（也即線性時(shí)間不變性）來(lái)解決這個(gè)問題，這樣就可以在不實(shí)現(xiàn)大狀態(tài)的情況下，進(jìn)行等效的「卷積模式」計(jì)算。

但這次CMU和普林斯頓研究者的方法是選擇性SSM，只能循環(huán)計(jì)算。

為了解決這個(gè)計(jì)算瓶頸，他們利用了其他高效的硬件感知算法（如FlashAttention）使用的技術(shù)。

需要注意的是，對(duì)于Mamba（和一般的SSM），這種方法只能在SRAM中實(shí)現(xiàn)擴(kuò)展?fàn)顟B(tài)，而不是在主存儲(chǔ)器中。

此外，scan實(shí)現(xiàn)比基本的PyTorch/JAX快30倍，當(dāng)序列長(zhǎng)度變長(zhǎng)時(shí)，比二次FlashAttention還要快幾個(gè)數(shù)量級(jí)。

而且，由于采用了固定大小的循環(huán)狀態(tài)（沒有KV緩存），Mamba的LM推理速度比Transformer快5倍。

從經(jīng)驗(yàn)上看，兩位作者取得的最重要的成果是在語(yǔ)言建模上，這也是以前的SSM所瞄準(zhǔn)的領(lǐng)域（比如H3，也即Mamba的前身）。

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

然而這時(shí)，自己的工作仍然不及Transformer。并且他表示，當(dāng)時(shí)沒有哪個(gè)模型能真正與精調(diào)后的Transformer相抗衡。

然而，驚喜忽然來(lái)了！

根據(jù)Chinchilla縮放定律進(jìn)行預(yù)訓(xùn)練時(shí)，Mamba的表現(xiàn)忽然就優(yōu)于一個(gè)非常強(qiáng)大的現(xiàn)代「Transformer++」模型（接近Llama模型）！

而在300B token上訓(xùn)練完成后，Mamba的性能，已經(jīng)大大優(yōu)于同類的開源模型。

最后，作者總結(jié)道：硬件感知思維可以開啟新的架構(gòu)設(shè)計(jì)。

展望未來(lái)，這種新架構(gòu)能否利用圍繞Transformers構(gòu)建的硬件/庫(kù)？它將如何改變其他領(lǐng)域（基因組學(xué)、音頻、視頻）的序列擴(kuò)展？

為此，作者還發(fā)布了一系列模型的權(quán)重（參數(shù)量最高可達(dá)2.8B，在300B token上訓(xùn)練），以及快速推理代碼。

項(xiàng)目地址：https://github.com/state-spaces/mamba

擊敗Transformer的架構(gòu)，是怎樣誕生的

現(xiàn)在的基礎(chǔ)模型，幾乎都是基于Transformer架構(gòu)和其中最核心的注意力模塊來(lái)構(gòu)建的。

為了解決Transformer在處理長(zhǎng)序列時(shí)的計(jì)算低效問題，學(xué)界開發(fā)了很多二次方時(shí)間復(fù)雜度的架構(gòu)，比如線性注意力、門控卷積和循環(huán)模型，以及結(jié)構(gòu)化狀態(tài)空間模型（SSM）。

然而，這些架構(gòu)在處理語(yǔ)言時(shí)，表現(xiàn)并不如傳統(tǒng)的注意力模型。

研究人員發(fā)現(xiàn)，這些模型的主要弱點(diǎn)在于它們難以進(jìn)行基于內(nèi)容的推理，并因此作出了幾項(xiàng)改進(jìn)：

首先，通過讓SSM參數(shù)成為輸入數(shù)據(jù)的函數(shù)，可以解決這類模型在處理離散數(shù)據(jù)類型時(shí)的不足。

這就使得模型能夠根據(jù)當(dāng)前的token在序列長(zhǎng)度的維度上選擇性地傳播或遺忘信息。

其次，盡管這樣的調(diào)整使得模型無(wú)法使用高效的卷積，但研究人員設(shè)計(jì)了一種適應(yīng)硬件的并行算法，并在循環(huán)模式下實(shí)現(xiàn)它。

研究人員將這種選擇性的SSM集成進(jìn)了一個(gè)簡(jiǎn)化的端到端神經(jīng)網(wǎng)絡(luò)架構(gòu)中，這種架構(gòu)不需要注意力機(jī)制，甚至也不需要MLP（多層感知器）模塊，這就是研究人員提出的Mamba。

Mamba在快速推理方面表現(xiàn)出色（比Transformers高5倍的處理速度），并且隨著序列長(zhǎng)度的增加，其性能線性增長(zhǎng)，在處理長(zhǎng)達(dá)百萬(wàn)長(zhǎng)度的序列時(shí)表現(xiàn)更佳。

作為一個(gè)通用的序列處理模型，Mamba在語(yǔ)言、音頻和基因組學(xué)等多個(gè)領(lǐng)域都獲得了最先進(jìn)的性能表現(xiàn)。

在語(yǔ)言建模方面，Mamba-3B模型在預(yù)訓(xùn)練和后續(xù)評(píng)估中性能達(dá)了兩倍參數(shù)量的Transformers模型性能。

通過實(shí)證研究，研究人員驗(yàn)證了Mamba在作為基礎(chǔ)模型（FM）的核心框架方面的巨大潛力。

這種潛力不僅體現(xiàn)在預(yù)訓(xùn)練的質(zhì)量上，還表現(xiàn)在特定領(lǐng)域任務(wù)的性能上，涵蓋了多種模態(tài)和環(huán)境：

- 合成任務(wù)

在重要的合成任務(wù)中，如復(fù)制和歸納等，Mamba不僅能輕松解決，還能推斷出無(wú)限長(zhǎng)（>100萬(wàn)個(gè)token）的解決方案。

- 音頻和基因組學(xué)

在音頻波形和DNA序列建模方面，Mamba的表現(xiàn)優(yōu)于SaShiMi、Hyena和Transformers等先前的SOTA模型，無(wú)論是在預(yù)訓(xùn)練質(zhì)量還是下游指標(biāo)方面（例如，在具有挑戰(zhàn)性的語(yǔ)音生成數(shù)據(jù)集上，F(xiàn)ID降低了一半以上）。

在這兩種情況下，它的性能隨著上下文長(zhǎng)度的增加而提高，最高可達(dá)百萬(wàn)長(zhǎng)度的序列。

- 語(yǔ)言建模

Mamba是首個(gè)線性時(shí)間序列模型，無(wú)論是在預(yù)訓(xùn)練復(fù)雜度還是在下游任務(wù)評(píng)估中，都能實(shí)現(xiàn)Transformer級(jí)別的性能。

將模型規(guī)模擴(kuò)大到10億參數(shù)后，研究人員證明Mamba的性能超過了Llama等大量基線模型。

Mamba語(yǔ)言模型與同體量的Transformer相比，具有5倍的生成吞吐量，而且Mamba-3B的質(zhì)量與兩倍于其規(guī)模的Transformer相當(dāng)（與Pythia-3B相比，常識(shí)推理的平均值高出4分，甚至超過了Pythia-7B）。

選擇性狀態(tài)空間模型

研究人員利用合成任務(wù)的直覺來(lái)激發(fā)他們的選擇機(jī)制，然后解釋如何將該機(jī)制合并到狀態(tài)空間模型中。由此產(chǎn)生的時(shí)變SSM無(wú)法使用卷積，這對(duì)如何有效地計(jì)算它們提出了技術(shù)挑戰(zhàn)。

研究人員通過利用現(xiàn)代硬件上的內(nèi)存層次結(jié)構(gòu)的硬件感知算法克服了這個(gè)問題。然后，研究人員描述了一個(gè)簡(jiǎn)單的SSM架構(gòu)，沒有注意力機(jī)制，甚至沒有MLP模塊。最后，研究人員討論選擇機(jī)制的一些附加屬性。

動(dòng)機(jī)：選擇作為壓縮手段

研究人員認(rèn)為序列建模的一個(gè)基本問題是將上下文壓縮成更小的狀態(tài)。他們從這個(gè)角度來(lái)看待流行序列模型的權(quán)衡（tradeoffs）。

例如，注意力在某些方面非常有效，但是在另一些方面又很低效，因?yàn)樗耆粔嚎s上下文。從這一點(diǎn)可以看出，自回歸推理需要顯式存儲(chǔ)整個(gè)上下文（即KV緩存），這直接導(dǎo)致Transformers的線性時(shí)間推理和二次時(shí)間訓(xùn)練緩慢。

另一方面，循環(huán)模型是高效的，因?yàn)樗麪顟B(tài)是有限的，這意味著推理時(shí)間是恒定的，并且訓(xùn)練的時(shí)間也將會(huì)是線性的。

然而，注意力的有效性受到這種狀態(tài)壓縮上下文的程度的限制。

為了理解這一原理，研究人員重點(diǎn)關(guān)注兩個(gè)合成任務(wù)的運(yùn)行示例（如下圖2）。

選擇性復(fù)制（Selective Copying）任務(wù)通過改變要記憶的標(biāo)記的位置來(lái)修改流行的復(fù)制任務(wù)。它需要內(nèi)容感知推理才能記住相關(guān)標(biāo)記（彩色）并過濾掉不相關(guān)標(biāo)記（白色）。

歸納頭（Induction Heads）任務(wù)是一種眾所周知的機(jī)制，以前的研究假設(shè)它可以解釋LLM的大多數(shù)情境學(xué)習(xí)能力。它需要上下文感知推理來(lái)知道何時(shí)在適當(dāng)?shù)纳舷挛模ê谏┲挟a(chǎn)生正確的輸出。

這些任務(wù)揭示了LTI模型的失效模式。從循環(huán)的角度來(lái)看，它們的恒定動(dòng)態(tài)（例如（2）中的（A，B）轉(zhuǎn)換）不能讓它們從上下文中選擇正確的信息，或者影響沿輸入相關(guān)的序列傳遞的隱藏狀態(tài)方式。

從卷積的角度來(lái)看，眾所周知，全局卷積可以解決普通復(fù)制任務(wù)，因?yàn)樗恍枰獣r(shí)間感知，但由于缺乏內(nèi)容意識(shí)，它們?cè)谶x擇性復(fù)制任務(wù)上有困難（如上圖）。

更具體地說，輸入到輸出之間的間距是變化的，并且不能通過靜態(tài)卷積核進(jìn)行建模。

總之，序列模型的效率與有效性權(quán)衡的特征在于它們壓縮狀態(tài)的程度：高效模型的狀態(tài)必須要小，而模型效果好必須要求這個(gè)小狀態(tài)要包含上下文中所有必要信息的狀態(tài)。

而相反，研究人員構(gòu)建的序列模型的基本原則是選擇性：或者是關(guān)注或過濾輸入到序列狀態(tài)的上下文感知能力。

特別是，選擇機(jī)制控制信息如何沿著序列維度傳播或交互。

通過選擇改進(jìn)SSM將選擇機(jī)制納入模型的一種方法是：讓影響序列交互的參數(shù)（例如 RNN 的循環(huán)動(dòng)態(tài)或 CNN 的卷積核）依賴于輸入。

算法1和2說明了研究者使用的主要選擇機(jī)制。

主要區(qū)別在于簡(jiǎn)單地使輸入的幾個(gè)參數(shù)Δ、B、C成為函數(shù)，以及整個(gè)張量形狀的相關(guān)更改。

需要注意，這些參數(shù)現(xiàn)在具有長(zhǎng)度維度 ，這意味著模型已從時(shí)不變（time-invariant）改為時(shí)變（time-varying）。

這就失去了與卷積的等價(jià)性，并影響了其效率。

簡(jiǎn)化的SSM架構(gòu)

與結(jié)構(gòu)化SSM一樣，選擇性SSM是獨(dú)立的序列轉(zhuǎn)換，可以「靈活地合并到神經(jīng)網(wǎng)絡(luò)中」。

H3架構(gòu)是最著名的SSM架構(gòu)的基礎(chǔ)，該架構(gòu)通常由受線性注意力啟發(fā)的塊與 MLP（多層感知器）塊交織組成。研究人員通過將這兩個(gè)組同質(zhì)堆疊件合并為一個(gè)組件來(lái)簡(jiǎn)化這一架構(gòu)（如下圖）。

之所以這么處理是受到門控注意力單元（GAU）的啟發(fā)。該架構(gòu)涉及通過可控?cái)U(kuò)展因子來(lái)擴(kuò)展模型維度。對(duì)于每個(gè)塊，大多數(shù)參數(shù)（3ED^2）位于線性投影中，而內(nèi)部SSM貢獻(xiàn)較少。SSM參數(shù)的數(shù)量相比起來(lái)要小的多。

研究人員重復(fù)了這個(gè)塊，與標(biāo)準(zhǔn)標(biāo)準(zhǔn)化和殘差連接交織，形成Mamba架構(gòu)。

在實(shí)驗(yàn)中，始終將x設(shè)為E=2，并使用塊的兩個(gè)堆棧來(lái)匹配Transformer交錯(cuò)MHA（多頭注意力）和MLP塊的122個(gè)參數(shù)。

研究人員使用SiLU / Swish激活函數(shù)，其動(dòng)機(jī)是使門控 MLP 成為流行的「SwiGLU」變體 。最后，研究人員還使用了一個(gè)可選的歸一化層，動(dòng)機(jī)是RetNet在類似位置使用歸一層。

選擇機(jī)制是一個(gè)更廣泛的概念，可以以不同的方式應(yīng)用，例如更傳統(tǒng)的RNN或CNN、不同的參數(shù)（例如算法2中的 A），或使用不同的變換。

實(shí)證評(píng)估

合成任務(wù)：選擇性復(fù)制

復(fù)制任務(wù)是用來(lái)測(cè)試序列模型，特別是循環(huán)模型記憶能力的經(jīng)典合成任務(wù)。

LTI SSM（線性遞歸和全局卷積）可以通過只關(guān)注時(shí)間而不是推理數(shù)據(jù)輕松地解決這個(gè)任務(wù)。例如，構(gòu)建一個(gè)長(zhǎng)度完全正確的卷積核（圖2）。

對(duì)此，選擇性復(fù)制任務(wù)則可以通過隨機(jī)改變token的間距，來(lái)阻止這種走捷徑的方法。

表1顯示，H3和Mamba等門控架構(gòu)只能部分提升性能，而選擇機(jī)制（即將S4改進(jìn)為S6）則可以輕松解決這一問題，尤其是與更強(qiáng)大的架構(gòu)相結(jié)合時(shí)。

合成任務(wù)：歸納頭

歸納頭是一個(gè)從機(jī)械可解釋性的角度出發(fā)相對(duì)簡(jiǎn)單的任務(wù)，卻意外地能夠預(yù)測(cè)大語(yǔ)言模型（LLMs）的上下文學(xué)習(xí)能力。

這項(xiàng)任務(wù)要求模型進(jìn)行關(guān)聯(lián)性回憶和復(fù)制動(dòng)作：比如，模型之前在一個(gè)序列中遇到過「Harry Potter」這樣的詞組，那么當(dāng)「Harry」再次出現(xiàn)在同一個(gè)序列時(shí)，模型應(yīng)能夠通過回顧歷史信息并預(yù)測(cè)出「Potter」。

表2顯示，Mamba模型，或者更準(zhǔn)確地說是它的選擇性SSM層，由于能夠選擇性地記住相關(guān)的token，同時(shí)忽略中間其他的token，因此能夠完美地完成任務(wù)。

并且，它還能完美地泛化到百萬(wàn)長(zhǎng)度的序列，也就是訓(xùn)練期間遇到的長(zhǎng)度的4000倍。相比之下，其他方法的泛化能力都無(wú)法超過2倍。

語(yǔ)言建模

研究人員將Mamba與標(biāo)準(zhǔn)的Transformer架構(gòu)（即GPT-3架構(gòu)），以及目前最先進(jìn)的Transformer（Transformer++）進(jìn)行了對(duì)比。

后者基于PaLM和LLaMa架構(gòu)，其特點(diǎn)包括旋轉(zhuǎn)嵌入（rotary embedding）、SwiGLU MLP、使用RMSNorm替換LayerNorm、取消線性偏置，并采用更高的學(xué)習(xí)率。

圖4顯示，在從≈1.25億到≈13億的參數(shù)規(guī)模中，Mamba是首個(gè)在性能上媲美最強(qiáng)Transformer架構(gòu)（Transformer++）的無(wú)注意力模型。

表3展示了Mamba在一系列下游zero-shot評(píng)估任務(wù)中的表現(xiàn)。

其中，Mamba在訓(xùn)練時(shí)使用了與Pythia和RWKV相同的tokenizer、數(shù)據(jù)集和訓(xùn)練長(zhǎng)度（300B token）。

需要注意的是，Mamba和Pythia訓(xùn)練時(shí)的上下文長(zhǎng)度為2048，而RWKV為1024。

DNA建模

隨著大語(yǔ)言模型的成功，人們開始探索將基礎(chǔ)模型的范式應(yīng)用于基因組學(xué)。

DNA由具有特定詞匯表的離散符號(hào)序列組成，還需要長(zhǎng)程依賴關(guān)系來(lái)建模，因此被比作語(yǔ)言。

研究者將Mamba作為預(yù)訓(xùn)練和微調(diào)的FM骨干進(jìn)行了研究，研究背景與最近DNA長(zhǎng)序列模型的研究相同。

在預(yù)訓(xùn)練方面，研究者基本上按照標(biāo)準(zhǔn)的因果語(yǔ)言建模（下一個(gè)token預(yù)測(cè)）設(shè)置。

在數(shù)據(jù)集方面，基本沿用了鬣狗DNA的設(shè)置，它使用了HG38數(shù)據(jù)集進(jìn)行預(yù)訓(xùn)練，該數(shù)據(jù)集由單個(gè)人類基因組組成，在訓(xùn)練分割中包含約45億個(gè)token（DNA堿基對(duì)）。

結(jié)果如圖5（左）顯示，Mamba的預(yù)訓(xùn)練困惑度隨著模型規(guī)模的增大而平穩(wěn)提高，并且Mamba的擴(kuò)展能力優(yōu)于 HvenaDNA和Transformer++。

例如，在最大模型規(guī)?！?0M參數(shù)時(shí)，曲線顯示，Mamba可以用少3到4倍的參數(shù)，與Transformer++和HvenaDNA模型相媲美。

另外，圖5（右）顯示，Mamba能夠利用更長(zhǎng)的上下文，甚至長(zhǎng)達(dá)1M的極長(zhǎng)序列，并且其預(yù)訓(xùn)練困惑度會(huì)隨著上下文的增加而提高。

另一方面，鬣狗DNA模型會(huì)隨著序列長(zhǎng)度的增加而變差。

從卷積的角度看，一個(gè)非常長(zhǎng)的卷積核正在聚合一個(gè)長(zhǎng)序列上的所有信息。

圖6是類人猿DNA的分類，顯示了使用相同上下文長(zhǎng)度的預(yù)訓(xùn)練模型對(duì)長(zhǎng)度2^10到2^20的序列進(jìn)行微調(diào)后的準(zhǔn)確度。

音頻建模與生成

在音頻波形處理領(lǐng)域，主要對(duì)比的是SaShiMi架構(gòu)。該模型包括：

1. 一個(gè)U-Net主干，通過兩個(gè)階段的池化操作，其中每個(gè)階段都將模型的維度D增加一倍，池化因子為p，

2. 每個(gè)階段都交替使用S4和MLP模塊。

針對(duì)長(zhǎng)上下文的自回歸式預(yù)訓(xùn)練，研究人員采用了標(biāo)準(zhǔn)鋼琴音樂數(shù)據(jù)集——YouTubeMix進(jìn)行評(píng)估。數(shù)據(jù)集包含了4小時(shí)的獨(dú)奏鋼琴音樂，采樣率為16000Hz。

圖7展示了在保持計(jì)算量不變的情況下，訓(xùn)練序列長(zhǎng)度從8192（2^13）增加到≈1000000（2^20）時(shí)的效果。

無(wú)論是Mamba還是SaShiMi（S4+MLP）基線模型，表現(xiàn)都隨著上下文長(zhǎng)度的增加而穩(wěn)步提升。其中，Mamba在整個(gè)過程中都更勝一籌，而且序列越長(zhǎng)優(yōu)勢(shì)越明顯。

在自回歸語(yǔ)音生成方面，則使用基準(zhǔn)語(yǔ)音生成數(shù)據(jù)集SC09進(jìn)行評(píng)估。它由時(shí)長(zhǎng)1秒的語(yǔ)音片段組成，采樣頻率為16000 Hz，包含數(shù)字「0」到「9」，特征多變。

表4展示了Mamba-UNet與一系列基準(zhǔn)模型的自動(dòng)評(píng)估結(jié)果，其中包括WaveNet、SampleRNN、WaveGAN、DiffWave以及SaShiMi。

可以看到，小規(guī)模的Mamba模型在性能上就已經(jīng)超越了那些更大、采用了最先進(jìn)的基于GAN和擴(kuò)散技術(shù)的模型。而同等參數(shù)規(guī)模的Mamba模型，在保真度方面的表現(xiàn)更是大幅領(lǐng)先。

表5采用的是小規(guī)模Mamba模型，并探究了在外部和中心階段不同架構(gòu)的組合效果。

研究發(fā)現(xiàn)，無(wú)論在外部塊還是中心塊，Mamba模型的表現(xiàn)都優(yōu)于S4+MLP架構(gòu)，而在中心塊的性能排名為Mamba > S4+MLP > MHA+MLP。

速度和顯存基準(zhǔn)測(cè)試

圖8展示了scan操作（狀態(tài)擴(kuò)展N = 16）速度，以及Mamba端到端推理吞吐量的基準(zhǔn)測(cè)試。

結(jié)果顯示，當(dāng)序列長(zhǎng)度超過2k時(shí)，高效的SSM scan比目前最優(yōu)秀的注意力機(jī)制——FlashAttention-2還要快。而且，比起PyTorch標(biāo)準(zhǔn)的scan實(shí)現(xiàn)，速度提升更是高達(dá)20到40倍。

由于沒有鍵值（KV）緩存，因此Mamba可以支持更大的批處理大小，從而使推理吞吐量比同等規(guī)模Transformer高了4到5倍。

舉個(gè)例子，一個(gè)未經(jīng)訓(xùn)練的69億參數(shù)的Mamba（Mamba-6.9B），在推理處理能力上可以超過僅有13億參數(shù)、規(guī)模小5倍的Transformer模型。

與大多數(shù)深度序列模型一樣，顯存使用量與激活張量的大小成正比。表15顯示，Mamba的顯存需求與經(jīng)過優(yōu)化的Transformer相當(dāng)。

125M模型在單張A100 80GB GPU上訓(xùn)練時(shí)顯存的需求

在論文最后，作者表示，選擇性狀態(tài)空間模型在為不同領(lǐng)域構(gòu)建基礎(chǔ)模的廣泛應(yīng)用性，太令人興奮了。

種種實(shí)驗(yàn)結(jié)果表明，Mamba很有可能成為通用序列模型的主流框架，甚至有潛力跟Transformer一搏。