Mamba: Linear-Time Sequence Modeling with Selective State Spaces

近年來，Transformer 架構已成為深度學習中基礎模型（Foundation Models）不可或缺的核心，尤其在自然語言處理、語音及基因資料處理等多種序列任務上展現卓越性能。然而，Transformer 本身的注意力機制（Attention）計算複雜度為 O(N²)，當序列長度 N 大幅增長時，對計算資源及記憶體的需求驟增，極大限制了其在超長序列上的實際應用。此外，為了克服此短板，學界曾嘗試多種子二次時間（subquadratic-time）架構，包括線性注意力（linear attention）、門控卷積（gated convolution）、遞迴模型以及結構化狀態空間模型（Structured State Space Models, SSMs）。雖然這些方法在計算效率上有顯著提升，但在語言等離散模態的表現仍難以超越標準的 Transformer，顯示現有模型在「內容基礎推理」（content-based reasoning）能力上的欠缺。

在此背景下，來自 Gu Dao 等人的論文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》（ICLR 2024，獲得 Outstanding Paper）提出了一套創新性框架 Mamba，旨在同時兼顧效率與性能，特別關注序列模型在處理長序列、離散模態（如語言）時的資訊選擇與遺忘機制，並大幅改進現有 SSM 架構的限制。

核心方法與創新

Mamba 的核心設計理念是基於「選擇性狀態空間模型」（Selective State Spaces），其主要創新包括：

1. 讓狀態空間模型參數成為輸入的函數：
   傳統 SSM 透過一組固定參數定義狀態轉移，但這種固定參數在處理離散且多變的序列時（如文字），往往無法靈活反映當前輸入的內容。Mamba 採用輸入相關的函數動態決定 SSM 參數，意味模型可根據當前 token 動態選擇是否「保留」或「遺忘」先前隱藏狀態資訊，增強了模型的內容敏銳度與調節能力。這種設計使得序列中長期依賴的記憶以更有策略且資訊相關的方式被維持或更新，提升了內容基礎的推理能力。
2. 硬體感知（hardware-aware）的平行演算法設計：
   雖然讓參數動態依賴於輸入破壞了可直接使用高效卷積計算的便利，但作者設計了一套在遞迴模式下可充分發揮現代硬體（如 GPU）並行能力的算法，在不犧牲效率的前提下實現線性時空複雜度（O(N)）。此舉在保持序列長度可線性放大的同時，提供了高吞吐量的推論速度，實現比 Transformer 還快 5 倍以上的效果。
3. 簡化且有效的端到端架構設計：
   Mamba 網絡架構摒棄了典型 Transformer 中的注意力模塊，甚至不採用多層感知器（MLP）塊，全網絡核心即為選擇性 SSM 模組，有效降低了模型架構的複雜度。此簡化後的架構不僅利於理解與優化，也使得整體計算更為高效。

主要實驗結果

在多模態序列任務中，Mamba 展現出令人驚豔的性能和計算效率：

• 語言建模任務上，作者訓練了 30 億參數規模的 Mamba-3B，在預訓練與下游評估階段均超越了同尺寸的 Transformer 模型，更與 60 億參數的雙倍尺寸 Transformer 比肩，顯示出極佳的參數效率及泛化能力。
• 進一步實驗涵蓋了音訊及基因組學（genomics）等超長序列數據，Mamba 在序列長度上能擴展至百萬級別，且推論速度線性增長，解決了傳統 Transformer 在處理超長序列時因計算資源爆炸而難以實用的痛點。
• 與現有子二次時間架構相比，Mamba 不僅在效能評比中取得領先，同時在語言等離散模態的內容理解與推理方面，有顯著改善，證明選擇性狀態轉移的策略有效彌補了過去 SSMs 及類似模型在「內容感知」上的缺憾。

對 AI 領域的深遠影響

Mamba 論文為序列模型設計提供了重要的新方向，特別在於：

1. 融合靈活參數與高效計算的策略創新：
   Mamba 展示了動態參數化如何顛覆傳統狀態空間模型僵化的結構，為長序列建模引入了更多「內容選擇」的元素，這種設計理念可被拓展到更多模型架構，促進序列理解向更為靈活且精細的推理能力演化。
2. 提升超長序列處理的可行性：
   透過硬體感知的遞迴平行算法使得超長序列在理論與實務上均可高效運算，這對於基因組資料分析、語言長文分析、以及連續音訊流的即時處理，乃至未來多媒態序列統合學習，都提供了關鍵的技術支撐。
3. 開啟序列模型結構簡化的新篇章：
   Mamba 的成功證明，去除注意力機制和 MLP 結構並不必然犧牲模型性能與泛化能力，有助於引導後續研究探索更具計算效率、結構簡潔且理論基礎扎實的序列模型。

綜合來說，Mamba 以其創新性的 Selective State Spaces 模型架構，成功平衡了速度與性能，且克服了長序列及離散模態建模的難題，為未來大規模序列模型的設計提供了啟示，並可能推動多領域基礎模型的突破。本論文獲得 ICLR 2024 Outstanding Paper 獎項，足以見得其在學術界與實務領域的雙重產出與潛力。

對於具備基礎 AI 知識的工程師與研究生而言，深入理解 Mamba 的方法、實現及優化策略，將有助於開發更高效且適用於超長序列的序列模型，拓展在語言理解、音訊處理與生命科學領域的人工智慧應用。

---

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