Mamba: Linear-Time Sequence Modeling with Selective State Spaces

在當前深度學習領域中，序列建模(sequence modeling)是自然語言處理、語音識別、時間序列分析等多項任務的基石。隨著 Transformer 架構的普及，尤其是自注意力機制的廣泛應用，模型在捕捉長距離依賴關係上獲得了顯著提升，但其內在計算複雜度也隨序列長度呈二次成長，導致在處理超長序列時計算資源及記憶體消耗劇增。面對此一挑戰，如何在維持長距離建模效能之餘，實現線性時間複雜度的序列建模，是近年研究重要的焦點。

本篇由 Gu Dao 等人於 ICLR 2024 發表並榮獲 Outstanding Paper 獎項的論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》提出了全新的思路與方法。作者從狀態空間模型(State Space Models, SSM)切入，利用「選擇性狀態空間」（Selective State Spaces）概念，成功設計出一種能以線性時間複雜度進行序列建模的架構 Mamba，不僅保有優異性能，且顯著降低運算負擔，具有高度工程實用性與理論創新價值。

一、研究背景與動機

傳統 RNN 及 LSTM 類模型因其循環結構天生支持序列數據處理，但在捕捉長距離依賴上仍受限，且訓練時存在難以並行化的瓶頸。Transformer 則依賴自注意力機制，能直接建模長距離依賴關係，並有效利用現代硬體的並行能力，性能與表現上有質的飛躍。惟其計算與記憶體需求為 O(N²)，N為序列長度，極大限制了可處理序列的長度。

此外，近年興起的狀態空間模型（SSM）以其結合物理系統動態特性及隱式狀態的方法，對於連續序列建模展現出理論上的優勢及應用潛力，相關工作如 S4、SSM Transformer 等皆證實在多種長序列任務中效果良好。遺憾的是，現有 SSM 方法通常受到模型設計靈活度與計算效率的制約，在特定情況下仍難兼顧性能與效率。

因此，Gu Dao 等人發問：能否設計一種具備靈活性、且能依序列狀態選擇性更新的狀態空間模型，實現令運算成本隨序列長度線性增長，並且不犧牲建模品質，從而達成高效且可擴展的序列模型？

二、核心方法與創新

Mamba 的關鍵創新在於提出「選擇性狀態空間」（Selective State Spaces, SSS）理念，結合並擴展過去 SSM 的結構，使模型可以視不同時間點的序列資訊，有選擇性地更新隱含狀態，而非全部狀態同步更新。

具體來說，Mamba 模型設計了兩大機制：

• 局部選擇性更新：透過引入一組強制稀疏的控制向量，模型自動判斷哪些狀態應當在當前時間步被更新，哪些可維持不變，避免重複無謂計算。
• 高效狀態空間核函數：配合數學上可證明的某些穩定性與稀疏性條件，設計專門的核函數架構，使狀態轉換與輸出映射均可高效計算，且保留非線性建模能力。

整體架構融合了卷積、隱馬爾可夫模型（HMM）與長短期依賴結構的優勢，並通過精巧的演算法設計完成低於 O(N²)的運算複雜度，理論上達到 O(N) 線性時間標準。

此方法不僅突破了傳統 SSM 常見的固有限制，也相較於 Transformer 的自注意力計算，顯著節省記憶體及計算資源，尤其在極長序列（如數萬至十萬級Token）建模任務中展現絕佳擴展性。

三、主要實驗結果

作者在多組序列模建常見數據集進行實驗，包括語言建模（Language Modeling）、訊號預測(Time-series forecasting)、長文本理解等多種場景。

• 語言建模：在 WikiText-103 及 OpenWebText 等長文本資料集，Mamba 不僅在 perplexity（困惑度）上達到與先進 Transformer 與 SSM 模型持平的水準，甚至超越部分基準，且在序列長度超過 10,000 時能保持穩定表現而無明顯性能損失。
• 計算效率：在序列長度達數萬 token 的情況下，Mamba 訓練與推論時間大幅縮短，相較於標準 Transformer 模型平均節省約 5 至 10 倍的時間，記憶體使用量也顯著減少。
• 泛化能力：從時間序列預測任務來看，Mamba 展示出強大的長期依賴捕捉能力與噪聲魯棒性，有效提升預測準確率，超越傳統 LSTM 與部分 SSM 變種。

整體分析指出，Mamba 透過選擇性狀態更新策略，同時結合嚴謹的數學理論與工程最佳實踐，在保持模型表達力的同時，實現了前所未有的線性時間序列建模，使其成為處理超長序列環境下最具競爭力的方法之一。

四、對 AI 領域的深遠影響

Mamba 所提出的 Selective State Spaces 框架帶來了序列建模理論與實踐的重大突破。首先，其高效且靈活的狀態更新機制為未來序列模型架構提供了新的設計範式，不必再依賴全量狀態同步計算，突破了傳統序列模型在效能上的瓶頸。

其次，Mamba 透過嚴謹的數學保證與實證表現，深化了狀態空間理論在深度學習中的應用，進一步推動連續時間動態系統理論與神經網絡模型的融合發展。這不僅利於長序列任務，也能被拓展到控制系統、動態系統建模等領域。

此外，考慮到現代產業對大規模數據處理的迫切需求，Mamba 所帶來的計算資源節省與速度優勢，將大幅增進訓練及推論效率，降低碳足跡，促進綠色 AI 的發展方向。特別是在雲端服務和邊緣計算場景中，具備顯著實用價值。

最後，作為一篇獲得 ICLR Outstanding Paper 認可的研究，Mamba 代表了序列建模領域前沿的理論突破與技術創新。它鼓勵學術界與產業界重新思考序列模型的未來發展路徑，推動對更高效、更可擴展但不損失性能模型的需求，為長期依賴建模問題提供切實且具潛力的解決方案。

總結而言，Mamba 以其 Selective State Spaces 的全新思路，成功實現了線性時間長序列建模，兼顧效率與表現，為序列模型發展帶來階段性飛躍，未來預期將深刻影響自然語言處理、語音識別、金融預測等多種長序列任務的技術路線。

---

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