Mamba: Linear-Time Sequence Modeling with Selective State Spaces

隨著深度學習基座模型（foundation models）的蓬勃發展，Transformer 架構與其核心的注意力機制（attention）已成為絕大多數關鍵應用的主流選擇。然而，Transformer 在面對極長序列時，計算複雜度的次平方（quadratic）成長限制了其在效能與效率上的擴展性，促使研究者們積極尋求更高效的替代架構。

本篇由 Gu Dao 等人發表於 ICLR 2024、並獲得 Outstanding Paper 獎的論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》提出一種創新的方法，名為 Mamba，它兼顧序列建模的運算效率與表現能力，針對 Transformer 於長序列推理與離散數據（如語言）上的弱點提出關鍵性突破，並在多重模態展現卓越成果。

研究背景與動機

Transformer 以注意力機制捕捉序列中元素間的關聯，表現上佳，但計算成本為O(N²)，N為序列長度，難以處理百萬級長度的序列數據。為了解決此問題，先前研究陸續提出低於次平方複雜度的模組，包括線性注意力（linear attention）、門控卷積（gated convolution）、循環神經網絡（RNN）或結構化狀態空間模型（Structured State Space Models, SSMs）。但令人遺憾的是，這些模型在重要的離散數據模態（像是文字）上尚未超越 Transformer，關鍵原因在於它們缺乏「基於內容的推理能力（content-based reasoning）」，無法依據序列中的上下文動態調整資訊的傳遞與忘記策略。

核心方法與創新

本論文的核心貢獻在於重新設計狀態空間模型（SSMs），讓模型能夠「選擇性地」據內容（token）改變狀態轉移參數，進而動態調控資訊流:

• 參數化的輸入依賴性： 傳統 SSM 使用靜態參數處理序列，導致對離散模態表達不足。Mamba 將 SSM 的關鍵參數設計為輸入的函數，使模型能根據當前 token 動態調整狀態遷移，類似於門控機制，達成選擇性傳播或忘記資訊，加強了對離散元素的敏感度與推理能力。
• 硬體感知並行計算方法：輸入依賴參數的調整使得原本可利用快速卷積優化的 SSM 失去直接適用，作者設計了一種硬體優化的並行演算法，在保持循環狀態更新邏輯的同時，充分發揮現代硬體的多核平行性，大幅提升運算效率。
• 簡化架構設計： Mamba 完全摒棄了 Transformer 的注意力層與 MLP 層，架構極度精簡，同時具備線性時間複雜度（O(N)）。如此設計不僅降低了計算與顯存壓力，也減少了訓練與推理時的瓶頸。

主要實驗結果

在多模態長序列建模任務中，Mamba 展現了非常出色的性能突破：

• 極端長序列測試：在處理序列長度達百萬級的數據（例如語言文本或基因序列）時，Mamba 的推理速度達到 Transformer 的 5 倍以上，且保持線性擴展，成功突破傳統 Transformer 的計算瓶頸。
• 語言建模表現：Mamba-3B 模型在語言建模任務上不僅超越了同規模的 Transformer，甚至能匹敵兩倍大小 Transformer 的水準，且在預訓練及下游應用評測中均有亮眼表現，展現其在離散自然語言處理上的強大能力。
• 跨模態泛用性：除了語言，Mamba 也在音訊和基因組等模態中取得最先進的性能，證明其作為通用序列建模骨幹架構的潛力和彈性。

對 AI 領域的深遠影響

Mamba 的提出標誌著序列模型設計的一大進展。長期以來，Transformer 在自然語言處理、語音辨識及基因序列分析等領域中扮演關鍵角色，但其二次方時間與空間複雜度成為大幅推廣與應用的絆腳石。Mamba 用選擇性狀態空間的核心思想，不僅解決了過去線性時間模型難以做到基於內容推理的問題，也成功將模型規模和應用場景極速擴展到百萬級長度序列。

此外，Mamba 拋棄了注意力與 MLP 的設計，提出更簡潔、高效的序列處理框架，有助於未來 AI 基座模型在硬體資源有限的情況下廣泛部署，降低能源消耗與延遲，推動深度學習走向真正的大規模普及與實時應用。

最後，Mamba 在多種模媒的優異表現，意義重大。它展示了選擇性且可學習狀態轉移參數的設計思路在語言、聲音與生物信息學等多領域的可拓展性與有效性，有望成為下一代大規模序列建模的基石架構。

總結

Gu Dao 等人在《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》一文中，創新地引入了可依據內容動態調整參數的狀態空間模型，突破了過去線性模型在離散內容推理上的瓶頸，並透過硬體感知平行計算加速設計大幅提升推理效率。該方法不僅在標準語言建模和其他模態中刷新最先進成績，亦為大規模序列建模開闢了全新途徑，具有高度理論價值與實際應用潛能，是跨世代序列模型技術躍進的典範。

---

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