★ FEATURED ARTICLE 1. 为什么 LLM 圈突然都在聊 SSM一次关于“注意力”的路线反动如果你在 2023 年之前问做大模型的同行序列建模应该用什么结构十有八九会得到同一个答案Transformer加注意力机制。这几乎成了行业共识以至于很多人直接把“大模型”等同于“Transformer 堆叠”。但 2023 年下半年开始格局悄悄变了——Mamba 论文放出来之后社区里讨论“状态空间模型”SSMState Space Model的声音越来越大到了今天已经没有人敢忽视这条技术路线。先说一个最直观的背景Transformer 的核心是自注意力它的信息交互能力和序列长度呈平方级关系。简单说Token 数量翻倍计算量和显存开销大约是原来的四倍。这是数学上写死的规律任何工程优化都只能缓解、没法消除。你可以在 4096、8192 的上下文长度下活得很滋润一旦想上 10 万、100 万 Token 的长文本注意力机制就会变成一座绕不开的大山。大家天天说的“长上下文”本质是一场军备竞赛而这场竞赛的尽头必然有人会去寻找注意力之外的替代方案。状态空间模型正是在这个背景下被推到台前。它来自经典控制理论几十年前就被用来描述动态系统核心思路是系统在每一时刻的输出由当前输入和历史状态共同决定。这个描述天然适合序列建模。更关键的是它绕开了“每个 Token 都要和其他所有 Token 做交互”的路线改用一种线性的、递推的方式处理历史信息计算复杂度直接降到序列长度的线性级别。我最早注意这条路线是因为业内一位朋友在做法律文书的超长文本分类上下文经常到十几万 Token用 Transformer 根本训不动。他换用基于 SSM 的模型做基座之后显存占用掉了一个数量级速度还快了不少。对我这种被注意力机制折磨过太多次的人来说这种实际案例比任何理论推导都有说服力。但要真正理解 SSM不能只停留在“它可以处理长文本”这种模糊印象。这篇文章会把 SSM 的原理、代表性架构、工程落地时的关键细节以及接下来的方向全部拆开讲清楚。无论你是刚入门的研究生还是在做推理优化、模型选型的工程师这套内容应该都能帮你建立起一张相对完整的地图。2. 从控制论到深度学习SSM 到底是怎么描述序列的2.1 一个方程看懂状态空间模型的核心思想很多人第一次接触 SSM 会被“状态空间”这四个字吓到其实它的内核非常朴素。想象你在跟踪一辆行驶中的汽车你不会每一秒都把汽车的历史轨迹全记下来你只需要维护一个“当前状态”——比如位置和速度——然后基于当前状态和新的观测去推算下一时刻的状态。这就是状态空间的核心用有限维度的“状态”压缩整个历史信息。用数学语言来描述一个连续时间的状态空间模型包含两个方程状态方程( x(t) A x(t) B u(t) )输出方程( y(t) C x(t) D u(t) )这里的 ( u(t) ) 是输入序列对应 NLP 里的 Token 向量( x(t) ) 是隐状态( y(t) ) 是输出。A、B、C、D 都是可学习的参数矩阵。别被符号吓到你把 A 想象成“上一时刻状态对当前时刻状态的影响方式”B 是“新输入如何写入状态”C 是“状态如何映射到输出”整个系统就通了。到了深度学习里我们需要处理的是离散的 Token 序列所以要先把连续方程离散化。标准做法是用零阶保持ZOH或欧拉法把连续参数 A、B 转换成离散参数 ( \bar{A} )、( \bar{B} )。离散化之后模型就变成一个递推过程( x_k \bar{A} x_{k-1} \bar{B} u_k )( y_k C x_k )这个递推形式是不是很像 RNN是的SSM 的本质就是 RNN 的一种高阶泛化形态。你每一时刻只需要维护一个固定维度的隐状态 ( x_k )然后基于它产生输出。2.2 两种计算模式循环推理与卷积训练的配合SSM 最巧妙的地方在于它同时拥有两种计算形态——循环和卷积可以根据场景自由切换。在推理阶段也就是真正跑模型、一个 Token 一个 Token 生成内容的时候SSM 可以用循环模式。这意味着模型只需要维护固定大小的隐状态不需要像 Transformer 那样存储前面所有 Token 的 Key-Value 缓存。这一点有多重要我举个例子如果 KV Cache 的大小是 1GB那 1000 个并发用户的推理服务就要吃掉 1TB 内存而且这个数字还会随上下文长度继续增长。SSM 直接把这种显存压力变成了一个固定常量。在训练阶段SSM 又可以使用卷积模式。因为线性时不变系统有一个性质整个输入序列的输出可以表示成输入序列和一组卷积核的卷积运算。这样一来训练时可以把整个序列一次性并行处理避开 RNN 那种逐 Token 串行导致的训练效率低下问题。S4 论文里正是利用了这一性质实现了并行化训练。2.3 S4 的登场解决长序列建模的原始版本2021 年底Albert Gu 等人提出了 S4Structured State Spaces for Sequence Modeling这是 SSM 真正进入深度学习领域的关键节点。S4 做了两件非常重要的事第一它证明了随机初始化的 A 矩阵很难学到远程依赖——梯度在长序列上容易消失或爆炸。他们用 HiPPOHigh-order Polynomial Projection Operators理论来初始化 A 矩阵使得模型可以持续记忆序列的历史信息。简单说HiPPO 是一种特殊的初始化方法它让状态能够在数学上保证“最优地记忆”过去一段时间的输入。当时 S4 在 LRALong Range Arena基准上直接把 Transformer 系模型打爆长序列分类任务上优势非常明显。第二它解决了一个工程问题A 矩阵如果是满秩的计算复杂度仍然很高。S4 把 A 矩阵结构化成对角加低秩的形式配合 Cauchy 核算法把计算复杂度降到了可接受的范围。这为后续一系列变体铺平了道路。从工程人员的角度理解 S4它相当于证明了“序列建模不一定非注意力不可”它提供了一套完整的理论框架和可落地的算法。而 Mamba 则是在这个框架上做了更激进、更实用的改良。3. Mamba 与选择性机制让状态空间模型学会“选择性遗忘”要说 SSM 真正引爆社区热度还得是 Mamba。如果说 S4 是理论奠基那 Mamba 就是把 S4 从一个“学术玩具”变成了“能打的模型结构”。3.1 固定权重的问题SSM 早期版本的“近视眼”S4 虽然能处理长序列但它有一个致命缺陷它的 A、B、C 参数在序列处理过程中是固定不变的。这意味着模型对历史信息的压缩方式无论输入是“今天天气不错”还是“请把张三的合同在明天上午前发给李四”都采用同一套机制。这种“一视同仁”在语言任务上显然不合理。用一个人的记忆来类比一个合格的助手应该知道老板交代的“明天上午交给李四”这个信息很重要可能需要在对话过程中反复引用而“今天天气不错”这种寒暄过去就过去了。S4 的问题在于它会把所有信息以同等的固定比例写进状态无法根据内容判断重要程度。3.2 Mamba 的三个关键改进选择性、硬件感知、简化架构Mamba 的论文标题叫《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》翻译过来就是“带选择性状态空间的线性时间复杂度序列建模”。它的核心改进集中在三点第一选择性扫描Selective Scan。Mamba 让 A、B、C 矩阵变成输入的函数也就是说每个 Token 都会动态决定“上一时刻的状态保留多少”“当前输入有多少写入状态”“状态如何映射为输出”。这就打破了对所有 Token 一视同仁的限制。当模型读到一个重要的实体名或指令时B 矩阵会放大写入权重A 矩阵会让状态保留得更久遇到“嗯”“啊”“今天天气”这类低价值 Token模型可以快速把状态覆盖掉。这正是“选择性遗忘”的含义。第二硬件感知设计。Mamba 不再使用标准的卷积表示而是回到了循环扫描。但为了避免循环扫描在训练时太慢他们特别针对 GPU 的内存层次做了优化——将状态常驻在 SRAM静态随机存取存储器里避免频繁读写 HBM高带宽内存。这听起来很工程化但在实际训练中这种优化能带来几个数量级的加速。Flash Attention 的成功已经证明了一件事在大模型时代算法的“硬件友好度”和算法本身的数学优雅性同样重要。第三把 S4 中的复杂结构简化。S4 用到了许多专门设计的矩阵结构和分解技巧Mamba 干脆做减法把 A 矩阵简化为对角矩阵把深度可分离卷积Depthwise Convolution拿来做输入投影并配合 Swish 激活和层归一化。最终 Mamba 的模块非常干净输入先过一个线性层扩大维度经过 1D 卷积提取局部信息再走选择性扫描最后线性映射回原维度。3.3 实测观察Mamba 在哪些任务上真的赢了我在长文本场景里跑过 Mamba 和同尺寸 Transformer 的对比例子说一下直观感受。在单卡 A100 上做 100K 长度序列的语言建模Mamba 的显存占用大约只有 Transformer 的几分之一训练速度也明显占优。在需要长程依赖的任务上比如 DOCSTRING 摘要、DNA 序列分类Mamba 的表现可以和同规模的 Transformer 打平甚至更好。但在短文本、复杂指令跟随、多步推理这类任务上Mamba 并没有明显优势甚至有部分场景略逊于同参数的 Transformer。这也合理选择性机制虽然高效但它的信息交互模式毕竟和注意力不同在处理“需要同时关注多个分散位置的 Token”这类任务时SSM 的表达能力仍有天花板。这里要特别说一句很多人对 Mamba 有“一定能全面超越 Transformer”的期待这是不现实的。Mamba 的真实价值在于它在特定场景下提供了一种更经济、更高效的替代方案而不是推翻注意力机制本身。4. 工程实战从选型到部署的完整落地经验4.1 你该在什么场景下用 SSM 类的模型以及不该用的场景根据自己的实际项目经验我整理了一张 SSM 擅长的任务清单和避雷区域供参考适合选用 SSM 的场景不适合选用 SSM 的场景超长文本超过 32K Token分类、摘要需要精确实体间关系的多跳问答音频波形、基因组序列等信号类数据短文本上的极高性能竞赛端侧部署受限的实时流式推理需要成熟的工具链和生态的场景高并发长上下文推理服务对已有代码库依赖很深、迁移成本高的项目一个很典型的例子是智能客服。客服系统的上下文经常会包含整段历史会话记录如果放到 Transformer 模型里KV Cache 会随会话长度线性膨胀最后可能出现一个高并发服务撑不住的情况。SSM 因为隐状态固定缓存占用不随上下文长度增长天生适合流式对话。我自己做过一个对比同样处理平均 80K 上下文的客服会话Mamba 架构的推理服务在相同显存下吞吐量大约是 Transformer 的 2 到 3 倍。反之如果你的任务是“从一篇文章里找出所有互相矛盾的事实”这类任务高度依赖全局注意力SSM 的反而不容易做好。因为 SSM 对信息的压缩是有损的长距离的精确关系容易在压缩过程中丢失。这类任务老老实实用 Transformer 或者混合架构更合适。4.2 部署 Mamba 时容易踩的三个坑第一CUDA 核函数兼容性。Mamba 官方实现依赖 mamba-ssm 包中的自定义 CUDA 核函数这些核函数对 GPU 的架构有要求。你在 A100/H100 上跑得好好的换到 V100 可能直接编译不过更老的 GPU 连 Flash Attention 都用不了。部署前一定要确认 GPU 型号和 PyTorch 版本。我见过有同事在 T4 上跑 Mamba折腾了半天才发现是代码根本没编译成功报错信息还不明显。第二状态维度的超参选择。SSM 的状态维度 ( N ) 决定了模型能记住多少历史信息。这个值太小长依赖能力不够太大则训练和推理耗时上升。根据我的经验对于文本任务( N ) 维设置在 16 到 64 之间比较合理音频或基因组等信号数据可以适当加大。这个问题官方文档说的比较少需要实测调参。第三与 RAG 架构的配合。现在很多应用是“LLM RAG”的模式也就是先把外部知识库检索出来拼到 Prompt 里再交给模型生成。如果你把 SSM 模型作为基座要注意它对 Prompt 中的上下文顺序非常敏感。我实测下来SSM 对“越靠近当前位置的信息越重要”的假设比 Transformer 更强所以检索到的文档排序策略需要额外优化——把和当前查询最相关的文档放到最后效果会明显好于按相关度从高到低排列。4.3 推理优化跳过 KV Cache 之后真正省钱的地方在 Transformer 推理中KV Cache 往往是显存瓶颈。SSM 的循环模式没有 KV Cache只有固定大小的隐状态这使推理服务的内存成本变得更可控。但“没有 KV Cache”不代表不需要优化。SSM 的推理瓶颈转移到了状态更新本身批量大小越大状态矩阵乘法的计算量就越可观。工程上的处理方式通常有这几个方向使用张量并行Tensor Parallelism切分状态矩阵降低单卡压力对 SSM 的状态转移矩阵做低秩分解在精度损失可控的前提下进一步减少计算量利用连续批处理Continuous Batching策略把不同请求的 SeqLen 对齐提高 GPU 利用效率。这里要提醒一下虽然很多推理框架开始支持 Mamba但生态成熟度相比 vLLM、TensorRT-LLM 对 Transformer 的支持还有很大差距。生产环境如果要上线 SSM 模型往往需要自己写一部分 CUDA 扩展或者推理调度逻辑“开箱即用”目前还不太现实。5. 前沿方向混合架构、硬件协同与下一代状态空间模型5.1 混合架构的兴起SSM 和注意力不必二选一2024 年有一个明显的趋势就是 SSM 和注意力机制从“竞争关系”走向“互补关系”。最典型的就是 Jamba它把 Mamba 层和 Transformer 层交错堆叠在一起论文效果看上去很有意思既保留了 SSM 的长序列效率和低成本推理又用注意力层弥补了 SSM 在复杂指令跟从和精准检索上的短板。还有 Zamba我在个人项目里试过这个架构。它在每个 Mamba 模块里加入一个共享的注意力层让模型能够对全局信息做周期性的“回顾”。相比同参数量的纯 Transformer 或纯 MambaZamba 在长上下文任务上表现更均衡训练效率也不错。对于工程团队来说混合架构是当前性价比最高的尝试方向。它在保证模型能力下限的同时大幅降低了长上下文推理的显存成本。如果你在做技术选型我建议不要把目光锁定在“纯 SSM”上而是要考虑“SSM 占多少比例、注意力层放在哪里”这个更精细的配置问题。5.2 硬件协同设计算法倒逼芯片架构的时代过去做 AI 芯片设计的人更多是围绕 Transformer 的矩阵乘法和注意力模式来设计加速器。但 SSM 的循环扫描模式对芯片架构提出了新挑战它的非线性递归路径是加速的难点而这块在传统 GPU 上并不是强项。我看到的一个前沿方向是在 FPGA 或专用芯片上为 SSM 设计“扫描加速器”把循环扫描操作做成硬件原语。这有点类似于 Transformer 时代 Flash Attention 针对 GPU 内存层次做优化。如果 SSM 继续发展硬件和算法的协同设计可能会成为一个重要的竞争方向。这对做系统架构和底层优化的工程师来说是一个机会窗口。5.3 Mamba-2、Mamba-3 与更长序列的探索Mamba-2 的发布进一步优化了状态空间模型的理论框架。它的核心技术是引入“半可分矩阵”Semiseparable Matrices把 SSM 和某种形式的线性注意力统一到了一个数学框架下。这不仅让理论更优雅也在 GPU 上跑得更快。如果你去看 Mamba-2 的代码会发现它在处理大隐藏层维度时矩阵分解的策略比第一代精细很多。Mamba-3 引入了“内存层”的概念进一步强化了 SSM 处理超长上下文时的记忆容量。它给我的感觉是SSM 的进化速度并不比 Transformer 慢每隔几个月就有新的结构改进出来。如果你想在这个方向深耕保持论文追踪的频率非常重要——尤其要留意 Albert Gu 团队和斯坦福相关课题组的最新工作他们基本主导了这个方向的研究节奏。5.4 SSM 在知识库应用中的潜力自然语言到本体的新入口标题里提到了 LLM wiki、ontology 这些热词结合我自己的观察SSM 在知识库方向有几个值得关注的结合点。传统 RAG 管线的瓶颈之一是检索到的文本片段再长进入 LLM 上下文后能被利用的信息仍然受限于模型的注意力能力。如果用 SSM 作为长文本编码器把整个知识库文档一次性编码成一个固定维度的状态向量检索或问答时直接对接这个状态向量就能绕开“把所有命中文本都塞进 Prompt”的做法。这种“先压缩后检索”的范式在处理大规模企业知识库时很有想象空间。另外SSM 也可以用于本体Ontology构建和关系抽取任务。由于它对超长上下文的建模能力更强可以一次性读入大量文档跨文档抽取实体关系这对于自动构建知识图谱是天然的优势。当然这些方向目前都还在探索期还没有特别成熟的工业级方案。但从趋势看SSM 不是只用来做语言模型的替代品它作为一种高效的序列建模工具会在更多 AI Infra 层面找到应用场景。6. 给入门者的实操建议最后聊点实在的如果你想自己动手体验 SSM应该怎么开始。别上来就啃理论论文我建议先跑通官方 Mamba 仓库的 Demo用一个小数据集微调一个模型观察它在长文本上的实际效果。然后再去读 S4 的 HiPPO 初始化和选择性扫描的源码把每个参数的含义弄清楚。之后再上手混合架构比如跑一下 Jamba 或 Zamba 的开源权重和同规模的 Transformers 做对比。环境配置上需要注意Python 版本建议 3.10 以上PyTorch 2.1 以上CUDA 11.8 或更高GPU 显存至少 16GB微调 1.4B 级别模型。我用的是 24GB 显存的 RTX 3090 / 4090亲测可以跑通 Mamba 1.4B 的推理和 LoRA 微调。更小的机器可以从 130M 参数级别的模型入手。一个小细节用 HuggingFace Transformers 加载 Mamba 模型时要确认 transformers 版本不要过低因为 Mamba 的支持是 4.39.0 之后才合入的。如果你在使用官方 mamba_ssm 包还需要手动编译 CUDA 扩展务必先跑一下官方自带的测试脚本确认环境没问题再开始正式实验。关于学习路径我个人的体会是先跑实验建立直觉再回头补数学效率远高于从论文开始啃。SSM 的理论框架有深度但真正驱动我理解它的还是亲眼看到它在长文本任务上的表现差异。理论与实践交替进行是掌握这条技术路线最快的路径。
阅读完成 · 觉得有帮助？
