大模型依赖注意力读懂长序列，却付出 O(n²) 的代价。Mamba 用状态空间模型（SSM）替代注意力，把序列建模压到线性复杂度。本文拆解它的数学直觉、选择性扫描（selective scan）原理，并给出可运行的 PyTorch 实现。
一、为什么注意力有 O(n²) 的天然瓶颈
Transformer 的注意力让每个 token 直接“看”序列里所有其他 token，这是它长程建模能力强的根源。但代价是：对长度为 n 的序列，要算 n×n 的注意力分数矩阵，计算量和显存都随 n 平方增长。当 n 从 2K 涨到 200K，开销直接翻一万倍——这正是长文本、基因组、音频等场景下注意力最吃力的地方。要缓解，要么像 MLA 多头潜在注意力 那样压缩 KV 缓存，要么干脆换掉注意力本身，而 Mamba 走的是后一条路。
SSM（State Space Model，状态空间模型）是一类比注意力古老得多的序列建模工具：它用一个固定大小的“隐藏状态”来描述系统的过去，新输入进来只更新这个状态，再生成输出。因为状态大小与序列长度无关，整条序列的计算量天然是线性的。
二、状态空间模型（SSM）是什么：用“状态”记住过去
2.1 连续时间的 SSM 直觉
最朴素的 SSM 是一组微分方程：输入信号 x(t) 经过一个状态 h(t) 的演化，再投影成输出 y(t)。状态 h 像是系统的“记忆条”，它会按矩阵 A 衰减、被输入 B 推动，而输出 C 从状态里读出来。和注意力“人人见面”不同，SSM 是“状态替你记着历史”，新来一个 token 只需读当前状态，复杂度与序列长度脱钩。
2.2 离散化：让 SSM 跑在 GPU 上
微分方程不能直接喂给神经网络，需要用步长 Δ 离散化，变成每步的递推。离散化后，状态按 A_bar 衰减、被 B_bar 推动，输出由 C 读出。这一步的“零阶保持”离散化，正是后面 Mamba 让 Δ 跟着输入走的关键伏笔。下面用 PyTorch 写一个最小可运行的 SSM 单步与整序列版本。
import torch
def ssm_step(x_t, h, A, B, C, D):
# x_t: 当前时间步输入 (d,)；h: 上一时刻状态 (n,)
h_new = A * h + B * x_t # 状态演进（线性、与 n 无关）
y_t = C @ h_new + D * x_t # 输出投影
return y_t, h_new
# 示意初始化：对角 A 取负保证稳定衰减
d, n = 16, 64
A = -torch.ones(n)
B = torch.randn(n)
C = torch.randn(n)
D = torch.zeros(d)
h = torch.zeros(n)
seq = torch.randn(200, d) # 一条 200 步的序列
out = []
for x_t in seq: # 逐时间步递推，O(n) 而非 O(n^2)
y_t, h = ssm_step(x_t, h, A, B, C, D)
out.append(y_t)
注意这段代码里每个时间步只做矩阵-向量乘法，状态 h 的大小固定为 n，与序列长度 200 无关。这正是 SSM 线性复杂度的来源：无论序列多长，内存里始终只有一个定长状态在滚动。
三、选择性状态空间：Mamba 的核心突破
3.1 为什么传统 SSM 不够
早期的线性 SSM（如 S4）把 A、B、C、Δ 都设成“与输入无关”的固定参数。这带来一个致命缺陷：它对所有输入一视同仁地做等权平滑，无法像注意力那样“看到某个关键词就重点关注”。换句话说，固定参数 SSM 记性是有的，但“选择性”为零——它不能根据内容决定记住什么、忽略什么，而语言建模恰恰极度依赖这种选择性。
3.2 让参数随输入变化（selective）
Mamba 的命名来自“选择性扫描”（Selective Scan）。它的做法很直白：把 B、C 和步长 Δ 改成由当前输入 x 通过一个小型线性层实时算出来，而不是写死。于是，遇到关键 token 时 Δ 变大、状态更新更“激进”；遇到无关填充时 Δ 变小、状态被平滑掉。这样 SSM 既保留了线性复杂度，又获得了注意力的“内容感知”能力。
import torch.nn.functional as F
def selective_step(x_t, h, A, D, sB, sC, sdt):
# sB, sC, sdt 都是“随输入变化”的小投影，这是 Mamba 的关键
dt = F.softplus(sdt(x_t)) # 步长由输入决定，Δ > 0
A_bar = torch.exp(dt * A) # 离散化后的状态衰减
B_bar = dt * sB(x_t) # 离散化后的输入投影
h_new = A_bar * h + B_bar * x_t # 选择性状态更新
y_t = sC(x_t) @ h_new + D * x_t # 选择性输出投影
return y_t, h_new
# 推理时状态 h 仍是定长，但每一步的“记多少”由输入说了算
这里的“选择性”可以用一个生活比喻理解：固定参数 SSM 像一台恒定转速的录音机，什么都录但都不清晰；Mamba 像一位会抓重点的速记员，听到要紧处就多记一笔（Δ 大），听到废话就略过（Δ 小）。正是这种按内容调节的记忆，让线性模型第一次在语言任务上追平甚至超过同等规模的注意力模型。
四、Mamba 架构：把 SSM 拼成深度网络
单个 SSM 层并不够，Mamba 把它组织成一个类 Transformer 的块：输入先过线性投影扩维，分两路——一路走选择性 SSM（含因果卷积 + 选择性扫描），另一路走 SiLU 门控，两路相乘后再投影回去，外面套残差连接。和 Transformer 块相比，它只是把“注意力子层”换成了“SSM 子层”，其余 LayerNorm、残差、MLP 几乎一致，因此工程上迁移成本很低。
五、与 Transformer / 线性注意力的取舍
| 维度 | Transformer 注意力 | Mamba（选择式 SSM） | 线性注意力 |
|---|---|---|---|
| 计算复杂度 | O(n²) | O(n) | O(n) |
| 显存随序列增长 | 平方增长 | 线性增长 | 线性增长 |
| 因果性 | 需 mask | 天然因果 | 天然因果 |
| 长程建模 | 强但昂贵 | 强且省 | 易退化 |
| 推理时携带的状态 | 需缓存全部 KV | 仅定长状态 | 需维护状态矩阵 |
三者都能做序列建模，但权衡不同：注意力长程最强但最贵；线性注意力把平方降成线性却常因“低秩假设”在长文本上丢信息；Mamba 用“选择性”弥补了普通线性模型的退化，又保住了定长状态带来的推理效率。需要强调的是，Mamba 并不是“全面取代”注意力——在很多中等长度、需要强双向交互的任务上，注意力依然更稳。
六、真实落地场景：长序列、基因组、音频
Mamba 的线性复杂度在“序列长到离谱”的场景里价值最大。典型包括：长文档与百万 token 级语料处理（推理时只需滚动一个定长状态，显存不再随长度爆炸）、DNA/蛋白质序列（动辄上万个 token，注意力根本喂不下）、音频与语音（采样率高、序列极长）、以及时间序列预测（传感器流数据）。在这些领域，Mamba 类模型往往能用更低的算力拿到媲美注意力的效果。
七、工程落地与避坑
| 坑 | 根因 | 对策 |
|---|---|---|
| 不用官方 CUDA 扫描就慢 | 朴素 Python 循环无法并行 | 用厂商提供的并行 selective scan 核 |
| 短序列反而更慢 | 固定开销 > 省下的 O(n²) | 短文本仍优先 Transformer |
| 双向信息缺失 | SSM 天然因果、只看过去 | 需要双向时叠反向 SSM 或回退注意力 |
| 训练不稳 | Δ 与 A 的指数离散化敏感 | 初始化加约束、用 stable 参数化 |
如果你打算把这类模型接进自己的服务，推理层可以复用 SGLang 推理引擎 或 vLLM 的连续批处理能力；本地先把模型跑起来，参考 DeepSeek 本地部署 的同类思路；要编排成智能体工作流，则可以用 Dify+Ollama 把模型封装成可调用节点。无论哪条路，Mamba 的价值都在于：当你被注意力的 O(n²) 卡住脖子时，它提供了一个已经被验证的“线性替代”选项。
八、小结
Mamba 的本质，是用选择性状态空间（SSM）把“注意力”替换为“定长状态 + 按内容选择性更新”：既拿到线性的计算与显存复杂度，又通过让 Δ、B、C 随输入变化，补回了普通线性模型缺失的“选择性”。它并非万能，但在超长序列、基因组、音频这类注意力吃不消的场景里，是一个已经被大量实验验证的务实选择。理解到这一步，你再去看任何“线性注意力/状态空间”架构，都能快速抓住它究竟在和 Transformer 的哪条短板较劲。
