Mamba-3: Improved Sequence Modeling using State Space Principles

Mamba-3:利用狀態(tài)空間原理改進序列建模

https://arxiv.org/pdf/2603.15569

摘要

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

1 引言

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

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

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

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

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

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

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

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

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

2 預備知識

2.1 符號

2.2 SSM 預備知識

2.3 結(jié)構(gòu)化掩碼表示與狀態(tài)空間對偶性

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

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

3 方法論

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

3.1 指數(shù) - 梯形離散化

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

3.1.1 指數(shù)調(diào)整離散化概覽

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

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

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

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

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

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

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

3.2 復數(shù)值狀態(tài)空間模型

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

3.2.1 具有指數(shù) - 歐拉離散化的復數(shù) SSM

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

3.2.2 具有指數(shù) - 梯形離散化的復數(shù) SSM

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

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

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

3.3 多輸入，多輸出

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

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

這緩解了參數(shù)的倍數(shù)增長，使其變?yōu)楦侠淼募臃▍?shù)數(shù)量增長。附錄 C 詳細說明了參數(shù)化細節(jié)，且我們論文中的所有 MIMO 變體均通過降低 MLP 寬度，使其參數(shù)數(shù)量與對應的 SISO 模型相匹配。

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

3.4 Mamba-3 架構(gòu)

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

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

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

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

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

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

4 實證驗證

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

4.1 語言建模

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

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

4.1.1 MIMO

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

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

4.1.2 檢索能力

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

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

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

4.3 推理效率與性能的權(quán)衡

4.4 快速 Mamba-3 內(nèi)核

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

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

5 相關工作

5.1 線性時間序列混合器

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

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

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

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

5.2 狀態(tài)追蹤與復數(shù)狀態(tài)空間模型

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

復數(shù)狀態(tài)空間模型。Mamba 之前的結(jié)構(gòu)化 SSM 經(jīng)常是復數(shù)值的，根植于傳統(tǒng) SSM 理論。它們通常也在視覺和音頻等領域表現(xiàn)出色，這些領域具有明確的基于頻率的信息內(nèi)容，而不是語言。雖然一些模型如 H3（Fu et al. 2023）、RetNet（Yutao Sun et al. 2023）和 Megalodon（Ma et al. 2024）在針對語言建模時保留了復數(shù)值 SSM，但它們?nèi)匀幻黠@遜于 Transformer。

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

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

5.3 多輸入，多輸出

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

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

5.4 狀態(tài)空間模型視角

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

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

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

6 結(jié)論與未來工作

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

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