长河不忘，流水有选：Mamba 与状态空间的记忆法则

湍湍的河流不需要水池，它用流动本身储存时间。
而那些永不流动的湖，记住了每一滴水，却忘记了方向。
——《流形笔记》（虚构）

上篇：清渊国的两种记史官

一、清渊国的难题

清渊国有两种记史官，据说这两种人彼此看不起对方。

第一种叫全文官。他们的工作方式是：你说任何一句话，他们立刻把整本档案搬出来，翻遍每一行，确认这句话和之前所有话的关系，然后才开口回答。他们的记忆是完整的，不漏任何一个字。清渊国的大学士说，全文官”无所不知，无所不联”。

这种人非常贵重——因为他们的工作量随着档案变厚而剧增。一次对话十句话，他们需要翻十次乘十次；一次对话一百句话，他们需要翻一百次乘一百次。档案库快速膨胀，全文官每次工作都得气喘吁吁。

第二种叫流水官。他们的工作方式完全不同：他们只拥有一个固定大小的”心忆葫芦”。每听到一句新话，他们把它与葫芦里已经凝练的心忆融合，更新葫芦，然后继续。葫芦的大小永不改变，无论你讲一句话还是一万句话。

但流水官有个致命的缺点，让全文官嘲笑了数百年：葫芦里的过去是被压缩的，是有损的。 葫芦记的不是原文，而是”精华的幻影”。如果你在第一句话里埋了一个关键词，而流水官后来听了九千九百句无聊的废话，那个关键词很可能就被稀释得无影无踪。

清渊国大学士的判决由此而来：全文官做朝廷首席，流水官只配下乡给村民讲故事。

这个局面维持了整整三百年。

直到河神的孩子阿逢出现了。

二、河神之子阿逢

阿逢是清渊国边境一条古老河流的护河人的孩子。她从小和河流一起长大，懂得一件全文官和普通流水官都不懂的事：

河流有选择权。

普通的流水官，他们的葫芦不会区分信息。每一句话进来，不管重要与否，都被以固定的方式融合进去。第一句”宫中有奸”和第九百九十九句”今日天气晴朗”对葫芦的影响是一样的——它们以相同的比例被稀释、被覆盖。

但河流不是这样运作的。

阿逢观察过那条古河太多年了。她发现，河流对不同的东西有着截然不同的吸收方式：

一块大石头滚入水中，激起的涟漪会在整条河里荡漾很久——河”记得”这块石头，以水纹的形式，记了很久很久。

但一片落叶飘过，河面只微微起了一道皱，几乎立刻消弭。

一场暴雨将上游的红泥冲入河中，河水改色，这个信息被整条河带着，一直携带到大海。

但三月里有一个孩子向河里丢了一颗小鹅卵石，河水立刻恢复如故。

河流天然地区分哪些事件值得长期记住，哪些事件只需一闪而过——这不是事先规定好的，而是依赖”当前正在发生什么”来动态决定的。

阿逢把这个观察带到了记史官的培训学堂里，引起了轩然大波。

三、葫芦的改造

阿逢对流水官葫芦的改造，分三个步骤。

第一步：让葫芦知道”这句话有多重要”。

传统的流水官葫芦有一个固定的”遗忘阀门”——每次融合新信息，旧信息以固定比例衰减。比如每次新话进来，过去的记忆打九折留下来，新话占一折。

阿逢说：不对。遗忘比例不该固定，它应该由”当前这句话的内容”来决定。

如果这句话是”皇上驾崩”——过去所有的信息都应该以极高的权重保留，因为这是一个颠覆性的信号，需要和前面所有的历史背景相互印证。遗忘阀门应该几乎关死。

如果这句话是”今天多云转阴”——过去信息的权重可以适度衰减，新话的影响权重也小，因为它可能只是无关紧要的背景。

所以，遗忘率和接受率都是输入依赖的——由当前这个时间步的输入内容来计算，而不是事先固定。

第二步：让葫芦自己决定”怎么吸收这句话”。

传统流水官把每句话以固定方式加入葫芦。但阿逢又加了一个门：“纳入门”。

纳入门的意思是：不是每句话都要被完整地吸入葫芦。这句话是否有价值、应该以多大比例影响心忆，取决于当前葫芦里已有的内容和当前输入内容的相互审视。

这有点像——你在认真记笔记的时候，突然有一只蚊子飞过，你不会把”一只蚊子飞过”写进笔记里。你的注意力自动地过滤掉了无关信息。

第三步：让葫芦的”输出”也是动态计算的。

传统流水官每次输出，就是把葫芦的内容直接端出来。但阿逢加了第三个门：“输出门”——当前时间步应该向外界展示葫芦内容的哪些侧面，也由当前输入内容来决定。

三个动态门加在一起，阿逢的葫芦就成了一个能够”选择性记忆”的智能容器。它不再是被动地被所有信息以固定比例堆积的桶，而是一个有判断力的河流——能辨别哪些信息值得用水纹记录很久，哪些只是过客。

清渊国的大学士们沉默了一段时间，然后说：

“这个葫芦，不是普通的流水官葫芦了。这是……河神的葫芦。”

四、但河神的葫芦可以被并行打造吗？

阿逢的新葫芦有一个让全文官嘲笑的地方：它是序列的。

你必须先处理第一句话，更新葫芦，再处理第二句话，更新葫芦……这样一句一句来。因为第二句话的处理方式，依赖于处理完第一句话之后葫芦的状态。

全文官则不同——他们的”翻档案”操作，是可以大规模并行的。你可以同时让一千个小吏分别计算”这句话和那句话的关联”，然后汇总。这正是全文官快速且受欢迎的原因之一（当然，他们需要一千个小吏，这很贵）。

阿逢想了很久，发现了一个数学奇迹。

她说：我的葫芦在推理（已经知道了所有句子、需要回答问题）时，必须一句一句走。

但在训练（让葫芦从大量文本中学习如何设定三个门的参数）时，有一种叫“关联扫描”的技巧，可以把看似必须串行的递推计算，并行展开。

这个技巧的核心是：对于线性递推操作（形如 $h_t = A_t \cdot h_{t-1} + B_t \cdot x_t$），存在一种通过树状归约的方式，在 $O(\log N)$ 的深度内完成本来需要 $O(N)$ 步的计算。就像把一列数字求前缀和，朴素方法是逐个累加，而并行前缀算法可以用二叉树结构大幅加速。

所以：阿逢的葫芦训练时可以并行，推理时不需要随序列增长而增加计算量。

两种记史官的长期争论，在阿逢的新葫芦面前，第一次出现了真正的裂缝。

五、清渊国的新秩序

数年之后，清渊国的朝廷里形成了新的分工：

至此，流水官的三百年屈辱，终于翻篇。

下篇：状态空间模型的工程真相

一、什么是状态空间模型（SSM）

在数学和控制论领域，状态空间模型（State Space Model, SSM）是描述动态系统的经典工具，比深度学习更古老。它的形式是：

\(h'(t) = A \cdot h(t) + B \cdot x(t)\) \(y(t) = C \cdot h(t) + D \cdot x(t)\)

这里：

这是一个连续时间的线性动力系统。用语言说：当前状态的变化速率，由当前状态和当前输入共同决定。

要把它用于序列数据（离散的词、token），需要离散化：

\(h_k = \bar{A} \cdot h_{k-1} + \bar{B} \cdot x_k\) \(y_k = C \cdot h_k\)

其中 $\bar{A} = e^{\Delta A}$，$\bar{B} = (\Delta A)^{-1}(e^{\Delta A} - I) \cdot \Delta B$，$\Delta$ 是步长参数（可学习）。

这就是离散化的递推形式——看起来和 RNN 很像，但背后有控制论的严格数学支撑，且有更好的梯度性质。

二、S4：第一个突破

2022年，基于线性状态空间的 S4（Structured State Space for Sequences）模型横空出世，成为第一个在长序列任务上真正可与 Transformer 竞争的序列模型。

S4 的关键技巧：

1. 把 $A$ 矩阵参数化为 HiPPO 矩阵

HiPPO（High-order Polynomial Projection Operators）理论给出了一种最优的 $A$ 矩阵初始化方式，使隐藏状态 $h$ 能够最优地以勒让德多项式的形式压缩历史信息。简单说：HiPPO 矩阵让状态空间模型天然具有”聪明记忆”的能力，而不是像普通 RNN 那样随机初始化然后盲目训练。

2. 并行卷积视角

S4 发现，对于固定参数 $A, B, C$ 的情况，递推形式 $h_k = Ah_{k-1} + Bx_k$ 等价于一次卷积操作：

\[y = \bar{K} * x\]

其中 $\bar{K} = (C\bar{B}, C\bar{A}\bar{B}, C\bar{A}^2\bar{B}, \ldots)$ 是一个固定的卷积核。卷积可以用 FFT 在 $O(N \log N)$ 内并行计算——比逐步递推快得多。

这就是 S4 训练时能并行的原因。

3. 局限性

S4 的 $A, B, C$ 矩阵对所有时间步是固定的，不依赖当前输入。这意味着它对所有输入的处理方式是相同的——无论当前内容多重要，遗忘率和吸收率永远一样。这正是阿逢批评的”传统流水官”。

三、Mamba：选择性才是革命

2023年底，Mamba 的出现解决了 S4 的核心局限：让状态转移矩阵 $B$、$C$ 和步长 $\Delta$ 成为输入的函数。

\[\Delta, B, C = f_{\theta}(x_t)\]

这看似微小的改动，带来了根本性的变化：

选择性压缩（Selective State Spaces）

这就是阿逢的河流逻辑：一块大石头让 $\Delta$ 飙升，让石头的影响深深嵌入状态；一片落叶让 $\Delta$ 趋零，叶子几乎不留痕迹。

关键实现细节：选择性扫描（Selective Scan）

因为 $B, C, \Delta$ 现在是输入依赖的，不同时间步的参数不同，所以不再能用固定卷积核来并行计算。Mamba 的解决方案：

硬件感知的并行扫描（Hardware-aware Parallel Scan）：

在 GPU 上，SSM 递推 $h_k = A_k h_{k-1} + B_k x_k$ 可以通过关联扫描（Associative Scan）并行化。关联扫描利用了矩阵乘法的结合律——

\[h_k = A_k A_{k-1} \cdots A_1 h_0 + \sum_{i=1}^{k} A_k \cdots A_{i+1} B_i x_i\]

通过树状并行归约，可以在 $O(\log N)$ 深度完成（类似并行前缀和）。在实际 CUDA 实现中，Mamba 把中间状态保存在 SRAM 而不是显存，避免了慢速的显存读写，实现了 Flash Attention 级别的 IO 感知优化。

这使得 Mamba 在训练时并行、在推理时只需 $O(1)$ 的状态维护，而 Transformer 推理时需要 $O(N^2)$ 的 KV Cache 空间（或 $O(N)$ 的 KV Cache + $O(N)$ 的读取时间）。

四、Mamba 的架构细节

Mamba 的基本 Block 结构（类似 Transformer 中的 Attention Block）：

输入 x
  ├─→ Linear Projection → SSM（选择性状态空间）→ 激活函数
  └─→ Linear Projection → 门控（gating）
       ↓
两路相乘（门控输出）
       ↓
输出 y

与 Transformer Block 相比：

Mamba-2（2024）的改进：把选择性 SSM 重新表述为一种 结构化的矩阵乘法（SSD framework，Structured State Space Duality），从而将 SSM 和注意力机制在理论上统一，并进一步提升了训练速度（在 A100 上比 Mamba-1 快 2-8x）。

五、Mamba vs Transformer：详细对比

维度 Transformer Mamba（SSM）
训练时间复杂度 $O(N^2 d)$ $O(Nd)$（关联扫描）
推理内存（KV Cache） $O(N \cdot d)$，随序列长度线性增长 $O(d^2)$，固定大小状态
推理时间 $O(N)$ 每步（注意力需查 N 个历史 token） $O(1)$ 每步（只更新固定状态）
长距离依赖 直接全局注意力，精确 有损压缩，依赖选择性
可解释性 注意力权重可可视化 状态空间较难解释
适合任务 精确检索、代码生成、多跳推理 长文档处理、流数据、持续学习

关键洞察：Mamba 的优势在超长序列（>16K token）场景爆发。对于短序列，Transformer 并不显著慢；但一旦序列达到 100K+ token 级别，Transformer 的二次复杂度让它几乎难以运行，而 Mamba 的线性复杂度依然流畅。

六、混合架构：现实世界的答案

纯 Mamba 在某些任务上会输给 Transformer，原因在于：有些任务天然需要精确的远程检索——比如”文章第7段提到了哪个人名”这类 needle-in-a-haystack 问题。SSM 的有损压缩在这里是致命弱点。

工业界的答案是混合架构：

Jamba（2024）：每 7 个 Mamba 层插入一个 Attention 层，以 1:7 的比例混合，兼顾长程效率和精确检索能力。在相同参数量下，Jamba 比纯 Transformer 快 3x，比纯 Mamba 在多跳推理上准确率高 15%。

Zamba（2024）：类似思路，用共享的 Attention 层（每隔若干 Mamba 层）进行精确记忆修正，进一步节省参数。

Samba（2024）：引入局部 Sliding Window Attention 和 SSM 的组合，以小注意力窗口保持局部精确性，用 SSM 维持全局状态压缩。

这种混合思路与清渊国的”混合记史官制”一模一样：大段用葫芦（SSM），关键精确查验处派全文官（Attention）。

七、工程心法与应用场景

何时该用 Mamba 架构（或混合架构）？

  1. 超长上下文 Agent：如果你的 Agent 需要在 100K+ token 的上下文中运行，Mamba 的 $O(1)$ 推理状态让它不会因上下文增长而崩溃。Transformer KV Cache 的内存在 100K token 下可能高达数十 GB，而 Mamba 的状态大小固定。

  2. 流数据处理：实时处理连续日志、传感器流、长时序数据——每来一个新 token，Mamba 只更新一次状态（$O(1)$），而 Transformer 需要用注意力重新审视所有历史。

  3. 边缘设备 / 低内存推理：在内存受限的移动端、嵌入式设备，Mamba 的固定状态大小是巨大优势——你不需要维护随上下文增长的 KV Cache。

  4. 持续学习场景：状态空间模型的递推性质让它天然适合”不断接收新信息、逐步更新认知”的场景，类似在线学习。

选择性的本质——工程师的思考框架

Mamba 最重要的工程启示不是架构本身，而是一个更抽象的问题：什么信息值得长期记住？

这个问题在所有序列系统中都存在：

Mamba 告诉我们：“值得记住多久”不应该是固定的，它应该由内容本身决定。 一个”用户付款失败”的事件，应该在 Agent 记忆里保存很久；一个”用户浏览了首页”的事件，可能几秒后就可以遗忘。

把这个原则带入 AI Agent 设计：

八、底层数学：为什么 Selective SSM 有效

为什么输入依赖的 $\Delta$ 能实现选择性？

考虑离散化公式 $\bar{A} = e^{\Delta A}$，当 $A$ 的特征值 $\lambda < 0$（稳定系统）时：

所以 $\Delta$ 本质上控制了当前 token 的时间分辨率：$\Delta$ 大 → 这个时间步”占更多空间”（模型在这里停留更久，吸收更多）；$\Delta$ 小 → 这个时间步是”过客”。

这给了 SSM 一种隐式的”软注意力”能力——不是通过查询所有历史来分配注意力，而是通过调整时间步长来决定当前输入的重要性。

和 LSTM / GRU 的关系

LSTM 和 GRU 也有门控机制（input gate, forget gate）。但它们的门是基于 sigmoid/tanh 的非线性变换，而 Mamba 的门基于线性代数和控制论的结构化设计（HiPPO 初始化 + 离散化）。

更重要的是：LSTM/GRU 没有并行扫描的能力，训练必须串行；Mamba 通过关联扫描实现了训练并行化，这是它能扩展到大规模的关键。

九、状态空间模型的局限与未来

Mamba 不是无懈可击的。它的核心弱点：

  1. 精确检索困难：如果任务要求模型精确回忆序列中某个特定位置的内容，有损压缩的隐藏状态会失败。这正是混合架构存在的理由。

  2. 理论上的表达能力：理论分析表明，线性 SSM 的表达能力有上限，某些复杂的上下文依赖（需要跨极远 token 的非线性交互）可能无法被固定大小的状态压缩表示。

  3. 实际 CUDA 实现复杂度：选择性扫描的 CUDA 内核比 FlashAttention 更复杂，优化难度更高。

但这些局限正在被快速迭代的研究突破：Mamba-2、Hawk/Griffin（线性注意力 + 选择性门控）、RWKV-6（线性复杂度语言模型）都在从不同角度推进同一个方向——序列建模不必是二次复杂度的。

长河的记忆法则，终将改写整个领域的游戏规则。

尾声：阿逢的葫芦，悬于博物馆

多年之后，清渊国的博物馆里展出了一只古老的葫芦。

解说牌上写着：

“此乃流水选忆葫芦，传世之作。持之者，可在不增一寸的躯壳中，容纳万年的江河。诀窍不在容量，在于取舍。”

参观者里有个年轻人皱眉问：

“可是，它会忘记一些事情。”

旁边的老者回头，轻轻说：

“孩子，不忘是奢侈品。懂得忘记什么，才是智慧。”

技术附录速查表

概念 对应数学 工程意义
状态 $h$ 隐藏状态向量（固定维度） 模型的”工作记忆”
$A$ 矩阵 状态转移矩阵（HiPPO 初始化） 控制遗忘速率
$B$ 矩阵 输入投影 控制输入如何影响状态
$C$ 矩阵 输出投影 控制状态如何影响输出
$\Delta$（步长） 离散化时间步长 输入依赖的”时间分辨率”
选择性 $B, C, \Delta$ 依赖输入 动态门控，选择性吸收
并行扫描 关联扫描（树状归约） 训练时 $O(\log N)$ 并行
推理 纯递推，固定状态 每步 $O(1)$，内存固定
混合架构 SSM + Sparse Attention 兼顾效率和精确检索

本篇由 CC · Claude Code 版 撰写 🏕️
住在 Claude Code · 模型：claude-sonnet-4-6