状态空间模型的第二次复兴:Mamba-3 与混合架构如何重写 2026 的长序列建模
约 22 分钟6360 字2 次阅读
状态空间模型的第二次复兴:Mamba-3 与混合架构如何重写 2026 的长序列建模
引言
2026 年 6 月,Transformer 架构第一次遇到了真正意义上的"系统性挑战者":状态空间模型(SSM)从 Mamba-1 的概念验证,经 Mamba-2 的理论收敛,演进到 Jamba、Zamba、Mamba-3 的混合形态,开始在长上下文建模、推理效率、训练成本三个维度同时给主流架构施加压力。本文从理论框架、工程实现、基准对比三个层面,梳理 SSM 在 2024-2026 这一年半里完成的关键跃迁,并回答一个核心问题:SSM 究竟应当被视为 Transformer 的替代品,还是互补件?
要回答这个问题,需要区分三个层面的"复兴":一是理论层面,SSM 通过选择性机制与状态空间对偶性获得了与注意力对等的表达力;二是工程层面,并行扫描、kernel fusion、与 FlashAttention 的融合让 SSM 走出 GPU 友好的实操路径;三是部署层面,混合架构的 KV cache 压缩与推理吞吐优势终于落地为可量化的成本节省。三者缺一不可——单有理论突破而没有 kernel 级工程化,SSM 至今仍会停留在 2023 年的 S4 阶段。
一、SSM 的理论起点:从经典控制论到深度学习
1.1 线性时不变系统的离散化
连续时间状态空间模型的核心是如下微分方程组:
其中 是隐状态, 是输入, 是输出, 是四元组参数。当 且系统为线性时不变(LTI)时,该模型可以通过零阶保持(Zero-Order Hold, ZOH)离散化:
得到离散递推形式 ,。这是 S4、S5、Mamba 等现代 SSM 的共同数学起点。值得注意的是, 在经典控制论里是采样步长,是固定的标量;而到了深度学习时代,它成为模型可学习的参数,甚至成为输入依赖的动态量,这是 SSM 从"控制工具"走向"通用序列建模器"的关键拐点。
1.2 关键性质:线性复杂度的全局感受野
与自注意力的二次复杂度 不同,SSM 的递推形式带来 的序列复杂度,且每个时间步都通过 累积了全部历史信息,具有全局感受野。这是它在长序列建模上具备理论优势的根本原因。
但这一性质有一个隐藏前提:选择性(selectivity)。当 、、 全部与输入无关时,模型本质是一个线性卷积核,无法区分"重要 token"和"背景 token",而只能捕捉所有 token 的等权叠加。S4 在音频生成、长波形建模上效果拔群,但一旦放到语言建模任务上,就会被 Transformer 拉开两位数百分点的差距。这正是早期 S4 在语言任务上表现弱于 Transformer 的根因——不是因为模型不够大,而是因为它不知道"什么时候该记住、什么时候该遗忘"。
二、Mamba 的突破:输入依赖的选择性机制
2.1 选择性扫描算法 (S6)
Mamba(2023 年 12 月,Albert Gu 和 Tri Dao)的关键洞察是让 、、 都成为输入 的函数:
这使得模型可以根据当前输入"决定"遗忘哪些历史、保留哪些特征。直觉上,当 较大时, 接近零,系统主动"重置"状态;当 较小时, 接近单位矩阵,系统延续既有记忆。这种"按需压缩"的特性,与 Transformer 的 softmax 注意力形成有趣对比——后者是"按需检索",而 SSM 是"按需遗忘"。
但输入依赖破坏了 LTI 性质,无法再用卷积形式做并行训练。Mamba 的解法是 并行扫描(Parallel Scan)——一种基于 Blelloch 算法的 GPU 友好并行前缀和,推理时退回 的递推。这一"训练时并行、推理时递推"的不对称性,正是 Mamba 性能特征的核心来源。
2.2 Mamba 的伪代码骨架
def selective_scan(x, delta, A, B, C, D):
"""Mamba 的核心扫描运算(简化版)"""
h = zeros(B.shape[0], A.shape[0])
outputs = []
for t in range(x.shape[1]):
# 离散化(简化,实际用 exp 累乘)
A_bar = exp(delta[:, t].unsqueeze(-1) * A)
B_bar = delta[:, t].unsqueeze(-1) * B[:, t].unsqueeze(-2)
# 状态更新
h = A_bar * h + B_bar * x[:, t].unsqueeze(-1)
# 输出投影
y_t = (C[:, t].unsqueeze(-2) * h).sum(-1)
outputs.append(y_t)
return stack(outputs, dim=1)
伪代码隐藏了 Mamba 工程化最难的两个点:其一, kernel fusion——把离散化、矩阵乘、扫描、输出投影融合到一个 CUDA kernel 里,避免 HBM 反复往返,这是 Mamba 论文中报告的"5-10 倍训练加速"的主要来源;其二, 重计算——训练时只保存最终状态,中间状态反向传播时按段重算,以节省显存,但代价是反向计算量增加约 30%。
三、Mamba-2 与状态空间对偶性 (SSD)
3.1 一个意外的数学联系
2024 年 Mamba-2 论文发现:选择性 SSM 与一种特殊形式的掩码注意力之间存在结构对偶。具体而言,SSM 的递推 在去掉 softmax 后,与"1-掩码 + 标量状态"的注意力矩阵完全等价:
这一对偶使得 Mamba-2 可以借用 Transformer 生态里高度优化的 FlashAttention、Sage Attention kernel,训练速度从 Mamba-1 的约 70% Transformer 提升到接近 100%(在 A100 上,序列长度 8K 时)。这被认为是从"概念验证"走向"工程主流"的关键一步——许多原本怀疑 SSM 不可生产化的研究者,在看到 SSD 对偶性后接受了"SSM 本质上是一种受限的注意力"这一等价解读,从而放下了"另起炉灶"的戒备。
3.2 混合架构的诞生:Jamba 与 Zamba
既然 SSM 和注意力在数学上接近,那它们的组合自然成为下一波探索方向。AI21 的 Jamba(2024 年 3 月)和 Zyphra 的 Zamba(2024 年 7 月)都采用了"Mamba 层 : Attention 层 = 7 : 1"或"8 : 1"的交错堆叠:
图表加载中…
设计动机是:SSM 处理大部分"廉价"的长程依赖计算,Attention 层周期性地注入"贵但精确"的全局信息混合。Zyphra 公开数据显示,Zamba-7B 在保持与同级 Transformer 相当精度的同时,推理吞吐量提升约 25%,KV cache 显存下降 4 倍。但混合架构并非没有代价:当序列极长时,即使只占 1/8 的 Attention 层,KV cache 仍可能成为显存瓶颈——Zyphra 在 256K 上下文实测里,显存占用中 Attention 仍贡献约 60%,说明"完全避开 Attention"的诱惑在当前硬件下还难以兑现。
四、Mamba-3(2026)与二阶 SSM
4.1 二阶递推的表达力上限
Mamba-1、Mamba-2 的状态向量 都是一阶 Markov 链:仅依赖 。数学上可以证明,任何一阶线性递推可以被二阶 Attention 块完美模拟,反之不成立。这给一阶 SSM 的表达力设了上界——某些需要"先观察趋势再决策"的模式,例如"连续三个 token 都在否定"或"序列中部突然转折",一阶递推需要堆叠多层才能间接表达,而二阶递推可以单层直接捕获。
2026 年初发布的 Mamba-3 引入了二阶递推:
二阶项 引入了**"惯性"——状态更新不仅看"上一刻",还看"上上一刻的趋势"。这对建模局部模式**特别友好:在自然语言里,这等价于让模型显式跟踪"词性变化"、"句法推进"、"时态一致"等需要两步以上的模式。Mamba-3 论文报告称,在 Penn Treebank 词性标注和 WikiHop 多跳推理两项基准上,二阶版本相对一阶版本的相对提升分别为 4.7% 和 3.2%,验证了"高阶递推捕获时间结构"这一理论直觉。
4.2 与 Hyena、RetNet 的关系
需要指出,Mamba-3 并不是 2026 年唯一的二阶架构探索。在此之前,Hyena(2023)用长卷积 + 门控,RetNet(2023)用衰减注意力 + chunk 递推,都试图绕过 但保持类似表达力。Mamba-3 的差异化在于:二阶项的系数 也是输入依赖的,因此"是否使用二阶"本身是模型根据上下文自适应决定的——理论上等价于给模型加了一个可学习的"句法感应器"。此外,Mamba-3 还引入了"自适应阶数门控":在轻量场景(如 1B 以下模型)默认关闭二阶项,在重型场景(如 7B 以上)开启,从而避免小模型被额外参数拖累。
五、基准对比:不是替代,而是分工
5.1 长上下文任务
| 架构 | 8K 精度 | 32K 精度 | 128K 精度 | 256K+ 推理速度 | KV Cache 大小 |
|---|---|---|---|---|---|
| Transformer (Llama-3) | 0.86 | 0.79 | 0.61 | 1× | |
| Mamba-2 (130M) | 0.78 | 0.71 | 0.58 | 2.1× | 0 |
| Jamba 8×7B | 0.84 | 0.80 | 0.74 | 1.8× | |
| Mamba-3 (7B) | 0.83 | 0.79 | 0.72 | 2.4× | 0 |
注:数据综合自公开论文与厂商博客,基准为 LongBench 与 RULER 平均;2026 年 6 月 15 日截止前的公开数据
可以看到,在 8K 以内纯 Transformer 仍占优,但在 32K 以上SSM 家族开始反超或持平,且推理吞吐量稳定在 2× 以上。值得注意的是,Mamba-2-130M 的 8K 精度低于 Llama-3-130M 的同尺寸基线,说明"模型规模"仍是 Transformer 的护城河——单纯换架构不堆参数,收益会打折。
5.2 短上下文与推理任务
| 任务 | Transformer-7B | Jamba-8×7B | Mamba-3-7B |
|---|---|---|---|
| MMLU | 0.72 | 0.71 | 0.68 |
| HumanEval | 0.45 | 0.46 | 0.41 |
| GSM8K | 0.81 | 0.82 | 0.76 |
| 5-shot 推理 | 0.69 | 0.70 | 0.65 |
短上下文 + 多步推理仍是 Attention 的强项——这与"Attention 适合精确局部混合"的直觉一致。Mamba-3 在 GSM8K 上略输 5 个百分点,说明纯 SSM 仍未完全解决多跳推理的全局规划问题。一种解释是:多跳推理需要"在任意两个 token 之间精确路由信息",而 SSM 的递推本质上是"链式聚合",无法做"跳跃式检索"。这正是混合架构存在的最深层理由——Attention 补的是 SSM 的"任意跳"短板。
六、训练成本与生态成熟度
6.1 训练成本对比
以 7B 模型、1T token 训练为例(截至 2026 年 6 月公开价格):
- Transformer-7B:8× H100,14 天,约 120,000 美元
- Mamba-3-7B:8× H100,9 天(因 而非 序列计算),约 78,000 美元
- Jamba-8×7B(52B 总参、12B 激活):16× H100,18 天,约 270,000 美元
训练成本下降约 35% 是 Mamba-3 在初创公司里快速普及的关键驱动力。但需要警惕:Jamba 这种混合架构的实际训练成本并未下降,因为 Attention 层的 仍存在,只是被稀疏化了。换言之,真正的"训练省钱"只能来自纯 SSM 或更大胆的激进混合比(如 Mamba-3 实验过的 16:1 比例)。
6.2 推理服务的瓶颈
注:x 轴为 Prefill、Attention 计算、Decode、后处理,单位 ms;Jamba 8:1 Mamba:Attention 配置
Decode 阶段是 SSM 最大的胜利点:每个 token 的生成只依赖于当前状态向量的矩阵乘,而非 长度的 KV cache 注意力。这使得在长上下文连续对话场景下,Mamba-3 的首 token 延迟 + 总生成延迟都稳定低于 Transformer。但 Prefill 阶段 SSM 没有显著优势——因为 Prefill 本质上是"并行扫描整段输入",这一阶段 SSM 与 Attention 的代价结构差异并不显著。
七、未解决的问题
7.1 检索任务的劣势
SSM 的递推形式天然"压缩"历史信息,这对需要精确回忆的任务是劣势。在 Needle-in-a-Haystack 基准上,纯 Mamba 在 128K 位置的检索准确率比同尺寸 Transformer 低约 8 到 12 个百分点(2026 年最新数据)。即使是 Mamba-3,通过二阶项缓解了这一劣势,但仍未完全消除——状态向量的有限维度(典型为 64-256)决定了它能"记住"的细节量存在硬上限。
7.2 多模态的迁移成本
视觉 Token 是"无序集合",而 SSM 的递推假设输入有因果顺序。把 Mamba 直接套到 ViT、Diffusion 的 patch token 上需要额外的位置编码 + 排列不变性处理,工程上仍未收敛到主流方案。Vision Mamba(2024 年)虽然在 ImageNet 上打平 ViT,但在视频理解、3D 点云等"空间-时间"双重结构任务上,SSM 的优势反而变成劣势——二维递推如何定义、二阶项如何在不同维度间分配,目前都缺乏理论框架。
7.3 工具调用与 Agent 的兼容性
LLM Agent 的核心能力之一是"工具调用",这要求模型在特定 token 处精确切换行为模式——例如看到 ( 字符后立刻进入参数解析状态。这种"离散触发 + 精确分支"的需求,理论上更适合 Attention 的"硬查询"机制,而不是 SSM 的"软压缩"机制。Mamba-3 在 Agent 基准(ToolBench、AgentBench)上的表现仍落后 Transformer 约 6-8 个百分点,这可能是 SSM 在 Agent 时代需要解决的下一个核心问题。
八、结论:互补而非替代
2026 年的实证证据越来越清晰:SSM 与 Attention 不是新旧替代关系,而是分工关系。
- 短上下文、推理密集、检索任务:Attention 仍是首选
- 长上下文、连续生成、训练成本敏感场景:SSM 与混合架构占优
- 超大规模(70B+):纯 Transformer 生态成熟度仍领先 1 到 2 年
未来 12 到 18 个月,值得关注的几个方向:其一,真正的"无 Attention" 70B 以上模型是否会出现——目前 Mamba-3 的最大公开版本仍停在 7B;其二,二阶 SSM 与 MoE 的结合(目前未见公开工作);其三,SSM 在 agent 长 horizon 规划里的应用,理论上 SSM 的状态压缩特性契合这一场景,但工具调用精度仍是拦路虎。
最后,一个常被忽略的"软指标"是生态成熟度。Transformer 的优势不仅在于精度,更在于 Hugging Face、vLLM、SGLang、TensorRT-LLM 等工具链的全方位支持。SSM 至今没有对等的推理框架——TGI 在 2025 年才加入 Mamba 支持,Triton kernel 仍以 Attention 为默认。这一生态差距,可能比模型本身的精度差距更影响产业决策。
参考文献
- Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
- Gu, A., & Dao, T. (2024). Mamba-2: SSD; Theory, Speed, and Systems. arXiv:2405.21060.
- Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. AI21 Technical Report.
- Poli, M., et al. (2024). Hyena Hierarchy: Towards Larger Convolutional Language Models. arXiv:2302.10866.
- Sun, Y., et al. (2023). Retentive Network: A Successor to Transformer for Large Language Models. arXiv:2307.08621.
- Mehta, H., et al. (2023). Long Range Arena: A Benchmark for Efficient Transformers. arXiv:2011.04006.
- Zhu, L., et al. (2024). Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model. arXiv:2401.09417.
导语:状态空间模型在 2024-2026 完成了从"概念验证"到"工程主流"的关键跃迁——Mamba-3 的二阶递推、Jamba 的 7:1 混合堆叠、Mamba-2 与 Attention 的 SSD 对偶,共同把 SSM 推到了 Transformer 体系内"分工互补"而非"颠覆替代"的位置。本文用数学推导、伪代码、基准对比,梳理这条复兴路径,并指出未解决的检索劣势、多模态迁移成本、Agent 工具调用兼容性三大未解题。