相信很多大家對Mamba-2新架構(gòu)出世一統(tǒng)江湖！普林斯頓CMU華人再出神作，性能狂飆8倍還不知道吧，今天菲菲就帶你們一起去了解一下~.~！

【新智元導(dǎo)讀】在開源社區(qū)引起「海嘯」的Mamba架構(gòu)，再次卷土重來!這次，Mamba-2順利拿下ICML。通過統(tǒng)一SSM和注意力機制，Transformer和SSM直接成了「一家親」，Mamba-2這是要一統(tǒng)江湖了?

年前，Mamba被頂會ICLR拒稿的消息曾引起軒然大波。

甚至有研究人員表示:如果這種工作都被拒了，那我們這些「小丑」要怎么辦?

這次，新一代的Mamba-2卷土重來、再戰(zhàn)頂會，順利拿下了ICML2024!

仍是前作的兩位大佬（換了個順序），仍是熟悉的配方:

論文地址:https://arxiv.org/pdf/2405.21060

開源代碼和模型權(quán)重:https://github.com/state-spaces/mamba

不同的是，作者在更高的視角上，統(tǒng)一了狀態(tài)空間模型（SSM）和注意力機制(Attention)，也就是文章標(biāo)題所說的「Transformers are SSMs」。

——這下咱們都是一家人了，不用動不動就「打生打死」了。

性能方面，Mamba-2采用了新的算法（SSD），比前代提速2-8倍，對比FlashAttention-2也不遑多讓，在序列長度為2K時持平，之后便一路遙遙領(lǐng)先。

在Pile上使用300B token訓(xùn)練出的Mamba-2-2.7B，性能優(yōu)于在同一數(shù)據(jù)集上訓(xùn)練的Mamba-2.8B、Pythia-2.8B，甚至是更大的Pythia-6.9B。

從理論上整合了SSM和Transformer，同等性能下，模型更小，消耗更低，速度更快。

更重要的是，能夠利用GPU的硬件資源（矩陣乘法單元），以及針對Transformer的一系列優(yōu)化。

——Mamba-2大有一統(tǒng)江湖之勢。

1代Mamba，爆發(fā)式占領(lǐng)AI社區(qū)

事實上，關(guān)于1代Mamba的各種研究一直在爆發(fā)性地增長，arxiv已經(jīng)被各種Mamba所占領(lǐng)，谷歌學(xué)術(shù)的引用量也達到了350多。

后續(xù)工作如雨后春筍一般冒出，包括視覺、基因組學(xué)、圖表等的直接應(yīng)用，以及回憶能力、上下文學(xué)習(xí)能力、形式語言表達能力等方面的研究。

作者興奮地表示:「我們多年來一直在追求的高效序列模型研究路線，真正引起了機器學(xué)習(xí)社區(qū)的共鳴。」

唯一遺憾的是，Mamba遭到ICLR拒稿，所以關(guān)于Mamba到底有沒有前途這個事也就被打上了問號。

現(xiàn)在，問題解決了，不但論文被接收了，而且還證明了Transformer和Mamba其實是一家人——

「你說我不行?那Transformer到底行不行?」

值得注意的是，之前很火的Vision Mamba以及另一篇關(guān)于Mamba的研究也殺入了ICML2024。

對于改進Mamba的初衷，作者表示，當(dāng)前AI社區(qū)的大家都在努力解決Transformer的問題，盡管SSM的特性和效果都相當(dāng)好，但卻跟社區(qū)的努力方向不一致。

這次的Mamba-2可以把針對Transformer的優(yōu)化都用上，不浪費大家的努力。

新架構(gòu)一統(tǒng)江湖

在介紹新架構(gòu)之前，小編先幫大家簡單理一下背景。

狀態(tài)空間模型SSM之所以如此令人著迷，是因為它們顯得如此之「基礎(chǔ)」。

比如，它們與序列模型的許多主要范式，都有著豐富的聯(lián)系。

它們似乎抓住了連續(xù)、卷積和循環(huán)序列模型的本質(zhì)，把所有這些元素都包含在了一個簡單優(yōu)雅的模型里。

不過，另一個主要的序列模型范式——注意力機制的變體，卻更加無所不在。

然而SSM卻總感覺和Attention是脫節(jié)的。

在這里，研究者們發(fā)出了「靈魂拷問」——SSM和注意力之間的概念聯(lián)系是什么?有無可能將二者結(jié)合起來?

那就要從公式說起了。

狀態(tài)空間模型SSM可以這么定義:

這是個微分方程，利用導(dǎo)數(shù)定義進行代換:

可以得到SSM的解:

這個東西就跟RNN一毛一樣了:

所以可以認(rèn)為SSM等價于RNN。

如果將RNN的遞歸結(jié)構(gòu)展開，那么它又可以等價于卷積:

此時，便可以利用卷積的特性進行并行訓(xùn)練，而進行推理時又可以享受RNN帶來的O（1）復(fù)雜度。

當(dāng)然，好事不能讓你全占了，這種結(jié)構(gòu)仍然逃不過固有的梯度爆炸（或消失），以及難以勝任選擇性復(fù)制和上下文學(xué)習(xí)等任務(wù)。

為此，Mamba在SSM的基礎(chǔ)上加入了能夠隨輸入變化的參數(shù)。

不過這樣做的代價是失去了固定kernel帶來的并行性，所以作者另辟蹊徑，使用前綴和的方式來加速RNN的訓(xùn)練。

不過，從計算角度來看，Mamba在硬件效率上仍然遠(yuǎn)不如注意力機制。

原因在于，目前常用的GPU、TPU等加速器，是為矩陣乘法進行過專門優(yōu)化的。

1代Mamba吃不到硬件矩陣運算單元的紅利，盡管推理時有速度優(yōu)勢，但訓(xùn)練時問題就大了。

所以作者就想，我能不能把Mamba的計算重構(gòu)成矩陣乘法呢?

于是，新一代的Mamba誕生了。

結(jié)構(gòu)化狀態(tài)空間對偶性:SSD

Mamba-2的核心，是結(jié)構(gòu)化狀態(tài)空間對偶性（State Space Duality，SSD）的概念:

1.SSD模型指的是一個特定的獨立層，比如注意力層或狀態(tài)空間模型（SSM），可以被整合到深度神經(jīng)網(wǎng)絡(luò)中;

2.SSD框架是一個用于推理該模型（以及更多理論連接）的通用框架;

3.SSD算法是一種比以前的SSM更高效地計算SSD層的算法。

SSD框架（紅色，藍(lán)色）:狀態(tài)空間模型(即半分離矩陣)和結(jié)構(gòu)化掩碼注意力涵蓋了大量高效的序列模型。它們的交集就是SSD模型(紫色)

原始的Mamba（或更準(zhǔn)確地說，其核心「S6」層）實際上是一個具有對角結(jié)構(gòu)的選擇性狀態(tài)空間模型(SSM)。

Mamba-2的SSD層只做了一個小改動:它進一步限制了對角矩陣

　　免責(zé)聲明：本文由用戶上傳，與本網(wǎng)站立場無關(guān)。財經(jīng)信息僅供讀者參考，并不構(gòu)成投資建議。投資者據(jù)此操作，風(fēng)險自擔(dān)。 如有侵權(quán)請聯(lián)系刪除！