Transformer 即為 SSM：透過結構化狀態空間對偶性的廣義模型與高效演算法

Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

原始論文：Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality 作者：Tri Dao, Albert Gu arXiv ID：2405.21060v1 日期：2024-05-31 標籤：SSM Mamba Mamba-2 Linear Attention State Space Duality Transformer Sequence Modeling ICML2024

摘要

雖然 Transformer 一直是深度學習在語言建模上成功的主要架構，但近期狀態空間模型 (State-Space Model, SSM) 例如 Mamba 已被證明在中小規模上能匹敵甚至超越 Transformer。我們證明這兩個模型家族其實密切相關，並建立一套豐富的理論框架，透過經典的結構化半可分矩陣 (Structured Semiseparable Matrix) 的各種分解方式，將 SSM 與多種注意力 (Attention) 變體連結起來。我們的狀態空間對偶性 (State Space Duality, SSD) 框架使我們能設計出新架構 Mamba-2，其核心層是 Mamba 選擇性 SSM (Selective SSM) 的精煉版本，速度快 2-8 倍，同時在語言建模上仍與 Transformer 競爭。

1. 引言

Transformer，特別是以因果方式處理輸入序列的僅解碼器模型 (Decoder-only Model)，例如 GPT 與 Llama，是現代深度學習成功的主要驅動力之一。許多研究嘗試近似其核心注意力層以解決效率問題，包括訓練時序列長度的二次方擴展、以及自迴歸生成 (Autoregressive Generation) 時所需的線性大小快取。

與此並行的是另一類序列模型——結構化狀態空間模型 (Structured State Space Model)。這類模型在訓練時對序列長度呈線性擴展，在生成時擁有常數狀態大小。它們在長距離任務 (如 S4) 上表現強勁，最近也在小到中等規模的語言建模上 (如 Mamba) 匹敵或超越 Transformer。然而，SSM 的發展似乎與 Transformer 社群的集體努力 (例如理論理解與現代硬體優化) 相互脫節。因此，相較於 Transformer，理解、實驗、以及高效訓練 SSM 仍然具有挑戰。

我們的主要目標是建立一套豐富的理論連結，將結構化 SSM 與注意力的各種變體聯繫起來。這將允許我們把原本為 Transformer 開發的演算法與系統優化轉移到 SSM 上，朝著建構超越 Transformer 並能更有效擴展序列長度的基礎模型 (Foundation Model) 之目標前進。一個里程碑式的貢獻是線性注意力 (Linear Attention, LA) 框架，它透過顯示二次核化注意力 (Kernelized Attention) 與特定線性遞迴 (Linear Recurrence) 的等價性，連結了自迴歸注意力與線性 RNN。本論文同樣以多重觀點，連結線性複雜度的 SSM 與二次複雜度的形式，結合 SSM 與注意力之優勢。

狀態空間對偶性 (State Space Duality)

我們連結結構化 SSM 與注意力變體的框架稱為結構化狀態空間對偶性 (Structured State Space Duality, SSD)，是透過結構化矩陣 (Structured Matrix)——具有次二次參數量與乘法複雜度的矩陣——的抽象來建立的。我們發展了兩種廣泛框架來表示序列模型：一是矩陣轉換 (Matrix Transformation)，二是張量縮併 (Tensor Contraction)，每種框架揭示對偶性的不同視角。

我們的技術貢獻包括：

• SSM 與半可分矩陣的等價性 (第 3 節)：我們證明狀態空間模型與一類經過充分研究的結構化矩陣——半可分矩陣 (Semiseparable Matrix)——等價。本論文的核心訊息是：計算狀態空間模型的不同方法都可被重新表述為對結構化矩陣的各種矩陣乘法演算法。

• 大幅改進線性注意力理論 (第 4 節)：我們先用張量縮併語言提供其遞迴形式的精準證明，然後將其推廣為新的結構化遮罩注意力 (Structured Masked Attention, SMA) 家族。

• 連結 SSM 與 SMA (第 5 節)：我們證明它們有大幅交集，彼此互為對偶，同時擁有 SSM 式線性形式與注意力式二次形式。我們也證明任何具備快速遞迴形式的核注意力 (Kernel Attention) 方法必為 SSM。

高效演算法 (Efficient Algorithms)

我們的框架揭示了計算 SSM 的新穎且易於實作的演算法 (第 6 節)。我們引入新的 SSD 演算法，基於半可分矩陣的區塊分解 (Block Decomposition)，同時利用線性 SSM 遞迴與二次對偶形式，在所有主要效率軸線上獲得最佳折衷。一個專屬的 SSD 實作比 Mamba 的優化選擇性掃描 (Selective Scan) 實作快 2-8 倍，同時允許更大的遞迴狀態大小 (是 Mamba 的 8 倍以上)。SSD 與 FlashAttention-2 的優化實作高度競爭，在序列長度 2K 處交叉，序列長度 16K 時快 6 倍。

架構設計

採用 SSM 等新架構的一個主要障礙是專為 Transformer 量身打造的生態系。我們的框架允許利用為注意力建立的成熟慣例與技術，建構 SSM 的架構設計詞彙。例如我們將多頭注意力 (Multi-head Attention, MHA) 中的「頭 (head)」概念引入 SSM，並指出 Mamba 架構是一個多輸入 SSM (Multi-Input SSM, MIS)，類比於多值注意力 (Multi-Value Attention, MVA)。

修改後的 Mamba 區塊結合 SSD 作為內部 SSM 層，產生 Mamba-2 架構。我們在與 Mamba 相同的設定下研究 Mamba-2 的 Chinchilla 縮放定律 (Scaling Law)，發現它在困惑度 (Perplexity) 與 wall-clock 時間上均 Pareto 主導 Mamba 與 Transformer++。例如，在 The Pile 資料集 300B tokens 上訓練的 2.7B Mamba-2 在標準下游評估上勝過 Mamba-2.8B、Pythia-2.8B，甚至 Pythia-6.9B。

系統優化 (Systems Optimizations)

SSD 框架連結 SSM 與 Transformer，使我們能利用為 Transformer 開發的系統優化 (第 8 節)。例如，張量平行 (Tensor Parallelism, TP) 是訓練大型 Transformer 的關鍵模型平行技術，我們設計 Mamba-2 為 TP 友善的，將每個區塊的同步點數量減半。對於非常長的序列，我們描述如何用序列平行 (Sequence Parallelism) 訓練 SSM。對於變長序列的微調，我們展示 Mamba-2 可以高效地處理變長序列而不需 padding。

模型程式碼與預訓練檢查點 (Checkpoint) 開源於 https://github.com/state-spaces/mamba。

2. 背景與概觀

2.1 結構化狀態空間模型 (Structured State Space Models)

結構化狀態空間序列模型 (S4) 是近期一類深度學習序列模型，與 RNN、CNN、以及經典狀態空間模型廣泛相關。它們受到一個連續系統啟發，將 1 維序列 $x \in \mathbb{R}^T \mapsto y \in \mathbb{R}^T$ 透過一個隱式潛在狀態 $h \in \mathbb{R}^{(T,N)}$ 對應起來。

結構化 SSM 的一般離散形式為：

$$ \begin{aligned} h_t &= A h_{t-1} + B x_t \quad &(1a) \quad &\quad& h_t &= A_t h_{t-1} + B_t x_t \quad &(2a) \\ y_t &= C^\top h_t \quad &(1b) \quad &\quad& y_t &= C_t^\top h_t \quad &(2b) \end{aligned} $$

其中 $A \in \mathbb{R}^{(N,N)}$, $B \in \mathbb{R}^{(N,1)}$, $C \in \mathbb{R}^{(N,1)}$。結構化 SSM 之名來自於控制時序動態的 $A$ 矩陣必須具有結構，以便能在深度神經網路中高效地計算這個序列到序列的轉換。最早引入的結構是對角加低秩 (Diagonal Plus Low-Rank, DPLR) 與對角結構 (Diagonal)。

連續時間模型 (Continuous-time Models)：原始結構化 SSM 起源於對函數 $x(t) \in \mathbb{R} \mapsto y(t) \in \mathbb{R}$ 的連續時間映射。從連續時間觀點來看，矩陣 $(A, B)$ 不直接學習，而是從底層參數 $(\mathring{A}, \mathring{B})$ 與一個參數化的步長 $\Delta$ 透過離散化規則 (Discretization Rule) $A = f_A(\Delta, \mathring{A})$ 與 $B = f_B(\Delta, \mathring{B})$ 產生。

遞迴模型 (Recurrent Models)：方程 (1) 與 (2) 是對輸入 $x$ 線性的遞迴形式。結構化 SSM 因此可被視為一類 RNN，但其線性使其可避免傳統 RNN 的循序計算。儘管簡化，SSM 作為序列轉換仍是完全表達 (Universal Approximation) 的。

卷積模型 (Convolutional Models)：當 SSM 動態隨時間恆定時，模型稱為線性時間不變 (Linear Time-Invariant, LTI)，等價於卷積。SSM 也可以被視為 CNN 的一種，但卷積核透過 SSM 參數隱式參數化，且通常是全域而非局部的。

選擇性狀態空間模型 (Selective State Space Models)：方程 (2) 中參數 $(A, B, C)$ 隨時間變化的形式由 Mamba 引入，稱為選擇性 SSM (Selective SSM)。相較於標準 LTI 形式 (1)，這個模型可以選擇性地專注於或忽略每個時間步的輸入。它已被證明在資訊密集領域 (如語言) 上比 LTI SSM 表現好得多，同時比 Transformer 更有效率地擴展。我們將 Mamba 中使用的選擇性 SSM 稱為 S6 模型 (Selective Scan + Selection mechanism + Structured SSM)。

結構化 SSM 作為序列轉換

定義 2.1：對任意輸入維度 $P$，序列轉換 (Sequence Transformation) 是一個對輸入 $X \in \mathbb{R}^{(T,P)}$ 操作，產生輸出 $Y \in \mathbb{R}^{(T,P)}$ 的參數化映射。

我們將 SSM 視為對 $P$ 個獨立通道操作的單頭序列轉換，每個通道有自己的 $(A, B, C)$ 矩陣。這是一個重要視角：SSM 處理輸入的每一通道是獨立的，這使我們可以透過增加狀態大小 $N$ 來提升表達能力。

2.2 注意力機制

注意力 (Attention) 的核心是將序列中各元素彼此關聯的機制。最常見的變體——softmax 自注意力 (Self-Attention)——可寫為：

$$ Y = \text{softmax}(QK^\top) \cdot V $$

其中 $Q, K, V \in \mathbb{R}^{(T,N)}$。其機制讓每個位置都能存取所有先前位置的資訊。其主要缺點是訓練時的二次擴展 ($T^2$ 成本) 以及推論時的線性大小快取。

線性注意力 (Linear Attention)：Katharopoulos 等人 (2020) 觀察到，若在查詢-鍵 (Query-Key) 計算中省略 softmax，並使用核特徵映射 (Kernel Feature Map) $\psi$，則注意力可被重寫為一個遞迴形式，達到線性時間複雜度。其形式為：

$$ Y_t = \frac{\psi(Q_t)^\top \sum_{s \le t} \psi(K_s) V_s^\top}{\psi(Q_t)^\top \sum_{s \le t} \psi(K_s)} $$

省略歸一化分母項即為核線性注意力，可視為一種帶有特定狀態結構的 RNN。

2.3 結構化矩陣

矩陣 $M \in \mathbb{R}^{T \times T}$ 通常以密集形式儲存與運算，需要 $O(T^2)$ 參數與 $O(T^2)$ 計算成本。結構化矩陣 (Structured Matrix) 是指：

1. 可以用次二次參數壓縮表示
2. 具有快速演算法 (例如矩陣-向量乘法)，其複雜度低於樸素的 $O(T^2)$

常見例子包括稀疏矩陣、低秩矩陣、Toeplitz 矩陣、Cauchy 矩陣、Vandermonde 矩陣，以及蝴蝶矩陣 (Butterfly Matrix)。

2.4 概觀：結構化狀態空間對偶性

我們連結 SSM 與注意力的結構化狀態空間對偶性 (SSD) 框架建立在兩個觀點上：

• 矩陣轉換觀點 (第 3 節)：將 SSM 重寫為一個 $T \times T$ 矩陣 $M$，這個矩陣是半可分的。
• 張量縮併觀點 (第 4-5 節)：將線性注意力推廣為結構化遮罩注意力 (SMA)，並識別出 SSM 與 SMA 之交集。

3. 狀態空間模型即為半可分矩陣

本節我們展示一個關鍵連結：所有結構化 SSM 都對應於一個半可分矩陣。

3.1 矩陣轉換形式

考慮具有時變參數 $(A_t, B_t, C_t)$ 的選擇性 SSM 方程 (2)，展開遞迴：

$$ y_t = \sum_{s=0}^{t} C_t^\top A_{t:s+1} B_s x_s $$

其中 $A_{t:s+1} = A_t A_{t-1} \cdots A_{s+1}$ (對 $s = t$ 取為單位矩陣 $I$)。整個 SSM 因此可表示為矩陣乘法 $y = Mx$，其中 $M$ 是下三角矩陣，元素為：

$$ M_{ji} = \begin{cases} C_j^\top A_{j:i+1} B_i, & i \le j \\ 0, & i > j \end{cases} $$

3.2 半可分矩陣

定義 3.1：一個 (下三角) 矩陣 $M$ 稱為 $N$-半可分 ($N$-Semiseparable) 若其下三角部分 (對角線及以下) 中任何包含於對角線下方或對角線本身的子矩陣的秩都至多為 $N$。

定義 3.2 (順序半可分 SSS 表示)：下三角矩陣 $M \in \mathbb{R}^{(T,T)}$ 具有 $N$-順序半可分 ($N$-Sequentially Semiseparable, $N$-SSS) 表示，若可寫為： $$M_{ji} = C_j^\top A_{j:i+1} B_i \quad \text{對於 } i \le j$$ 其中 $B_0, \ldots, B_{T-1}, C_0, \ldots, C_{T-1} \in \mathbb{R}^N$ 且 $A_0, \ldots, A_{T-1} \in \mathbb{R}^{(N,N)}$。

定理 3.5：一個 (下三角) 矩陣 $M$ 是 $N$-半可分的，若且唯若它具有 $N$-SSS 表示。

這個定理確立了狀態空間模型與半可分矩陣的精確等價：

定理 3.7 (SSM 即為半可分矩陣)：方程 (2) 的狀態空間模型轉換等價於透過 $N$-SSS 矩陣 $M$ 進行矩陣乘法。

換言之：計算 SSM 的不同演算法，例如線性遞迴 (1a) 或樸素的全展開矩陣乘法，都對應於對同一個半可分矩陣的不同矩陣乘法演算法。

1-半可分矩陣 (1-SS Matrices)：當 $N = 1$ 時，$A_t$ 退化為純量 $a_t \in \mathbb{R}$，矩陣 $M$ 可寫為：

$$ M_{ji} = a_j a_{j-1} \cdots a_{i+1} = \prod_{k=i+1}^{j} a_k $$

這稱為 1-SS 矩陣，等價於純量 SSM 掃描或 cumprodsum 運算子。我們將 1-SS 矩陣記為 $L_a$，作為純量遞迴的矩陣化形式。

4. 結構化遮罩注意力

4.1 從注意力到結構化遮罩注意力

回想一下因果線性注意力 (Causal Linear Attention) 的核心結構：

$$ Y = (Q K^\top \odot L) V $$

其中 $L$ 是因果遮罩 (Causal Mask)，是一個下三角全 1 矩陣。等價地，這可寫為一個 4 路張量縮併：

$$ Y_{tn} = \sum_{s, p} L_{ts} Q_{tn} K_{sn} V_{sp} $$

不同的縮併順序對應不同的計算演算法：

• 二次形式 (Quadratic Form)：先計算 $G = QK^\top \in \mathbb{R}^{T \times T}$，再 $G \odot L$，再乘以 $V$。成本 $O(T^2N)$。
• 線性形式 (Linear Form)：透過遞迴沿著序列累積 $K^\top V$。成本 $O(TN^2)$。

線性注意力的關鍵洞察是：因果遮罩 $L$ 的特殊結構 (全 1 下三角) 使得遞迴形式可行。

定義 4.2 (結構化遮罩注意力, SMA)：給定結構化矩陣 $L$，結構化遮罩注意力定義為帶有遮罩 $L$ 的線性注意力： $$Y_{tp} = \sum_{s, n} L_{ts} Q_{tn} K_{sn} V_{sp}$$ 其效率取決於 $L$ 的結構。

定理 4.1：對於任何結構化矩陣 $L$，SMA 的計算成本由矩陣-向量乘法 $L \cdot v$ 的成本決定。若 $L$ 允許 $O(T)$ 時間矩陣乘法，則對應 SMA 也具有線性時間遞迴形式。

這提供了線性注意力的一個強力推廣：透過選擇不同的結構化遮罩 $L$，我們可以得到不同效能特性的 SMA 變體。

5. 狀態空間對偶性

現在我們連結 SSM 與 SMA。關鍵觀察是：當 $L$ 是 1-SS 矩陣時，SMA 與 SSM 重合。

5.1 純量恆等 SSM 的對偶性

考慮一個特殊情況：$A_t = a_t I$ (純量乘以單位矩陣)。在這情況下，$A$ 軸的縮併簡化為與 1-SS 矩陣 $L_a$ 的矩陣乘法。SSM 的矩陣形式變為：

$$ M = L_a \odot (C B^\top) $$

(其中 $C, B \in \mathbb{R}^{(T, N)}$，$L_a$ 由純量序列 $a_t$ 產生。)

這正是帶有 1-SS 遮罩的 SMA 形式！我們得到：

定理 (狀態空間對偶性)：純量恆等結構的選擇性 SSM 等價於 1-SS 結構化遮罩注意力。

換言之，選擇性 SSM 可以同時用兩種形式計算：

• 線性形式 (SSM 遞迴)：成本 $O(TN^2)$，但需要使用純量操作而非矩陣乘法。
• 二次形式 (1-SS SMA / 二次注意力)：成本 $O(T^2 N)$，但完全基於矩陣乘法，硬體效率更高。

這個對偶性是我們命名 狀態空間對偶性 (State Space Duality, SSD) 的核心。我們將同時擁有這兩種形式的模型稱為 SSD 模型。

5.2 SSD 與 Mamba 的關係

Mamba 使用的核心 S6 層具有對角 $A_t$，而 SSD 進一步限制為純量 $A_t = a_t I$。這個略微更受限的形式換來顯著的硬體效率優勢，因為純量恆等情況使對偶二次形式成為簡單的矩陣乘法，可以充分利用張量核 (Tensor Core)。

5.3 SSD 即為矩陣轉換

從矩陣觀點看，SSD 矩陣具有特殊的「遮罩外積」結構：

$$ M = L \odot (C B^\top) $$

其中：

• $L \in \mathbb{R}^{(T,T)}$ 是 1-SS 矩陣 (純量遞迴)
• $C, B \in \mathbb{R}^{(T,N)}$ (注意力中對應 $Q, K$)
• 輸入 $X \in \mathbb{R}^{(T,P)}$ 透過 $Y = M \cdot X$ 對映到輸出

這個矩陣形式同時是 1-SS SMA 與半可分矩陣 (兩者交集)。

6. 硬體高效的 SSD 演算法

雖然 SSD 同時擁有線性與二次形式，但兩種純粹形式各有缺點：純線性遞迴無法充分利用矩陣乘法單元，而純二次形式對長序列來說成本過高。我們提出一個 區塊分解演算法 (Block Decomposition Algorithm)，將兩者結合，達到最佳折衷。

區塊分解觀點

將序列分為長度為 $Q$ 的塊 (chunk)，總共 $T/Q$ 塊。SSD 矩陣 $M$ 可分解為：

• 對角塊 (Diagonal Blocks)：每個 $Q \times Q$ 塊，表示塊內輸入到輸出的影響
• 下三角區塊 (Off-diagonal Blocks)：表示前一個塊到當前塊的影響，可低秩分解

定理 6.1 (SSD 區塊分解)：SSD 矩陣 $M$ 可以分解為一個塊對角矩陣加上低秩矩陣，其中低秩部分由三個因子組成 (右因子 $B$、中間因子 $A$、左因子 $C$)。

6.1 對角塊

對角塊容易處理，它們是更小規模的自相似問題。第 $j$ 個塊計算 $\text{SSM}(A_R, B_R, C_R)(x_R)$，範圍 $R = jQ : (j+1)Q$。對於小塊長度 $Q$，使用對偶二次 SMA 形式更有效。塊可以平行計算。

6.2 低秩塊

低秩分解有 3 個項，相應地有 3 部分計算：

• 右因子 (Right Factors)：稱為 $B$-塊因子，計算每塊的最終狀態 (假設初始狀態為 0)。形狀為 $(N, P)$ 的張量，與擴展隱狀態 $h$ 相同維度。
• 中心因子 (Center Factors)：稱為 $A$-塊因子，是長度 $T/Q$ 的純量 SSM 掃描，計算考慮所有先前輸入後的真實塊邊界隱狀態。
• 左因子 (Left Factors)：稱為 $C$-塊因子，計算每塊的輸出 (考慮先前塊的影響，但塊內輸入為 0)。

6.3 計算成本

設 $N = P = Q$，總計算成本為：

• 總 FLOP：$O(TN^2)$
• 總記憶體：$O(TN)$
• 主要工作為形狀 $(N, N)$ 的矩陣乘法

與純 SSM 和注意力的比較：

SSD 與線性 SSM 有相同的 FLOP 計數，但避免了純量操作，能利用矩陣乘法單元。與二次注意力相比，SSD 具有可控的狀態擴展 $N$，將歷史壓縮到固定大小，而注意力需要快取整個歷史。

完整 PyTorch 實作 (Listing 1)：

def segsum(x):
    """Naive segment sum calculation. exp(segsum(A)) produces a 1-SS matrix,
       which is equivalent to a scalar SSM."""
    T = x.size(-1)
    x_cumsum = torch.cumsum(x, dim=-1)
    x_segsum = x_cumsum[..., :, None] - x_cumsum[..., None, :]
    mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0)
    x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
    return x_segsum

def ssd(X, A, B, C, block_len=64, initial_states=None):
    """
    Arguments:
        X: (batch, length, n_heads, d_head)
        A: (batch, length, n_heads)
        B: (batch, length, n_heads, d_state)
        C: (batch, length, n_heads, d_state)
    Return:
        Y: (batch, length, n_heads, d_head)
    """
    assert X.dtype == A.dtype == B.dtype == C.dtype
    assert X.shape[1] % block_len == 0

    # Rearrange into blocks/chunks
    X, A, B, C = [rearrange(x, "b (c l) ... -> b c l ...", l=block_len) for x in (X, A, B, C)]

    A = rearrange(A, "b c l h -> b h c l")
    A_cumsum = torch.cumsum(A, dim=-1)

    # 1. 計算每塊內的對角輸出 (對角區塊)
    L = torch.exp(segsum(A))
    Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, L, X)

    # 2. 計算每塊內的狀態 (低秩分解的右因子, B 項)
    decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)

    # 3. 計算塊間 SSM 遞迴 (低秩分解的中間因子, A 項)
    if initial_states is None:
        initial_states = torch.zeros_like(states[:, :1])
    states = torch.cat([initial_states, states], dim=1)
    decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
    new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
    states, final_state = new_states[:, :-1], new_states[:, -1]

    # 4. 計算每塊的狀態到輸出轉換 (低秩分解的左因子, C 項)
    state_decay_out = torch.exp(A_cumsum)
    Y_off = torch.einsum('bclhn,bchpn,bhcl->bclhp', C, states, state_decay_out)

    # 加總塊內與塊間項 (對角與離對角塊)
    Y = rearrange(Y_diag+Y_off, "b c l h p -> b (c l) h p")
    return Y, final_state

[Figure 5: SSD 演算法示意圖。透過將狀態空間模型寫為半可分矩陣的矩陣轉換觀點，我們透過區塊分解矩陣乘法演算法發展出更具硬體效率的 SSD 模型計算。對角塊代表塊內計算，離對角塊代表透過 SSM 隱狀態進行的塊間計算。]

7. Mamba-2 架構

透過連結 SSM 與注意力，SSD 框架使我們能為兩者發展共享的詞彙與技術庫。本節我們討論利用 Transformer 原本概念來理解與修改 SSD 層的範例，得到 Mamba-2 架構。

7.1 區塊設計

平行參數投影 (Parallel Parameter Projections)：Mamba-1 從以 SSM 為中心的觀點出發，將選擇性 SSM 層視為從 $X \mapsto Y$ 的映射，SSM 參數 $A, B, C$ 是 $X$ 的函數。在 Mamba-2 中，SSD 層被視為從 $A, X, B, C \mapsto Y$ 的映射。因此，在區塊開始時用單一投影平行產生 $A, X, B, C$ 是合理的。注意這類比於標準注意力，其中 $X, B, C$ 對應於平行創建的 $Q, K, V$ 投影。

採用平行投影對 $A, B, C, X$ 略微減少參數，更重要的是更適合張量平行化。

額外正規化 (Extra Normalization)：在初步實驗中，我們發現大型模型容易產生不穩定。我們透過在區塊最後輸出投影前加入額外正規化層 (例如 LayerNorm、GroupNorm 或 RMSNorm) 來緩解這點，類似 NormFormer 架構。

[Figure 6: Mamba-2 架構。Mamba-2 區塊簡化了 Mamba 區塊，移除循序線性投影；SSM 參數 $A, B, C$ 在區塊開始就產生，而不是作為 SSM 輸入 $X$ 的函數。額外正規化層改善穩定性。$B$ 與 $C$ 投影只有單一頭，由 $X$ 的多頭共享，類比於多值注意力 (MVA)。]

7.2 序列轉換的多頭模式

回顧 SSM 定義為一個序列轉換，其中 $A, B, C$ 參數有狀態維度 $N$。一個多頭序列轉換由 $H$ 個獨立頭組成，總模型維度 $D = $ d_model。

多頭模式 (Multihead Patterns)：

命題 7.2：Mamba 架構的選擇性 SSM (S6) 層可以視為具有頭維度 $P = 1$ (每個通道有獨立 SSM 動態 $A$) 與多輸入 SSM (MIS) / 多值注意力 (MVA) 頭結構：$B, C$ 矩陣 (對應於注意力中的 $K, Q$) 在輸入 $X$ 的所有通道間共享 (對應於注意力中的 $V$)。

我們發現 MVA 模式在實證上表現最佳，這個觀察並非由總狀態大小解釋。

分組頭模式 (Grouped Head Patterns)：類似於 GQA (Grouped-Query Attention)，可將 Mamba-2 中的 MIS 頭模式擴展為分組輸入 SSM (Grouped-Input SSM, GIS)，或同義的分組值注意力 (Grouped-Value Attention, GVA)。

7.3 來自線性注意力的其他 SSD 擴展

核注意力近似於 softmax 注意力：我們在 Mamba-2 中加入靈活的核特徵映射，並應用於 $B$ 與 $C$ 分支。預設我們選擇 $\psi$ 為元素式 Swish/SiLU 函數。

包含正規化 (分母) 項：透過將 $X$ 加上一個額外的全 1 欄，可以一併計算分母項。然而消融實驗顯示 (第 9.4.3 節)，這些核近似方法並未顯著改善效能，因此 Mamba-2 預設不採用它們。

8. SSM 系統優化

我們描述 SSM (尤其是 Mamba-2) 的多項系統優化，用於大規模高效訓練與推論。重點放在大規模訓練的張量平行與序列平行，以及高效微調與推論的變長序列。

8.1 張量平行 (Tensor Parallel)

張量平行 (Tensor Parallelism, TP) 是一種模型平行技術，將每層 (例如注意力、MLP) 拆分到多個加速器上 (例如 GPU)。它最初為 Transformer 開發，要套用到其他架構並不直接。

Mamba-1 的 TP 挑戰：在 Mamba-1 中，由於 $\Delta, B, C$ 是 $x_c$ 的函數，需要在計算 $\Delta, B, C$ 之前進行額外的 all-reduce，最終一個區塊需要兩次 all-reduce，將通訊時間加倍。

Mamba-2 的 TP 設計：我們將 $\Delta, B, C$ 直接從 $u$ 投影 (而非從 $x_c$)，使這些投影矩陣可以拆分。我們在每個區塊內使用 GroupNorm，群組數可被 TP 度數整除，使 TP 群組內 GPU 在區塊內無需通訊。如此每個區塊只需一次 all-reduce，類似 Transformer 中注意力與 MLP 區塊。

[Figure 7: Mamba-2 區塊的平行化。(左：張量平行) 拆分輸入投影矩陣 $W^{(x)}, W^{(z)}$ 與輸出投影矩陣 $W^{(o)}$。每個 SSM 頭住在單一裝置上。選擇 GroupNorm 作為最終正規化避免額外通訊。(右：序列/上下文平行) 類似於 SSD 演算法，用多個裝置可以沿序列維度拆分。每個裝置計算其序列的狀態，然後將狀態傳給下一個 GPU。]

8.2 序列平行 (Sequence Parallel)

對於非常長的序列，可能需要將輸入與激活沿序列長度維度拆分到不同 GPU。有兩種主要技術：

1. 殘差與正規化的序列平行 (SP)：將 TP 中的 all-reduce 分解為 reduce-scatter 與 all-gather。Mamba-2 架構使用相同的殘差與正規化結構，因此 SP 可直接套用。
2. token-mixing 操作的序列平行 (即上下文平行 CP)：對 SSM，每個 worker 取一個初始狀態，計算其輸入的 SSM，回傳最終狀態並傳給下一個 worker。通訊頻寬與 worker 數量呈線性，比注意力的二次擴展好得多。這個分解與 SSD 演算法中的區塊分解完全相同。

8.3 變長序列

在微調或推論時，模型可能需要處理不同長度的輸入。對 SSM 與 Mamba 來說，我們可以將整個批次視為一個長序列，避免在個別序列間傳遞狀態。這等價於將一個序列尾端的 token $t$ 設定 $A_t = 0$，防止資訊流向下一個序列的 token $t+1$。

9. 實證驗證

我們在合成召回任務上 (第 9.1 節)、標準語言建模上 (第 9.2 節) 對 Mamba-2 進行實證評估，並驗證 SSD 演算法比 Mamba-1 高效得多 (第 9.3 節)，最後消融 Mamba-2 架構的各種設計選擇 (第 9.4 節)。

9.1 合成任務：聯想召回

合成聯想召回 (Associative Recall) 任務常用於測試語言模型查找上下文資訊的能力。多查詢聯想召回 (Multi-Query Associative Recall, MQAR) 任務要求模型記憶多個鍵值對。

[Figure 8: MQAR 結果。聯想召回任務對 SSM 來說很有挑戰，因為它們必須將所有相關資訊記憶到其遞迴狀態。SSD 層結合改進的架構允許 Mamba-2 採用更大的狀態大小，比 Mamba-1 甚至原版注意力表現顯著更好。]

結果顯示：Mamba-1 在這項任務上掙扎，而 Mamba-2 在所有設定下都表現良好。即使在控制狀態大小 ($N=16$) 時，Mamba-2 也顯著優於 Mamba-1。增加狀態大小 ($N$ 從 16 到 64 到 256) 持續改善 MQAR 表現，驗證了狀態大小的重要性。

9.2 語言建模

我們遵循 LLM 標準協定，在 The Pile 資料集上訓練並評估 Mamba-2 架構。模型大小 (深度與寬度) 遵循 GPT-3 規格，從 125m 到 2.7B。

9.2.1 縮放定律

[Figure 9: 縮放定律。約 125M 到 1.3B 參數的模型，在 The Pile 上訓練。Mamba-2 匹敵或超過 Mamba 與強力的 Transformer++ recipe。相較於 Transformer 基準，Mamba-2 在效能 (困惑度)、理論 FLOPs 與實際 wall-clock 時間上都 Pareto 主導。]

9.2.2 下游評估

Table 1 (Zero-shot Evaluations)：

對於每個模型大小，Mamba-2 都優於 Mamba，並通常匹敵兩倍大小的 Pythia。

9.2.3 混合模型：將 SSD 層與 MLP 和注意力結合

最近與並行的工作顯示，SSM 層與注意力層的混合架構可改善 Transformer 或純 SSM 模型的品質。我們發現約 10% 的注意力層比例表現最佳。

Table 2 (Combining SSD and Attention Blocks)：350M 模型 48 層，不同數量的注意力層。

Table 3 (Hybrid 評估，2.7B 規模 300B tokens)：

加入僅僅 6 層注意力就能顯著改善純 Mamba-2 (以及 Transformer++)。我們假設 SSM 層作為一般序列到序列映射運作良好，而注意力層作為快速參考序列中先前 token 的檢索機制，避免強制模型將所有上下文壓縮到記憶 (SSM 狀態)。

9.3 速度基準

[Figure 10: 效率基準。(左) SSD 在大狀態擴展 ($N=64$) 時比 Mamba 融合掃描快 2-8 倍，從序列長度 2k 開始比 FlashAttention-2 快。(右) 序列長度 4K：增加狀態擴展使 Mamba 優化掃描實作線性減速。SSD 可處理更大的狀態擴展因子而無太多減速。]

SSD 由於使用矩陣乘法作為子常式，能利用 GPU 上專門的矩陣乘法單元 (張量核)。因此它比 Mamba 的融合聯想掃描快 2-8 倍，從序列長度 2K 起就比 FlashAttention-2 快。

9.4 架構消融

Table 4 (Mamba-2 區塊消融)：

平行投影節省參數且略勝循序投影，更重要的是有利於更大模型的張量平行。額外正規化層也略微改善表現，更重要的是在更大規模時改善訓練穩定性。

Table 5 (多頭結構消融)：所有模型 $N=64$, $P=64$，訓練至 Chinchilla 縮放定律 token 數量。

驚人地，我們發現 MVA 與 MQA 或 MKA 頭模式之間有顯著差異，儘管它們看似非常相似。MVA 模式 (Mamba 原本使用) 表現最好。

核近似消融 (Table 6, 7)：我們嘗試的核近似方法 (cosFormer、Random Feature Attention、Performer、Based、ReBased) 沒有比簡單的點對點非線性激活函數明顯改善。Mamba-2 預設使用 $\psi(x) = \text{Swish}(x)$ 以遵循 Mamba-1。

10. 相關工作與討論

10.1 狀態空間模型

結構化狀態空間模型可以沿以下軸線特徵化：

1. 是時間不變還是時間變化
2. 系統的維度
3. 遞迴轉移 $A$ 的結構

SSD 可以描述為具有 SISO 維度與純量恆等結構的選擇性 SSM。

時間變化 (選擇性)：原始結構化 SSM (S4) 是線性時間不變 (LTI) 系統。SSD 是時間變化的結構化 SSM，也稱為選擇性 SSM (在 Mamba 中引入)。選擇性 SSM 與 RNN 的閘控機制密切相關，包括傳統 RNN (LSTM、GRU) 與更現代的變體 (QRNN、SRU、RWKV、HGRN、Griffin)。

維度與狀態擴展：SSD 與其前身 SSM (S4、H3、Mamba) 的重要特徵是它是單輸入單輸出 (Single-Input Single-Output, SISO) 系統，輸入通道獨立處理。這帶來更大的有效狀態大小 $ND$。Mamba 表明在資訊密集領域 (如語言) 中狀態擴展至關重要。SSD 的主要優勢之一是允許更大的狀態擴展因子而不減慢模型。

結構：與先前結構化 SSM 相比，SSD 對狀態轉移 $A_t$ 的表達能力有所限制。我們假設可能可以精煉結構化矩陣演算法以改進到一般對角 SSM 案例。

10.2 結構化矩陣

狀態空間對偶性的第一個觀點採用這些模型作為矩陣序列轉換或「矩陣混合器 (Matrix Mixer)」的觀點：可以表示為沿序列維度 $T$ 的矩陣乘法的序列轉換。

之前提出的矩陣混合器包括 MLP-Mixer (非結構化矩陣)、FNet (傅立葉變換矩陣)、M2 (Butterfly/Monarch 矩陣)、Toeplitz 矩陣等。SSD 框架的核心結果是將 SSM 視為具有特定結構 (半可分矩陣) 的矩陣混合器。

10.3 (線性) 注意力

與標準 (因果) 注意力相比，SSD 只有兩個主要差異：

1. SSD 不使用標準注意力的 softmax 激活，這是給予注意力其二次複雜度的原因。
2. SSD 將 logits 矩陣乘以一個輸入相關的 1-半可分遮罩。這個遮罩可以視為替代標準注意力中的 softmax。

這個半可分遮罩也可被視為提供位置資訊。元素 $a_t$ 在 RNN 意義上扮演「閘」，或「選擇」機制，其累積乘積 $a_{j:i}$ 控制位置 $i$ 與 $j$ 之間允許的互動量。我們可以將 SSD 的 1-SS 遮罩視為相對位置嵌入 (Relative Positional Embedding) 的更原則性形式 (與 sinusoidal、AliBi、RoPE 等位置嵌入相比)。

我們強調 SSD 並未推廣標準 softmax 注意力，或任何沒有有限特徵映射 $\psi$ 的注意力核變換。與一般注意力相比，SSD 的優勢在於擁有可控的狀態擴展因子 $N$ 來壓縮歷史，而不是二次注意力對整個歷史的快取 (隨序列長度 $T \gg N$ 擴展)。

10.4 相關模型

我們強調一系列正在發展的相關工作，這些工作與 Mamba 和 Mamba-2 非常相似：

• RetNet 與 TransNormerLLM 用衰減項而非累積和推廣線性注意力，提出對偶平行/遞迴演算法以及混合「分塊 (chunkwise)」模式。
• GateLoop 並行提出使用輸入相關衰減因子 $A_t$，並發展與 SSD 相同的對偶二次形式 (它們稱為「替代注意力」)。
• Gated Linear Attention (GLA) 提出帶有資料相關閘的線性注意力變體，並有計算分塊模式的高效演算法。
• HGRN 引入帶有輸入相關閘的 RNN，HGRN2 改進並加入狀態擴展。
• Griffin 與 RecurrentGemma 顯示帶有輸入相關閘控與本地注意力的 RNN 可與強力現代 Transformer 競爭。
• Jamba 顯示將 Mamba 與少量注意力層結合在語言建模上表現非常好。
• xLSTM 採用狀態擴展與其他閘控、正規化和穩定化技術改進 LSTM。
• RWKV-4 是基於不同線性注意力近似 (attention-free Transformer) 的 RNN。最近改進為 RWKV-5/6 (Eagle 與 Finch) 架構，採用了選擇性與狀態擴展概念。

11. 結論

我們提出一個基於充分研究的結構化矩陣類別的理論框架，彌合了 SSM 與注意力變體之間的概念差距。這個框架揭示了近期 SSM (例如 Mamba) 在語言建模上為何能與 Transformer 表現相當的洞見。此外，我們的理論工具提供了改進 SSM (以及潛在地改進 Transformer) 的新想法，連結了兩邊的演算法與系統進展。作為展示，這個框架引導我們設計了一個位於 SSM 與結構化注意力交集處的新架構 (Mamba-2)。

致謝

感謝 Angela Wu 對於如何以數值穩定方式高效計算 $\Delta$ 梯度的建議。感謝 Sukjun Hwang 與 Aakash Lahoti 協助 MQAR 實驗。

References

[1] Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints". arXiv:2305.13245 (2023).

[2] Yaroslav Aksenov et al. "Linear Transformers with Learnable Kernel Functions are Better In-Context Models". arXiv:2402.10644 (2024).

[3] Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. "In-Context Language Learning: Architectures and Algorithms". ICML 2024.

[4] Ameen Ali, Itamar Zimerman, and Lior Wolf. "The Hidden Attention of Mamba Models". arXiv:2403.01590 (2024).

[5] Simran Arora, Sabri Eyuboglu, Aman Timalsina, et al. "Zoology: Measuring and Improving Recall in Efficient Language Models". ICLR 2024.

[6] Simran Arora, Sabri Eyuboglu, Michael Zhang, et al. "Simple Linear Attention Language Models Balance the Recall-Throughput Tradeoff". ICML 2024.

[7] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. "Neural Machine Translation by Jointly Learning to Align and Translate". ICLR 2015.

[9] Maximilian Beck et al. "xLSTM: Extended Long Short-Term Memory". arXiv:2405.04517 (2024).

[10] Stella Biderman et al. "Pythia: A Suite for Analyzing Large Language Models across Training and Scaling". ICML 2023.

[14] Aleksandar Botev et al. "RecurrentGemma: Moving Past Transformers for Efficient Open Language Models". arXiv:2404.07839 (2024).

[16] James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. "Quasi-recurrent Neural Networks". arXiv:1611.01576 (2016).

[18] Tom Brown et al. "Language Models are Few-shot Learners". NeurIPS 2020.

[19] Krzysztof Choromanski et al. "Rethinking Attention with Performers". ICLR 2021.

[20] Aakanksha Chowdhery et al. "PaLM: Scaling Language Modeling with Pathways". JMLR 2023.

[23] Tri Dao. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning". ICLR 2024.

[25] Tri Dao, Daniel Y. Fu, et al. "Hungry Hungry Hippos: Towards Language Modeling with State Space Models". ICLR 2023.

[29] Soham De et al. "Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models". arXiv:2402.19427 (2024).

[35] Leo Gao et al. "The Pile: An 800GB Dataset of Diverse Text for Language Modeling". arXiv:2101.00027 (2020).

[39] Albert Gu. "Modeling Sequences with Structured State Spaces". PhD thesis, Stanford University, 2023.

[40] Albert Gu and Tri Dao. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces". arXiv:2312.00752 (2023).

[42] Albert Gu, Karan Goel, and Christopher Ré. "Efficiently Modeling Long Sequences with Structured State Spaces". ICLR 2022.

[43] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. "On the Parameterization and Initialization of Diagonal State Space Models". NeurIPS 2022.

[49] Sepp Hochreiter and Jürgen Schmidhuber. "Long Short-Term Memory". Neural Computation 9.8 (1997).

[50] Jordan Hoffmann et al. "An Empirical Analysis of Compute-Optimal Large Language Model Training". NeurIPS 2022.

[51] Samy Jelassi et al. "Repeat After Me: Transformers Are Better Than State Space Models at Copying". ICML 2024.

[52] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention". ICML 2020.

[53] Tobias Katsch. "GateLoop: Fully Data-Controlled Linear Recurrence for Sequence Modeling". arXiv:2311.01927 (2023).

(完整參考文獻共 100+ 條，請見原論文)

術語對照表