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

隨著深度學習技術的迅猛發展，Transformer 架構因其強大的序列建模能力，已成為基礎模型（foundation models）中的主流方案，廣泛應用於語言、視覺、語音、基因序列等多種模態。然而，Transformer 的核心組件──自注意力機制（self-attention）在長序列處理時計算成本呈現二次方增長，使得處理非常長的序列時面臨計算和記憶體瓶頸。為了降低此類負擔，業界與學術界一直在探索次二次時間複雜度(sub-quadratic)的序列建模架構，如線性注意力（linear attention）、門控卷積（gated convolution）、循環模型以及結構化狀態空間模型（Structured State Space Models, SSMs），但這些替代方案在處理語言等離散模態時，往往在性能上仍不及傳統 Transformer。

一、研究背景與動機

傳統 SSMs 雖然在長序列的計算效率上具有理論優勢，能做到線性時間複雜度，且理論上具備捕捉長距離依賴的能力，但實務上卻存在一個關鍵瓶頸──缺乏基於內容（content-based）的選擇性推理能力。換句話說，這些模型在序列中以統一且固定的狀態轉移參數傳遞信息，無法針對當前輸入 Token 決定保留或遺忘哪些訊息，這在語言等離散資料的序列建模上造成明顯劣勢。這一問題直接影響了 SSM 在重要應用場景中的表現，限制了其廣泛應用。

Gu Dao 等人在其論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》（ICLR 2024，Outstanding Paper）中，針對上述限制提出全新解決方案。他們主張：透過讓 SSM 的狀態轉移參數可根據輸入資料動態調整，賦予模型選擇性「記憶」與「遺忘」的能力，便能大幅提升 SSM 在離散模態與長序列建模上的表現。

二、核心方法與技術創新

1. 選擇性狀態空間模型（Selective State Spaces）
論文的核心貢獻是提出一種稱為「Selective State Spaces」的機制，突破傳統 SSM 固定參數的限制。具體來說，他們設計了一套可依據當前輸入 Token 動態調整狀態轉移矩陣和輸出矩陣的參數生成函數。這使得模型可以根據內容決定哪些歷史信息應該被保留、哪些應被忘卻，類似 Transformer attention 的「聚焦」效應，但又保持了線性時間複雜度。

這種設計的挑戰在於，動態參數化使得模型無法再直接利用高效的傅立葉變換卷積等演算法，但作者巧妙設計了針對硬體友好（hardware-aware）的並行遞歸演算法，在不犧牲計算效率的前提下，仍能維持快速推論的能力。

2. Mamba 架構設計
基於上述 Selective SSM，作者進一步構建了一個極簡化的端到端神經網路架構──Mamba。這個架構完全省略傳統 Transformer 中的自注意力層與多層感知器層（MLP），只藉由多層堆疊的 Selective SSM 模組實現序列建模。

Mamba 的特色包括：

• 線性時間複雜度與序列長度擴展性，適合處理百萬級長度的序列
• 硬體優化且並行遞歸計算方案，使推論吞吐量達到傳統 Transformer 五倍以上
• 一體化的動態狀態調節能力，提升對離散與連續模態的適應力

三、主要實驗與結果分析

為驗證 Mamba 的效果，作者在多個序列建模核心任務上進行了嚴謹實驗，涵蓋語言建模、語音識別、以及基因序列分析等多模態領域。

1. 語言模型性能
在最具代表性的語言建模任務中，尺度達 3 億參數的 Mamba-3B 模型，不但在預訓練階段達成超越相同規模 Transformer 模型的效果，更驚人的是，其性能相當於傳統 Transformer 6 億參數模型（約兩倍規模），展示了顯著提升的參數效率。

2. 長序列推論效率
Mamba 以線性時間複雜度處理長序列，可輕鬆擴展至百萬字元長度的序列，並在推論速度上實現超過 5 倍於 Transformer 的吞吐量。這對於需要長距離上下文信息建模的任務極具價值，突破以往 Transformer 在長序列時段計算瓶頸。

3. 多模態泛化能力
除此之外，Mamba 在語音和基因組資料的建模上均取得了當前最先進的競賽結果，證明其所設計的選擇性狀態空間不僅限於語言模態，也具有廣泛的實用性。

四、對 AI 領域的深遠影響與未來展望

Mamba 所展現的創新，對序列模型領域具有突破性的啟發意義：

• 突破線性模型的表現瓶頸：過去 SSM 及其他線性時間模型難以挑戰 Transformer 是因為缺乏內容敏感的記憶調控機制。Mamba 以 Selective State Spaces 成功填補這一缺口，為後續研究打開新的思路。
• 計算效率與模型容量平衡的典範：Mamba 兼具極高的計算效率及強大表現，這對於超大規模基礎模型尤其重要，隨著模型尺寸及資料集變得更龐大，Mamba 提供了可行的技術路徑。
• 多模態泛化能力：論文展示此架構在語言、音訊、基因組等多種序列模態均有亮眼成績，表明 Selective SSM 模型極具泛用性，有望成為未來多模態基礎模型的重要組成。
• 硬體與軟體協調的設計範例：作者設計的硬體友好並行演算法彰顯了深度學習理論與實務硬體優化結合的必要性，未來 AI 加速器與模型設計更緊密結合的趨勢也被再次強調。

整體而言，Mamba 不僅是對現有 Transformer 及其變種的強有力挑戰，更為序列建模技術的理論與實踐提供了豐富的創新方案。未來，我們可期待基於 Selective State Spaces 的進一步擴展，結合更多領域的專屬知識，催生更高效且強大的序列模型架構，在自然語言處理、多媒體理解、生物資訊等領域持續推動 AI 技術的革新。

---

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