Transformer 的自注意力機制在序列變長時,記憶體與運算量會隨長度平方成長,這是它最貴的一筆帳。狀態空間模型(State Space Model,SSM)以 Mamba 為代表,提出用一個固定大小的「狀態向量」沿時間遞迴地壓縮歷史資訊,把推論成本壓到線性等級。到了 2026 年,純 SSM 已不再是實驗品,而是與注意力層交錯組裝成「混合架構」(Hybrid Architecture),在多家實驗室的正式產品中出貨。今天就把這條 Transformer 之外的路線,從原理到工程權衡講清楚。
📖 學(核心)
SSM 為什麼能做到線性推論
一般注意力在解碼每個新 token 時,都要回頭掃過前面所有 token(存在 KV Cache 裡),因此記憶體與長度成正比、運算與長度平方成正比。SSM 換了一套思路:它維護一個固定維度的隱藏狀態 h,每讀進一個 token 就用遞迴式更新這個狀態,再從狀態輸出結果。關鍵在於狀態的大小不會隨序列變長,所以無論輸入是一千還是十萬個 token,每一步的成本都一樣,推論時間對長度呈線性,而且不需要無限膨脹的 KV Cache。
Mamba 的核心貢獻是「選擇性」(Selective)狀態空間:它讓狀態轉移的參數變成隨輸入內容而動態調整,等於讓模型自己決定「哪些資訊該留在狀態裡、哪些該遺忘」。這解決了早期 SSM(如 S4)參數固定、無法根據語意做內容篩選的弱點,使 SSM 首度在語言建模上逼近 Transformer 的品質。
遞迴與注意力其實是一體兩面
Mamba-2 提出的「狀態空間對偶」(State Space Duality,SSD)框架揭示了一個漂亮的事實:SSM 的時間遞迴可以被向量化,寫成一個帶遮罩的矩陣形式。在這個框架下,只要令 Q:=C、K:=B、V:=X,並把狀態的衰減視為一個資料相依的遮罩 L,遞迴式的 SSM 就能被表示成平行的類注意力運算;最簡單的情況下,SSD 退化成因果線性注意力(causal linear attention),而 L 就是那個因果三角遮罩。
這個對偶關係的實務意義很大:它讓 SSM 能像 Transformer 一樣在訓練時平行運算、吃滿 GPU,而在推論時又切換回省記憶體的遞迴模式,同時提高了算術強度(arithmetic intensity),改善原本 SSM 在解碼階段偏向記憶體受限(memory-bound)的硬體利用率問題。
Mamba-3:補上狀態追蹤與硬體利用的短板
發表於 ICLR 2026 的 Mamba-3,針對線性時間模型的兩大老問題做了系統性改良。第一,線性模型在「狀態追蹤」類任務(例如判斷奇偶、逐步算術)上表現不佳;第二,解碼時硬體利用率低。Mamba-3 引入三項技術:用梯形離散化(trapezoidal discretization)取代較粗糙的歐拉法,提升數值近似品質;採用多輸入多輸出(MIMO)的公式來提高算術強度、讓 GPU 更忙得有效率;並在理論上把複數值 SSM 與資料相依的旋轉位置編碼(RoPE)連了起來。這代表 SSM 這條路線已從概念驗證走向成熟精修。
為什麼贏家是「混合」而非「取代」
純 SSM 有個天生弱點:狀態是固定大小的壓縮摘要,面對需要精準回憶某個遙遠 token 的「大海撈針」檢索任務時,它不如能逐一回看全歷史的注意力。業界的答案不是二選一,而是交錯堆疊。AI21 的 Jamba 是第一個大規模的 Transformer-Mamba-MoE 混合模型,以大約 1:7 的比例交錯注意力層與 Mamba 層,並每隔幾個區塊插入 MoE 層。直覺是:少量的完整注意力層負責提供精準檢索能力,大量的 Mamba 層以線性成本承擔序列處理的主體工作。成果是 KV Cache 比同級 Transformer 小約 8 倍,在 256K 上下文下僅需約 4GB KV Cache,吞吐量顯著提升。到 2026 年,這類混合 SSM-注意力設計已成為長上下文模型的主流做法之一。
🧠 記
- SSM 用固定大小的遞迴狀態壓縮歷史,推論成本對序列長度呈線性,且不需隨長度膨脹的 KV Cache。
- Mamba 的關鍵是「選擇性」:狀態轉移參數隨輸入內容動態變化,能自主決定記住或遺忘什麼。
- SSD 對偶框架把遞迴與線性注意力統一起來,兼得訓練平行化與推論省記憶體。
- Mamba-3(ICLR 2026)用梯形離散化、MIMO 公式與 RoPE 連結,補強了狀態追蹤能力與硬體利用率。
- 主流做法是「混合架構」:少量注意力層負責精準檢索,大量 SSM 層負責線性長序列處理,如 Jamba。
✍️ 實踐
- 先量化你的痛點:用 profiler 觀察現有 Transformer 服務在長上下文下的 KV Cache 記憶體與延遲曲線,確認瓶頸是否來自序列長度的平方成長。
- 找一個已釋出的混合模型(如 Jamba 系列)實測,在你的長文件任務上比較吞吐量、記憶體佔用與答案品質,而非只看排行榜分數。
- 特別設計「大海撈針」檢索測試:在超長輸入中埋入關鍵事實,驗證混合模型的注意力層是否足以支撐你需要的精準回憶。
- 讀通 Mamba-2 的 SSD 對偶推導,理解 Q/K/V 與 SSM 參數的對應關係,這是判斷各家混合設計取捨的基礎。
- 評估工程成本:確認你的推論框架與量化工具鏈是否已支援 SSM/混合層,避免落地時卡在缺乏 kernel 支援。
- 對照注意力比例:實驗不同的注意力層對 SSM 層比例(如 1:7 到 1:5),觀察它對檢索能力與效率的影響,再決定自訓或選型方向。
🔗 延伸學習
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces(原始論文)
- Mamba-3: Improved Sequence Modeling using State Space Principles(ICLR 2026)
- Jamba: A Hybrid Transformer-Mamba Language Model
- Attention was never enough: Tracing the rise of hybrid LLMs(AI21)
💬 問 AI
我正在為 {任務類型,例如長文件摘要/程式碼理解/RAG} 評估是否從純 Transformer 換到混合 SSM-注意力架構(如 Mamba/Jamba)。
我的關鍵限制是 {上下文長度需求} 與 {硬體/記憶體預算}。
請幫我:
1. 比較純注意力、純 SSM、混合架構在我這個任務上的優劣與風險;
2. 針對「精準檢索 vs 線性效率」的取捨,建議合適的注意力層比例;
3. 列出上線前應做的三個實測(含大海撈針測試的設計),並說明各自要看什麼指標。