摘要核心问题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 之前序列建模主要有三条路线各有死穴模型家族训练并行性推理复杂度长程依赖致命短板RNNLSTM/GRU❌ 必须串行O(n)⚠️ 易梯度消失无法并行训练扩不动CNNConv1d/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_inputC状态到输出的映射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 线性时不变LTISSM 的第一性传统 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 会快速遗忘长程依赖照样学不到。HiPPOHigh-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 聚合进固定维度记忆是同一哲学。维度线性注意力经典 SSMLTI记忆载体固定维度聚合矩阵状态向量 h内容相关性✅ 依赖输入φ 映射❌ 时间不变A/B/C 与输入无关训练方式直接矩阵乘卷积 / FFT长程能力取决于 φ 设计依赖 HiPPO 初始化关键差异经典 SSM 是时间不变的——无论输入是什么它对如何压缩历史的策略是固定的。这让它高效但也导致一个致命缺陷无法做内容相关的、选择性地关注例如复制某个特定词“跳过无关段落”。而这是 Transformer 注意力最擅长的。Mamba 的全部创新都围绕如何把选择性还给 SSM展开。五、Mamba 的突破选择性状态空间Selective SSM5.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 带宽几十倍低于 SRAMI/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 四类序列模型对比特性RNNCNN(Conv)TransformerMamba (SSM)训练并行❌ 串行✅ 卷积✅ 全并行✅ 并行扫描推理复杂度O(n)O(n·k)O(n²)O(n)显存(长序列)O(1)状态O(k)O(n) KVO(1)状态内容选择性❌隐式❌✅ 全局✅ 选择性长程依赖⚠️ 易忘⚠️ 受限✅✅硬件友好⚠️ 串行慢✅✅✅感知扫描一句话总结 Mamba 的定位在保持线性复杂度与恒定显存的同时补齐了经典 SSM 缺失的内容选择性使它第一次在语言建模上能正面对抗同规模 Transformer。7.3 实测表现代表性结果Mamba-3B 在 Pile 上的困惑度 ≈ 同规模 Transformer / LLaMA 但序列越长推理吞吐优势越大8K→1M 上下文几乎线性 Mamba-2.8B / 在下游任务(常识/QA)追平或略超 Pythia-2.8B 混合模型 JambaMamba MoE Attention证明二者可互补八、Mamba-2 与结构化状态空间对偶SSD2024 年的Mamba-2进一步发现选择性 SSM 与一个特殊结构的结构化矩阵乘法半可分矩阵在数学上对偶——而后者可以映射成 GPU 高度优化的矩阵乘法。核心洞见SSD, State Space Duality 选择性扫描串行 ⇔ 半可分矩阵乘法并行 → 用矩阵乘替代扫描训练再快 2-8 倍 → 同时统一了SSM 视角与注意力视角 意义 ✓ 训练效率大幅提升 ✓ 理论上把 SSM 与线性注意力拉到同一框架 ✓ 为后续注意力-SSM 混合架构扫清障碍这标志着序列建模从Transformer vs RNN vs SSM 三方混战走向用统一数学语言描述记忆与计算的新阶段。九、实战跑通 Mamba9.1 用 Hugging Face 加载 Mamba 生成文本# 需要: pip install transformers torchfromtransformersimportAutoModelForCausalLM,AutoTokenizer model_idstate-spaces/mamba-130m-hf# 也可换 mamba-2.8b-hftokenizerAutoTokenizer.from_pretrained(model_id)modelAutoModelForCausalLM.from_pretrained(model_id,trust_remote_codeTrue)promptThe future of artificial intelligence isinputstokenizer(prompt,return_tensorspt)# Mamba 自回归生成无 KV Cache 膨胀问题长序列友好outputsmodel.generate(**inputs,max_new_tokens64,temperature0.8,top_p0.9,)print(tokenizer.decode(outputs[0],skip_special_tokensTrue))注意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) 由输入实时算出选择性参数dtdt_proj(x_k)# (d_inner,) 门控步长BkB_proj(x_k)# (d_state,)CkC_proj(x_k)# (d_state,)# 2) 离散化一阶近似训练稳定dAtorch.exp(dt[:,None]*A)# (d_inner, d_state)dBdt[:,None]*Bk[None,:]# (d_inner, d_state)# 3) 状态更新选择性dt 大则刷新dt 小则保留h_newdA*h_prevdB*x_k[:,None]# (d_inner, d_state)# 4) 输出 跳跃连接y_k(h_new*Ck[None,:]).sum(dim-1)# (d_inner,)returny_kD*x_k,h_new这段代码浓缩了 Mamba 的精髓参数随输入变化选择性 固定维度状态O(n)。十、常见误区纠正误区1Mamba 会完全取代 Transformer。 事实混合架构Jamba、Zamba、Samba把 SSM 与注意力互补使用—— SSM 处理长程线性压缩注意力处理局部精确交互往往更优。 误区2SSM 不能做长上下文。 事实SSM 天然 O(n) 且状态维度固定反而比 Transformer 更适合超长序列 Mamba 论文展示过 1M token 级别的语言建模。 误区3选择性 SSM 和注意力是一回事。 事实都做内容相关加权但注意力显式存全部 KVO(n) 显存 SSM 把历史压进固定状态O(1) 显存没有全局注意力矩阵 表达力与归纳偏置都不同。 误区4Mamba 只是 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 S4HiPPO LTI-SSM 卷积训练 长序列建模突破 ↓ 2022 GSS / S5 扩展 SSM 表达力 ↓ 2023 Mamba选择性 SSM 硬件感知扫描 语言建模正面对抗 Transformer ↓ 2024 Mamba-2SSD 对偶 训练提速 理论统一 ↓ 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 ProjectionsGu Dao (2023)—Mamba: Linear-Time Sequence Modeling with Selective State SpacesDao 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 ModelKatharopoulos 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 模块在合成复制任务上验证选择性