摘要
- 核心问题：Transformer 的 Self-Attention 是 O(n²) 复杂度，序列越长越昂贵，且自回归生成受限于串行依赖。有没有一种序列模型，既能像 RNN 一样线性复杂度，又能像 Transformer 一样并行训练？
- 主要贡献：系统拆解状态空间模型（SSM）的连续系统本质与离散化、HiPPO 长程记忆初始化、Mamba 的选择性状态空间（Selective SSM）、硬件感知并行扫描，并厘清 SSM 与线性注意力的血缘关系。
- 阅读收获：① 看懂 SSM 的"连续系统 → 离散递推"数学框架；② 理解为什么传统 SSM 做不到内容相关推理、而 Mamba 可以；③ 弄清 Mamba 如何在 GPU 上既保持 O(n) 又接近卷积的并行速度；④ 能亲手用 Hugging Face 跑通一次 Mamba 文本生成。
一、引言：Transformer 的"算力墙"
回顾本系列前 18 篇，Transformer 几乎是我们所有讨论的基石：Self-Attention、位置编码、MoE、KV Cache、长上下文……它强大的根源，在于注意力机制能让任意两个位置直接交互——这种全局、内容相关的建模能力，正是大模型"涌现"的温床。
但能力有代价。注意力有一个绕不开的瓶颈：时间与显存复杂度都是 O(n²)（n 为序列长度）。
序列长度 n 带来的代价（注意力部分）： 注意力矩阵大小 = n × n → 长序列时爆炸 KV Cache 显存 ∝ n → 推理越聊越占显存（见第14篇） 单步生成 依赖前一步 → 自回归天生串行，难并行加速 实际体感： n = 2K → 注意力矩阵 400 万项，尚可 n = 32K → 10 亿项，显存吃紧 n = 1M → 1 万亿项，单卡无法承载更深层的问题在于：注意力其实是在"显式存储"全部历史（通过 n×n 的注意力矩阵 / KV Cache），而人类阅读时并不需要把整本书抄在草稿纸上逐字比对——我们靠"理解后的记忆"来处理长文。
这就引出了一个核心问题：
能否设计一种序列模型： ✓ 训练时可并行（像 Transformer） ✓ 推理时复杂度 O(n)（像 RNN） ✓ 能建模长程依赖（像 Transformer，而非退化 RNN）答案就是本文的主角——状态空间模型（State Space Model, SSM），以及它的革命性实现Mamba。
二、序列建模的三条路线
在 SSM 之前，序列建模主要有三条路线，各有死穴：
| 模型家族 | 训练并行性 | 推理复杂度 | 长程依赖 | 致命短板 |
|---|---|---|---|---|
| RNN（LSTM/GRU） | ❌ 必须串行 | O(n) | ⚠️ 易梯度消失 | 无法并行训练，扩不动 |
| CNN（Conv1d/WaveNet） | ✅ 可卷积 | O(n·k)（k=核宽） | ⚠️ 感受野受限 | 全局依赖需堆很深/很宽 |
| Transformer | ✅ 全并行 | O(n²) | ✅ 全局 | 序列越长越贵 |
SSM 的野心是集三家之长：拿到 RNN 的 O(n) 推理、Transformer 的并行训练，同时用"连续状态"天然建模长程依赖。
直观理解：RNN 用隐藏状态 h 压缩历史，但 h 是"定死的计算图"难并行；Transformer 用注意力"重新读取"历史，但代价是 O(n²)；SSM 则把"状态压缩历史"这件事建立在连续系统数学之上，从而既能被离散化用于推理，又能被转换成卷积用于并行训练。
三、状态空间模型（SSM）基础
3.1 连续系统视角
SSM 的出发点不是离散序列，而是一个连续时间的线性动力系统：
状态方程： h'(t) = A · h(t) + B · x(t) 输出方程： y(t) = C · h(t)x(t)：输入信号（连续）h(t)：状态向量（维度 d_state，通常 16/64/128），是系统对历史的"压缩记忆"y(t)：输出A：状态转移矩阵（d_state × d_state）——决定记忆如何随时间演化/遗忘B：输入到状态的映射（d_state × d_input）C：状态到输出的映射（d_output × d_state）
关键直觉：h(t) 是一个"有损但连续"的历史摘要。当前输出 y(t) 只通过 h(t) 依赖过去，而不需要回看全部 x(τ)。这正是 O(n) 复杂度的来源。
3.2 离散化：从连续到递推
计算机只能处理离散的 Token 序列，因此要把连续系统离散化。最常用的方式是零阶保持（Zero-Order Hold, ZOH）：假设在一个时间步 Δ 内，输入 x 保持恒定。
离散化后得到递推形式：
h_k = Ā · h_{k-1} + B̄ · x_k y_k = C · h_k其中离散化参数为：
Ā = exp(Δ · A) # 状态转移（矩阵指数） B̄ = (Δ · A)^{-1} · (exp(Δ · A) − I) · Δ · B # 输入投影工程上常用一阶近似
Ā ≈ I + ΔA、B̄ ≈ ΔB（当 Δ 很小时），更简单稳定。
这个递推式长得和 RNN 一模一样：当前状态 = 转移矩阵·旧状态 + 输入映射·当前输入。区别在于 A、B 来自连续系统的严谨离散化，而非手动设计。
3.3 线性时不变（LTI）：SSM 的"第一性"
传统 SSM 有一个重要性质：A、B、C 与具体输入无关（时间不变）。这意味着：
对于任意输入序列 x，输出 y 满足： y = (C 卷积 Ā) * x （* 表示卷积）也就是说，LTI 的 SSM 等价于一个固定卷积核的 1D 卷积！这带来一个巨大好处：可以用 FFT / 卷积算法在O(n log n)内并行训练，而不必串行扫描。
这解释了为什么 SSM 早期（如 S4 模型）能在长序列上击败 Transformer——它把"历史压缩"变成了一个可并行计算的卷积。
3.4 HiPPO：让状态真正"记住"长历史
光有卷积还不够。早期 SSM 的 A 矩阵若随便初始化，状态 h 会快速遗忘，长程依赖照样学不到。
HiPPO（High-order Polynomial Projection Operators）给出了一个优雅的初始化方案：它让状态 h(t) 始终逼近"截至当前时刻的历史输入 x(τ), τ≤t"在一组正交多项式基上的投影系数。
HiPPO 的核心思想： 用 d_state 个正交基函数（如 Legendre 多项式）拟合历史 状态 h 存的是"拟合系数"而非原始信号 → 少量维度即可高保真压缩任意长历史 → 长程依赖不再因维度固定而丢失HiPPO 初始化是 S4 模型成功的关键，也是后续所有 SSM（包括 Mamba）的标配。
四、SSM 与线性注意力的"血缘关系"
理解 SSM，就不能不提线性注意力（Linear Attention）——它们解决的是同一个问题：把 O(n²) 的注意力降为线性。
标准注意力： Attn = softmax(QK^T / √d) · V # O(n²) 因为 QK^T 是 n×n 线性注意力： Attn = φ(Q) · (φ(K)^T · V) # 先聚合 KV 到固定维度 = φ(Q) · (Σ_k φ(k_k) v_k^T) # O(n·d²) 与 n 线性相关而 SSM 的递推h_k = Ā·h_{k-1} + B̄·x_k本质上也是在把历史信息压缩进固定维度的状态 h——这与线性注意力把 KV 聚合进固定维度记忆是同一哲学。
| 维度 | 线性注意力 | 经典 SSM（LTI） |
|---|---|---|
| 记忆载体 | 固定维度聚合矩阵 | 状态向量 h |
| 内容相关性 | ✅ 依赖输入（φ 映射） | ❌ 时间不变（A/B/C 与输入无关） |
| 训练方式 | 直接矩阵乘 | 卷积 / FFT |
| 长程能力 | 取决于 φ 设计 | 依赖 HiPPO 初始化 |
关键差异：经典 SSM 是时间不变的——无论输入是什么，它对"如何压缩历史"的策略是固定的。这让它高效，但也导致一个致命缺陷：无法做内容相关的、选择性地关注（例如"复制某个特定词"“跳过无关段落”）。而这是 Transformer 注意力最擅长的。
Mamba 的全部创新，都围绕"如何把选择性还给 SSM"展开。
五、Mamba 的突破：选择性状态空间（Selective SSM）
5.1 问题：时间不变 = 无法"选择性关注"
经典 SSM 的 A、B、C 是固定的。想象你在读一段文字：
"请把上面那句【重要结论】抄写下来，忽略中间的水词。"人类会选择性地把"重要结论"记进工作记忆，把水词忘掉。但 LTI-SSM 对"重要结论"和"水词"一视同仁——它的遗忘率是预先定死的，不随内容变化。于是它擅长"平滑地压缩一切"，却不擅长"精准地挑出某一点"。
论文里用一个简单的**"抓特定token"任务**证明：LTI 模型（包括 S4）在这个任务上几乎失败，而 Transformer 轻松通过。
5.2 解法：让参数依赖输入
Mamba 的核心洞见极简：让 Δ、B、C 成为输入 x 的函数。
传统 SSM： Ā, B̄, C = 常数（与 x 无关） Mamba SSM： Δ_k, B_k, C_k = f(x_k) （每个 token 不同！）具体做法：用 tiny 的线性层（投影到 d_state / d_inner）从当前输入 x_k 实时算出 Δ_k、B_k、C_k。
这带来"选择性"的本质效果：
Δ_k 很大 → "这一步很重要，快更新记忆" → Ā_k = exp(Δ_k·A) 接近 0 矩阵 → 旧记忆被快速覆盖/刷新 Δ_k 很小 → "这一步是噪声，忽略" → Ā_k ≈ I → 旧记忆几乎原样保留换句话说，Δ 充当了注意力的"门控"：类似 Transformer 的 softmax 权重，但作用在"状态更新"上而非"全局相似度"上。这让 Mamba 既能选择性记忆、又能保持 O(n)。
妙处：选择性让 SSM 从"只能平滑压缩"升级为"可内容感知推理"，补上了与 Transformer 之间的能力鸿沟，同时复杂度仍是线性的。
5.3 代价：卷积并行性被破坏
但"参数依赖输入"意味着 SSM 不再是 LTI，也就无法再用卷积并行训练。Mamba 必须在"并行"和"选择性"之间做技术权衡——这正是下一节硬件算法的用武之地。
六、硬件感知并行扫描（Hardware-aware Parallel Scan）
选择性让经典卷积失效，但朴素的串行递归（RNN 式）在 GPU 上又慢得离谱（GPU 擅长大规模并行矩阵乘，不擅长一步步串行）。Mamba 的解法是硬件感知并行扫描（Parallel Scan），其核心是减少 HBM（显存）与 SRAM（片上缓存）之间的来回读写。
朴素递归的瓶颈： 每步都要从 HBM 读状态 h、写回 h → HBM 带宽几十倍低于 SRAM，I/O 成为瓶颈 Mamba 的做法（GPU 友好的扫描）： 1. 把序列分块（chunk），每块载入 SRAM 2. 在 SRAM 内做并行扫描（associative scan，利用结合律） 3. 只在块边界与 HBM 交互 → 显存 I/O 从 O(n) 降到接近 O(n / chunk_size) → 既保持 O(n) 计算，又逼近卷积的吞吐数学上依赖扫描的结合律：
(h_1 op x_1) op (h_2 op x_2) = 可任意分组并行计算这使得"选择性扫描"能在 GPU 上高效并行，实测训练速度接近同规模 Transformer，推理随序列长度线性扩展——真正拿到了 RNN 的复杂度 + Transformer 的并行性。
七、Mamba 架构与复杂度全景
7.1 Mamba 块结构
Mamba 用"SSM 替换注意力"的方式接入标准 LLM 主干（类似 LLaMA，但把 Self-Attention 换成 Selective SSM）：
输入 x │ ├─→ RMSNorm → 线性扩展(expand) → SiLU │ → 因果 1D 卷积 → SiLU │ → 选择性 SSM（核心） │ → 线性压缩 │ → 残差相加 └─→ (跳跃连接) 特点： ✓ 无注意力 → 无 n×n 矩阵、无 KV Cache 膨胀 ✓ 推理时状态 h 是固定维度 → 显存随 n 线性、不平方7.2 四类序列模型对比
| 特性 | RNN | CNN(Conv) | Transformer | Mamba (SSM) |
|---|---|---|---|---|
| 训练并行 | ❌ 串行 | ✅ 卷积 | ✅ 全并行 | ✅ 并行扫描 |
| 推理复杂度 | O(n) | O(n·k) | O(n²) | O(n) |
| 显存(长序列) | O(1)状态 | O(k) | O(n) KV | O(1)状态 |
| 内容选择性 | ❌（隐式） | ❌ | ✅ 全局 | ✅ 选择性 |
| 长程依赖 | ⚠️ 易忘 | ⚠️ 受限 | ✅ | ✅ |
| 硬件友好 | ⚠️ 串行慢 | ✅ | ✅ | ✅（感知扫描） |
一句话总结 Mamba 的定位：在保持线性复杂度与恒定显存的同时，补齐了经典 SSM 缺失的内容选择性，使它第一次在语言建模上能正面对抗同规模 Transformer。
7.3 实测表现（代表性结果）
Mamba-3B 在 Pile 上的困惑度 ≈ 同规模 Transformer / LLaMA 但：序列越长，推理吞吐优势越大（8K→1M 上下文几乎线性） Mamba-2.8B / 在下游任务(常识/QA)追平或略超 Pythia-2.8B 混合模型 Jamba（Mamba + MoE + Attention）证明二者可互补八、Mamba-2 与结构化状态空间对偶（SSD）
2024 年的Mamba-2进一步发现：选择性 SSM 与一个特殊结构的结构化矩阵乘法（半可分矩阵）在数学上对偶——而后者可以映射成 GPU 高度优化的矩阵乘法。
核心洞见（SSD, State Space Duality）： 选择性扫描（串行） ⇔ 半可分矩阵乘法（并行） → 用矩阵乘替代扫描，训练再快 2-8 倍 → 同时统一了"SSM 视角"与"注意力视角" 意义： ✓ 训练效率大幅提升 ✓ 理论上把 SSM 与线性注意力拉到同一框架 ✓ 为后续"注意力-SSM 混合架构"扫清障碍这标志着序列建模从"Transformer vs RNN vs SSM 三方混战"，走向"用统一数学语言描述记忆与计算"的新阶段。
九、实战：跑通 Mamba
9.1 用 Hugging Face 加载 Mamba 生成文本
# 需要: pip install transformers torchfromtransformersimportAutoModelForCausalLM,AutoTokenizer model_id="state-spaces/mamba-130m-hf"# 也可换 mamba-2.8b-hftokenizer=AutoTokenizer.from_pretrained(model_id)model=AutoModelForCausalLM.from_pretrained(model_id,trust_remote_code=True)prompt="The future of artificial intelligence is"inputs=tokenizer(prompt,return_tensors="pt")# Mamba 自回归生成（无 KV Cache 膨胀问题，长序列友好）outputs=model.generate(**inputs,max_new_tokens=64,temperature=0.8,top_p=0.9,)print(tokenizer.decode(outputs[0],skip_special_tokens=True))注意：Mamba 没有 KV Cache，但推理仍是一个 token 一个 token 产出（自回归本质不变）。它的优势在于长序列时显存稳定、吞吐更高。
9.2 最小选择性 SSM 推理示意（伪代码）
importtorchdefselective_ssm_step(x_k,h_prev,A,dt_proj,B_proj,C_proj,D):""" x_k : 当前输入 (d_inner,) h_prev : 上一时刻状态 (d_state, d_inner) A : 固定的状态矩阵 (d_state, d_inner) —— 通常每个通道一份 """# 1) 由输入实时算出"选择性参数"dt=dt_proj(x_k)# (d_inner,) 门控步长Bk=B_proj(x_k)# (d_state,)Ck=C_proj(x_k)# (d_state,)# 2) 离散化（一阶近似，训练稳定）dA=torch.exp(dt[:,None]*A)# (d_inner, d_state)dB=dt[:,None]*Bk[None,:]# (d_inner, d_state)# 3) 状态更新（选择性：dt 大则刷新，dt 小则保留）h_new=dA*h_prev+dB*x_k[:,None]# (d_inner, d_state)# 4) 输出 + 跳跃连接y_k=(h_new*Ck[None,:]).sum(dim=-1)# (d_inner,)returny_k+D*x_k,h_new这段代码浓缩了 Mamba 的精髓：参数随输入变化（选择性）+ 固定维度状态（O(n)）。
十、常见误区纠正
误区1："Mamba 会完全取代 Transformer。" 事实：混合架构（Jamba、Zamba、Samba）把 SSM 与注意力互补使用—— SSM 处理长程线性压缩，注意力处理局部精确交互，往往更优。 误区2："SSM 不能做长上下文。" 事实：SSM 天然 O(n) 且状态维度固定，反而比 Transformer 更适合超长序列 （Mamba 论文展示过 1M token 级别的语言建模）。 误区3："选择性 SSM 和注意力是一回事。" 事实：都做内容相关加权，但注意力显式存全部 KV（O(n) 显存）， SSM 把历史压进固定状态（O(1) 显存），没有全局注意力矩阵， 表达力与归纳偏置都不同。 误区4："Mamba 只是 RNN 换皮。" 事实：关键在"可在 GPU 上并行扫描训练" + "选择性"， 既有 RNN 的 O(n)，又有接近 Transformer 的训练吞吐， 这是纯 RNN 做不到的。十一、总结与展望
本文核心论断
1. SSM 的本质：用连续动力系统的"状态"压缩历史， 离散化后等价于 RNN 递推，LTI 时又等价于卷积（可并行训练）。 2. HiPPO 初始化让状态能高保真压缩任意长历史—— 这是 SSM 长程能力的基础。 3. 经典 SSM 是时间不变的，无法内容选择性关注； Mamba 让 Δ/B/C 依赖输入，补上了这块短板。 4. 选择性破坏了卷积并行，但硬件感知并行扫描在 GPU 上 既保 O(n) 又逼近卷积吞吐——工程上真正可用。 5. Mamba-2 揭示 SSM 与结构化矩阵乘对偶， 统一了 SSM 与（线性）注意力的数学语言。技术演进脉络
1990s 状态空间模型（控制论） ↓ 2019 S4（HiPPO + LTI-SSM + 卷积训练） 长序列建模突破 ↓ 2022 GSS / S5 扩展 SSM 表达力 ↓ 2023 Mamba（选择性 SSM + 硬件感知扫描） 语言建模正面对抗 Transformer ↓ 2024 Mamba-2（SSD 对偶） 训练提速 + 理论统一 ↓ 2025+ Jamba / 混合架构 与注意力、MoE 深度融合下期预告（下次发布）：《阶段总结：技术原理篇核心知识回顾》—— 把第9~19篇的 11 讲核心公式、架构图、代码范式串成一张"技术原理篇知识地图"，并给出针对不同目标的复习路径。
参考资料
- Gu et al. (2021)—Combining Recurrent, Convolutional, and Continuous-time Models with the Linear State Space Layer (S4)
- Gu et al. (2022)—HiPPO: Recurrent Memory with Optimal Polynomial Projections
- Gu & Dao (2023)—Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- Dao & Gu (2024)—Transformers are SSMs: Generalized Models and Efficient Parallel Inference for Structured State Space Models (Mamba-2 / SSD)
- Lieber et al. (2024)—Jamba: A Hybrid Transformer-Mamba Language Model
- Katharopoulos et al. (2020)—Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention（线性注意力对偶）
延伸讨论
思考题：
- 既然 Mamba 有选择性，为什么很多落地场景仍选择"Transformer + 长上下文"而非纯 Mamba？
- 如果让你设计一个"注意力 + SSM"混合块，你会把 SSM 放在浅层还是深层？为什么？
实践作业：
- 用
state-spaces/mamba-130m-hf生成一段文本，再换成同规模EleutherAI/pythia-160m对比困惑度与生成速度 - 把 9.2 节的伪代码补全成可运行的 PyTorch 模块，在合成"复制任务"上验证选择性
