Mamba: Linear-Time Sequence Modeling with Selective State Spaces

在現代深度學習的大規模基礎模型中，Transformer架構憑藉其強大的注意力機制，成為了語言處理、語音、視覺等多種序列任務的核心。但Transformer的計算複雜度是二次方（quadratic）隨序列長度增長，導致當輸入序列極長時計算成本和記憶體消耗巨大，阻礙了其在如全文本檢索、基因組學、長音頻等長序列處理上的應用。為了解決此瓶頸，先前研究提出多種亞二次方時間（subquadratic-time）方法，包括線性注意力（linear attention）、門控卷積（gated convolution）、循環模型與結構化狀態空間模型（Structured State Space Models, SSMs）等，但這些替代方案普遍在處理重要模態——尤其是自然語言時，難以匹敵 Transformer 的性能，成為推廣應用的瓶頸。

在本篇於 ICLR 2024 獲得 Outstanding Paper 獎項的論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》中，Gu Dao 則提出了一種嶄新的結構化狀態空間模型（SSM）變形，稱為「Selective State Spaces」，同時設計專為此架構量身打造的硬體友好型並行演算法，並將其整合入一個無須傳統Transformer注意力或甚至MLP結構的簡化神經網路框架Mamba。該方法突破性地達成線性時間複雜度，並在多種序列模態上取得超越大型Transformer的表現。

研究動機與背景

Transformer 雖然是現今序列建模的主流架構，關鍵瓶頸在於計算與記憶體使用呈現 O(N²) 複雜度，Ｎ為序列長度，對於長序列問題造成極大限制。過去嘗試改進方法為縮減注意力矩陣的計算，或者引入基於狀態空間系統的模型（SSM），後者能以連續時間的狀態演化方程來表達序列依賴，有潛力以線性時間處理長距離依賴。然而，現有 SSM 模型普遍無法有效處理離散模態（如語言），核心問題是它們缺乏「基於內容」的推理能力，即無法根據當前輸入動態調整訊息傳遞與遺忘的策略，導致通用性與靈活性不足，表現在語言任務上的效果遠遜 Transformer。

核心方法與創新

作者團隊的主要創新在於提出「Selective State Spaces」，其核心概念是將 SSM 的參數設計為輸入的函數，使模型能根據當前序列元素「選擇」性地更新或遺忘狀態。具體而言，傳統 SSM 參數固定，一旦訓練完成即無法隨序列內容改變，引入動態參數後，每一步狀態更新將依據該token的特徵微調，有如 Transformer 的注意力使能依據內容調控訊息流通，帶來更強的語境敏感性。

然而，這種基於輸入的動態調整破壞了原有 SSM 可依賴的快速卷積運算結構，理論上增大了計算負擔。為此，論文貢獻其獨特的「硬體感知並行演算法」，允許在循環模式下仍可高效並行計算，最大限度利用現代硬體的多核心及向量化特性，開啟了選擇性 SSM 的實用前景。

為驗證方法可行，作者構建了一個名為 Mamba（無注意力的簡化端對端架構）的神經網路框架，完全去除傳統 Transformer 的注意力層與MLP區塊，僅以選擇性 SSM 作為序列核心模組。這種簡化不僅大幅減少架構複雜度，更由於線性時間與高度並行化設計，實現推理速度較同尺寸 Transformer 快5倍。

主要實驗結果

Mamba 在多個重要序列模態展示強健性能。語言建模實驗中，3億參數的 Mamba 模型不僅優於同參數的 Transformer，甚至可媲美2倍規模 Transformer，不論是在預訓練還是下游任務評估，皆展現極具競爭力的結果。這反映出Selective State Spaces不僅改善了 SSM 處理離散語言的能力，還有效提升了整體序列理解與生成能力。

在其他領域，如音頻與基因組序列，Mamba 同樣展現穩定且優越表現，尤其是在處理長達百萬級序列長度時，依旧保持良好的推理效率與精度，這是現存大多數 Transformer 或類 Transformer 架構無法實現的。

對 AI 領域的深遠影響

Mamba 的提出對序列模型領域帶來多重啟示。首先，它證明了結構化狀態空間模型在引入動態、內容依賴參數後，具備與 Transformer 競爭的能力，彌補了過去 SSM 在離散模態上的缺陷。其次，硬體感知的並行演算法設計，提供了一條線性時間序列模型實作的可行道路，為大規模序列資料處理提供了理論與工程支持。

此外，Mamba擺脫了注意力及MLP塊，顯示即便在非Transformer架構中，只要序列更新機制足夠靈活與高效，同樣能在多元模態取得一流水準，這挑戰了Transformer看似「不可或缺」的地位，引領學界思考更輕量卻高效的序列建模新範式。

最後，Mamba在大規模長序列建模中具備線性伸縮性與高效推理能力，為如全文本搜索、語音識別、基因組分析等長序列場景提供了嶄新解決方案，推動AI技術向處理極端尺度數據邁進，對推動各行業基於長序列的智能應用具有深遠意義。

總結而言，Mamba憑藉Selective State Spaces的創新設計及高效實作，成功架構了一條不同於Transformer但性能不讓的線性時間序列模型新路徑。未來隨著模型規模與資料規模持續增長，其在神經序列模型設計指引和長序列實務應用上的價值，將持續被重視與探究。

---

論文資訊
📄 Mamba: Linear-Time Sequence Modeling with Selective State Spaces
👥 Gu, Dao
🏆 ICLR 2024 · Outstanding Paper
🔗 arxiv.org/abs/2312.00752