利用狀態空間原理改進序列建模：Mamba-3

Mamba-3: Improved Sequence Modeling using State Space Principles

Mamba-3:利用狀態空間原理改進序列建模

https://arxiv.org/pdf/2603.15569

摘要

擴展推理時計算已成為大型語言模型（LLM）性能的重要驅動因素，使得推理效率成為與模型質量并列的模型設計核心焦點。雖然當前基于 Transformer 的模型提供了強大的模型質量，但其二次方計算和線性內存使得推理成本高昂。這激發了具有降低的線性計算和恒定內存需求的次二次方模型的發展。然而，許多最近的線性模型為了算法效率犧牲了模型質量和能力，在狀態追蹤等任務上失敗。此外，它們理論上線性的推理在實踐中仍然硬件效率低下。在推理優先視角的指導下，我們引入了三個受線性模型的狀態空間模型（SSM）視角啟發的核心方法改進。我們結合了：(1) 源自 SSM 離散化的更具表現力的遞歸，(2) 能夠實現更豐富狀態追蹤的復數值狀態更新規則，以及 (3) 在不增加解碼延遲的情況下提高模型性能的多輸入多輸出（MIMO）公式。結合架構改進，我們的 Mamba-3 模型在檢索、狀態追蹤和下游語言建模任務上取得了顯著增益。在 1.5B 規模上，與下一個最佳模型（Gated DeltaNet）相比，Mamba-3 將平均下游準確率提高了 0.6 個百分點，Mamba-3 的 MIMO 變體進一步將準確率提高了 1.2 個百分點，總增益為 1.8 個百分點。在狀態大小實驗中，盡管使用其前身一半的狀態大小，Mamba-3 實現了與 Mamba-2 相當的困惑度。我們的評估證明了 Mamba-3 推進性能 - 效率帕累托前沿的能力。

1 引言

測試時計算已成為大型語言模型進步的關鍵驅動因素，思維鏈推理和迭代細化等技術表明，推理時擴展可以解鎖新能力（Snell et al. 2024; Wu et al. 2025）。并行、智能體工作流的迅速興起僅加劇了對此類模型高效推理和部署的需求（Anthropic 2026; OpenAI 2026）。這種范式轉變使得推理效率（Kwon et al. 2023; Li et al. 2024）至關重要，因為 AI 系統的實際影響現在關鍵取決于它們在部署期間執行大規模推理的能力。模型架構設計在確定推理效率方面起著根本作用，因為架構選擇直接決定了生成期間的計算和內存需求。雖然基于 Transformer 的模型（Vaswani et al. 2017）是當前的行業標準，但它們從根本上受限于通過 KV 緩存線性增加的內存需求和通過自注意力機制二次方增加的計算需求。這些缺點激發了最近關于次二次方模型的工作路線，例如狀態空間模型（SSM）和線性注意力，它們在保持恒定內存和線性計算的同時，實現了與其 Transformer 對應模型相當或更好的性能。這些模型已進入主流，諸如 Mamba-2（Dao and Gu 2024）和 Gated DeltaNet（GDN）（Schlag, Irie, and Schmidhuber 2021; S. Yang, B. Wang, Y. Zhang, et al. 2025）等層最近被納入大規模混合模型中，這些模型以高得多的效率匹配純 Transformer 替代方案的性能（Kimi Team et al. 2025; NVIDIA et al. 2025; Tencent Hunyuan Team et al. 2025; A. Yang et al. 2025）。

盡管線性模型取得了成功，但在提高其性能方面仍有重大進展空間，特別是在推進模型質量和推理效率之間的帕累托前沿方面。例如，Mamba-2 旨在提高相對于 Mamba-1（Gu and Dao 2024）的訓練速度和簡單性，但犧牲了一些表現力，因此在推理匹配的模型上表現更差。此外，它們已被證明缺乏某些能力，例如糟糕的狀態追蹤能力，即簡單地確定比特序列的奇偶性（Grazzi, Siems, Zela, et al. 2025; Sarrof, Veitsman, and Hahn 2024）。最后，盡管這些次二次方模型因其理論上高效的推理而受到重視并因此被廣泛采用，但它們的推理算法在硬件上效率不高。特別是，因為這些算法是從訓練視角開發的，它們的解碼階段具有低算術強度（FLOPs 與內存流量的比率），導致大部分硬件保持空閑。

為了從推理優先范式開發性能更高的模型，我們在 Mamba-2 之上引入了三個核心方法變更，受次二次方模型的以 SSM 為中心的視角影響。

指數 - 梯形離散化。我們提供了一種離散化時變、選擇性 SSM 的簡單技術。通過我們的框架，我們可以推導幾種新的離散化方法。我們的一個實例化，稱為“指數 - 歐拉”，形式化了 Mamba-1 和 Mamba-2 的啟發式離散化，此前缺乏理論依據。我們新的“指數 - 梯形”實例化是“指數 - 歐拉”的更具表現力的泛化，其中遞歸可以展開以揭示應用于 SSM 輸入的隱式卷積。結合顯式 B , C偏置項，Mamba-3 可以在經驗上替換語言模型架構中的短因果卷積，此前假設這對于遞歸模型是必不可少的。

復數值狀態空間模型。通過將 Mamba-3 的底層 SSM 視為復數值，我們實現了比 Mamba-2 更具表現力的狀態更新。這種更新規則的變更旨在使訓練和推理輕量級，克服了許多當前線性模型中缺乏狀態追蹤能力的問題。我們表明，我們的復數值更新規則等價于數據依賴的旋轉嵌入，并且可以高效計算（Su et al. 2023），并在經驗上證明了其解決先前線性模型能力之外的合成任務的能力。

多輸入多輸出（MIMO）SSM。為了提高解碼期間的 FLOP 效率，我們從基于外積的狀態更新切換到基于矩陣乘法的狀態更新。從 SSM 的信號處理基礎視角來看，這種轉換恰好與從單輸入單輸出（SISO）序列動力學到多輸入多輸出（MIMO）的泛化一致。在這里，我們發現 MIMO 特別適合推理，因為額外的表現力使得在解碼期間內存受限的狀態更新期間能夠進行更多計算，而不增加狀態大小和損害速度。

綜上所述，這些改進構成了我們Mamba-3層的核心。在方法論上，我們注意到這些都自然地從以 SSM 為中心的視角產生，但從現代線性層的其他流行視角（如線性注意力或測試時回歸）來看并非顯而易見；我們在第 5 節進一步討論這些聯系。在經驗上，我們在一系列合成狀態追蹤和語言建模任務上驗證了我們新模型的能力和功能。

• 更好的質量。在 1.5B 規模上，Mamba-3 (MIMO) 相比 Transformer 將下游語言建模準確率提高了+2.2，相比 Mamba-2 提高了+1.9 個百分點，相比 GDN 提高了，而 Mamba-3 (SISO) 相比下一個最佳模型 GDN 提高了+0.6 個百分點。此外，在狀態大小實驗中，狀態大小為 64 的 Mamba-3 (MIMO) 與狀態大小為 128 的 Mamba-2 的困惑度相匹配，有效地以一半的延遲實現了相同的語言建模性能
• 新能力。Mamba-3 對 SSM 狀態的復數化使其能夠解決 Mamba-2 無法解決的合成狀態追蹤任務。我們在經驗上證明，高效的類 RoPE 計算能夠近乎完美地解決算術任務，而沒有 RoPE 的 Mamba-3 和 Mamba-2 的表現并不比隨機猜測更好。
• 推理效率。Mamba-3 (MIMO) 提高了硬件利用率。在固定狀態大小下，相比 Mamba-2，它將解碼 FLOPs 提高了高達4×，同時保持相似的掛鐘解碼延遲，并同時改善困惑度和下游性能。我們發布了 Mamba-3 的快速訓練和推理內核。1

Mamba-3 (SISO) 相比之前的線性模型提高了質量和能力，而 Mamba-3 (MIMO) 相比 Mamba-3 (SISO) 和其他強基線進一步提升了性能，同時保持了與 Mamba-2 匹配的推理速度。我們的兩種 Mamba-3 變體通過其強大的建模能力和硬件高效設計，推進了性能 - 延遲帕累托前沿。

2 預備知識

2.1 符號

2.2 SSM 預備知識

2.3 結構化掩碼表示與狀態空間對偶性

Mamba-2 表明，一大類 SSM 允許采用矩陣形式，該形式將時間步遞歸向量化。通過狀態空間對偶性（SSD）框架，遞歸 SSM 可以在并行形式內表示，該形式結合了元素級掩碼來模擬狀態轉移衰減。

SSD 為線性遞歸與可并行化（基于矩陣乘法）計算形式之間的對偶性提供了一個通用框架。

3 方法論

我們介紹了 Mamba-3，一種具有三個新創新的狀態空間模型：用于更具表現力動態的“指數 - 梯形”離散化（第 3.1 節），用于狀態追蹤的復數值狀態空間（第 3.2 節），以及用于提高建模能力和推理時硬件利用率的多輸入多輸出（MIMO）（第 3.3 節）。這些進展解決了當前次二次方架構在質量、能力和效率方面的局限性。我們在第 3.4 節將這些結合到一個更新后的 Mamba 架構塊中。

3.1 指數 - 梯形離散化

離散化方法在經典控制理論中已得到充分研究，在早期深度學習中的 SSM 工作中使用了幾種規范公式（Gu, Goel, and Ré 2022; Gu, Gupta, et al. 2022; Smith, Warrington, and Linderman 2023）。這些機制傳統上被陳述并應用于線性時不變（LTI）系統，它們的推導不直接適用于線性時變（LTV）系統。此外，雖然 Mamba-1 在沒有證明的情況下將零階保持（ZOH）方法適應于 LTV 系統，但與選擇性 SSM 相關的復雜性促使使用額外的啟發式近似，這種近似缺乏理論依據，并且不對應任何已建立的離散化技術。在以下小節中，我們通過我們的離散化框架形式化了當前 LTV SSM 中使用的先前啟發式方法，并利用它提出了一種更具表現力的離散化方案。

3.1.1 指數調整離散化概覽

我們引入了一種簡單的推導，它導出了針對 LTV（線性時變）狀態空間模型的一類新離散化方法。該方法可以通過多種方式實例化；我們展示了其中一種實例化產生了 Mamba-1/2 中使用的啟發式方法，從而在理論上證明了它的合理性（指數 - 歐拉）。我們還介紹了一種用于 Mamba-3 的更強大的離散化方法（指數 - 梯形）。

我們將公式 (4) 稱為指數 - 歐拉離散化方法，源于指數積分后接歐拉近似。這一推導證明了 Mamba-1/-2 實現中使用的公式的合理性。

指數 - 梯形（Mamba-3）。然而，歐拉法則僅提供狀態輸入積分的一階近似，其局部截斷誤差按縮放。相比之下，我們引入了一種廣義梯形法則，它提供了積分的二階精確近似，比歐拉法則具有更高的精度。具體而言，它用依賴于數據的、兩個區間端點的凸組合來近似積分。這一泛化擴展了經典的梯形法則（Süli and Mayers 2003），后者簡單地平均區間端點（圖 1）。

命題 1（指數 - 梯形離散化）。通過廣義梯形法則近似公式 (16) 中的狀態輸入積分，得到遞歸關系，

這種并行形式使得針對 SSM 輸出的硬件高效、以矩陣乘法（matmul）為中心的訓練計算成為可能。

我們注意到，Mamba-3 的卷積連接也可以通過這種并行對偶形式看出，其中公式 (7) 中乘以 2-帶狀矩陣代表了權重為 β , γ的卷積。在附錄 A.1 中，我們利用 SSD 張量收縮機制證明了該并行形式等價于帶有狀態 - 輸入卷積的基礎 SSM。

備注 5。Mamba-3 的結構化掩碼可以被視為對 Mamba-2 的泛化，后者不使用 2-帶狀矩陣，而是僅擁有包含 γ t
的對角矩陣 (3)。

3.2 復數值狀態空間模型

現代 SSM 的設計以效率為核心目標，這是由擴展到更大模型和更長序列的需求所驅動的。例如，連續的架構逐漸簡化了狀態轉移矩陣：S4 (Gu, Goel, and Ré 2022) 使用了復數值正規加低秩（NPLR）矩陣，Mamba (Gu and Dao 2024) 將其簡化為實數對角線，而 Mamba-2 (Dao and Gu 2024) 進一步將其簡化為單個縮放單位矩陣。盡管這些簡化在很大程度上保持了語言建模性能，但最近的研究 (Grazzi, Siems, Zela, et al. 2025; Merrill, Petty, and Sabharwal 2025; Sarrof, Veitsman, and Hahn 2024) 表明，限制于實數、非負特征值轉移會降低模型在簡單狀態追蹤任務上的能力——這里主要指可解群機制（TC?），例如奇偶性校驗——這可以通過單層 LSTM 解決。這一局限性在 (Grazzi, Siems, Schrodi, et al. 2024) 的定理 1 中被形式化，它源于將轉移矩陣的特征值限制在實數上，這無法表示“旋轉”隱藏狀態動態。例如，考慮定義在二進制輸入 { 0 , 1 }
上的奇偶函數，定義為。此任務可以使用更新公式執行：，其中 R ( ? )
是一個 2-D 旋轉矩陣。這種旋轉動態無法用實特征值表示。

3.2.1 具有指數 - 歐拉離散化的復數 SSM

為了恢復這種能力，我們從復數 SSM (8) 開始，它能夠表示狀態追蹤動態。我們表明，在離散化下（命題 5），復數 SSM 可以表述為具有由 2 × 2 旋轉矩陣組成的塊對角轉移矩陣的實數 SSM（命題 2）。隨后我們表明，這等同于分別在輸入和輸出投影 B , C
上應用數據依賴旋轉嵌入。這一結果建立了復數 SSM 與數據依賴 RoPE 嵌入之間的理論聯系（命題 3）。最后，Su et al. (2023) 中使用的“RoPE 技巧”允許高效實現復數值狀態轉移矩陣，與實數值 SSM 相比，其計算開銷極小。

3.2.2 具有指數 - 梯形離散化的復數 SSM

在推導了具有指數 - 歐拉離散化的復數 SSM 的遞歸后，推廣到指數 - 梯形離散化是相似的。命題 4 提供了 Mamba-3 帶有 RoPE 技巧的完整遞歸。

命題 4（具有指數 - 梯形離散化的旋轉嵌入等價性）。用指數 - 梯形法則（命題 1）離散化復數 SSM 產生遞歸

證明見附錄 B.3。 我們實證驗證了，我們通過數據依賴 RoPE 實現的復數 SSM，能夠解決帶有和不帶有標準 RoPE 的實數值 SSM 無法解決的狀態追蹤任務（表 5b），支持了理論主張。

3.3 多輸入，多輸出

擴展測試時計算開啟了模型能力的新前沿，例如智能體工作流，其中推理占據了整體計算預算越來越大的份額。這使得人們對語言模型的推理效率重新關注，并促進了 SSM 和次二次方層的采用，這些層具有固定大小的隱藏狀態，因此提供更低的計算和內存需求。盡管這些新層與 Transformer 相比具有更低的掛鐘時間，但它們的解碼嚴重受限于內存，導致硬件利用率低。在本節中，我們利用 SSM 視角引入了一種對 Mamba-3 遞歸的方法論改進，允許在不增加解碼掛鐘時間的情況下增加模型 FLOPs，從而在相同解碼速度下獲得更好的模型。

解碼算術強度。為了提高硬件效率，我們需要考慮 token 生成的算術強度，定義為給定操作的 FLOPs 除以輸入 - 輸出字節數。由于 SSM 解碼用空閑計算飽和了內存帶寬（即受限于內存），我們希望增加其算術強度，以有效地將計算與內存 I/O 重疊。更具體地說，Mamba 中單次生成的算術強度約為每字節 2.5 次操作（表 2a），而對于 NVIDIA H100-SXM5，bfloat16 矩陣乘法的算術強度約為每字節 295 次操作（NVIDIA 2022）。因此，SSM 解碼遠未達到計算受限狀態，而且此外，尚不清楚如何調整 Mamba 中的現有參數來緩解硬件效率的缺乏。我們注意到，這一觀察結果通常適用于其他次二次方模型，例如因果線性注意力。

這緩解了參數的倍數增長，使其變為更合理的加法參數數量增長。附錄 C 詳細說明了參數化細節，且我們論文中的所有 MIMO 變體均通過降低 MLP 寬度，使其參數數量與對應的 SISO 模型相匹配。

備注 6。為簡單起見，本節的所有討論均針對較簡單的 2 項遞歸（例如由指數 - 歐拉離散化產生的遞歸）；推廣到 3 項指數 - 梯形遞歸的情況是類似的。

3.4 Mamba-3 架構

整體架構遵循 Llama（Grattafiori et al. 2024），交替使用 Mamba-3 和 SwiGLU 塊，并采用預歸一化（pre-norm）。Mamba-3 塊保留了其前身的整體布局，同時引入了幾個關鍵修改。

更新的 SSM 遞歸。SSD 層被替換為命題 4 中定義的更具表現力的復數值指數 - 梯形 SSM。Mamba-3 默認采用 SISO SSM，以便與其他類 SISO 模型進行公平比較，但其 MIMO 變體可以訓練和部署為基線 Mamba-3 的更強替代方案（表 3）。我們的 SSM A A 是復數的，包含由數據依賴投影產生的實部和虛部。結合圖 2，這被劃分為實數值 A A 和虛數值 Θ Θ；前者像 Mamba-2 中一樣傳入 SSD 黑盒，而后者通過 RoPE 技巧計算。

BC / QK 歸一化。在 B , C
投影之后添加了 RMS 歸一化，這模仿了現代 Transformer（Henry et al. 2020; Wortsman et al. 2023）和其他近期線性模型（Hu et al. 2025; S. Yang, Kautz, and Hatamizadeh 2025）中常用的 QKNorm。我們將其稱為 BC 歸一化（BCNorm）或 QK 歸一化（QKNorm），兩者可互換使用。我們發現 BCNorm 也能穩定大規模運行，從而使得在我們的純 Mamba-3 模型中移除了門后 RMSNorm 層（該層在 Mamba-2 中引入用于穩定性）。然而，在混合模型中，移除的 RMSNorm 層對于長上下文外推至關重要（表 4）。

B , C偏置。類似于 Yu and Erichson (2025)，其證明了在 Mamba-1 的分塊變體中向 B 添加通道特定偏置賦予了通用逼近能力，Mamba-3 在 BCNorm 之后將可學習的、頭特定的、通道維度的偏置整合到了 B 和 C 分量中。

我們假設這些偏置也在模型中誘導了類似卷積的行為。具體而言，向 B B 和 C C添加偏置將數據獨立組件引入 SSM 中，使其功能更類似于卷積。關于偏置參數化的消融實驗位于附錄 F。

數據獨立偏置參數的組合，加上指數 - 梯形離散化（其本身在狀態輸入上誘導了卷積），在經驗上能夠免除 Mamba-2 和大多數現代遞歸模型中存在的短因果卷積及其伴隨的激活函數（第 4.2 節）。

4 實證驗證

我們通過 Mamba-3 模型在一系列合成和現實世界任務上，實證驗證了我們以 SSM 為中心的方法論變更。第 4.1 節在語言建模和基于檢索的任務上評估 Mamba-3。第 4.2 節消融了我們新 SSM 組件（如離散化和復數轉移）的效果。第 4.3 節探討了 Mamba-3 家族的推理效率以及 MIMO Mamba-3 相比 SISO 變體在固定推理計算下的優勢，第 4.4 節基準測試了我們 Mamba-3 訓練和推理內核的性能。

4.1 語言建模

所有模型均使用 FineWeb-Edu 數據集（Penedo et al. 2024）的 1000 億 token 進行預訓練，使用 Llama-3.1 分詞器（Grattafiori et al. 2024），上下文長度為 2K，采用相同的標準訓練協議。訓練和評估細節可在附錄 D 中找到。

在所有四個模型規模上，Mamba-3 在各種下游任務上優于流行的基線（表 3）。我們強調，Mamba-3 不使用外部短卷積，該卷積已被經驗性地確定為許多高性能線性模型中的重要組件（Allen-Zhu 2025; Gu and Dao 2024; S. Yang, Kautz, and Hatamizadeh 2025）。

4.1.1 MIMO

我們旨在通過在相同設置下訓練秩 R = 4
的 MIMO 模型來調查其語言建模能力，從而進一步驗證 MIMO 的收益。為了確保總參數數量與基于 SISO 的模型相當，我們降低了 MIMO 模型中 MLP 層的內部維度，以補償由于 MIMO 投影導致的增加。例如，在 1.5B 參數模型中，MLP 內部維度僅減少了 6.6%，從 4096 減少到 3824。詳見附錄 C。

在驗證困惑度和我們的語言評估任務套件（表 3）上，我們看到當 Mamba-3 模型從 SISO 轉向 MIMO 時有顯著增益。即，我們在 1.5B 模型上實現了 0.11 的顯著困惑度增益，圖 3 說明了我們驗證損失的向下偏移。在語言評估方面，與 SISO 相比，我們在大多數任務上看到增益，導致相比 SISO 平均增益 1.2 個百分點。

4.1.2 檢索能力

除了標準語言建模外，線性模型的一個重要衡量標準是它們的檢索能力——它們能從序列早期回憶信息的程度如何（A. Arora et al. 2025; S. Arora, Eyuboglu, et al. 2025）。與注意力模型不同，注意力模型可以通過增長的 KV 緩存自由回顧過去上下文，線性模型必須將上下文壓縮到固定大小的狀態中。這種權衡反映在 Transformer 基線顯著更強的檢索分數上。為了在此視角下評估 Mamba-3，表 4 使用我們第 4.1 節中的預訓練 1.5B 模型，在現實世界和合成大海撈針（NIAH）任務（Hsieh et al. 2024）上將其與基線進行比較。我們將任務序列長度限制為 2K token 以匹配訓練設置，并遵循 S. Arora, Eyuboglu, et al. (2025) 和 S. Arora, Timalsina, et al. (2024)，對我們的現實世界任務采用完形填空風格格式，以反映下一個 token 預測目標。

Mamba-3 在現實世界關聯回憶和問答（TQA, SQuAD）上具有競爭力，但在從半結構化或非結構化數據中提取信息時表現掙扎（SWDE, FDA）。然而，在合成 NIAH 任務上，Mamba-3 在大多數情況下超過或匹配基線，并且值得注意的是，相比其前身 Mamba-2，表現出明顯更好的分布外檢索能力。

使用混合模型改進檢索。由于固定狀態大小的自然基于檢索的弱點，我們預測線性層將主要用在混合架構中，該架構通過二次方自注意力層減輕這一缺點。為了評估 Mamba-3 在此架構范式內的表現，我們以交錯方式訓練相同規模的混合模型，線性層與 NoPE 自注意力（B. Yang et al. 2025）的比例為 5:1。正如先前工作（Waleffe et al. 2024）所見，混合模型優于 Transformer 基線。我們發現，將預輸出投影 RMSNorm（表 4 中的 pre-gate, grouped RMSNorm）重新引入 Mamba-3 層，提高了長度泛化檢索能力，代價是輕微的上下文內現實世界檢索任務，并且當與自注意力混合時，作為線性序列混合骨干網絡具有高度競爭力。然而，由于相互競爭的權衡（附錄 E，表 9），理想的歸一化類型（分組 vs 默認）及其放置位置（門前 vs 門后）仍不清楚，因為我們發現混合模型及其確切特征和動態是復雜且通常反直覺的，最近的工作如 Cabannes et al. (2025) 也呼應了這一點。

4.3 推理效率與性能的權衡

4.4 快速 Mamba-3 內核

我們通過優化的內核補充了 Mamba-3 的方法論進展，這些內核在實際場景中提供快速推理。我們為 Mamba-3 實現了一系列新的推理內核——使用前向（預填充）路徑的 Triton 和解碼的 CuTe DSL——并在表 6 中將它們的每 token 解碼延遲與 Mamba-2 和 GDN 發布的 Triton 內核進行比較。該評估在單個 H100 上以批量大小 128 測量單個解碼步驟，針對 FP32 和 BF16 數據類型；模型為 1.5B 參數，模型維度為 2048，狀態維度 ∈ {64, 128}。在所有配置中，SISO 在基線中實現了最低的延遲。MIMO 憑借其更高的算術強度，在不顯著增加解碼運行時間的情況下增加了解碼 FLOPs。我們的基準測試表明，我們的 CuTe DSL 解碼實現具有競爭力，且 Mamba-3 的額外組件（指數 - 梯形更新、復數值狀態和 MIMO 投影）是輕量級的。這支持了我們整體的推理優先視角：Mamba-3 允許簡單、低延遲的實現，同時提供強大的實證性能。

表 7 基準測試了不同解碼序列長度下的端到端延遲以及相同序列長度的預填充時間。解碼時間與表 6 一致，其中 Mamba-3 (SISO) 最快；Mamba-3 (MIMO) 與 Mamba-2 相當；并且隨著序列長度增長，所有線性方法都比優化注意力更快。我們還看到，正如第 3.3 節所討論的，MIMO 為預填充帶來了適度的開銷。基準測試的詳細信息見附錄 G。

5 相關工作

5.1 線性時間序列混合器

越來越多的工作尋求用線性運行時間的替代方案替換基于二次方 softmax 的注意力機制（Bahdanau, Cho, and Bengio 2014; Vaswani et al. 2017）。主要方法可以分為三大框架：線性注意力、測試時訓練和狀態空間模型。

許多新興的線性注意力（LA）模型旨在通過核特征圖近似 softmax 注意力（Choromanski et al. 2022; Katharopoulos et al. 2020），而最近的模型已經拋棄了特征圖，轉而使用查詢和鍵之間的原始點積，并由衰減或掩碼調節（Yutao Sun et al. 2023; S. Yang, B. Wang, Shen, et al. 2024）。最近，用鍵值對調節狀態內存的快速權重編程器 Schlag, Irie, and Schmidhuber (2021) 也被歸入“線性注意力”這一統稱之下。S. Yang, Kautz, and Hatamizadeh (2025) 和 S. Yang, B. Wang, Y. Zhang, et al. (2025) 源于這一工作路線，并通過用 delta 規則遞歸替換加法內存更新來增強傳統線性注意力。這進一步激發了一系列工作，以提高基于 delta 規則構建的線性模型的效率和能力（Hu et al. 2025; Kimi Team et al. 2025）。

一條并行的測試時訓練（TTT）或測試時回歸（TTR）工作線將序列建模視為推理期間的在線學習任務。在這里，遞歸狀態代表過去輸入的壓縮摘要，遞歸步驟更新狀態以記憶新信息（Yu Sun et al. 2025; Tandon et al. 2025; T. Zhang et al. 2025）。等價地，這些方法可以被視為全局回歸目標的優化，遞歸狀態更新代表迭代優化過程，例如梯度下降的變體（K. A. Wang, Shi, and Fox 2025）。

結構化狀態空間模型（SSM）是受經典信號處理和動態系統啟發的現代遞歸模型的另一種視角。早期版本的 SSM 如 S4（Gu, Goel, and Ré 2022; Gupta, Gu, and Berant 2022; Smith, Warrington, and Linderman 2023）使用具有結構化狀態轉移矩陣（例如對角或低秩加對角）的線性時不變（LTI）層，以促進長上下文任務的高效計算和穩定學習（Gu, Goel, and Ré 2022; Gupta, Gu, and Berant 2022; Smith, Warrington, and Linderman 2023）。Mamba-1（Gu and Dao 2024）中向 SSM 引入時變、輸入依賴的選擇性，減少了自注意力和線性模型在信息密集模態（尤其是語言建模）上的差距。隨后，Mamba-2（Dao and Gu 2024）通過結構化狀態空間對偶性（SSD）形式化了 SSM 與（線性）注意力之間的聯系，我們在此工作中以此為基礎。

5.2 狀態追蹤與復數狀態空間模型

表現力與狀態追蹤。最近的工作刻畫了遞歸、恒定內存混合器可以維持的狀態類型，揭示了先前基于 SSM 的模型中的算法缺陷。Merrill, Petty, and Sabharwal (2025) 表明，在有限精度下，實際 SSM 坍縮為 TC0，導致在諸如 上的排列組合等任務上失敗，除非擴展原語。類似地，Yu and Erichson (2025) 證明了單層 Mamba 不是通用逼近器。已經提出了幾種修改來提高表現力。例如，同一項工作表明，塊偏置變體僅需細微更改（通過塊分解或通道特定偏置）即可恢復通用逼近屬性。允許負特征值或非三角轉移使得線性 RNN——包括對角和 Householder/DeltaNet 形式——能夠捕捉奇偶性，并在溫和假設下捕捉正則語言（Grazzi, Siems, Zela, et al. 2025）。復數值參數化提供了另一條增強表現力的途徑。

復數狀態空間模型。Mamba 之前的結構化 SSM 經常是復數值的，根植于傳統 SSM 理論。它們通常也在視覺和音頻等領域表現出色，這些領域具有明確的基于頻率的信息內容，而不是語言。雖然一些模型如 H3（Fu et al. 2023）、RetNet（Yutao Sun et al. 2023）和 Megalodon（Ma et al. 2024）在針對語言建模時保留了復數值 SSM，但它們仍然明顯遜于 Transformer。

此外，由于這些模型是 LTI 的，并且使用與現代選擇性 SSM（如 Mamba）非常不同的算法計算（特別是卷積或顯式遞歸），它們通常不使用 RoPE 技巧來處理復數部分。一個例外是 RetNet，它引入了一種介于線性注意力和 Mamba-2 之間的模型，使用常數標量衰減（不同于 LA 中的無衰減和 Mamba-2 中的數據依賴衰減），并通過 RoPE 實現額外的常數復數相位。

一般來說，經驗發現復數對語言建模沒有幫助，因此在 Mamba-1 及其后繼版本中被淘汰，包括線性注意力和測試時訓練方面的并行工作線。Mamba-3 代表了第一個具有復數值狀態轉移的現代遞歸模型，引入它是為了增加表現力和狀態追蹤能力的特定目的。通過結合 RoPE 技巧，據我們所知，這代表了基于理論動機的數據依賴 RoPE 的首次使用。

5.3 多輸入，多輸出

S4（Gu, Goel, and Ré 2022）是一個單輸入單輸出（SISO）LTI 系統，其中輸入的每個維度都被分配了自己獨立的 SSM。此類 SISO 模型具有比經典 RNN 大得多的遞歸狀態，并且需要更復雜的數學機制來高效計算它們。為了簡化模型，S5（Smith, Warrington, and Linderman 2023）和 LRU（Orvieto et al. 2023）用直接應用于整個向量化輸入的多輸入多輸出（MIMO）SSM 替換了一組 SISO SSM。這一改變降低了有效狀態容量，但通過直接使用并行掃描計算遞歸，啟用了一條替代計算路徑。雖然這種狀態容量和建模性能之間的權衡在 LTI 模型中不太明顯，但 Mamba-1 (S6)（Gu and Dao 2024）和 Mamba-2（Dao and Gu 2024）由于大狀態大小在時變設置中的重要性，回到了 SISO 系統。與增加的狀態大小相關的計算瓶頸通過 Mamba-1 的硬件感知并行掃描算法和 Mamba-2 的基于矩陣乘法的算法得到解決。

將 MIMO 引入 Mamba-3 與先前的工作顯著不同。與旨在簡化訓練算法但以略微降低表現力為代價的先前 MIMO 模型不同，Mamba-3 的 MIMO 結構旨在提高建模能力同時保持推理效率。因此，其狀態擴展保持在 Mamba-1/-2 水平，以維持建模能力，同時權衡額外的訓練計算。

5.4 狀態空間模型視角

雖然現代遞歸模型有幾個 largely converge 的不同觀點（第 5.1 節），但每個框架都有略微不同的解釋和動機，可能導致不同的設計空間和擴展。特別是，線性注意力和測試時訓練聯系更緊密，或許可以在關聯記憶框架下歸為一類，該框架明確旨在通過“鍵值”存儲記憶輸入數據；要么通過 LA 中對規范 KV 方法（即二次方注意力）的近似，要么通過最小化 TTT 中的軟優化目標。另一方面，狀態空間模型有不同的淵源，這反映在術語（例如， A , B , C , X 而不是 Q , K , V
）及其自然擴展中。值得注意的是，Mamba-3 的方法論改進都特別與 SSM 觀點相關，并且較少由關聯記憶框架驅動。

1. 指數 - 梯形離散化。SSM 觀點需要對 governing 系統的連續 ODE 進行離散化；我們的指數 - 梯形離散化源于改進的離散化方法。由于關聯記憶方法不使用離散化，如何在替代觀點下解釋諸如指數 - 梯形之類的 3 項遞歸并不明顯。
2. 復數值狀態轉移。復數 SSM 長期以來一直是動態系統的主要內容，將復數值視為選擇性 SSM 的擴展是很自然的。另一方面，關聯記憶框架將 A A 狀態轉移解釋為目標函數的系數，例如對應于優化目標中 L2 正則化（或權重衰減）項的權重（K. A. Wang, Shi, and Fox 2025）。然而，作為回歸目標的系數，復數值是沒有意義的；因此，Mamba-3 在這些框架內并不明顯可解釋。
3. 多輸入，多輸出。MIMO 是狀態空間模型文獻中的經典概念，并不自然出現在關聯記憶（線性注意力或測試時訓練）框架中。然而，我們確實注意到，本文介紹的 MIMO 公式并不直接局限于 SSM 理論——而是由計算視角驅動——并且我們的技術也可以適應其他現代遞歸模型。

線性時間序列模型的開發繼續取得蓬勃進展，這里的討論僅捕捉了其中的一部分。隨著這些模型的不斷發展，我們預計將出現一個不斷增長的統一框架空間、改進的理解和新的泛化。

6 結論與未來工作

我們介紹了 Mamba-3，這是一種狀態空間模型，相比先前的 SSM 具有幾項方法論改進：通過指數 - 梯形離散化實現更強大的遞歸；通過復數值狀態轉移提高表現力；以及通過 MIMO 公式實現更高的推理效率和建模能力。Mamba-3 的基礎 SISO 版本提供了強大的語言建模結果，無論是獨立使用還是在交錯混合架構中，并且在性能 - 效率權衡的帕累托前沿上超越了先前的線性序列模型。MIMO 版本以較慢的訓練速度換取了更強的建模能力，同時與 Mamba-2 相比保持了具有競爭力的推理效率。總的來說，Mamba-3 中的技術展示了從狀態空間模型視角出發的簡單且有理論依據的改進，并為高效序列模型開辟了新的方向和設計原則。

原文鏈接：https://arxiv.org/pdf/2603.15569