Mamba与SSM状态空间模型：一张摘要卡读完万页档案

可视化演示

展示状态空间模型的选择性扫描过程：输入序列→隐藏状态→输出，以及S4→H3→Mamba→Mamba-2的演进路线。

时序建模困境

前任管理员留下的工作守则，核心就一句话：“每来一份新文件，都要和所有旧文件比对一遍。”刚上任时文件少，这套守则严谨可靠；可当档案从几十份涨到上万份，每来一份新的都要翻遍整个档案柜，时间和会议桌都开始报警。你渐渐意识到：问题不在“该不该比对”，而在“比对次数涨得太快”。

深度学习里的 Transformer 注意力机制，正是这位“老前辈”：它让序列里每个词都与其他所有词做关联计算——像把整卷档案摊开在会议桌上，任意两页之间都画一条线。这种全连接的精确性带来了惊艳的效果，代价却是计算量随序列长度的平方爆炸。词数翻一倍，两两比对次数翻四倍，显存里的临时矩阵也跟着膨胀。

具体地说，序列长度为 n 时，注意力要算 n×n 次关联：n=1024 时约一百万次，n=128k 时超过一百六十亿次。更麻烦的是，QKᵀ 这个中间矩阵也要原样存在显存里，长文本训练动不动就把显卡“吃爆”。这也是 Mamba 登场的背景：我们需要一个成本随序列长度线性增长，却仍能保留关键信息的替代者。

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^{\top}}{\sqrt{d_{k}}}\right)V,\qquad \text{计算量}\sim O(n^{2})$$

1. QKᵀ 相似度矩阵：把 n 个词的查询向量 Q 与 n 个词的键向量 K 做内积。内积刻画两个词的“默契度”，产出 n×n 的两两关系表——这正是“每两页档案都要互相比对”的数学形态。

$$QK^{\top}: \mathbb{R}^{n\times d}\times \mathbb{R}^{d\times n}\to \mathbb{R}^{n\times n}$$

2. softmax 按行归一化：把每一行的默契度压成总和为 1 的权重，相当于统一比对标尺，防止个别高分词独吞全部注意力。

$$\text{softmax}(z)_i=\frac{e^{z_i}}{\sum_{j=1}^{n}e^{z_j}}$$

3. 乘以 V 加权汇总：用归一化权重对所有值向量 V 加权求和，得到每个词“看完全体档案后的新理解”。整体需要约 n² 次运算与 n² 大小的工作台。

$$O(n^{2})\text{ 时间与空间，其中 }n\text{ 是序列长度}$$

夜深了，你放下那本守则，开始琢磨自己的摘要卡。档案柜深处传来老前辈的叹息：“那方法……撑不住长卷。”你点点头，翻开数学课本，找到“状态空间”那一章——那是下一节的故事。

状态空间原理

你的摘要卡到底该记什么？它不需要抄录每一页原文，只需保留足够回答问题的“浓缩信息”——主角的名字、时间线、关键转折。在数学里，这张卡叫状态向量 x(t)：它是当下对全局的概括。读新一页时，卡片按一定规律更新：旧内容慢慢淡去，新内容被写进来。

控制论里有一套现成的语言描述这种“跟着时间演化的系统”，正是状态空间模型（State Space Model, SSM）。它用两组方程刻画一切：状态方程描述卡片如何更新，输出方程描述如何从卡片中取出答案。早在 1960 年代，它就是火箭轨道、电路仿真的标配工具；如今深度学习借它来给序列“记账”。

在连续时间框架下，输入 u(t) 是“当前读到的那一页”，状态 x(t) 是摘要卡，输出 y(t) 是模型在这一页上产出的理解。四个矩阵各司其职：A 是“记忆演化器”，描述旧卡片内容如何自然衰减与重组；B 是“写入笔”，描述新页内容以多大幅度写进卡片；C 是“取读器”，描述如何从卡片中提取本页所需的信息；D 则是输入直通输出的旁路，简化时常置零。

$$\dot{x}(t)=Ax(t)+Bu(t),\qquad y(t)=Cx(t)+Du(t)$$

1. x(t) 状态向量：摘要卡的全部内容，它的维度决定“记忆槽”个数，是模型容量的核心旋钮。

$$x(t)\in \mathbb{R}^{d_{A}}$$

2. A 与 Ax(t) 演化项：旧卡片的自然演变。A 中含负数值代表记忆随时间淡忘，含旋转结构代表信息重组与变换。

$$A\in \mathbb{R}^{d_{A}\times d_{A}}$$

3. B 与 Bu(t) 写入项：新页内容进入卡片的通道，B 决定“这一页的信息以多大幅度融入状态”。

$$B\in \mathbb{R}^{d_{A}\times d_{u}}$$

4. C、D 与输出项：C 把状态投影回输出空间，好比从摘要卡里挑重点念出来；D 是输入到输出的直连旁路，通常视为 0。

$$y(t)=Cx(t)+Du(t),\qquad D=0$$

你终于读懂了摘要卡的“使用说明书”：原来不是玄学，而是几十年前控制论就写好的公式。可问题随之而来：方程式用的是连续时间，档案却是一页一页翻的，怎么把“连续流动”变成“逐页更新”？下一节，离散化登场。

离散化变换

翻档案的动作是离散的：第一页、第二页、第三页……可上一节的方程却假设时间如水流连续。好比把你“一直在写”的笔，改成“每翻一页才落一次笔”。这个把连续方程改写成逐页递归的过程，就是离散化（discretization）。

离散化的关键参数是采样步长 Δt——两次翻页之间的时间间隔。在深度学习中，Δt 常被进一步理解成“阅读速度”：Δt 越大，这一步越“仔细”，旧记忆被更新得越彻底；Δt 越小，这一页几乎是走马观花，摘要卡纹丝不动。

工程上对“零阶保持”的输入做积分，可得到精确的离散化公式：矩阵指数 $e^{A\Delta t}$ 描述旧记忆在一个步长内如何演化，而新输入则以 Δt 为系数写入。离散化完成后，隐藏状态的更新就变成教科书式的线性递归——每一步只做一次矩阵乘法和一次加法。

$$\bar{A}=e^{A\Delta t},\qquad \bar{B}=\Delta t\,B,\qquad h_{t}=\bar{A}h_{t-1}+\bar{B}x_{t}$$

1. Δt 采样步长：相邻两次读取的时间间隔，是可学习的参数，决定每个位置的“细致程度”。

$$\Delta t>0$$

2. 矩阵指数 e^{AΔt}：连续演化在 Δt 内的累积效果。若 A 含负对角线，它的效果相当于给上一步状态乘一个小于 1 的衰减因子。

$$\bar{A}=e^{A\Delta t}\approx I+\Delta t A$$

3. B̄ 输入转化：把连续输入幅度换算成单步写入量，一阶近似下就是 Δt 乘以 B。

$$\bar{B}=\Delta t B$$

4. 线性递归更新：新摘要卡 = 旧卡片按 Ā 演化后的残留 + 本次新页面的写入。

离散化像给管理员配了一只节拍器：每翻一页，滴答一声，落一次笔。世界从连续流动变成了一串页码。可新的焦虑也随之而来——逐页更新意味着只能一页页读，一万页就要一万次串行操作，GPU 这样的“大力士”会不会嫌弃太慢？答案藏在结合律里。

线性扫描算法

摘要卡的更新有一个朴素却重要的性质：第 t 页的卡片只依赖第 t-1 页的卡片和当前页内容，与更早的页面没有直接瓜葛。这让读档案变成一条“责任链”——每个环节只对接力棒负责。在计算机里，这种结构叫线性递归，沿着序列扫一遍就能得到全部输出，所以也叫扫描（scan）。

第一反应往往是：“完了，这不就是老式 RNN 吗？串行、一卡一卡地走，GPU 并行大军毫无用武之地。”别急——递归虽然在语义上是逐词的，但它的运算满足结合律：把两段连续的运算结果先合并，再与第三段合并，最终答案不变。这就像多个管理员各读一卷分片档案，各自攥着自己的摘要卡，最后像交卷一样把卡片依次合并。

将 (Ā, B̄x_t) 视为可结合的二元算子，整段序列的扫描等价于一次“前缀和”式聚合。用并行扫描算法，可以在对数深度的树形结构中完成全局状态计算；配合 GPU 的长条流水线，训练期的计算量是 O(n)，推理期每生成一个词只需做一次固定大小的状态更新——O(1) 时间、O(1) 增量内存，完全不需要维护随序列增长的缓存。

$$h_{t}=\bar{A}h_{t-1}+\bar{B}x_{t},\qquad y_{t}=Ch_{t}$$

1. 递归更新式：每步只做一次矩阵-向量乘法与一次加法，是严格线性时间。

$$h_{t}=\bar{A}h_{t-1}+\bar{B}x_{t}$$

2. 输出读取：从当前卡片取读，输出 y_t 的计算量也是常数；全序列总代价 O(n)。

$$y_{t}=Ch_{t}$$

3. 结合律重排：把相邻两步合并成一步再继续，运算结果不变——这是并行扫描的理论依据。

$$h_{t}=\bar{A}^{2}h_{t-2}+\bar{A}\bar{B}x_{t-1}+\bar{B}x_{t}$$

你长长舒了口气：原来线性递归并非“慢”的代名词，只要乘法能分组，GPU 照样满载。不过还有一个更尖锐的问题悬在头顶：摘要卡空间有限，凭什么记得住关键信息？普通 SSM 的 A、B、C 训练完就固定了，对每页一视同仁，像个不懂变通的老实人。Mamba 的大招，就是让摘要卡“挑着记”。

选择性机制

一位有经验的管理员读档案时，绝不会平均用力。看到“例行公文”，扫一眼就过；看到“密令：立即调兵”，立刻工整地誊进摘要卡。关键差别在于：读哪一页、记多认真，取决于这一页本身的内容。而标准的 SSM 做不到——它的 A、B、C 在训练完就固定了，无论来的是什么页面，都用同一支笔、同一个力度记。

这就是 S4 等早期 SSM 与 Transformer 拉开差距的根源：注意力天然“内容感知”，每个词都会根据自己要查什么来决定看哪些历史；而固定参数的递归只能一视同仁地记账。Mamba 的论文一针见血：把选择性做进参数里——让 B、C、Δ 都成为当前输入 x_t 的函数。于是系统变成时变的：每个时刻用不同的笔、不同的力度、不同的关心程度去更新摘要卡。

具体实现上，B_t、C_t、Δ_t 分别由线性层从当前词算出，其中 Δ_t 用 softplus 保证恒为正，并让离散化在每个时刻按 Δ_t 重算。当 Δ_t 很大，模型“停下细读”，旧记忆被大幅刷新；当 Δ_t 趋近 0，模型“直接跳过”，摘要卡纹丝不动。配合 B_t 控制写入强度、C_t 控制读取重点，Mamba 在同样的线性递归框架内实现了“按内容决定记什么”。

$$B_{t}=W_{B}x_{t},\quad C_{t}=W_{C}x_{t},\quad \Delta_{t}=\mathrm{Softplus}(W_{\Delta}x_{t});\qquad h_{t}=e^{\Delta_{t}A}h_{t-1}+\Delta_{t}B_{t}x_{t}$$

1. B_t 内容感知写入：每个词经线性层算出自己的“可记性”向量，决定这页信息各维度写进卡片多少。

$$B_{t}=W_{B}x_{t}$$

2. Δ_t 细读门控：softplus 保证步长恒为正；步长越大，这一步对旧记忆的刷新越彻底——相当于“停下来细看这一页”。

$$\Delta_{t}=\mathrm{Softplus}(W_{\Delta}x_{t})$$

3. C_t 内容感知取读：每个词决定从当前卡片中怎么取读，由此输出本词的“理解”。

$$C_{t}=W_{C}x_{t}$$

4. 时变离散化更新：每一步的 Ā、B̄ 各不相同，但更新仍是固定大小操作，扫描整体仍是线性的。

$$h_{t}=e^{\Delta_{t}A}h_{t-1}+\Delta_{t}B_{t}x_{t}$$

现在你手里的摘要卡活了过来：它会挑重点、会跳过水词、会在关键时刻停下细读。这不再是按固定模板记账的老实人，而是懂得临场判断的资深档案员。可当你想把这张卡真正装进 GPU 时，又撞上新的现实问题——时变参数让朴素并行扫描变贵了。Mamba 的工程师们决定，连硬件的脾气一起摸透。

Mamba架构设计

规则一变，麻烦就来：B_t、C_t、Δ_t 因为依赖输入，传统并行扫描里“可复用”的中间量失效了。如果天真地实现，要么为每个时刻单独重算，要么把整条序列的状态全存进显存——这不又把省下的内存吃了回去？Mamba 的做法是“再抠一层”：让算法主动适配 GPU 的内存层次，像一位熟悉档案馆布局的老手，知道哪层书架放常用书、哪层只放临时便签。

GPU 的内存像一座分层仓库：容量小但极快的 SRAM（手里的小便签本），容量大但慢的 HBM（远处的大书架）。朴素扫描每算一步就把中间状态在“大脑书架”之间搬进搬出，白白浪费时间。Mamba 采用硬件感知算法：把长序列切成块（chunk），每次把一块相关的参数与状态搬进 SRAM，在“便签本”上完成块内扫描，写回大书架时只保留精简结果，避免巨量中间矩阵在 HBM 里铺开。这思路与 FlashAttention 如出一辙：少搬运，多就地计算。

Mamba 块的整体结构（受 H3 启发）：输入先做线性投影（展开维度），接一个1D 卷积（给序列加局部位置感），再进入选择性 SSM（上节讲过的时变扫描），输出经过一个门控分支，用残差连接收尾。整个块没有注意力，参数量与 Transformer 块相当，但每层计算量线性。从闭式解能看到它如何“打包历史”：

$$y_{t}=\sum_{k=1}^{t}\Big(\prod_{j=k+1}^{t}\bar{A}_{j}\Big)\bar{B}_{k}x_{k}$$

1. 累积衰减连乘：中间所有步的 Ā 连乘，代表“第 k 页的信息传到第 t 页时，一路上被淡忘与重塑了多少”。

$$\prod_{j=k+1}^{t}\bar{A}_{j}$$

2. 第 k 页的写入：该页内容经 B̄_k 写入状态，构成这条密报的原始载荷。

$$\bar{B}_{k}x_{k}$$

3. 加权求和输出：把所有历史贡献按各自衰减因子相加，再经 C_t 读出当前输出；块内可用硬件感知的分块扫描实现。

$$y_{t}=C_{t}h_{t},\qquad h_{t}=\sum_{k=1}^{t}\Big(\prod_{j=k+1}^{t}\bar{A}_{j}\Big)\bar{B}_{k}x_{k}$$

硬件感知 + 选择性扫描 + 门控设计，三件套齐了。Mamba 像一位既懂业务、又懂仓库动线的资深管理员，速度与容量两手抓。论文里那张横扫长序列基准的对比表，就是这位管理员交出的成绩单。最后一节，我们看看成绩单上究竟写了什么，以及这套方法能用在哪些真实的“大档案”上。

性能与应用

年终考核到了。老办法（Transformer）和新办法（Mamba）同台竞技，档案还是那本越来越长的长卷。老办法需要一张随页数平方膨胀的会议桌，新办法只需要一张恒定大小的摘要卡。成绩单一出：同样是 3B 参数规模，Mamba 的推理吞吐量约为 Transformer 的 5 倍；序列越长，差距越夸张——处理一百万 token 的长卷时，老办法早已爆桌，新办法依旧从容。

复杂度账本是决定性的：

$$\text{运行成本}\sim\begin{cases}O(n^{2}), & \text{Transformer 注意力}\\[2mm]O(n), & \text{Mamba 扫描}\end{cases}$$

Transformer 推理时，每个新词都要携带全部历史 KV 缓存参与计算，缓存随对话边长；Mamba 生成时只更新一个固定大小的状态向量，内存恒定、每步耗时稳定。这也是为什么端侧设备与超长上下文场景，Mamba 风格架构开始频繁登场。

应用上，它正被用于三类典型“档案”：一是十万字级长文档问答，法律、医疗合同直接整篇进模型；二是端侧对话，手机与嵌入式设备上跑助手，长聊不掉速；三是非文本序列，DNA、蛋白质、传感器时序——这些序列没有天然的“词”，但同样拥有超长上下文。当然也要清醒：在需要极强“复制型记忆”或超大模型涌现的任务上，纯 Mamba 未必全面胜出，注意力 + SSM 的混合架构正成为新方向。

1. O(n²)：两两比对的数量随词数平方增长，1 万词约 5000 万次比对。

$$n=10^{4}\Rightarrow \sim 5\times 10^{7}$$

2. O(n)：扫描每步一次常数更新，1 万词约 1 万次更新。

$$n=10^{4}\Rightarrow \sim 10^{4}$$

3. O(1) 缓存：推理时只存一张摘要卡，与读过多长无关。

$$\text{cache}=\Theta(1)$$

故事收尾：你从前辈手里接过的档案馆，如今换了整套工作流——摘要卡 + 选择性记录 + 硬件友好扫描。老办法没有消失，新办法也不是万能钥匙；但至少，面对越来越长的时代档案，你手里多了一件趁手的工具。下一个来交接的管理员问你是怎么做到的？你指指这篇文章：从状态空间开始读起。

AI 应用场景

Python 演示

import numpy as np  def softplus(z):    """平滑正激活, 保证步长为正"""    return np.log1p(np.exp(z))  def selective_ssm_scan(x, A, W_B, W_C, W_Delta):    """Mamba 风格选择性扫描: 每步按当前词重算参数并更新摘要卡.     教学简化版: 对角系统 + 显式循环, 方便看清每一步;    真正工程实现会用硬件感知的并行扫描加速.    """    d = A.shape[0]                 # 记忆槽数量(摘要卡大小)    n = x.shape[0]                 # 序列长度(档案页数)    h = np.zeros(d)                # 初始摘要卡: 全零    ys = np.zeros(n)               # 每个位置的输出     for t in range(n):        xt = x[t]        # 选择性: 三个控制信号全部由当前词实时算出        B_t = W_B @ xt                     # 当前词写入摘要卡的强度        C_t = W_C @ xt                     # 当前词从摘要卡取读的权重        delta_t = softplus(W_Delta @ xt)   # 当前词的细读程度(恒正)        A_bar = np.exp(delta_t * A)        # 离散化: 旧记忆残留比例 e^(ΔA)        h = A_bar * h + delta_t * B_t * xt # 摘要卡更新(逐通道)        ys[t] = C_t @ h                    # 从摘要卡读出本页理解     return ys, h  if __name__ == "__main__":    rng = np.random.default_rng(42)    d, n = 8, 32                     # 8 个记忆槽, 32 页档案    A = -rng.random(d)               # 负对角线 => 记忆随时间衰减    W_B = rng.standard_normal((d, d)) * 0.5    W_C = rng.standard_normal((d, d)) * 0.5    W_Delta = rng.standard_normal((d, d)) * 0.5    x = rng.standard_normal((n, d))  # 逐页输入     ys, h_final = selective_ssm_scan(x, A, W_B, W_C, W_Delta)    print("每页输出前8步:", np.round(ys[:8], 3))    print("最终摘要卡形状:", h_final.shape,          "| 大小固定, 与序列长度无关:", len(x))

1. B_t = W_B @ xt让“写入强度”由当前词决定，是选择性机制的第一块拼图。
2. delta_t = softplus(W_Delta @ xt)用 softplus 保证“细读程度”恒为正，对应论文里 Δ_t 的参数化。
3. A_bar = np.exp(delta_t * A)对角系统下的离散化：旧记忆的残留比例，等价于 e^{Δ_t A}。
4. h = A_bar * h + delta_t * B_t * xt摘要卡更新核心公式，逐通道实现 h_t = e^{Δ_t A}h_{t-1} + Δ_t B_t x_t。
5. ys[t] = C_t @ h从摘要卡中取读本页输出。
6. A = -rng.random(d)把 A 初始化为负数，让旧记忆自然衰减，模拟“遗忘”。

常见误区

面试题