來自FlashAttention作者
現在ChatGPT等大模型一大痛點:
處理長文本算力消耗巨大
,背後原因是
Transformer架構中注意力機制的二次複雜度
。
FlashAttention作者Tri Dao
參與提出的新架構,成爲有力挑戰者,引起大量關注:
曼巴
(曼巴,一種蛇),在語言任務上擊敗/匹配Transformer性能,具有
線性複雜度
和
5倍推理吞吐量
。
具體來說,Mamba在語言、音頻、DNA序列模態上都實現SOTA。
在最受關注的語言任務上,
Mamba-3B超越同等規模的Transformer,與兩倍大的Transformer匹敵
。
並且相關代碼、預訓練模型checkpoint都已開源。
兩位作者的解讀都獲得大量轉發。
有網友發現,連在線預測平臺上的“Transformer在2027年還是SOTA嗎?”都在這一天出現明顯下降。
有選擇處理信息+硬件感知算法。
Mamba是一種
狀態空間模型
(SSM,狀態空間模型)。
建立在更現代的適用於深度學習的
結構化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縮放定律下預訓練時,語言任務優於同類開源模型。
對比對象中的
變壓器+ +
爲標準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é。
艾伯特固
現在是CMU助理教授,多年來一直推動SSM架構發展。
他曾在DeepMind 工作,目前是
Cartesia的
的聯合創始人及首席科學家。
三刀
,以FlashAttention、FlashDecoding系列工作聞名,現在是普林斯頓助理教授,和Together AI首席科學家,也在Cartesia AI擔任顧問。
Cartesia AI公司介紹中提到致力於基於新架構構建下一代基礎模型,現在看來主要就是指創新的SSM架構。
聯創及CEO
卡蘭高爾
同爲斯坦福博士畢業,也是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