現在 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
— 完 —
點這裏關注我，記得标星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
