Mamba-3：推理优先的状态空间模型

概述

Mamba-3 是由卡内基梅隆大学、普林斯顿大学、Cartesia AI 和 Together AI 的研究团队于 2026 年 3 月发布的下一代状态空间模型（State Space Model, SSM）。与前身 Mamba-2 专注于训练效率不同，Mamba-3 是首个推理优先（Inference-First） 的状态空间模型，旨在解决现代 LLM 部署中日益重要的推理效率问题。

核心洞见：Mamba-3 通过三个核心方法论改进——指数-梯形离散化、复数值状态更新规则和多输入多输出（MIMO） formulation——在保持推理延迟不变的同时，显著提升了模型质量与能力。¹

问题背景：为什么需要推理优先的设计？

从训练优先到推理优先的范式转变

过去两年中，LLM 的计算范式已从训练优先转向推理优先：

阶段	核心瓶颈	代表模型
2022-2024	预训练效率	Mamba-2（2-8× 训练加速）
2024-2026	后训练与部署	Mamba-3（推理效率优化）

后训练方法（如 RLVR）的扩展以及智能体工作流（agentic workflows）的兴起，使得推理需求急剧增长。Codex、Claude Code 等系统需要持续高强度地调用推理端点。

Mamba-2 的遗留问题

Mamba-2 为加速预训练，对底层 SSM 机制进行了大量简化：

• 将对角转移矩阵简化为标量-单位阵结构
• 每个 token 更新执行的计算量相对于内存移动量过小
• 解码阶段内存-bound：GPU 张量核心处于空闲状态，主要时间花在数据搬运上

┌─────────────────────────────────────────────────────┐
│                    计算范式对比                      │
│                                                     │
│  训练阶段：                                         │
│  ┌─────────────────────────────────────────┐       │
│  │         计算密集型（Compute-bound）      │       │
│  │         硬件持续执行运算                  │       │
│  └─────────────────────────────────────────┘       │
│                       ↓                            │
│  解码阶段：                                         │
│  ┌─────────────────────────────────────────┐       │
│  │         内存密集型（Memory-bound）       │       │
│  │         张量核心空闲，等待数据搬运        │       │
│  └─────────────────────────────────────────┘       │
└─────────────────────────────────────────────────────┘

核心问题

线性模型的核心优势在于：计算量随序列长度线性增长（得益于固定大小的状态）。然而，这也意味着所有历史信息必须压缩到一个固定表示中——这与 Transformer 通过不断增长的 KV 缓存存储所有历史信息的方式截然相反。

核心矛盾：如果不能增长状态，如何让固定状态完成更多工作？²

Mamba-3 的三大核心改进

Mamba-3 从经典控制理论和状态空间理论中汲取灵感，提出三个核心改进方向：

改进方向	技术手段	效果
(1) 更具表达力的递推	指数-梯形离散化	状态动力学更丰富
(2) 更强的状态追踪能力	复数值状态空间	解决奇偶校验等任务
(3) 更高的硬件利用率	MIMO formulation	不增加延迟下提升质量

---

1. 指数-梯形离散化（Exponential-Trapezoidal Discretization）

1.1 连续时间状态空间模型

状态空间模型最初是一个简单的连续时间常微分方程（ODE）：

$$\begin{array}{rl}{h}^{\prime }(t)& =A(t)h(t)+B(t)x(t)\\ y(t)& =C^{\top }(t)h(t)\end{array}$$

其中：

• $x(t)\in \mathbb{R}$ 是输入
• $y(t)\in \mathbb{R}$ 是输出
• $h(t)\in {\mathbb{R}}^N$ 是大小为 $N$ 的隐藏状态（也称状态大小）
• $A\in {\mathbb{R}}^{N\times N}$ 是状态转移矩阵
• $B\in {\mathbb{R}}^N$、$C\in {\mathbb{R}}^N$ 是输入/输出投影

1.2 零阶保持（ZOH）离散化

为获得可用的递推形式，需要对连续 ODE 进行离散化。Mamba-1 和 Mamba-2 使用零阶保持（Zero-Order Hold, ZOH）离散化：

$$\begin{array}{rl}{h}_t& =e^{{\Delta }_t{A}_t}{h}_{t-1}+A_t^{-1}(e^{{\Delta }_t{A}_t}-I)B_t{x}_t\\ y_t& =C_t^{\top }{h}_t\end{array}$$

其中 ${\Delta }_t$ 是时间步长。

1.3 指数-欧拉（Exponential-Euler）方法

Mamba 实际实现中采用的是指数-欧拉方法的启发式近似：

$$\begin{array}{rl}{\bar{A}}_t& =e^{{\Delta }_t{A}_t}\\ {\bar{B}}_t& ={\Delta }_t{B}_t\end{array}$$

这对应于将积分 ${\int }_{{\tau }_{t-1}}^{{\tau }_t}\exp (({\tau }_t-\tau )A_t)B(\tau )x(\tau )d\tau$ 用欧拉法则近似。

1.4 新的指数-梯形离散化

Mamba-3 提出了更一般的离散化框架。核心思想是：

1. 使用积分因子（integrating factor） $e^{-{\int }_0^{t}A(s)ds}$ 调整系统动力学
2. 将状态转移和状态输入积分解耦处理

具体推导过程如下：

对系统 $h^{\prime }(t)=A(t)h(t)+B(t)x(t)$ 应用积分因子 $z(t):=e^{{\int }_0^t-A(s)ds}$：

$$\begin{array}{rl}z(t)h^{\prime }(t)& =z(t)A(t)h(t)+z(t)B(t)x(t)\\ \frac{d}{dt}(z(t)h(t))& =z(t)B(t)x(t)\end{array}$$

在区间 $[{\tau }_{t-1},{\tau }_t]$ 上积分：

$$z({\tau }_t)h({\tau }_t)-z({\tau }_{t-1})h({\tau }_{t-1})={\int }_{{\tau }_{t-1}}^{{\tau }_t}z(\tau )B(\tau )x(\tau )d\tau$$

重新整理得到：

$$h({\tau }_t)=\exp \left({\int }_{{\tau }_{t-1}}^{{\tau }_t}A(s)ds\right)h({\tau }_{t-1})+{\int }_{{\tau }_{t-1}}^{{\tau }_t}\exp \left({\int }_{\tau }^{{\tau }_t}A(s)ds\right)B(\tau )x(\tau )d\tau$$

1.5 梯形规则（Trapezoidal Rule）

欧拉法则只使用区间的一个端点，是一阶近似（局部截断误差为 $O({\Delta }_t^2)$）。Mamba-3 采用梯形规则，使用区间两端点的凸组合来近似积分，是二阶近似（局部截断误差为 $O({\Delta }_t^3)$）：

$${\int }_{{\tau }_{t-1}}^{{\tau }_t}z(\tau )B(\tau )x(\tau )d\tau \approx {\Delta }_t\left[(1-{\lambda }_t)\exp ({\Delta }_t{A}_t)B_{t-1}{x}_{t-1}+{\lambda }_t{B}_t{x}_t\right]$$

其中 ${\lambda }_t\in [0,1]$ 是可学习的参数（门控机制）。

1.6 Mamba-3 的递推形式

最终的指数-梯形离散化递推形式为：

$$h_t={\alpha }_t{h}_{t-1}+{\beta }_t{B}_{t-1}{x}_{t-1}+{\gamma }_t{B}_t{x}_t$$

其中：

• ${\alpha }_t=e^{{\Delta }_t{A}_t}$
• ${\beta }_t=(1-{\lambda }_t){\Delta }_t{e}^{{\Delta }_t{A}_t}$
• ${\gamma }_t={\lambda }_t{\Delta }_t$

关键发现：这个三项递推等价于在状态输入上应用宽度为 2 的隐式数据依赖卷积！

# Python 伪代码：指数-梯形离散化
def mamba3_step(h_prev, x_t, x_tm1, dt, A, B, lambda_):
    alpha = torch.exp(dt * A)           # 状态衰减
    beta = (1 - lambda_) * dt * alpha   # 前一时刻贡献
    gamma = lambda_ * dt                # 当前时刻贡献
    
    # 隐式卷积：结合前一刻和当前输入
    h_t = alpha * h_prev + beta * (B @ x_tm1) + gamma * (B @ x_t)
    return h_t

1.7 并行表示与 SSD 框架

与 Mamba-2 的 SSD 框架 类似，Mamba-3 的递推形式可以展开为并行矩阵形式：

$$Y=(L_1\odot C^{\top }B)X+(L_2\odot C^{\top }B)X$$

其中 $L_1$ 是 1-半可分矩阵，$L_2$ 是 2-带矩阵。这一分解是硬件高效训练算法的基础。

---

2. 复数值状态空间（Complex-Valued SSM）

2.1 状态追踪的挑战

当前线性模型的一个著名局限是缺乏状态追踪能力。例如，最简单的任务——奇偶校验（判断 0 和 1 序列的和是否为偶数）——对于 Mamba 模型来说，在常数层数内是不可解的。

理想解决方案需要隐藏状态追踪当前累加和是奇数还是偶数，并根据下一个输入交替，这建模了一个简单的双状态自动机。然而，Mamba 模型将转移矩阵限制为 ${\bar{A}}_t\in [0,1]$，强制模型学习朴素解法：将所有值相加后取模 2。

2.2 旋转的威力

奇偶校验和其他模 $m$ 任务可以通过旋转解决。可视化方法是：有一个 2D 向量绕原点旋转。整个可能的角分布 $[0,2\pi ]$ 被划分为 $m$ 个区间，向量旋转 $\frac{2\pi }{m}$ 来对齐当前的模 $m$ 余数。

复数特征值的引入使得隐藏状态能够在潜在空间中进行旋转，这对奇偶校验或计数等任务至关重要。

2.3 复数值 SSM 的实数表示

在计算机系统中处理复数值非常麻烦，因为其乘法交互很复杂。幸运的是，对角复数值连续 SSM 可以表示为离散实数值 SSM（不引入额外的近似损失）。

核心思想是将大小为 $N$ 的复数值状态 SSM 扩展为大小为 $2N$ 的实数值状态：

$$h_t^{\mathbb{R}}=\left[\begin{array}{c}\text{Re}(h_t)\\ \text{Im}(h_t)\end{array}\right]$$

在指数-欧拉离散化下，得到以下递推：

$$h_t=e^{{\Delta }_t{A}_t}\underset{R_t}{\underbrace{\left[\begin{array}{cc}\cos ({\Delta }_t{\theta }_t)& -\sin ({\Delta }_t{\theta }_t)\\ \sin ({\Delta }_t{\theta }_t)& \cos ({\Delta }_t{\theta }_t)\end{array}\right]}}{h}_{t-1}+{\Delta }_t{B}_t{x}_t$$

对于更大的状态，旋转矩阵 $R_t$ 是块对角的。

2.4 与数据依赖 RoPE 的等价性

关键定理：复数值 SSM 等价于数据依赖的旋转位置嵌入（RoPE）。³

设输出为：

$$y_t=C_t^{\top }{\bar{B}}_t+\cdots +C_t^{\top }(\bar{A}R{)}_{t\cdots 1}^{\times }{\bar{B}}_0$$

由于 $\bar{A}$ 是缩放单位矩阵，我们可以忽略 $\bar{A}$ 项。结果项 $C_i^{\top }{R}_i\cdots R_{j+1}{\bar{B}}_j$ 可以表示为：

$$(R_i\cdots R_0{C}_i{)}^{\top }(R_j\cdots R_0{\bar{B}}_j)$$

因此，旋转可以直接嵌入到 $B$ 和 $C$ 项中，而不是直接调整转移矩阵！

这正是 RoPE trick：将旋转应用到 $B,C$ 投影上，利用高效的 RoPE 内核实现复数动态。

# RoPE trick：高效实现复数旋转
def rope_bc(B, C, theta_cumsum, dt):
    # 计算累积角度
    angle = theta_cumsum * dt
    
    # 应用旋转到 B 和 C
    cos_a = torch.cos(angle)
    sin_a = torch.sin(angle)
    
    # 旋转后的 B 和 C
    B_rot = B * cos_a + rotate(B) * sin_a
    C_rot = C * cos_a + rotate(C) * sin_a
    
    return B_rot, C_rot

2.5 状态追踪能力对比

实验表明，复数值 SSM 能够解决先前 Mamba 模型无法完成的状态追踪任务：

模型	奇偶校验	模运算	括号匹配
Mamba-2（实数值）	✗	✗	✗
Mamba-3（复数值）	✓	✓	✓

---

3. 多输入多输出（MIMO） formulation

3.1 内存-bound 问题的本质

训练阶段是计算密集型（compute-bound），而部署（尤其是解码）是内存密集型（memory-bound）。在解码时，硬件的计算单元大部分时间处于空闲状态，等待数据在不同内存层级间移动。

算术强度（Arithmetic Intensity） 是衡量硬件利用率的指标：执行的计算操作数与内存移动量的比值。

对于 Mamba-2 单头 SSM：

$$\begin{array}{rl}{\mathbf{h}}_t& =a_t{\mathbf{h}}_{t-1}+B_t{\mathbf{x}}_t\\ {\mathbf{y}}_t& =C_t^{\top }{\mathbf{h}}_t\end{array}$$

其中 $a_t$ 是标量衰减，${\mathbf{x}}_t,{\mathbf{y}}_t\in {\mathbb{R}}^P$，${\mathbf{h}}_t\in {\mathbb{R}}^{N\times P}$。

• 内存移动量：约 $2(1+2N+P+NP)$ 字节
• FLOPs：约 $5NP-P$
• 算术强度：约 2.5（远低于 H100 的 300 ops/byte 阈值）

3.2 MIMO 系统定义

从信号处理角度看，从 SISO（单输入单输出）到 MIMO（多输入多输出）的转换正是 SSM 的泛化形式。

MIMO SSM 的定义：

$$\begin{array}{rl}{\mathbf{h}}_t& =a_t{\mathbf{h}}_{t-1}+{\mathcal{B}}_t{\mathbf{x}}_t^{\top }\\ {\mathbf{y}}_t& ={\mathcal{C}}_t^{\top }{\mathbf{h}}_t\end{array}$$

其中：

• ${\mathcal{B}}_t\in {\mathbb{R}}^{N\times R}$、${\mathcal{C}}_t\in {\mathbb{R}}^{N\times R}$
• ${\mathbf{x}}_t\in {\mathbb{R}}^{P\times R}$、${\mathbf{y}}_t\in {\mathbb{R}}^{P\times R}$
• $R$ 是系统的秩（rank）

3.3 关键洞察：保持状态大小，增加计算量

MIMO 的核心思想是：在保持状态大小不变的情况下，增加每步的计算量。

指标	SISO	MIMO ($R=4$)
状态大小	$N\times P$	$N\times P$（不变）
内存移动	$O(NP)$	$O(NP+NRP)$
FLOPs	$O(NP)$	$O(NPR)$
算术强度	~2.5	$\sim 2.5R$

当 $R\ll P,N$（通常 $P=64,N=128,R=4$）时，算术强度随 $R$ 线性增长！

3.4 分块训练算法（Chunked Training）

MIMO 变体将 SISO 的表达能力扩展 $R^2$ 倍，但训练计算量仅增长 $R$ 倍！

这是通过分块训练算法实现的：

1. 序列被划分为大小为 $C$ 的块
2. 块内使用二次（并行）算法计算输出
3. 块间使用线性（顺序）算法聚合状态

对于 MIMO，通过将块大小减小到 $\frac{C}{R}$，总 FLOPs 仅增长 $R$ 倍。

3.5 MIMO 的实际效果

模型规模	困惑度提升	下游任务平均提升	解码延迟
Mamba-3 SISO vs Mamba-2	-0.03	+0.6pp	持平
Mamba-3 MIMO vs SISO	-0.11	+1.2pp	持平
Mamba-3 MIMO 总计	-0.14	+1.8pp	持平

---

4. 硬件效率与实现

4.1 硬件友好的架构设计

Mamba-3 的三大改进都经过精心设计，以确保硬件效率：

改进	实现方式	硬件友好性
指数-梯形离散化	SSD 框架扩展	兼容现有并行算法
复数值 RoPE	利用现有 RoPE 内核	无需新内核
MIMO	TileLang 实现	专用 CUDA 内核

4.2 延迟对比（1.5B 规模）

在 1.5B 规模下，Mamba-3 SISO 在所有序列长度的 prefill + decode 延迟上均优于 Mamba-2、Gated DeltaNet，甚至超过 Llama-3.2-1B（Transformer）！

序列长度    Mamba-2    Mamba-3 SISO    Transformer    Mamba-3 MIMO
   128      基准        更快 (~20%)       更慢           持平
   512      基准        更快 (~25%)       更慢           持平
  2048      基准        更快 (~30%)       更慢           持平

4.3 开源实现

Mamba-3 的实现采用多种技术栈的混合：

• Triton：SISO 版本的预填充
• TileLang：MIMO 版本的高效实现
• CuTe DSL：底层 CUDA 内核优化

# 安装 Mamba-3
MAMBA_FORCE_BUILD=TRUE pip install --no-cache-dir --force-reinstall \
    git+https://github.com/state-spaces/mamba.git --no-build-isolation

---

5. 数学公式汇总

5.1 连续时间 SSM

$$\begin{array}{rl}{h}^{\prime }(t)& =A(t)h(t)+B(t)x(t)\\ y(t)& =C^{\top }(t)h(t)\end{array}$$

5.2 指数-梯形离散化递推

$$h_t={\alpha }_t{h}_{t-1}+{\beta }_t{B}_{t-1}{x}_{t-1}+{\gamma }_t{B}_t{x}_t$$

其中：

• ${\alpha }_t=e^{{\Delta }_t{A}_t}$
• ${\beta }_t=(1-{\lambda }_t){\Delta }_t{e}^{{\Delta }_t{A}_t}$
• ${\gamma }_t={\lambda }_t{\Delta }_t$

5.3 复数值状态转移（含旋转）

$$h_t=e^{{\Delta }_t{A}_t}{R}_t{h}_{t-1}+{\Delta }_t{B}_t{x}_t$$

其中 $R_t$ 是旋转矩阵：

$$R_t=\left[\begin{array}{cc}\cos ({\Delta }_t{\theta }_t)& -\sin ({\Delta }_t{\theta }_t)\\ \sin ({\Delta }_t{\theta }_t)& \cos ({\Delta }_t{\theta }_t)\end{array}\right]$$

5.4 MIMO SSM

$$\begin{array}{rl}{\mathbf{h}}_t& =a_t{\mathbf{h}}_{t-1}+{\mathcal{B}}_t{\mathbf{x}}_t^{\top }\\ {\mathbf{y}}_t& ={\mathcal{C}}_t^{\top }{\mathbf{h}}_t\end{array}$$

5.5 算术强度

$${\text{AI}}_{\text{SISO}}\approx \frac{5NP-P}{2(1+2N+P+NP)}\approx 2.5$$ $${\text{AI}}_{\text{MIMO}}\approx 2.5R$$

---

6. 性能-效率前沿（Pareto Frontier）

Mamba-3 重新定义了亚二次模型在固定推理预算下的性能前沿：

                    性能
                      ↑
                      │    ★ Mamba-3 MIMO
                      │   ╱
                      │  ╱  ★ Mamba-3 SISO
                      │ ╱
                      │╱  ☆ Mamba-2
                      │╱ ☆
                      │╱☆
                     ─┼──────────────────→ 推理延迟
                      │
                      最低

核心成就：

• 在 1.5B 规模下，下游任务平均准确率比 Gated DeltaNet 高 0.6pp
• MIMO 变体再提升 1.2pp，总计 1.8pp
• 在仅使用 Mamba-2 一半状态大小的情况下，达到相当的困惑度

---

7. 与 Mamba-2 的关系

Mamba-3 是对 Mamba-2 SSD 框架的直接扩展：

特性	Mamba-2	Mamba-3
离散化	零阶保持（启发式）	指数-梯形（正式推导）
状态类型	实数值	复数值（RoPE trick）
输入输出	SISO	SISO/MIMO 可选
设计目标	训练效率	推理效率
核心创新	SSD 对偶性	更丰富的动力学

Mamba-3 将 SSD 层替换为更表达力的复数值指数-梯形 SSM，同时保持与现有 SSD 硬件算法的兼容性。

---

参考资料

相关链接

• Mamba-2 状态空间对偶性理论
• Mamba 表达力理论
• Mamba 硬件算法
• 状态空间模型
• Transformer 与注意力机制

---

本文档遵循 Mamba-3 官方论文与作者博客的技术细节，日期：2026-05-03

Footnotes

1. Aakash Lahoti et al., “Mamba-3: Improved Sequence Modeling using State Space Principles”, arXiv:2603.15569, 2026. https://arxiv.org/abs/2603.15569 ↩

2. Cartesia AI Blog, “Mamba-3: An Inference-First State Space Model”, 2026. https://blog.cartesia.ai/p/mamba-3/ ↩

3. Tri Dao, “Mamba-3 Part 2 - Methodological Deep Dive”, 2026. https://tridao.me/blog/2026/mamba3-part2/ ↩