Mamba: Linear-Time Sequence Modeling with Selective State Spaces

隨著序列建模在自然語言處理、時間序列分析以及生物資訊等多項應用領域中扮演關鍵角色，如何在計算效率與模型表達能力之間取得平衡，一直是學術與工業界的挑戰。經典的序列模型如長短期記憶網路（LSTM）與門控循環單元（GRU）雖然在捕捉長距依賴方面有明顯優勢，但因為序列長度增加，計算時間成長趨近線性甚至更高。另一方面，以Transformers為代表的自注意力機制在捕捉全局依賴關係上表現優秀，卻因為自注意力的計算複雜度為二次方，伴隨長序列時難以擴展。針對此問題，近期「狀態空間模型」（State Space Models, SSM）成為新的研究熱點，因其能在理論上合併遞歸結構與頻域分析，提供線性時間複雜度的序列建模。

本文由Gu Dao等人發表於ICLR 2024的獲獎論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》即是針對狀態空間模型潛力進行深化與突破的代表作。作者認為現有的SSM變體仍面臨兩大關鍵困境：

• 一方面，現有的SSM通常在選擇模型結構（例如核函數形式和狀態維度）時缺乏靈活調整，導致模型難以兼顧多樣化的序列特性與複雜度需求。
• 另一方面，雖然SSM理論支持線性時間運算，但實際在長序列上實現高效且穩定的計算仍具挑戰，尤其是在保持模型表現力與精確近似間的平衡。

核心方法與創新
Mamba主要透過「Selective State Spaces」的設計，引入一種多元且可選擇的狀態空間結構，用以靈活捕捉序列中的多尺度特徵與動態變化。具體而言，Mamba構建了一組可切換的基態空間（基於不同核函數及狀態維度），並透過學習機制根據序列特性動態分配與調整每個狀態空間的貢獻。這種選擇性策略不僅提升模型對序列變異的適應能力，也能有效避免模型過度參數化所帶來的計算負擔。

技術細節上，Mamba採用了下列關鍵創新：

• 多樣化狀態空間組合：不再將單一型態的狀態空間應用於整個序列，而是透過一組不同頻域特性與記憶能力的基態空間集成，讓模型能在不同位置靈活使用最適合的子空間。
• 稀疏選擇機制：設計了一種輕量的稀疏選擇篩選器，使得在每一步的序列處理中，能快速且有效地挑選出部分狀態空間進行狀態更新，降低計算量且保證線性時間複雜度。
• 穩定且高效的數值運算：Mamba巧妙利用數值策略提升狀態空間求解的穩定性，特別是在長序列及高狀態維度下依舊保持計算效率與精度，避免常見的梯度消失或爆炸問題。

整體而言，Mamba不僅是對狀態空間模型架構的結構性升級，也在訓練效率與泛化表現間取得良好折衷。

主要實驗結果
作者在多個公開序列建模基準上對Mamba進行了嚴格評估，包括語音識別、語言建模、時間序列預測等挑戰性任務。實驗顯示，Mamba不僅在模型準確度上超越現有最先進的SSM模型與部分Transformer變體，同時在序列長度成長時依舊維持明顯的線性時間加速。

具體數據方面，Mamba在一些語言建模任務的困難長序列設定中取得了約5%相對降低的困惑度（Perplexity），且訓練速度比傳統SSM提升約2~3倍。此外，模型在多尺度時間序列預測中的表現，展現出其狀態空間選擇策略對捕捉不同時間依賴性的優異適應性。實驗還驗證了Mamba在模型參數量控制下，能夠出色地平衡計算效率與預測精度，適合用於對計算資源有限的工業應用。

對 AI 領域的深遠影響
Mamba的提出標誌著狀態空間模型序列建模研究的一大進展。傳統SSM在理論上的優勢，經過Mamba的創新設計轉化為實際可用且高效的工具，有望改變長序列資料處理的算法生態。具體來說：

• 突破效率瓶頸：Mamba確立了在線性時間內實現靈活高效多尺度序列建模的技術路徑，為處理數百萬規模的長序列數據提供可能，這對語音、基因序列分析到極端時間軸的金融數據均具實際意義。
• 促進新一代序列模型設計：透過選擇性狀態空間的框架，Mamba啟示未來序列模型不必拘泥於單一固定結構，而是可動態融合多種內建結構，提升模型普適性與泛化力。
• 與現有架構相容： 基於狀態空間的結構具有良好的模組化特性，便於與Transformer、RNN等方法混合設計，催生更多跨範式融合模型。
• 開展理論與實務橋接：Mamba成功將狀態空間理論的深層次數學結構與深度學習優化、硬體效率考量緊密結合，是AI理論與系統實作融合的典範。

綜合而言，《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》為序列建模領域帶來一套具備前瞻性與可操作性的全新模型架構，不僅推動了線性時間序列建模的研究，也為實際應用中長序列數據的高效處理提供了嶄新范式。對於目前依賴Transformer龐大計算資源的AI工程師與研究人員而言，Mamba提供了值得深入學習與探索的重要技術路徑。

---

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