Mamba: Linear-Time Sequence Modeling with Selective State Spaces 深度解析

隨著深度學習快速發展，基礎模型（foundation models）已成為推動自然語言處理、語音辨識、基因組分析等多領域突破的核心支柱。這些模型幾乎都建立在 Transformer 架構及其注意力（attention）模組之上，顯示出非凡的擷取長距序列中關鍵訊息的能力。然而，Transformer 因注意力機制的計算複雜度呈現二次方增長（O(N²)）的瓶頸，使其在處理極長序列時計算資源需求與時間成本急劇飆升，難以有效延展到百萬級別的序列長度。

為解決此一限制，研究界陸續提出許多次二次時間複雜度的架構，例如線性注意力（linear attention）、門控卷積（gated convolution）、循環模型（recurrent models）以及結構化狀態空間模型（Structured State Space Models, SSMs）。這些方法在理論上降低了計算負擔，但在核心表現上，尤其是語言模型等須依賴內容驅動推理的應用中，整體表現仍未達到 Transformer 的水準。Gu Dao 在 ICLR 2024 發表的論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》即聚焦於此問題，提出一套兼具計算效率與表現力的序列建模新架構 Mamba，並榮獲 Outstanding Paper 獎項肯定。

研究背景與動機

Transformer 的核心在於自注意力機制（self-attention），它透過計算序列中每對元素間的關聯權重來捕捉長距依賴，實現精確的內容導向推理。然而，此過程計算量隨序列長度平方級成長，限制了模型應用於長序列資料，如長篇文本、音訊訊號及基因序列等。此外，近期受到重視的結構化狀態空間模型（SSMs）在理論上能以線性時間複雜度處理序列，但在應對離散模態（如語言文字）及複雜推理任務時，依舊力不從心。

透過實驗觀察，作者指出問題癥結在於既有模型缺乏足夠的內容感知能力，也就是無法根據目前輸入的 token 動態選擇性地保留或遺忘先前狀態，導致模型表現受限。因此，如何在保持計算效率的前提下增強模型的內容選擇性（content-based reasoning）成為突破關鍵。

核心方法與創新

Mamba 模型的首要創新在於引入「選擇性狀態空間」（Selective State Spaces），透過讓 SSM 的參數成為輸入的函數，使參數得以根據當前序列中的 token 動態變化。此設計打破以往 SSM 參數固定不變的限制，使模型能針對不同內容決定信息在序列長度維度上的傳遞或捨棄，類似 Transformer 中根據內容計算注意力權重的機制。

然而，這樣的動態參數變化也阻礙了使用高效卷積算法，因為過去許多 SSM 透過固定參數得以快速批次化並行計算。為克服此挑戰，作者提出一套硬體感知（hardware-aware）的平行算法，利用循環模式（recurrent mode）有效並行處理計算流程，大幅提升運算吞吐量與效率。

完整架構上，Mamba 擺脫了傳統 Transformer 典型的注意力層與多層感知器（MLP）層，整合簡化為端對端的純 Selective SSM 神經網絡。此設計不僅具備線性時間複雜度，推論速度更達 Transformer 的約五倍，且支援序列長度隨意延展至百萬級，無明顯性能退化。

主要實驗結果

在多模態數據測試中，Mamba 展現出強大的泛化能力。尤其在語言建模任務中，Mamba-3B 版本不僅超越了相同參數量的 Transformer，甚至匹敵具有兩倍參數量的 Transformer 模型，於預訓練與下游評估皆獲得一致優異表現。

除了語言外，Mamba 在音訊處理與基因組分析等其他序列資料上也達到最先進水平，顯示其 Selective SSM 機制能有效捕捉多種序列資料中複雜的結構性與依賴關係。此外，效能評估證明模型推論延展性極佳，能流暢處理數以百萬計的序列長度，遠超 Transformer 受限於計算瓶頸的表現。

對 AI 領域的深遠影響

Mamba 以其創新的 Selective State Spaces 及硬體感知平行算法，突破了序列建模領域長久以來拉扯於效率與效果之間的瓶頸。這為處理超長序列，特別是在自然語言等離散內容驅動應用中，提供了一條可行且高效的替代方案，有望成為未來基礎模型架構的重要基石。

更廣泛地，Mamba 開啟了一條讓參數動態適應輸入內容的序列模型新思路，結合狀態空間模型的理論嚴謹與 Transformer 的內容敏感性，打破傳統固定參數多層結構的框架限制。此種彈性結構可能促成多模態、多任務跨域學習模型的發展，加速深度學習模型在實際應用中面對龐大數據與超長依賴挑戰的效能革新。

最後，論文所提出的硬體感知算法設計，亦強調了算法與計算架構協同優化的重要性，對於未來 AI 系統在推論速度與能效上的提升具指標性意義，有助推動端到端可擴展深度學習模型向更廣泛產業落地邁進。

總結而言，Mamba 不僅是序列建模技術上的重大突破，更是 AI 領域在探索高效能、大規模推理架構道路上的一盞重要明燈，值得業界與學界深入研讀及後續跟進。

---

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