Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

論文 · arXiv 2312.00752

Mamba：具選擇性狀態空間的線性時間序列塑模

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

發表於 2023-12-01cs.LGcs.AI arXiv PDF

用 XYZNOT 閱讀（全文翻譯與 AI 講解）

三句話總結

🎯 Transformer 運算量隨序列長度呈二次方增長，且自回歸推論需要龐大的 KV 快取；而先前如線性非時變（LTI）狀態空間模型等次二次方模型，在語言塑模等離散內容推理任務上則難以發揮良好表現。
🔧 使 SSM 參數具備輸入相依性（選擇性），以動態過濾或傳遞資訊；並搭配硬體感知的融合平行結合掃描與重運算演算法，避免將擴展後的隱藏狀態實例化至 GPU HBM 中。
📈 Mamba-3B 在下游零樣本（zero-shot）任務上達到或超越兩倍大小的 Transformer（例如平均準確率 63.3%，相較於 Pythia-2.8B 的 59.1% 與 Pythia-6.9B 的 61.7%），同時提供高達 5 倍的推論吞吐量。

摘要

現今驅動深度學習眾多亮眼應用的基礎模型，幾乎普遍建立在 Transformer 架構及其核心的注意力機制之上。許多具備次二次方時間複雜度的架構（如線性注意力、門控卷積與遞歸模型，以及結構化狀態空間模型，即 SSM）被相繼提出，旨在解決 Transformer 在長序列上的運算低效問題；然而，它們在語言等重要模態上的表現仍不及注意力機制。我們指出這類模型的一大核心缺陷在於缺乏基於內容的推理能力，並為此提出數項改進。首先，僅需讓 SSM 的參數成為輸入的函數，即可克服其在離散模態上的劣勢，使模型能依據當前權元（token）在序列長度維度上選擇性地傳遞或遺忘資訊。其次，儘管此項改動阻礙了高效卷積的使用，我們在遞歸模式下設計了一套具硬體感知的平行演算法。我們將這些選擇性 SSM 整合至一個簡化的端到端神經網路架構中，完全不包含注意力甚至 MLP 區塊（即 Mamba）。Mamba 享有極快的推論速度（吞吐量比 Transformer 高 5 倍）以及隨序列長度線性擴展的特性，且其效能隨著真實資料長度增加可擴展至百萬級序列。作為通用的序列模型骨架，Mamba 在語言、音訊與基因體學等多種模態上皆達到了尖端（state-of-the-art）水準。在語言塑模方面，我們的 Mamba-3B 模型在預訓練與下游評估中，均超越了同等規模的 Transformer，並能與其兩倍大小的 Transformer 旗鼓相當。

英文原文摘要

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5 $\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

論文解讀

問題與動機

基礎模型極度依賴 Transformer 架構與自注意力機制。雖然稠密注意力能在序列上下文視窗中有效路由資訊，卻帶來兩個根本限制：相對於序列長度 $L$ 的二次方計算複雜度 $O (L^{2})$ ，以及需要儲存過往鍵與值（KV 快取），導致每步推論需線性時間並引發記憶體瓶頸。

先前具備次二次方複雜度的方法——例如線性注意力、遞歸架構以及如 S4 等結構化狀態空間模型（SSM）——其計算時間隨序列呈線性或擬線性擴展。然而，結構化 SSM 歷來在文字等離散且資訊密集的模態上，表現始終無法媲美 Transformer。作者指出了根本原因：先前的 SSM 皆為線性非時變（LTI）系統。由於其轉移狀態的動態特性不隨輸入內容改變，因此無法選擇性地壓縮上下文、過濾無關雜訊，或在不同間距間靈活調整記憶留存。

核心概念

選擇機制即上下文壓縮（Selection Mechanism as Context Compression）： 序列塑模的本質是將歷史資訊壓縮為狀態。先前的 LTI SSM 使用固定的轉移參數，迫使其對所有時間步一視同仁。透過使 SSM 矩陣根據當前輸入權元動態參數化，模型便具備了依內容過濾或選擇性記憶資訊的能力。
突破無法卷積的瓶頸（Overcoming the Non-Convolving Bottleneck）： 讓參數具有輸入相依性，打破了線性遞歸與全域卷積之間的數學等價性。少了卷積支援，樸素的遞歸做法必須在 GPU 高頻寬記憶體（HBM）中實例化大小為 $(B, L, D, N)$ 的擴展潛在狀態。Mamba 透過將離散化、平行結合掃描（parallel associative scan）與輸出投影融合於快速的 GPU 片上 SRAM 中解決此問題，將 HBM 的存取頻寬需求降至最低。
統一架構（Unified Architecture）： Mamba 以單一均質區塊取代了多頭注意力（MHA）與 MLP 區塊交錯的設計，將 H3 的門控設計與 SwiGLU 風格的門控 MLP 結合在一起。

方法

連續狀態空間模型透過隱式潛在狀態 $h (t) \in R^{N}$ ，將輸入 $x (t) \in R$ 映射至 $y (t) \in R$ ： $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

利用步長 $Δ$ 進行零階保持（ZOH）離散化，連續系統可轉化為離散遞歸參數： $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

選擇機制（S6）

Mamba 不再讓 $(Δ, B, C)$ 隨時間恆定，而是使其依賴於輸入：

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ ，其中 $τ_{Δ} = softplus$ 且 $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

這推廣了經典的 RNN 門控機制：當 $N = 1, A = - 1, B = 1$ 時，遞歸直接簡化為門控更新 $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ ，其中 $g_{t} = σ (Linear (x_{t}))$ 。

硬體感知的狀態擴展（Hardware-Aware State Expansion）

為了在不將 $(B, L, D, N)$ 狀態實例化至 HBM 的情況下維持高吞吐量：

參數 $(Δ, A, B, C)$ 直接從 HBM 載入至片上 SRAM。
離散化為 $(\overline{A}, \overline{B})$ 以及具工作效率的平行結合掃描完全在 SRAM 內部完成計算。
潛在狀態透過 $C$ 降維縮減至 $(B, L, D)$ 形狀後，再寫回 HBM。
反向傳播期間的中間狀態直接在 SRAM 中重新計算，而非儲存在 HBM 中。

實驗結果

合成資料外推能力： 在歸納頭（Induction Heads）任務中，於序列長度 256 下訓練的 Mamba 能完美外推（100% 準確率）至長度 $2^{20} = 1, 048, 576$ （長度擴展 $4000 \times$ ），而注意力基準模型在超過 $2^{14}$ 時便失效或發生記憶體不足（OOM）。在選擇性複製任務上，Mamba 達到 99.8% 的準確率，相較之下 H3 為 57.0%，S4 僅為 18.3%。
語言塑模： 在 The Pile 資料集上以 Chinchilla 擴展法則評估至 13 億（1.3B）參數，Mamba 的表現達到甚至超越 Transformer++（包含旋轉位置嵌入、SwiGLU、RMSNorm）。在下游零樣本任務（如 LAMBADA、PIQA、HellaSwag、WinoGrande、ARC）中，Mamba-2.8B 獲得 63.3% 的平均準確率，超越了 Pythia-2.8B（59.1%）與 Pythia-6.9B（61.7%），並匹敵或超越多數 6B–7B 參數模型。
基因體學與音訊： 在序列長度為 100 萬（ $2^{20}$ ）的大猿 DNA 分類任務中，Mamba 分別達到 71.67%（1.4M 參數）與 81.31%（7M 參數）的準確率，大幅優於 HyenaDNA 的 54.87%。在 SC09 語音生成任務中，Mamba 將 FID 從 1.99（SaShiMi）降至 0.94（6.1M 參數）與 0.67（24.3M 參數）。
運算效率： 融合掃描核心的執行速度比標準 PyTorch 掃描實作快 20–40 倍，且在 32K 上下文長度下比 FlashAttention-2 快達 7 倍。推論吞吐量相較於同等規模的 Transformer 最高可提升 5 倍。

限制與開放問題

連續訊號表現： 雖然輸入相依的選擇機制改善了離散推理（如文字、DNA），但對原始音訊波形預訓練的消融實驗顯示，純線性非時變（LTI）模型（如標準 S4）在均勻取樣的連續資料上表現仍然較佳。
模型規模： 評估的模型規模上限約為 30 億參數與 3000 億預訓練權元；在 7B–70B+ 及更大規模下的實證表現仍有待進一步驗證。
下游生態系統： Transformer 擁有針對提示工程（prompt engineering）、情境操作、KV 快取量化與轉接器微調（adapter tuning）的成熟工具鏈。SSM 如何適應並融入這些互動模式仍是一個活躍的研究方向。

關鍵術語

Structured State Space Models (SSMs)（結構化狀態空間模型）: 一類源自連續線性動態系統的深度序列模型，透過隱式潛在狀態來映射一維輸入。
Linear Time Invariance (LTI)（線性非時變性）: 一種數學特性，指系統轉移參數在所有時間步中保持靜態不變；這雖能實現快速的全域卷積表示，卻限制了依內容推理的能力。
Selection Mechanism（選擇機制）: 將 SSM 的轉移系統動態（Δ, B, C）參數化為輸入權元的函數，使模型能在序列長度維度上動態傳遞、重置或忽略上下文。
S6: 透過結合性平行掃描而非卷積計算的選擇性結構化狀態空間模型的縮寫。
Parallel Associative Scan（平行結合掃描）: 一種具工作效率的前綴和計算演算法，可在加速器硬體上沿序列長度維度平行計算線性遞歸。
Hardware-Aware State Expansion（硬體感知狀態擴展）: 一種 GPU 核心融合與重運算策略，在快速的片上 SRAM 中執行離散化與掃描，避免在往返高頻寬記憶體（HBM）時產生龐大的中間狀態。
Induction Heads（歸納頭）: 一種合成的聯想回想機制，用以評估模型是否能在遇到先前出現過的上下文權元時，準確重現其後續跟隨的權元。
Mamba Block（Mamba 區塊）: 一種簡化且均質的序列模型區塊，結合了輸入投影、一維卷積、選擇性 SSM 與相乘門控，不需額外的注意力或 MLP 層。

常見問題

為什麼選擇性 SSM 無法使用快速 FFT 卷積進行訓練？

快速全域 FFT 卷積要求系統動態（矩陣 A、B 與 C）在時間維度上保持恆定（即具備線性非時變性）。由於 Mamba 使 B、C 與 Δ 在每個時間步都成為當前輸入權元的函數，卷積核在每個位置皆不相同，從而打破了卷積等價性，因此必須改用遞歸的平行掃描。

Mamba 在訓練期間如何避免遞歸狀態擴展帶來的記憶體暴增？

具有狀態擴展維度 N 的樸素遞歸需要在 GPU HBM 中實例化形狀為 (B, L, D, N) 的中間狀態。Mamba 透過核心融合技術，僅將低維輸入載入 SRAM，在高速 SRAM 內部完成離散化與平行掃描，僅將縮減後的 (B, L, D) 輸出寫回 HBM，並在反向傳播期間於 SRAM 中重新計算中間活化值。

Mamba 的 Δ 參數與遞歸門控之間的關係為何？

步長參數 Δ 扮演著廣義門控機制的角色。在 N=1 且 A=-1 的零階保持（ZOH）離散化條件下，選擇性遞歸在數學上直接簡化為標準 RNN 門控：較大的 Δ 會促使模型關注當前輸入並重置狀態；較小的 Δ 則會保留歷史狀態並忽略當前輸入。

Mamba 在生成過程中會使用 KV 快取嗎？

不會。因為 Mamba 在生成過程中是以自回歸狀態空間遞歸運作，它只需維持固定大小的隱藏狀態（每層維度為 D × N），而不必快取所有過往權元的鍵與值，因此每步推論均可在常數時間 O(1) 內完成。

本頁解讀由 AI 根據論文生成，可能有錯誤；細節請以原文為準。

用 XYZNOT 閱讀（全文翻譯與 AI 講解）