Mamba: Linear-Time Sequence Modeling with Selective State Spaces

隨著深度學習基礎模型迅速發展，Transformer 架構因其靈活的注意力機制，成為處理序列資料（如語言、音訊、基因組等）最主流且效果卓越的架構。然而，Transformer 在長序列上的計算複雜度達到二次方等級，造成計算資源消耗及推論延遲顯著增加，限制了其在超長序列上的應用。面對這一瓶頸，學界發展了多種亞二次方複雜度的架構，包括線性注意力（linear attention）、門控卷積（gated convolution）、循環模型及結構化狀態空間模型（structured state space models, SSMs）等，但這些方法在具挑戰性的離散序列（如自然語言）上的表現仍未達到 Transformer 級別，難以全面取代。

本篇於 ICLR 2024 獲獎的論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》由 Gu Dao 等人提出了一種嶄新的序列建模架構——Mamba。作者察覺到現有 SSMs 及類似模型的核心缺陷在於缺乏內容為本（content-based）的推理能力，使得模型無法根據當前輸入動態選擇性地保存或遺忘訊息，尤其在處理符號化或切割明確的離散資料時效果不佳。基於此，Mamba 重新設計了 SSM 參數，使之成為輸入的函數，賦予模型動態調節信息流的能力，實現選擇性記憶與遺忘，提升對內容的響應能力。

研究背景與動機

Transformer 依賴的自注意力機制在序列長度增加時，計算與記憶消耗呈平方級別增長（O(N²)），這是其無法有效處理超長序列的主要限制。為改善此一困境，研究者嘗試多樣方法，希望在保有 Transformer 精準建模能力同時，將複雜度降低。例如，線性注意力透過近似注意力分數計算減少複雜度，SSM 使用連續時間狀態空間理論，藉由求解微分方程描述長序列的長距依賴。然而，這些替代模型多依賴固定且與輸入無關的參數，缺少依據當下內容動態調節記憶的能力，因此在離散且結構複雜的資料（語言、程式碼等）上表現不佳，這個缺陷成為推廣這些亞二次方模型的關鍵瓶頸。

核心方法與創新

本論文的首要創新為引入「選擇性狀態空間模型（Selective SSMs）」。具體做法是，讓 SSM 模型中的參數不再是固定設定，而是透過神經網路學習輸入依賴的函數。換言之，模型對每個時間步的輸入 token 進行判斷，動態調節狀態轉換矩陣及輸出映射，從而有能力「選擇性地」保留或刪除序列資訊。

然而，引入參數依賴性後，原本條件下能應用高效率捲積加速的計算路徑被打破，不再能使用快速傅里葉等方法直接並行計算。為此，作者設計了「硬體感知並行演算法」，在循環模式下仍能充分利用現代硬體（如 GPU、TPU）的向量化與平行運算優勢，達成運算效率的大幅提升。

基於此，研究團隊打造出全新架構—Mamba。Mamba 完全去除 Transformer 的注意力模組及多層感知機（MLP）層，僅基於選擇性 SSM 組成極簡終端對端網路。此架構在保持序列建模能力的同時，以線性時間複雜度處理超長序列，且預先訓練與下游任務均展現卓越性能。

主要實驗結果

研究團隊在多模態真實數據集上做了嚴謹驗證，包括語言模型訓練、音訊處理及基因序列分析。結果顯示：

• Mamba 推論速度比傳統 Transformer 快約 5 倍，在千萬級序列長度的推理效率與表現上具有顯著優勢。
• 在語言建模任務中，Mamba-3B（參數約 30 億）超越了尺寸相同的 Transformer，並且表現不亞於規模達 60 億參數的 Transformer 模型，無論在預訓練還是微調任務皆如此。
• 跨音訊與基因序列資料，Mamba 也取得或競爭目前最先進的結果，展示其多模態泛化能力與強勁的泛用性。

對 AI 領域的深遠影響

Mamba 不僅展示了一條可行且高效的超長序列建模路徑，更具體化了讓模型具備「動態選擇記憶」的可能，大幅提升了結構化狀態空間模型在離散模態中的表現瓶頸，為序列建模引入了新的範式。

此項研究推動亞二次方複雜度架構從理論走向實務，解決了 Transformer 在超長序列任務上的痛點。其提出的硬體感知並行演算法及依輸入動態調整參數的設計，在效率與效果間取得絕佳平衡，為未來神經網路在計算機硬體上的最佳化提供重要啟示。

此外，Mamba 架構的極簡設計及領先性能，將促使更多基礎模型架構擺脫「注意力—MLP」的傳統套路，挖掘狀態空間理論與參數動態調節的潛力，對自然語言處理、音訊分析及基因資料等多領域均可產生深遠影響。

綜合而言，Mamba 代表了序列模型設計上的一大躍升：在處理長序列時擁有線性時間複雜度，保持甚至超越 Transformer 的強建模能力，同時兼顧高效推論與跨模態泛化，為未來基礎模型開發及應用鋪陳了嶄新的技術藍圖。對正在追求高效且可擴展序列建模的工程師與研究者而言，Mamba 提供了值得深入研究與借鑒的最新前沿成果。

---

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