Transformer挑戰者出現！斯坦福CMU聯合團隊，開源模型及代碼_風聞

　　夢晨 發自 凹非寺

　　量子位 | 公眾號 QbitAI

　　現在ChatGPT等大模型一大痛點：

　　處理長文本算力消耗巨大，背後原因是Transformer架構中注意力機制的二次複雜度。

　　FlashAttention作者Tri Dao參與提出的新架構，成為有力挑戰者，引起大量關注：

　　Mamba（曼巴，一種蛇），在語言任務上擊敗/匹配Transformer性能，具有線性複雜度和5倍推理吞吐量。

　　具體來説，Mamba在語言、音頻、DNA序列模態上都實現SOTA。

　　在最受關注的語言任務上，Mamba-3B超越同等規模的Transformer，與兩倍大的Transformer匹敵。

　　並且相關代碼、預訓練模型checkpoint都已開源。

　　兩位作者的解讀都獲得大量轉發。

　　有網友發現，連在線預測平台上的“Transformer在2027年還是SOTA嗎？”都在這一天出現明顯下降。

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

　　Mamba是一種狀態空間模型（SSM，State Space Model）。

　　建立在更現代的適用於深度學習的結構化SSM（S4, Structured SSM）基礎上，與經典架構RNN有相似之處。

　　在先前研究的Mamba主要有三點創新：

　　對輸入信息有選擇性處理

　　硬件感知的算法

　　更簡單的架構

　　選擇性狀態空間模型

　　作者認為，序列建模的一個基礎問題是把上下文壓縮成更小的狀態。

　　從這個角度來看，注意力機制雖然高性能但低效率，需要顯式地存儲整個上下文（也就是KV緩存），直接導致訓練和推理消耗算力大。

　　類RNN的循環神經網絡具有有限的狀態，高效，但性能受到對上下文壓縮程度的限制。

　　Mamba的解決辦法，是讓模型對信息有選擇性處理，可以關注或忽略傳入的內容，即使狀態大小固定也能壓縮上下文。

　　一個直觀的類比：

　　Transformer就像人類每寫一個字之前，都把前面的所有字+輸入都複習一遍，所以寫的慢。

　　RNN每次只參考前面固定的字數，寫的快，但容易忘掉更前面的內容。

　　Mamba每次參考前面所有內容的一個概括，越往後寫對前面內容概括得越狠，丟掉細節保留大意。

　　在其前身結構化狀態空間模型（S4）中，四個參數A、B、C、∆都是固定的，不隨輸入變化。

　　在Mamaba中，作者讓這些參數B、C、∆成為輸入的函數，讓模型能夠根據輸入內容自適應地調整其行為。

　　硬件感知的狀態擴展

　　為了讓古老的SSM在現代GPU上也能高效計算，Mamba中使用了FlashAttention同款技術。

　　核心思想是利用內存的不同層級結構處理SSM的狀態，減少高帶寬但慢速的HBM內存反覆讀寫這個瓶頸，具體來説：

　　在更高速的SRAM內存中執行離散化和遞歸操作，再將輸出寫回HBM。

　　通過並行掃描算法實現並行化。

　　當輸入從HBM加載到SRAM時，中間狀態不被保存，而是在反向傳播中重新計算。

　　簡化的SSM架構

　　將大多數SSM架構的基礎塊，與現代神經網絡中普遍存在的門控MLP相結合，組成新的Mamba塊。

　　重複這個塊，與歸一化和殘差連接結合，構成Mamba架構。

　　實驗結果

　　Mamba在Chinchilla縮放定律下預訓練時，語言任務優於同類開源模型。

　　對比對象中的**Transformer++**為標準GPT-3架構加上谷歌PaLM和Meta Llama中的改進方案，也就是已知最強Transformer配方。

　　下游任務上，每個規模尺寸的Mamba都是同類最佳，並且通常與兩倍規模的基線性能匹配。

　　特別是當序列長度增加到512k時，相比使用FlashAttention-2的Transformer快幾個數量級，而且不會內存不足。

　　Transformer的下一步？

　　最終，Mamba是第一個真正實現匹配Transformer性能的線性時間序列模型，無論是在預訓練困惑度還是下游任務評估方面。

　　並且在音頻和DNA序列建模上也優於之前的SOTA模型，表現出一定的通用性。

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

　　Stability AI創始人當即表示關注。

　　英偉達科學家Jim Fan也對Transformer的挑戰者出現感到興奮。

　　論文兩位作者Albert Gu和Tri Dao，博士都畢業於斯坦福大學，導師為Christopher Ré。

　　Albert Gu現在是CMU助理教授，多年來一直推動SSM架構發展。

　　他曾在DeepMind 工作，目前是Cartesia AI的聯合創始人及首席科學家。

　　Tri Dao，以FlashAttention、FlashDecoding系列工作聞名，現在是普林斯頓助理教授，和Together AI首席科學家，也在Cartesia AI擔任顧問。

　　Cartesia AI公司介紹中提到致力於基於新架構構建下一代基礎模型，現在看來主要就是指創新的SSM架構。

　　聯創及CEO Karan Goel同為斯坦福博士畢業，也是Mamba的前身S4論文作者之一。

　　對於Mamba的下一步，在論文中有提到“探索新架構是否能適用於Transformer已建立起的豐富大模型生態”。

　　其中包括微調、自適應、提示學習、上下文學習、指令微調、RLHF、量化……也就是要把基礎模型發展成GPT-3.5、Llama同類的助手模型了。

　　但作者也提到，目前的實驗規模較小，要全面評估SSM是否能與Transformer和其他架構如RWKV、微軟RetNet競爭，至少還需要驗證7B規模。

　　在擴展SSM的過程中，還會遇到新的工程挑戰和對模型的調整，論文中沒有涉及。

　　最後，Albert Gu還分享了為什麼把新架構起名為一種毒蛇的名字：

　　速度快、對序列建模問題很致命、前身S4是SSSS（嘶嘶嘶嘶）。

　　論文：

　　https://arxiv.org/abs/2312.00752

　　參考鏈接：

　　[1]https://twitter.com/_albertgu/status/1731727672286294400

　　[2]https://twitter.com/_albertgu/status/1731727672286294400