Mamba-3 方法论深度解析

概述

本文深入解析 Mamba-3 的三项核心方法论改进，这些改进均源自经典状态空间理论。与 Mamba-2 的 SSD 框架 一脉相承，Mamba-3 在推理优先（Inference-First）思想的指导下，对底层 SSM 数学进行了系统性升级。¹²

核心问题：如何让固定大小的状态执行更多工作，从而在保持推理效率的同时提升模型表达能力？³

Mamba-3 通过以下三条路径解决这一问题：

1. 更表达性的递推关系：通过指数-梯形（Exponential-Trapezoidal）离散化方案实现
2. 更丰富的状态追踪能力：通过复数值 SSM（等价于数据依赖的 RoPE）实现
3. 更高的硬件利用效率：通过多输入多输出（MIMO）结构实现

1. 升级的离散化方案

1.1 连续时间状态空间模型

状态空间模型最原始的形式是一个连续时间常微分方程（ODE）：

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

其中：

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

在传统控制理论中，这些系统是**线性时不变（LTI）**的，所有参数都是常数。

1.2 零阶保持（ZOH）离散化回顾

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}$$

其中离散化后的参数为：

• ${\bar{A}}_t=e^{{\Delta }_t{A}_t}$
• ${\bar{B}}_t=A_t^{-1}(e^{{\Delta }_t{A}_t}-I)B_t$

一个关键发现：Mamba 实际使用的离散化方法混合了 ZOH 和 Euler 两种方案，其实现为 ${\bar{A}}_t=\exp ({\Delta }_t{A}_t)$、${\bar{B}}_t={\Delta }_t{B}_t$。这种”混合启发式”方法虽然缺乏理论支撑，但在实践中表现优异。¹

1.3 指数调整离散化框架

为建立统一的离散化理论，Mamba-3 引入了**指数调整（Exponential-Adjusted）**方法。

1.3.1 积分因子法

对于简单 ODE $f^{\prime }(t)=Af(t)$，其解析解为 $f(t)=e^{tA}f(0)$。由于参数 $A$ 直接影响变化率，当 $A$ 参数化变化剧烈时，显式方法（如 Euler）被迫使用小步长，限制系统表达能力。

为解决这一问题，对系统 $h^{\prime }(t)=A(t)h(t)+B(t)x(t)$ 施加积分因子 $e^{-{\int }_0^{t}A(s)ds}$：

$$\begin{array}{rl}{e}^{{\int }_0^t-A(s)ds}{h}^{\prime }(t)& =e^{{\int }_0^t-A(s)ds}A(t)h(t)+e^{{\int }_0^t-A(s)ds}B(t)x(t)\\ (e^{{\int }_0^t-A(s)ds}h(t){)}^{\prime }& =e^{{\int }_0^t-A(s)ds}B(t)x(t)\end{array}$$

因为 $(e^{{\int }_0^t-A(s)ds}{)}^{\prime }=-A(t)e^{{\int }_0^t-A(s)ds}$。

1.3.2 离散化推导

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

$$\begin{array}{rl}{\int }_{{\tau }_{t-1}}^{{\tau }_t}\frac{d}{d\tau }(z(\tau )h(\tau ))d\tau & ={\int }_{{\tau }_{t-1}}^{{\tau }_t}z(\tau )B(\tau )x(\tau )d\tau \\ 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 \end{array}$$

其中 $z(t):=e^{{\int }_0^t-A(s)ds}$。整理得：

$$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$$

在 LTV 情形下，对 $A(s)$ 使用右侧保持假设（Right-Hold）：$\forall s\in [{\tau }_{t-1},{\tau }_t],A(s)=A({\tau }_t)=A_t$，则：

$$h_t\approx \exp ({\Delta }_t{A}_t)h_{t-1}+{\int }_{{\tau }_{t-1}}^{{\tau }_t}\exp \left(({\tau }_t-\tau )A_t\right)B(\tau )x(\tau )d\tau$$

1.3.3 恢复 prior Mamba 离散化

ZOH 恢复：假设 $B(\tau )x(\tau )$ 在区间右端点采样：

$$\begin{array}{rl}{\int }_{{\tau }_{t-1}}^{{\tau }_t}\exp \left(({\tau }_t-\tau )A_t\right)B({\tau }_t)x({\tau }_t)d\tau & =B_t{x}_t{\int }_{{\tau }_{t-1}}^{{\tau }_t}\exp \left(({\tau }_t-\tau )A_t\right)d\tau \\ & =A_t^{-1}\left(\exp ({\Delta }_t{A}_t)-I\right)B_t{x}_t\end{array}$$

指数-Euler 恢复：使用 Euler 规则近似积分，并在右端点保持 $B,x$：

$${\int }_{{\tau }_{t-1}}^{{\tau }_t}\exp \left(({\tau }_t-\tau )A_t\right)B({\tau }_t)x({\tau }_t)d\tau \approx {\Delta }_t{B}_t{x}_t$$

这正是 Mamba-1/2 实际使用的离散化方案！

1.4 指数-梯形离散化

1.4.1 梯形规则的二阶精度

Euler 规则仅一阶精度，局部截断误差为 $O({\Delta }_t^2)$。Mamba-3 使用广义梯形规则，提供二阶精度近似：

$${\int }_{{\tau }_{t-1}}^{{\tau }_t}f(\tau )d\tau \approx {\Delta }_t\left[(1-{\lambda }_t)f({\tau }_{t-1})+{\lambda }_{t}f({\tau }_t)\right]$$

其中 ${\lambda }_t\in [0,1]$ 是数据依赖的标量。当 ${\lambda }_t=1/2$ 时恢复经典梯形规则；Mamba-3 发现数据依赖的 ${\lambda }_t$ 效果更好。

1.4.2 最终递推公式

将梯形规则应用于状态-输入积分：

$$\begin{array}{rl}{h}_t& =\exp ({\Delta }_t{A}_t)h_{t-1}+(1-{\lambda }_t){\Delta }_t\exp ({\Delta }_t{A}_t)B_{t-1}{x}_{t-1}+{\lambda }_t{\Delta }_t{B}_t{x}_t\\ & =:{\alpha }_t{h}_{t-1}+{\beta }_t{B}_{t-1}{x}_{t-1}+{\gamma }_t{B}_t{x}_t\end{array}$$

其中：

• ${\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$

1.4.3 隐式数据依赖卷积

关键观察：新的递推公式在状态-输入项上引入了结构化的时间混合。它等价于对 SSM 的状态-输入 $v_t=B_t{x}_t$ 执行宽度为 2 的数据依赖卷积。

这意味着 Mamba-3 不再需要外部的短因果卷积层——离散化本身就隐式地实现了卷积功能！

1.5 并行表示与 SSD 扩展

根据 SSD 框架，展开递推关系（设 $h_{-1}=0$）：

$$\begin{array}{rl}{h}_0& ={\gamma }_0{B}_0{x}_0\\ h_1& =({\alpha }_1{\gamma }_0+{\beta }_1)B_0{x}_0+{\gamma }_1{B}_1{x}_1\\ h_2& ={\alpha }_2({\alpha }_1{\gamma }_0+{\beta }_1)B_0{x}_0+({\alpha }_2{\gamma }_1+{\beta }_2)B_1{x}_1+{\gamma }_2{B}_2{x}_2\\ & ⋮\end{array}$$

输出可表示为矩阵运算：

$$\left[\begin{array}{c}{y}_0\\ y_1\\ y_2\\ ⋮\end{array}\right]=\left(\left[\begin{array}{cccc}{\gamma }_0& & & \\ ({\gamma }_0{\alpha }_1+{\beta }_1)& {\gamma }_1& & \\ {\alpha }_2({\gamma }_0{\alpha }_1+{\beta }_1)& ({\gamma }_1{\alpha }_2+{\beta }_2)& {\gamma }_2& \\ ⋮& & & \ddots \end{array}\right]\odot \left[\begin{array}{cccc}{C}_0^{\top }{B}_0& & & \\ C_1^{\top }{B}_0& C_1^{\top }{B}_1& & \\ C_2^{\top }{B}_0& C_2^{\top }{B}_1& C_2^{\top }{B}_2& \\ ⋮& & & \ddots \end{array}\right]\right)\left[\begin{array}{c}{x}_0\\ x_1\\ x_2\\ ⋮\end{array}\right]$$

该形式可进一步分解为 1-半可分矩阵（Mamba-2 的衰减掩码）和 2-带矩阵的组合。

2. 复数值状态空间模型

2.1 动机：实数值 SSM 的表达能力局限

先前为提升效率而对 SSM 进行简化，导致了状态追踪能力的丧失。大量研究表明，当前线性 RNN 类模型在状态追踪任务上存在理论限制：

• 缺乏时间步间的非线性：限制了模型对复杂依赖关系的建模
• 结构化矩阵转移：如 Mamba-2 将对角转移简化为标量-单位阵

2.1.1 奇偶性（Parity）任务示例

奇偶性任务是检验状态追踪能力的经典基准：判断一串 0 和 1 的和是否为偶数。

理想解法需要隐藏状态追踪当前和是奇数还是偶数，然后在遇到下一个输入时交替状态——这等价于一个简单的两状态自动机。然而，Mamba 模型将转移限制为 ${\bar{A}}_t\in [0,1]$，这强制模型学习”求和后取模”的朴素解法：

$${\bar{A}}_t\in [0,1]\Rightarrow \text{无法实现符号交替}$$

这在短序列上可行，但当序列长度超过状态容量时便失效。

2.2 旋转视角：复数值的优势

2.2.1 旋转解决模运算问题

旋转可以优雅地解决模运算问题。直观理解：将一个 2D 向量绕原点旋转，整个可能的角分布 $[0,2\pi ]$ 被划分为 $m$ 个区间，向量根据当前模余数旋转 $\frac{2\pi }{m}$ 角度。

2.2.2 复数 SSM 的表示

核心结论：对角复数值连续 SSM 可以表示为离散化后的实数值 SSM，不引入任何额外近似损失。

对于状态 $N=2$ 的情况，使用指数-Euler 离散化得到：

$$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$ 是旋转矩阵。对于更大状态，$R_t$ 是块对角的。

2.3 RoPE 等价性：高效实现

2.3.1 从旋转状态到旋转投影

问题：直接旋转隐藏状态需要重新实现内核，涉及大量矩阵乘法。

关键洞察：由于 $\bar{A}$ 是缩放-单位矩阵，我们可以将旋转吸收到 $B,C$ 中。

时间步 $t$ 的输出可建模为：

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

忽略 $\bar{A}$ 项（可吸收到 $C$），得到 $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)$$

2.3.2 RoPE trick

命题（复数 SSM 与数据依赖 RoPE 的等价性）：对 $B,C$ 应用旋转等价于在 SSM 上实现复数值转移。

具体实现：

1. 计算累积角度：${\theta }_t=\text{cumsum}({\Delta }_t\cdot \theta )$
2. 使用 RoFormer 论文中的高效旋转乘法实现

这使得复数值 SSM 可以直接复用高效的 RoPE 内核，无需重新实现！

# 简化的 RoPE 等价实现
def apply_rope_trick(C, theta_cumsum, headdim):
    """RoPE trick: 将旋转应用到投影矩阵"""
    cos_vals = torch.cos(theta_cumsum)
    sin_vals = torch.sin(theta_cumsum)
    
    # 将 C 按维度分为前半和后半
    C_real, C_imag = C[..., :headdim//2], C[..., headdim//2:]
    
    # 应用旋转：C' = C_real * cos - C_imag * sin
    C_rotated = torch.cat([
        C_real * cos_vals - C_imag * sin_vals,
        C_real * sin_vals + C_imag * cos_vals
    ], dim=-1)
    
    return C_rotated

2.3.3 状态追踪能力验证

实验证明，复数值 SSM（通过 RoPE trick 实现）能够解决先前的 Mamba 模型无法完成的状态追踪任务，包括奇偶性等模运算问题。

3. 多输入多输出（MIMO）系统

3.1 问题：推理时的低算术强度

3.1.1 训练 vs 推理的计算特性

阶段	特性	说明
训练	计算密集型（Compute-bound）	硬件持续执行运算
推理/解码	内存密集型（Memory-bound）	计算单元大部分时间空闲，等待数据搬运

3.1.2 SSM 解码的算术强度分析

算术强度定义：执行运算次数 / 内存传输字节数。

典型 Mamba-2 SSM 结构（标量衰减 $a_t$，head 维度 $P$）：

$$\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}$$

其中 ${\mathbf{x}}_t,{\mathbf{y}}_t\in {\mathbb{R}}^P$，${\mathbf{h}}_t\in {\mathbb{R}}^{N\times P}$。

内存流量：$2(1+2N+P+NP)$ 字节（使用 2 字节数据）

FLOPs：约 $5NP-P$（包含 $a_t{\mathbf{h}}_{t-1}$ 缩放、$B_t{\mathbf{x}}_t$ 外积、$C_t^{\top }{\mathbf{h}}_t$ 矩阵乘法）

算术强度：$\approx 2.5$ ops/字节

作为对比，H100 的矩阵乘法算术强度约为 300 ops/字节。这意味着默认 SSM 解码是彻头彻尾的内存密集型操作！

3.2 MIMO 公式

3.2.1 从 SISO 到 MIMO

核心思想：在保持状态大小不变的同时，增加每次状态更新的 FLOPs。

将 ${\mathcal{B}}_t,{\mathcal{C}}_t$ 扩展为 $N\times R$，${\mathbf{x}}_t,{\mathbf{y}}_t$ 扩展为 $P\times R$：

$$\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}$$

其中 $R$ 是系统的秩。

3.2.2 算术强度提升

总 FLOPs：$4NPR+NP-PR=O(R)$ 增加

内存流量：基本保持不变（对于小 $R$）

结果：算术强度随 $R$ 线性增长！当 $R\ll P,N$（通常 $P=64,N=128,R=4$）时效果最佳。

3.3 参数高效实例化

3.3.1 朴素方案的参数爆炸

直接扩展投影大小会导致 $R\times$ 参数增长，因为 $x,B,C$、输出门 $Z$、输出投影都需要调整。

3.3.2 利用多值注意力结构

由于 $B,C$ 在所有 head 间共享（多值注意力结构），可以以较小开销扩展：

$$DN\to DNR(\text{约可忽略})$$

对于 $x,y,Z$（每 head 独有，是主要参数来源），保持原始投影不变，然后使用数据无关的可学习向量对每个维度进行元素级缩放：

$$\text{参数减少}:DPR\to DP+PR$$

3.4 训练 vs 推理的不对称性

3.4.1 训练成本增长 $R$ 倍，而非 $R^2$ 倍

直觉上，MIMO 输出需要 $R^2$ 个 SISO SSM。但分块训练算法使得 FLOPs 仅增长 $R$ 倍：

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

对于 MIMO：

• 块间输出计算增加 $R$ 倍
• 块内输出计算增加 $R^2$ 倍

通过将块大小减小到 $C/R$，总 FLOPs 仅增加 $R$ 倍。

3.4.2 为什么推理延迟不变？

推理时，GPU 张量核心在 SISO 解码中处于空闲状态。MIMO 的额外 FLOPs 可以被空闲的计算单元吸收，而不增加实际延迟：

$$\text{延迟}\approx \text{max}(\text{计算时间},\text{内存访问时间})$$

由于原始 SISO 是内存密集型，MIMO 增加的计算时间仍在内存访问时间范围内，因此总延迟保持不变！

4. 与 Mamba-2 的对比分析

4.1 架构层面的变化

组件	Mamba-2	Mamba-3
离散化	指数-Euler（一阶）	指数-梯形（二阶）
状态转移	标量 $\times$ 单位阵	复数值（旋转）
输入结构	SISO	SISO / MIMO（可选）
短卷积	外部独立层	隐式融入递推
归一化	RMSNorm	RMSNorm + QKNorm

4.2 性能对比

4.2.1 延迟基准（H100, batch=128）

Prefill 延迟（秒）：

模型	n=512	1024	2048	4096	16384
vLLM (Llama-3.2-1B)	0.26	0.52	1.08	2.08	12.17
Gated DeltaNet	0.51	1.01	2.01	4.00	16.21
Mamba-2	0.51	1.02	2.02	4.02	16.22
Mamba-3 SISO	0.51	1.01	2.02	4.01	16.22
Mamba-3 MIMO (R=4)	0.60	1.21	2.42	4.76	19.44

Prefill + Decode 延迟（秒）：

模型	n=512	1024	2048	4096	16384
vLLM (Llama-3.2-1B)	4.45	9.60	20.37	58.64	976.50
Gated DeltaNet	4.56	9.11	18.22	36.41	145.87
Mamba-2	4.66	9.32	18.62	37.22	149.02
Mamba-3 SISO	4.39	8.78	17.57	35.11	140.61
Mamba-3 MIMO (R=4)	4.74	9.48	18.96	37.85	151.81

关键发现：

• Mamba-3 SISO 在所有序列长度上实现最快的 prefill + decode 延迟
• Mamba-3 MIMO 速度与 Mamba-2 相当，但性能显著更强

4.2.2 状态大小 vs 性能

在受控实验中分析状态大小（推理速度的代理指标）与性能的关系：

• Mamba-3 在相同的困惑度下，可以使用 Mamba-2 一半的状态大小
• 这意味着更快的推理速度，同时保持模型质量

4.3 训练效率

Mamba-3 与 Mamba-2 在架构形状（模型维度、状态大小等）上完全兼容，训练时间相当。

MIMO 变体需要更长的训练时间（因为 FLOPs 增加 $R$ 倍），但推理延迟不变。

5. 实现细节

5.1 内核技术栈

Mamba-3 使用三层域特定语言（DSL）实现高效内核：

DSL	用途	特点
Triton	Prefill（SISO）	标准开发语言，可控分块与融合，平台无关
TileLang	Prefill（MIMO）	显式共享内存控制，策略性内存层级操作
CuTe DSL	Decode	Python 接口生成 CUDA 级内核，细粒度 tensor 布局控制

5.2 融合结构

Mamba-3 的算法设计使得这些不同层次的实现成为可能：

• 指数-梯形递推：简单的三项结构，易于融合
• RoPE trick：复用现有高效 RoPE 内核
• MIMO：矩阵-矩阵乘法（matmul）而非外积，可利用快速 tensor core

6. 总结与展望

6.1 核心贡献

Mamba-3 的三项方法论改进均源自经典状态空间理论：

1. 指数-梯形离散化：二阶精度递推，隐式实现卷积，移除外部短卷积层
2. 复数值 SSM：通过 RoPE trick 高效实现，解决状态追踪任务
3. MIMO 结构：在不增加推理延迟的情况下，提升模型表达能力和硬件利用率

6.2 更广阔的图景

Mamba-3 代表了 SSM 设计的范式转变：从训练优先到推理优先。这一转变反映了 LLM 部署的现实需求：

• 后训练方法（RLVR）需要大量推理 rollouts
• Agentic 工作流持续高强度调用推理端点
• 实时交互场景要求极低延迟

6.3 相关阅读

• Mamba-3：推理优先的状态空间模型 — 模型概述与应用场景
• Mamba-2 状态空间对偶性理论 — SSD 框架的数学基础
• Cartesia Blog: Mamba-3 — 官方产品视角解读³
• Goomba Lab: Mamba-3 Part 2 — 原论文方法论详解¹
• arXiv: Mamba-3 — 完整论文²

---

Footnotes

1. Mamba-3 Part 2 - Methodological Deep Dive, Goomba Lab, 2026. https://goombalab.github.io/blog/2026/mamba3-part2/ ↩ ↩² ↩³

2. Lahoti, A., Li, K. Y., Chen, B., Wang, C., Bick, A., Kolter, J. Z., Dao, T., & Gu, A. (2026). Mamba-3: Improved Sequence Modeling using State Space Principles. arXiv:2603.15569. https://arxiv.org/abs/2603.15569 ↩ ↩²

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