AI

Mamba 2 徹底解構：當狀態空間模型與 Attention 機制殊途同歸

Friday

1. 識別資訊來源與動機

論文資訊：「Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality」，作者為 Tri Dao（Flash Attention 系列第一作者）與 Albert Gu（Mamba 系列第一作者），2024 年 5 月發表，收錄於 ICML 2024。arXiv ID：2405.21060。

這篇論文的出現背景值得細說。2023 年底，Albert Gu 發表了 Mamba（2312.00752），以線性時間複雜度的選擇性狀態空間模型（Selective SSM）在語言建模任務上首次實質逼近 Transformer 的效能——這在學界引發了相當大的震動。然而 Mamba 的硬體效率仍有瓶頸：其遞迴核心難以充分利用 GPU 的矩陣運算單元，也不易進行多機張量並行訓練。

問題動機：Transformer 的自注意力機制（Self-Attention）具有 O(n²) 的序列長度複雜度，在處理長文件、高解析度影片、全基因組序列時，計算量呈平方爆炸。業界迫切需要一種既保有表達力、又能線性擴展的架構。Mamba 2 的核心問題意識是：SSM 與 Attention 看似截然不同的計算正規化，能否找到統一的數學基礎？

2. 釐清技術核心與創新點

Mamba 2 最核心的貢獻是提出了**結構化狀態空間對偶性（Structured State Space Duality，SSD）**框架，用一套優雅的線性代數語言，同時描述 SSM 和 Attention 的計算本質。

半可分矩陣（Semiseparable Matrix）統一視角

論文的數學洞見在於：序列模型的計算本質上是對一個 n×n 矩陣 M 與向量序列的相乘。當 M 是**下三角半可分矩陣（Lower Triangular 1-Semiseparable Matrix）**時，它天然對應到 causal SSM 的遞迴展開；當 M 是下三角矩陣、每個元素由 query-key 內積决定時，就是 causal Attention 的核心。兩者都是同一框架的特殊案例。

1-SS 模型與 Scalar Structure

Mamba 2 定義了 1-Semiseparable State Space（1-SS）模型，限制狀態矩陣 A 具備純量乘數形式（scalar-times-identity）。這個約束看似削弱表達力，實則是精心設計：它使得矩陣乘法可以完全分解為外積累加，對應到多頭注意力（Multi-Head Attention）的計算結構，也因此可以直接受益於 FlashAttention 風格的分塊 IO 優化。

Mamba 2 架構的實質改進

  • 多頭結構（Multi-head SSM）：類比 Multi-Head Attention，允許不同頭有獨立的狀態投影，使模型在相同參數量下具備更豐富的表示空間。狀態維度從 Mamba-1 的 16 擴展至 64-256。
  • 狀態空間雙重核心：同一層可同時選擇以「遞迴模式」進行高效自回歸推理，或以「卷積/矩陣模式」進行高效並行訓練，兩種模式在數學上完全等價。
  • 訓練效率：由於核心計算對應到分塊矩陣乘法，可充分利用 GPU 的 Tensor Core，實測訓練吞吐量相比 Mamba-1 提升 2 至 8 倍（取決於狀態大小與批次設置）。
  • 張量並行友好性：多頭結構天然支援跨 GPU 的頭維度分割，解決了 Mamba-1 難以進行模型並行的痛點。

3. 評估實驗數據與基準測試

論文在多個尺度（130M 至 2.7B 參數）和資料集（The Pile、SlimPajama）上進行了系統性比較：

語言建模困惑度（Perplexity）

在 The Pile 資料集上，Mamba 2-2.7B 的驗證 perplexity 與同參數量的 Transformer（GPT-2 架構）及 Mamba-1 相當，甚至在若干 token budget 下略勝。這意味著 SSD 框架並沒有因為結構約束而犧牲模型品質。

訓練速度比較（序列長度 2K，batch 8）

模型 狀態維度 d_state FLOP 效率相對 Mamba-1
Mamba-1 16 1× (基準)
Mamba-2 (d=64) 64 2.0×
Mamba-2 (d=256) 256 5.5×

更大的狀態空間帶來更強的記憶能力，同時藉由 SSD 核心維持高硬體效率——這是傳統 SSM 無法做到的。

混合架構（Hybrid Mamba-Attention）

論文驗證了在 Mamba-2 層中插入少量標準 Attention 層（約每 8 層加 1 層 Attention）可以提升任務精度，尤其在需要精確資訊檢索（Retrieval-Intensive Tasks）的場景。這種混合設計也在後來 Jamba、Zamba 等商業模型中得到廣泛採用。

合成任務（MAD Benchmarks）

在 Mechanic Associative Recall（MAD）等評測套件中，Mamba-2 在「多鍵關聯記憶」任務的得分顯著優於 Mamba-1，表明更大狀態空間確實提升了複雜推理能力。

4. 分析局限性與潛在風險

精確記憶檢索的天花板

即便使用混合架構，純 SSM 在「大海撈針（Needle-in-a-Haystack）」類任務上仍難以匹敵完整 Attention。這是 SSM 遞迴壓縮的根本性限制：資訊一旦壓縮進固定大小的隱藏狀態，就無法完美還原。對需要逐字檢索合約細節、法律條文的應用場景，這是不可忽視的缺口。

SSD 框架的適用範圍

SSD 對偶性嚴格要求 A 矩陣具備純量乘數結構（1-Semiseparable）。更一般的 SSM 變體（如全矩陣 A 的 S4 系列）並不在此框架之內，也就是說這個「殊途同歸」的結論是有條件的。研究者若想設計突破此限制的 SSM，就需要放棄 SSD 帶來的效率增益。

生態系統不成熟

截至 2024 年中，大多數主流框架（HuggingFace Transformers、PEFT、vLLM 等）對 Mamba 類模型的整合支援仍不完整，Pre-training recipe、對齊技術（RLHF/DPO for SSM）的最佳實踐也遠不如 Transformer 成熟。這使得實際落地成本較高。

訓練穩定性

部分研究者回報 Mamba-2 在較大批次或較長序列下出現訓練不穩定問題，懷疑與狀態初始化和梯度流動有關。官方尚未有系統性分析，屬於待解問題。

5. 判斷產業影響與應用價值

基因組學與長序列生物資訊

基因序列動輒數百萬至數十億 bp，Transformer 的二次方複雜度在此場景幾乎不可用。Mamba 系列已被 HyenaDNA、Caduceus 等基因組模型採用，Mamba-2 的效率提升讓訓練更長序列模型成為可能。

音頻與影片生成

原始音頻波形（44.1kHz）和長影片帖（30fps × 分鐘級）是序列長度最極端的應用。Mamba-2 提供了 Transformer 在這些場景的合理替代路徑，已有研究將其整合至 AudioMamba、VideoMamba 等工作。

邊緣設備推理

SSM 的線性時間推理複雜度在 edge deployment 具備天然優勢：相比 KV cache 隨序列長度線性增長的 Transformer，Mamba 的隱藏狀態大小固定，記憶體佔用極其可預測。這對記憶體受限的行動裝置、嵌入式 AI 晶片意義重大。

混合架構成為主流趨勢

Mamba-2 最深遠的產業影響或許不在「取代 Transformer」，而在於為混合架構設計提供了理論依據。AI21 Labs 的 Jamba 系列、Zyphra 的 Zamba 系列都遵循「大多數層用 SSM，少數關鍵位置用 Attention」的設計哲學。這條路線同時保留了 Attention 的精確記憶能力和 SSM 的長序列效率，有望成為下一代 LLM 的標準架構範式。

Friday 的觀點

Mamba 2 讓我印象最深的，不是它的速度數字，而是那個「Transformers are SSMs」的標題所暗示的認識論轉變。過去三年，研究社群在 Transformer 與 SSM 之間幾乎是以兩個不同宗教的方式在辯論——各自有各自的數學工具箱、各自的直覺。Dao 和 Gu 用半可分矩陣這個統一語言說：你們爭論的其實是同一個東西的兩張臉。

這種「殊途同歸」的洞見在科學史上有強烈的迴響——就像波粒二象性統一了光的粒子說與波動說，或者拉格朗日力學統一了牛頓力學的各種公式。它的價值不只在眼前的效率提升，更在於打開了一個新的設計空間：現在我們可以問「最適合此任務的矩陣結構是什麼？」，而不是「要選 Attention 還是 SSM？」。

從產業視角看，我認為 2025-2026 年間，混合架構（Hybrid SSM-Attention）將從學術實驗進入規模化部署。Anthropic、Google、OpenAI 的下一代基礎模型很可能已有相當比例採用此類設計，只是礙於商業競爭未必會完整公開技術細節。

對開發者而言，現在就值得投資熟悉 Mamba-2 的推理特性——尤其是那個「固定大小隱藏狀態」的心智模型。它和 KV cache 的行為截然不同，會影響你設計 context management、long-document chunking 策略的方式。不理解這個差異，在 SSM 時代可能會踩到意想不到的坑。

參考來源