Mamba-3: Improved Sequence Modeling using State Space Principles

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

https://arxiv.org/pdf/2603.15569

摘要

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

1 引言

測(cè)試時(shí)計(jì)算已成為大型語言模型進(jìn)步的關(guān)鍵驅(qū)動(dòng)因素，思維鏈推理和迭代細(xì)化等技術(shù)表明，推理時(shí)擴(kuò)展可以解鎖新能力（Snell et al. 2024; Wu et al. 2025）。并行、智能體工作流的迅速興起僅加劇了對(duì)此類模型高效推理和部署的需求（Anthropic 2026; OpenAI 2026）。這種范式轉(zhuǎn)變使得推理效率（Kwon et al. 2023; Li et al. 2024）至關(guān)重要，因?yàn)?AI 系統(tǒng)的實(shí)際影響現(xiàn)在關(guān)鍵取決于它們?cè)诓渴鹌陂g執(zhí)行大規(guī)模推理的能力。模型架構(gòu)設(shè)計(jì)在確定推理效率方面起著根本作用，因?yàn)榧軜?gòu)選擇直接決定了生成期間的計(jì)算和內(nèi)存需求。雖然基于 Transformer 的模型（Vaswani et al. 2017）是當(dāng)前的行業(yè)標(biāo)準(zhǔn)，但它們從根本上受限于通過 KV 緩存線性增加的內(nèi)存需求和通過自注意力機(jī)制二次方增加的計(jì)算需求。這些缺點(diǎn)激發(fā)了最近關(guān)于次二次方模型的工作路線，例如狀態(tài)空間模型（SSM）和線性注意力，它們?cè)诒３趾愣▋?nèi)存和線性計(jì)算的同時(shí)，實(shí)現(xiàn)了與其 Transformer 對(duì)應(yīng)模型相當(dāng)或更好的性能。這些模型已進(jìn)入主流，諸如 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）。

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

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

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

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

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

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

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

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

2 預(yù)備知識(shí)

2.1 符號(hào)

2.2 SSM 預(yù)備知識(shí)

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

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

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

3 方法論

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

3.3 多輸入，多輸出

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

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

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

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

3.4 Mamba-3 架構(gòu)

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

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

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ī)模運(yùn)行，從而使得在我們的純 Mamba-3 模型中移除了門后 RMSNorm 層（該層在 Mamba-2 中引入用于穩(wěn)定性）。然而，在混合模型中，移除的 RMSNorm 層對(duì)于長(zhǎng)上下文外推至關(guān)重要（表 4）。

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

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

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

4 實(shí)證驗(yàn)證

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

4.1 語言建模

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

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

4.1.1 MIMO

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

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

4.1.2 檢索能力

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

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

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

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

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

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

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

5 相關(guān)工作

5.1 線性時(shí)間序列混合器

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

許多新興的線性注意力（LA）模型旨在通過核特征圖近似 softmax 注意力（Choromanski et al. 2022; Katharopoulos et al. 2020），而最近的模型已經(jīng)拋棄了特征圖，轉(zhuǎn)而使用查詢和鍵之間的原始點(diǎn)積，并由衰減或掩碼調(diào)節(jié)（Yutao Sun et al. 2023; S. Yang, B. Wang, Shen, et al. 2024）。最近，用鍵值對(duì)調(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)存更新來增強(qiáng)傳統(tǒng)線性注意力。這進(jìn)一步激發(fā)了一系列工作，以提高基于 delta 規(guī)則構(gòu)建的線性模型的效率和能力（Hu et al. 2025; Kimi Team et al. 2025）。

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

結(jié)構(gòu)化狀態(tài)空間模型（SSM）是受經(jīng)典信號(hào)處理和動(dò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)移矩陣（例如對(duì)角或低秩加對(duì)角）的線性時(shí)不變（LTI）層，以促進(jìn)長(zhǎng)上下文任務(wù)的高效計(jì)算和穩(wěn)定學(xué)習(xí)（Gu, Goel, and Ré 2022; Gupta, Gu, and Berant 2022; Smith, Warrington, and Linderman 2023）。Mamba-1（Gu and Dao 2024）中向 SSM 引入時(shí)變、輸入依賴的選擇性，減少了自注意力和線性模型在信息密集模態(tài)（尤其是語言建模）上的差距。隨后，Mamba-2（Dao and Gu 2024）通過結(jié)構(gòu)化狀態(tài)空間對(duì)偶性（SSD）形式化了 SSM 與（線性）注意力之間的聯(lián)系，我們?cè)诖斯ぷ髦幸源藶榛A(chǔ)。

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

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

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

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

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

5.3 多輸入，多輸出

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

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

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

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

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

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

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

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

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