1607 字
8 分钟
SSM:替代 Transformer 的线性架构

长上下文的计算成本#

Self-Attention 的代价:每个 token 要关注序列中所有其他 token → 计算量随上下文长度平方级增长。翻倍上下文 = 四倍计算。

上下文长度KV Cache 占用(7B 模型)
4K~1-2 GB
128K30-40 GB VRAM

随着上下文增长,计算量和缓存需求都会增加。

SSM 直觉:像记笔记一样处理序列#

可以把 SSM 理解为不断更新的摘要:新信息到来时,决定保留什么、舍弃什么。

SSM 维护固定大小的隐藏状态(hidden state),每一步更新这个状态,而不是回看所有历史 token。整个序列历史被压缩进这个固定大小的状态中。

核心公式#

  • 状态更新:h_t = A·h_{t-1} + B·x_t(压缩历史 + 吸收新输入)
  • 输出生成:y_t = C·h_t(从状态中提取输出)

序列处理时间随长度线性增长,递推状态大小固定。

工业界的混合架构#

公司模型做法效果
NVIDIANemotron-H (2025)92% attention 替换为 Mamba-2推理吞吐比 LLaMA-3.1 高 3x
IBMBamba-9B混合 SSM与 LLaMA-3.1 8B 相当,显存占用大幅降低
MicrosoftPhi-4-mini-flash-reasoningSambaY(Mamba + 滑动窗口 attention + GMU)吞吐比前代高 10x

这些模型将 SSM 与 Attention 结合,用于降低推理成本。

SSM 的演化谱系#

S4 (2021) — 奠基#

Albert Gu (Stanford)。证明 SSM 能比 RNN 更好地处理长程依赖。参数固定,不依赖输入。

Mamba (2023) — 突破#

Albert Gu & Tri Dao。关键创新:让 SSM 参数 (A, B, C) 依赖输入(输入选择性)。模型可以选择性地决定记住什么、遗忘什么。Mamba-3B 在语言 benchmark 上超越同尺寸 Transformer,推理吞吐 5x。

Mamba-2 (2024) — 理论统一#

Dao & Gu (ICML)。证明了 Transformer 和 SSM 通过结构化半可分矩阵在数学上是相关的。SSD 层比 Mamba-1 快 2-8x。

Mamba-3 (2025)#

进一步提升序列建模能力,使用状态空间原理改进。

中间架构#

  • H3:最早尝试混合 SSM + Attention
  • Hyena:尝试用长卷积完全替代 Attention(未成为主流)
  • S5:改进训练稳定性

混合架构:Attention + SSM#

纯 SSM 的弱点:所有信息压缩进固定隐藏状态 → 精确回忆困难。问纯 Mamba “8000 token 前用户说的原话是什么?” 可能答不上来。Transformer 可以直接访问历史 token 的表示,更适合这类精确检索。

类比大脑:

  • 工作记忆(Attention):正在积极思考的东西,精确回忆
  • 长期记忆(SSM):总结和压缩一切其他东西

Jamba 架构:少数 Attention 层处理精确回忆,大量 SSM 层承担主要的序列处理。比例不是 50-50——通常每 8 层中有 1 层用 Attention。

关键发现:在混合环境下,Mamba-1 + Attention 的组合优于 Mamba-2 + Attention。组件之间的交互比各自单独表现更重要。

记忆对比#

维度TransformerSSM
记忆类型情景记忆(episodic)语义记忆(semantic)
存储方式KV Cache,显式存储所有历史 token压缩隐藏状态,丢失精确文本
精确回忆✅ 强❌ 弱
长序列成本O(N²),随长度爆炸O(N),恒定显存
220K tokens / 24GB无法运行✅ 畅跑

基准测试结果(2025)#

  • 短序列(<8K):Transformer 快 1.9x
  • 长序列(~57K):SSM 快 4x,节省 64% 显存

这些测试表明,序列变长后,SSM 的效率优势更明显。

SSM 如何用于 Agent 记忆#

1. 上下文窗口不是天花板#

Agent 运行的是极长序列——工具调用、观察日志、多步推理追踪……这些堆得很快。SSM 以固定、可预测的显存成本处理长上下文,有助于控制大量并发 Agent 会话的状态存储成本。

2. 用状态保留任务进展#

Agent 做 50 步任务时会累积大量运行上下文。如果把所有信息都放进上下文窗口,计算成本会增加,无关内容也可能干扰推理。

SSM 的隐藏状态结构上就是那个”写入-管理-读取”循环:每个 token 更新状态 → 无关信息消退 → 任务相关结构持续累积。模型随任务推进更新状态,必要时再配合外部检索。

Apple Research 结论:当 SSM 被授予与外部工具的交互能力时,尽管状态大小固定,它们能在算术、推理和编程任务上实现长度泛化,处理比训练时更长的任务。SSM 是在工具化和 agentic 场景下实用的高效替代方案。

3. 持续更新上下文表示#

哈佛论文指出:人脑的语言处理更接近 SSM 的方式,而非 Transformer。大脑维护的是压缩的、不断演进的上下文表征,不会重放每一个历史词来生成下一个词。

Agent 做长周期任务是最接近人类持续认知的 AI 模拟。人类管理多天项目时不需要每次会议的完整记录,只需要一个状态——事情进展如何、什么紧急、什么已解决。

SSM 的隐藏状态为这种持续记忆提供了一种归纳偏置(inductive bias)。

项目:Cerebra — 大脑仿生的 Agent 架构#

作者正在开源的 agent 项目,用大脑分区设计 agent 系统,不同区域用不同复杂度:

脑区功能复杂度
Cerebrum(大脑皮层)高阶认知:规划、推理、语言理解、工具整合O(N²) — LLM
Hippocampus(海马体)记忆系统:情景/语义/工作记忆O(N) — SSM/文件
Cerebellum(小脑)程序记忆、习惯执行、响应缓存O(N) — 线性
Amygdala(杏仁核)优先级与安全过滤O(N) — 线性
Thalamus(丘脑)中央路由、上下文门控O(N) — 线性
Hypothalamus(下丘脑)系统稳态:token 预算、速率限制、成本监控O(N) — 线性
Brainstem(脑干)自主运行循环:编排、生命周期管理O(N) — 线性
Corpus Callosum(胼胝体)LLM(平方) 与 SSM(线性) 子系统的桥接边界层
Basal Ganglia(基底节)动作选择与奖励追踪O(N) — 线性
Pineal Gland(松果体)调度与后台任务(系统休眠周期)O(N) — 线性

核心结论#

SSM 通过固定大小的递推状态控制历史存储开销,适合探索长期运行的 Agent 记忆。不过,压缩会损失细节,精确回忆仍可能需要 Attention 或外部检索。

SSM:替代 Transformer 的线性架构
https://blog.lpkt.cn/posts/concepts/state-space-models-agent-memory/
作者
lollipopkit
发布于
2026-05-06
许可协议
CC BY-NC-SA 4.0