Transformer挑戰(zhàn)者出現(xiàn)！斯坦福CMU聯(lián)合團(tuán)隊(duì)，模型代碼都開(kāi)源

夢(mèng)晨 2023-12-05 13:32:12 來(lái)源：量子位

夢(mèng)晨 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

現(xiàn)在ChatGPT等大模型一大痛點(diǎn)：

處理長(zhǎng)文本算力消耗巨大，背后原因是Transformer架構(gòu)中注意力機(jī)制的二次復(fù)雜度。

FlashAttention作者Tri Dao參與提出的新架構(gòu)，成為有力挑戰(zhàn)者，引起大量關(guān)注：

Mamba（曼巴，一種蛇），在語(yǔ)言任務(wù)上擊敗/匹配Transformer性能，具有線(xiàn)性復(fù)雜度和5倍推理吞吐量。

具體來(lái)說(shuō)，Mamba在語(yǔ)言、音頻、DNA序列模態(tài)上都實(shí)現(xiàn)SOTA。

在最受關(guān)注的語(yǔ)言任務(wù)上，Mamba-3B超越同等規(guī)模的Transformer，與兩倍大的Transformer匹敵。

并且相關(guān)代碼、預(yù)訓(xùn)練模型checkpoint都已開(kāi)源。

兩位作者的解讀都獲得大量轉(zhuǎn)發(fā)。

有網(wǎng)友發(fā)現(xiàn)，連在線(xiàn)預(yù)測(cè)平臺(tái)上的“Transformer在2027年還是SOTA嗎？”都在這一天出現(xiàn)明顯下降。

有選擇處理信息+硬件感知算法。

Mamba是一種狀態(tài)空間模型（SSM，State Space Model）。

建立在更現(xiàn)代的適用于深度學(xué)習(xí)的結(jié)構(gòu)化SSM（S4, Structured SSM）基礎(chǔ)上，與經(jīng)典架構(gòu)RNN有相似之處。

在先前研究的Mamba主要有三點(diǎn)創(chuàng)新：

• 對(duì)輸入信息有選擇性處理
• 硬件感知的算法
• 更簡(jiǎn)單的架構(gòu)

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

作者認(rèn)為，序列建模的一個(gè)基礎(chǔ)問(wèn)題是把上下文壓縮成更小的狀態(tài)。

從這個(gè)角度來(lái)看，注意力機(jī)制雖然高性能但低效率，需要顯式地存儲(chǔ)整個(gè)上下文（也就是KV緩存），直接導(dǎo)致訓(xùn)練和推理消耗算力大。

類(lèi)RNN的循環(huán)神經(jīng)網(wǎng)絡(luò)具有有限的狀態(tài)，高效，但性能受到對(duì)上下文壓縮程度的限制。

Mamba的解決辦法，是讓模型對(duì)信息有選擇性處理，可以關(guān)注或忽略傳入的內(nèi)容，即使?fàn)顟B(tài)大小固定也能壓縮上下文。

一個(gè)直觀的類(lèi)比：

Transformer就像人類(lèi)每寫(xiě)一個(gè)字之前，都把前面的所有字+輸入都復(fù)習(xí)一遍，所以寫(xiě)的慢。

RNN每次只參考前面固定的字?jǐn)?shù)，寫(xiě)的快，但容易忘掉更前面的內(nèi)容。

Mamba每次參考前面所有內(nèi)容的一個(gè)概括，越往后寫(xiě)對(duì)前面內(nèi)容概括得越狠，丟掉細(xì)節(jié)保留大意。

在其前身結(jié)構(gòu)化狀態(tài)空間模型（S4）中，四個(gè)參數(shù)A、B、C、?都是固定的，不隨輸入變化。

在Mamaba中，作者讓這些參數(shù)B、C、?成為輸入的函數(shù)，讓模型能夠根據(jù)輸入內(nèi)容自適應(yīng)地調(diào)整其行為。

硬件感知的狀態(tài)擴(kuò)展

為了讓古老的SSM在現(xiàn)代GPU上也能高效計(jì)算，Mamba中使用了FlashAttention同款技術(shù)。

核心思想是利用內(nèi)存的不同層級(jí)結(jié)構(gòu)處理SSM的狀態(tài)，減少高帶寬但慢速的HBM內(nèi)存反復(fù)讀寫(xiě)這個(gè)瓶頸，具體來(lái)說(shuō)：

• 在更高速的SRAM內(nèi)存中執(zhí)行離散化和遞歸操作，再將輸出寫(xiě)回HBM。
• 通過(guò)并行掃描算法實(shí)現(xiàn)并行化。
• 當(dāng)輸入從HBM加載到SRAM時(shí)，中間狀態(tài)不被保存，而是在反向傳播中重新計(jì)算。

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

將大多數(shù)SSM架構(gòu)的基礎(chǔ)塊，與現(xiàn)代神經(jīng)網(wǎng)絡(luò)中普遍存在的門(mén)控MLP相結(jié)合，組成新的Mamba塊。

重復(fù)這個(gè)塊，與歸一化和殘差連接結(jié)合，構(gòu)成Mamba架構(gòu)。

實(shí)驗(yàn)結(jié)果

Mamba在Chinchilla縮放定律下預(yù)訓(xùn)練時(shí)，語(yǔ)言任務(wù)優(yōu)于同類(lèi)開(kāi)源模型。

對(duì)比對(duì)象中的Transformer++為標(biāo)準(zhǔn)GPT-3架構(gòu)加上谷歌PaLM和Meta Llama中的改進(jìn)方案，也就是已知最強(qiáng)Transformer配方。

下游任務(wù)上，每個(gè)規(guī)模尺寸的Mamba都是同類(lèi)最佳，并且通常與兩倍規(guī)模的基線(xiàn)性能匹配。

特別是當(dāng)序列長(zhǎng)度增加到512k時(shí)，相比使用FlashAttention-2的Transformer快幾個(gè)數(shù)量級(jí)，而且不會(huì)內(nèi)存不足。

Transformer的下一步？

最終，Mamba是第一個(gè)真正實(shí)現(xiàn)匹配Transformer性能的線(xiàn)性時(shí)間序列模型，無(wú)論是在預(yù)訓(xùn)練困惑度還是下游任務(wù)評(píng)估方面。

并且在音頻和DNA序列建模上也優(yōu)于之前的SOTA模型，表現(xiàn)出一定的通用性。

作者在結(jié)論中提出，Mamba是通用序列模型骨干的有力候選者。

Stability AI創(chuàng)始人當(dāng)即表示關(guān)注。

英偉達(dá)科學(xué)家Jim Fan也對(duì)Transformer的挑戰(zhàn)者出現(xiàn)感到興奮。

論文兩位作者Albert Gu和Tri Dao，博士都畢業(yè)于斯坦福大學(xué)，導(dǎo)師為Christopher Ré。

Albert Gu現(xiàn)在是CMU助理教授，多年來(lái)一直推動(dòng)SSM架構(gòu)發(fā)展。

他曾在DeepMind 工作，目前是Cartesia AI的聯(lián)合創(chuàng)始人及首席科學(xué)家。

Tri Dao，以FlashAttention、FlashDecoding系列工作聞名，現(xiàn)在是普林斯頓助理教授，和Together AI首席科學(xué)家，也在Cartesia AI擔(dān)任顧問(wèn)。

Cartesia AI公司介紹中提到致力于基于新架構(gòu)構(gòu)建下一代基礎(chǔ)模型，現(xiàn)在看來(lái)主要就是指創(chuàng)新的SSM架構(gòu)。

聯(lián)創(chuàng)及CEO?Karan Goel同為斯坦福博士畢業(yè)，也是Mamba的前身S4論文作者之一。

對(duì)于Mamba的下一步，在論文中有提到“探索新架構(gòu)是否能適用于Transformer已建立起的豐富大模型生態(tài)”。

其中包括微調(diào)、自適應(yīng)、提示學(xué)習(xí)、上下文學(xué)習(xí)、指令微調(diào)、RLHF、量化……也就是要把基礎(chǔ)模型發(fā)展成GPT-3.5、Llama同類(lèi)的助手模型了。

但作者也提到，目前的實(shí)驗(yàn)規(guī)模較小，要全面評(píng)估SSM是否能與Transformer和其他架構(gòu)如RWKV、微軟RetNet競(jìng)爭(zhēng)，至少還需要驗(yàn)證7B規(guī)模。

在擴(kuò)展SSM的過(guò)程中，還會(huì)遇到新的工程挑戰(zhàn)和對(duì)模型的調(diào)整，論文中沒(méi)有涉及。

最后，Albert Gu還分享了為什么把新架構(gòu)起名為一種毒蛇的名字：

速度快、對(duì)序列建模問(wèn)題很致命、前身S4是SSSS（嘶嘶嘶嘶）。

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

參考鏈接：
[1]https://twitter.com/_albertgu/status/1731727672286294400
[2]https://twitter.com/_albertgu/status/1731727672286294400