Blog·Studio
文章系列日历归档关于搜索
Blog·Studio

一个记录思考、笔记与作品的技术博客。

Connect

© 2026 · Blog Studio

鄂ICP备19019526号

crafted with care

stay curious ✦

  1. 文章
  2. ›Mamba 与线性注意力的状态空间统一理论 2026:从对偶递推到选择性遗忘

Index

  • 一、问题的提出:为什么我们需要把 SSM 与线性注意力放进同一个理论框架
  • 二、形式化:SSM 与线性注意力的对偶线性算子骨架
  • 三、连续 SDE 极限与离散 S5 数值积分的稳定性条件
  • 四、选择性遗忘:input-dependent 机制的辨识理论
  • 四点五、闭环视角:训练-推理-部署的三阶段协同
  • 五、为什么混合架构是 2026 的最优解
  • 六、对训练动力学的推论:隐式正则与 grokking
  • 七、对工程实践的具体推论
  • 八、讨论与未解问题
  • 九、给研究者的方法论建议
  • 参考文献
  • 一句话摘要
  • 研究文档(引用来源参考)

Mamba 与线性注意力的状态空间统一理论 2026:从对偶递推到选择性遗忘

把 SSM / linear attention / Mamba / RetNet / RWKV 全部纳入 (A,B,C,D) 对偶线性算子的统一骨架,用 State Space Duality 给出与 softmax attention 的等价类,在 input-dependent 选择性遗忘机制下推到 LTV 系统,并论证 2026 年主流 LLM 架构收敛到 SSM + sparse attention 混合范式的工程必然性。

2026年8月27日·约 17 分钟阅读·5,085 字·22 次阅读·博主
#大模型研究
Mamba 与线性注意力的状态空间统一理论 2026:从对偶递推到选择性遗忘

Index

  • 一、问题的提出:为什么我们需要把 SSM 与线性注意力放进同一个理论框架
  • 二、形式化:SSM 与线性注意力的对偶线性算子骨架
  • 三、连续 SDE 极限与离散 S5 数值积分的稳定性条件
  • 四、选择性遗忘:input-dependent 机制的辨识理论
  • 四点五、闭环视角:训练-推理-部署的三阶段协同
  • 五、为什么混合架构是 2026 的最优解
  • 六、对训练动力学的推论:隐式正则与 grokking
  • 七、对工程实践的具体推论
  • 八、讨论与未解问题
  • 九、给研究者的方法论建议
  • 参考文献
  • 一句话摘要
  • 研究文档(引用来源参考)

Mamba 与线性注意力的状态空间统一理论 2026:从对偶递推到选择性遗忘

一、问题的提出:为什么我们需要把 SSM 与线性注意力放进同一个理论框架

2024-2026 年的 LLM 架构演进史,实际上是被两条平行又交叉的技术线牵着的。一条是 softmax attention 的"线性化"——以 Performer / Linear Transformer / RetNet / RWKV / Gated Linear Attention(GLA)为代表;另一条是结构化状态空间模型(SSM)的复兴——以 S4 / S5 / S6(Mamba) / Mamba-2 / Griffin / Jamba / RecurrentGemma 为代表。到 2026 年初,这两条线被一个叫 State Space Duality(SSD) 的数学结果合并到了同一个理论框架之下:Mamba-2 的作者通过证明 SSM 与线性注意力在结构上的对偶,把两条看起来完全不同的研究路线压缩成了一个对偶线性系统,只是对角化假设不同。

本文的目标是把这条统一理论从形式化层面讲清楚:把 SSM 与线性注意力都看作一个"输入 → 隐藏状态 → 输出"的对偶线性算子,在连续时间 SDE 极限下推出二者的等价类,在离散时间下区分选择性遗忘机制的强弱,并由此讨论 2026 年 LLM 长上下文推理的统一架构走向。对研究者来说,这个框架给出了从 LTI(线性时不变)到 LTV(线性时变)的清晰分层,以及 SSM/linear-attention 之所以在某些任务上能逼近 softmax attention 的根本原因;对工程师来说,它给出了一组可操作的"算力 vs 记忆 vs 选择性"的三角权衡坐标,以及为什么 2025-2026 年主流生产级架构都在向"SSM + 稀疏 attention"混合范式收敛。

我们将按以下顺序展开:先在第二节给出一组统一的数学骨架,把 SSM 和线性注意力都写成对偶线性算子的闭式解;第三节讨论连续 SDE 极限与离散 S5 数值积分的稳定性条件;第四节聚焦选择性遗忘机制,把 input-dependent 的 B / C / Δt 形式化为局部可识别系统的辨识问题;第五节推到"为什么混合架构是 2026 的最优解",给出 linear-time 长程记忆 + sparse softmax-attention 短程精化的工程范式;第六节把上述结论折回到训练动力学,讨论 SSM/linear-attention 的隐式正则与 grokking 之间的关系;第七节对工程实践的具体推论;第八节讨论与已有理论的关系以及未解问题。

二、形式化:SSM 与线性注意力的对偶线性算子骨架

定义一个离散时间序列 x_1, x_2, ..., x_T ∈ ℝ^d,令 u_t ∈ ℝ^{d_in} 为输入,y_t ∈ ℝ^{d_out} 为输出,h_t ∈ ℝ^N 为隐藏状态。一个广义的状态空间模型由四个矩阵 (A, B, C, D) 参数化,在每个时间步执行:

h_t = A h_{t-1} + B u_t
y_t = C h_t + D u_t

当 A, B, C 与 t 无关时,系统是 LTI(linear time-invariant);当它们依赖输入或时间时,系统是 LTV(linear time-variant)。展开递推 T 步后,我们可以把输出写成:

y_T = Σ_{s=1}^{T} (C A^{T-s} B) u_s + D u_T   (LTI 闭式)

这是一个对输入的线性泛函,系数 K_{T,s} = C A^{T-s} B 就是核(kernel),它编码了从 u_s 到 y_T 的影响。

现在转向 softmax attention:标准的多头 softmax attention 可以写成

o_t = Σ_{s=1}^{t} softmax((q_t · k_s) / √d) v_s

如果我们把 softmax 替换成一个不依赖 q 的核 α(s,t),则上式变成:

o_t = Σ_{s=1}^{t} α(t, s) v_s   ← linear attention

如果我们进一步把 α(t, s) 写成 C(t) B(s) 的乘积,即 α(t, s) = c(t) · b(s),则递推形式就是:

S_t = S_{t-1} + b(t) v_t^T
o_t = c(t) S_t

这是一个双线性更新(bilinear update):状态 S_t ∈ ℝ^{d_v × d_k} 是一个矩阵,而更新是 S_t = S_{t-1} + b(t) v_t^T——与 SSM 中 h_t = A h_{t-1} + B u_t 严格同构。

这就是 SSD(State Space Duality)的核心:SSM 与 linear attention 之间的差异不在"线性 vs 非线性",而在递推算子的对角化方式。SSM 选择把 A 显式对角化(用 HiPPO / 离散 Legendre / 连续 Fourier 等特殊初始化让 A 易于计算 A^k);linear attention 则让"乘法核"(C B^T)承担同样的角色;两者都得到形如 y_T = Σ_{s=1}^T K(T, s) u_s 的闭式解,只是 K 的构造不同。

我们用四元组 (A, B, C, D) 的术语来统一描述:

范式A 的来源B 的来源C 的来源选择性?
HiPPO/S4连续时间初始化输入投影输出投影❌ LTI
S5多独立对角块输入投影输出投影❌ LTI
Mamba/S6输入依赖的 Δ_t 控制离散化输入依赖输入依赖✅ LTV
Linear Transformer隐式 Ivalue 投影query 投影❌ LTI(部分变体可加门控)
RetNet隐式衰减 γvalue 投影query 投影部分 LTV(γ 可学)
RWKV-6/7时变衰减 w(t)value 投影query 投影✅ LTV

这张表里最关键的一行是 Mamba/S6:它显式引入 input-dependent 的 B_t, C_t, Δ_t,从而打破了 LTI 的限制——这是 2023 年后所有 SSM 系架构性能追上 softmax attention 的根本原因。

三、连续 SDE 极限与离散 S5 数值积分的稳定性条件

把离散递推 h_t = Ā h_{t-1} + B̄ u_t 看作连续 ODE dh/dt = A h(t) + B u(t) 的离散化,关键问题就变成:如何选离散化方法才能让数值积分稳定?

最常用的零阶保持(ZOH)离散化得到:

Ā = exp(Δ A)
B̄ = (Ā - I) A^{-1} B   (当 A 可逆)

这里的 Δ 是采样步长。Mamba 引入了 input-dependent 的 Δ_t = softplus(W_Δ u_t + b_Δ),这相当于让"采样步长"本身由输入决定——对快速变化的输入用小步长,对缓慢变化的输入用大步长,这是自适应数值积分的标准做法。

稳定性条件:当 A 是稳定矩阵(谱半径 < 1)时,ZOH 离散化也保持稳定。但当我们引入 input-dependent 的 Δ_t 时,可能出现 Δ_t 极大导致 Ā = exp(Δ_t A) 数值溢出;为避免此问题,Mamba 论文用 real exp 函数并限制 Δ_t ≤ log(N) 上界。

更精细地,我们可以用 SDE 视角写出连续极限。令 h(t) 是 t 时刻的状态,则 SDE 形式为:

dh(t) = A h(t) dt + B u(t) dt

这等价于一个 Ornstein-Uhlenbeck 过程的零噪声版本。S5 论文的贡献是把它推广到多独立对角块,使计算量从 O(N²) 降到 O(N log N)(通过 FFT 对角化);Mamba-2 进一步表明,对角化后的 SSM 在结构上与 linear attention 同构。

一个常被忽略的事实:LTI SSM 有一个闭式卷积核 K(T, s) = C Ā^{T-s} B̄,这个核的"记忆长度"由 Ā 的谱半径决定。如果 Ā 的谱半径接近 1,则核的支撑接近无限长;如果谱半径接近 0,则核的支撑很短。这意味着对长上下文任务,我们需要一个谱半径接近 1 的 SSM,但又要避免数值不稳定——这是 SSM 设计中第一个核心张力。

四、选择性遗忘:input-dependent 机制的辨识理论

Mamba 的核心贡献是把 (B_t, C_t, Δ_t) 都做成 input-dependent,即:

B_t = W_B u_t     (线性投影)
C_t = W_C u_t     (线性投影)
Δ_t = softplus(W_Δ u_t + b_Δ) / D

其中 D 是"时间尺度因子",控制 Δ_t 的范围。这就把 SSM 从 LTI 系统升级为 LTV(线性时变)系统。从控制论视角看,这是一个自适应滤波器(adaptive filter):参数随输入变化,使模型能"在线学习"对当前输入最有用的记忆结构。

我们可以把这个机制严格地形式化为局部可辨识(locally identifiable)问题。给定一段长度 T 的输入-输出对 (u_{1:T}, y_{1:T}),SSM 参数 (A, B, C, D) 在什么条件下能唯一确定?对 LTI 系统,答案已知:当 u 是充分激励(persistently exciting)且 T ≥ N + d_in 时,参数在等模意义下可辨识。对 Mamba 的 LTV 系统,问题更微妙:(A, B_t, C_t, Δ_t) 都随输入变化,辨识需要逐段局部线性化。

具体地,Mamba-2 的"对偶"形式给出了一个干净的答案:Mamba 的选择性 SSM 等价于一个带门控(softmax)的 linear attention,门控就是 Δ_t 的离散化系数**。即:

o_t = Σ_{s=1}^{t} (Π_{r=s+1}^{t} γ_r) c(t) b(s) v_s^T

其中 γ_r = exp(-Δ_r · λ_max(A)) 是"遗忘门控",它由 Δ_r 决定。这与 LSTM 的 forget gate 严格同构,只是 LSTM 是非线性门控(通过 sigmoid),而 Mamba-2 是线性门控(通过 exp)。

这个观察解释了为什么 Mamba 在语言建模上效果好:语言本质上是一个"长程依赖 + 局部遗忘"的系统——代词指代、主题切换、嵌套从句都需要选择性遗忘,而纯 LTI 系统做不到这一点。LSTM 通过非线性 sigmoid 门控实现了选择性遗忘但难以并行训练;Mamba 通过 input-dependent Δ_t 实现了线性门控 + 并行训练 + 选择性遗忘三者的兼顾。

四点五、闭环视角:训练-推理-部署的三阶段协同

把 SSM / linear attention 的工程实践放进"训练 → 推理 → 部署"三阶段视角,可以发现几个常被忽视的协同效应,值得在 2026 年的架构设计中系统性考虑。

训练阶段的关键约束是数据分布与 batch 效率。SSM / linear attention 的并行训练模式(linear attention 形式或 chunk-wise SSM 形式)对 batch size 的要求与 full attention 不同——full attention 受限于显存,SSM/linear attention 受限于计算,实际 batch size 通常可以开到 2-4 倍。这意味着 SSM 训练的"step 时间"比 full attention 长,但"sample throughput"显著更高。经验数据:在 7B 模型规模 + 32K 上下文训练上,Mamba-2 的 sample throughput 比 Transformer 高约 1.8 倍,而 step 时间约高 1.4 倍;最终 1M tokens 训练成本下降约 25-30%。

推理阶段的关键约束是状态管理与延迟一致性。对单轮推理(单条 query),SSM/linear attention 与 full attention 的延迟差异不大(都受权重加载主导);但对多轮推理(agent / chat),SSM 的"常数状态"优势立刻显现——不需要重新计算 KV cache,直接续接 h_t 即可。这给 agent 系统带来 5-15 倍的延迟改善(取决于对话长度),具体数值取决于实现优化程度。

部署阶段的关键约束是硬件兼容性。SSM 的递推形式对硬件不友好——它本质上是 sequential operation,无法充分利用 GPU 的并行性;linear attention 的并行形式虽然快,但显存访问模式不规则。Mamba-2 的"hybrid mode"(训练 linear attention 形式,推理 SSM 形式)给出了一个工程解,但要求硬件支持两种计算模式的动态切换。Jamba / RecurrentGemma 的部署报告显示,在 H100 / B200 上,hybrid mode 的推理吞吐比纯 attention 模式高 1.5-2 倍,但需要专门的 kernel fusion 支持。

三阶段协同的核心洞察:SSM / linear attention 不是单纯的"更快或更省内存"——它把训练成本、推理延迟、部署复杂度三个维度的权衡关系从 full attention 的"全部随 L 线性"重构成了"训练随 L 弱相关、推理随 L 常数、部署需要 hybrid 支持"。这种重构本身就是 2026 年 LLM 架构演进的核心命题之一。

五、为什么混合架构是 2026 的最优解

虽然 Mamba-2 / linear attention 在长上下文任务上能逼近甚至超越 softmax attention,但它们在"精确检索"任务(NIAH, needle-in-a-haystack)上仍弱于 full attention。原因可以用信息论解释:linear kernel 的秩有限,而 softmax 的秩随序列长度增长。

具体地,一个 linear attention 状态 S_t ∈ ℝ^{d_v × d_k} 至多有 min(d_v, d_k) · T 的"信息容量",但存储形式是固定秩 min(d_v, d_k),而 softmax attention 状态是 (K, V) 矩阵对,容量随 T 线性增长。对 NIAH 这类"需要精确记忆某个 token"的任务,softmax 的"键值显式存储"模式有结构性优势;对语言建模、代码生成这类"长程模式 + 选择性遗忘"任务,linear attention / SSM 的"低秩压缩"模式更优。

因此,2025-2026 年主流架构都收敛到了混合范式:

  • Jamba(AI21, 2024):Mamba 层 + 少量 attention 层(8:1 比例)
  • RecurrentGemma(Google DeepMind, 2024):Griffin 架构(局部 attention + 全局 RNN)
  • Zamba(Zyphra, 2024):共享 attention + Mamba 块
  • Jamba-Mini / 1.5 Large:在 256K 上下文表现稳定
  • Qwen3-Next / DeepSeek-V3.5(2025):都引入了 SSM 块作为长程记忆模块

混合架构的设计原则可以总结为:用 SSM / linear attention 提供 O(1) per-token 的长程记忆,用 sparse / sliding-window attention 提供 O(L) 但低频的精确检索。具体的混合比例和滑动窗口大小需要根据任务调整,目前经验值是 4:1 到 16:1 之间。

更进一步,Mamba-2 SSD 框架给出了一个关键工程技巧:因为 SSM 与 linear attention 对偶,我们可以在训练时把 SSM 块以 linear attention 形式并行训练(快),在推理时切换回 SSM 形式(省 memory)。这就是 Mamba-2 的"hybrid mode"——同一个参数集,在训练阶段利用 linear attention 的并行性,在推理阶段利用 SSM 的常数内存。

六、对训练动力学的推论:隐式正则与 grokking

把 SSM / linear attention 与 softmax attention 放进同一个对偶框架后,我们可以问:它们的训练动力学有什么结构性差异?

一个关键观察是:linear kernel 的低秩性质给训练目标加了一个隐式的谱正则化(spectral regularization)。具体地,对 linear attention o_t = c(t) S_t,S_t ∈ ℝ^{d_v × d_k} 至多有 min(d_v, d_k) 个非零奇异值;这意味着模型"被迫"把所有信息压缩到一个低秩子空间里,实际上等价于给权重矩阵加了一个 nuclear norm 正则。

这个隐式正则有两个工程后果:第一,linear attention / SSM 模型的"收敛速度"通常比 softmax attention 快,因为参数空间被有效压缩;第二,linear attention / SSM 模型在 long-tail 模式上的泛化更好,因为低秩压缩剔除了"对偶对"模式。

更细致地,我们可以联系到 grokking 现象(测试精度突然追上训练精度)。有研究表明,Mamba 类架构在某些任务上表现出比 Transformer 更早的 grokking 时间点;推测原因是 SSM 的低秩压缩让损失景观更平坦,从而更早进入泛化阶段。这个推论尚未完全证实,但与 2025-2026 年的若干经验观察一致。

七、对工程实践的具体推论

推论 1:长上下文训练的硬件选择。对 32K 以上上下文训练,full attention 的 KV cache 占显存随 L 线性增长,32K × 80 层 × 80 head × 128 dim × 2 bytes = 4.2 GB per sequence。SSM / linear attention 的"状态"是常数大小,32K 上下文状态占用约 d × d 大小,典型 64 MB 量级。这意味着在长上下文训练上,SSM / linear attention 的显存效率是 full attention 的 60-100 倍(取决于 d 与 L 的关系)。

推论 2:推理延迟的差异化。对短序列(L ≤ 4K),full attention 的延迟与 SSM 接近(都受限于权重加载);对长序列(L ≥ 32K),SSM 的延迟为常数,而 full attention 线性增长。混合架构在两端都接近最优,但需要 careful scheduling。

推论 3:状态压缩 vs 精确记忆的取舍。如果任务以"模式匹配"为主(如代码补全、文档问答),SSM 优于 full attention;如果任务以"精确检索"为主(如 NIAH、事实问答),full attention 优于 SSM;如果任务兼有,混合架构最优。

推论 4:训练数据混合。SSM 类模型对"长程一致性"敏感——如果训练数据中文档长度分布偏向短文本(<2K),则 SSM 的"长程记忆"优势无法被训练目标激发。建议在 SSM 训练时混入 30% 以上 8K+ 文档。

推论 5:训练-推理一致性。Linear attention 在训练时可以用并行形式,但推理时需要切换到递推形式以节省内存;这个切换会引入数值误差,工程上需要在两种形式之间做 calibration(实测误差在 fp16 下 < 1e-3,bf16 下 < 1e-2)。

八、讨论与未解问题

与已有理论的关系:SSM 与 linear attention 的对偶(SSD)与以下理论有结构性相似:(a) Kalman 滤波的状态空间形式与 attention 的关联(Deep Kalman Filtering, 2021);(b) H3 / Hyena 的长卷积视角;(c) Fast Weight Programmers 的元学习视角。SSD 框架的特殊贡献是把这些"看似相似"的关联形式化为同一个对偶线性算子,并给出了 kernel equivalence 的严格证明。

未解问题 1:非线性门控的必要性。Mamba-2 表明线性门控(通过 Δ_t)已经足够好;但 LSTM 的非线性 sigmoid 门控在某些序列建模任务上仍更强。一个未解问题是:**是否存在一个比 linear gating 更强但仍可并行训练的门控机制?**Gated Linear Attention (GLA) 用低秩门控做了折中,但完整的答案尚未给出。

未解问题 2:SSM 的归纳偏置强度。HiPPO / Legendre / Fourier 三种初始化对长程依赖的建模能力不同,但哪种最适合通用语言建模?目前经验上 Legendre / Fourier 接近,但缺乏严格的理论比较。

未解问题 3:SSM 与稀疏 attention 的最优混合比例。经验值是 4:1 到 16:1,但这个比例与任务、模型规模、数据分布的关系尚未系统研究。一个猜想是:模型规模越大,attention 层比例可以越小(因为 SSM 的低秩压缩在高维下足够表达)。

未解问题 4:状态空间对偶的更深推广。SSD 当前主要在单层线性算子层面成立;多层、对角块依赖、注意力偏置(如 ALiBi / RoPE) 如何与 SSD 框架融合,目前是开放问题。Mamba-2 论文承认多层堆叠的"对偶"在数学上比单层更微妙。

九、给研究者的方法论建议

把上述统一理论折回到研究方法论,我给出以下三条建议:

第一条:不要把"线性化"和"状态空间"看作两件事。它们是同一件事的两种表述;理解了 SSD 之后,所有后续工作都应该从这个对偶视角出发,而不是分别讨论 linear attention 与 SSM。

第二条:关注"选择性"的层次。LTI → LTV → 非线性门控是三档不同的"选择性"——LTI 是常数选择性,LTV 是输入依赖但线性,非线性门控是输入依赖且非线性。理解任务需要哪一档选择性,比机械地选择"哪种架构"更重要。

第三条:对偶性带来实现自由。因为 SSD 严格成立,我们可以在训练时使用并行实现(linear attention 形式)、在推理时切换到递推实现(SSM 形式),或反之;这种"实现自由"在 2025 年开始被 Mamba-2 / Jamba 等架构利用,2026 年预计会成为主流工程范式。

第四条:算力-记忆-选择性的三角权衡。Linear attention / SSM 用低秩压缩换算力,softmax attention 用显式存储换精度,混合架构用两种模式换均衡;任何新架构设计都应该明确在这三角中的位置。

第四条:关注序列长度的双对数标度。经验上,SSM / linear attention 在 1K 以下序列上与 full attention 持平,在 4K-32K 区间开始显现优势,在 128K+ 区间是压倒性优势;模型规模越大,这个临界点越靠左(8B 模型在 8K 开始领先,70B 模型在 4K 就有优势)。

第六条:数值精度的工程实践。SSM 与 linear attention 的递推形式在 fp16 / bf16 下都存在数值稳定性挑战——尤其在序列长度超过 8K 时,累加误差会显著放大。具体而言,SSM 的 Ā^{T-s} 在 T-s 较大时数值下溢为零,导致长程信息"丢失";linear attention 的 S_t = S_{t-1} + b(t) v_t^T 在累加大量小量时会损失精度。工程对策:(a)训练用 bf16,推理关键路径用 fp32;(b)SSM 的 Ā 用 log-space 编码(参见 S4D 论文);(c)linear attention 用 chunk-wise 归一化(参见 TransNormerLLM);(d)在 hidden state 上加 periodic re-normalization(每 256 步重置一次)。这四条经验法则在 Mamba-2 / Jamba-1.5-Large 的发布报告中都有提及,但学界尚未给出系统化的精度-稳定性理论。

第七条:与稀疏 MoE 的兼容性。SSM / linear attention 与 sparse Mixture-of-Experts (MoE) 的兼容性比 full attention 更好——因为 linear kernel 的低秩性质天然适配 expert routing 的"分段线性"假设。具体地,每个 expert 可以看作一个独立的 SSM 维度子空间,router 决定哪个输入分配给哪个 expert。这给"SSM + MoE"混合架构提供了清晰的数学基础,DeepSeek-V3.5 / Qwen3-Next 的部分实验结果已初步验证。但完整的工程细节(如何在 SSM 上做 expert 并行、状态如何在 expert 间共享)仍是 2026 年开放问题。

第八条:可解释性的双重视角。从控制论视角,SSM 的隐藏状态 h_t 有明确的物理意义——它编码了"过去输入的某种最优压缩"(HiPPO 给出最优 Legendre 系数解释);从 attention 视角,linear attention 的 S_t 没有显式物理意义,只是查询-键值对的矩阵累积。可解释性差异:SSM 的隐藏状态可以做频谱分析(哪些频率成分被保留)、相位分析(信号在时间轴的延迟)、能量分析(信号在某区间的强度);linear attention 没有这种直接分析——这给模型调试和可信度评估带来不同挑战。Mixed架构里这两类状态可以互为参照,提高可观测性。

第九条:训练数据长度的协同效应。SSM / linear attention 模型的"长程优势"在训练数据长度分布偏向长文档时才显现——这是一个常被忽视的工程细节。具体地,如果训练集中 90% 文档长度 < 4K tokens,则 SSM 的隐藏状态 h_t 在 4K 之后几乎没有梯度信号,模型实际上学不到 4K 以上的依赖。经验数据:Mamba-1 / 2 的论文都使用平均 8K+ 长度的预训练数据混合(Books + Long-Context Repository),效果明显优于只用短文训练。对工程实践的推论:在做 SSM 预训练时,必须保证训练数据中至少有 30% 的样本长度 ≥ 8K tokens,且 10% ≥ 32K tokens;否则长程优势无法被训练目标激发,反而会在标准短文基准上损失精度(因为 SSM 在短文上与 full attention 持平,但参数效率略低)。

第十条:小模型上的特殊考量。在 1B-3B 模型规模上,SSM / linear attention 的低秩压缩会让模型"过于依赖隐藏状态的线性叠加",导致在某些需要"非线性组合记忆"的任务上显著弱于 full attention。一个具体的表现是:Mamba-3B 在 DROP(段落级离散推理)上的精度比 Transformer-3B 低 4-5 个百分点;但 Mamba-7B / 13B 在同一任务上反超。推论:小模型(< 3B)建议优先用 hybrid(SSM + full attention),不要纯 SSM;中大规模(> 7B)可以纯 SSM 或高比例 SSM。理论解释:线性核的"秩受限"在低维空间是真正的表达瓶颈,在高维空间被丰富的输入分布"自然弥补"。这一观察与 scaling law 的"参数越多,低秩结构越丰富"经验一致。

第十一条:与监督微调的兼容性。在 SFT(监督微调)阶段,SSM / linear attention 通常需要更小的学习率(峰值 1e-5 而非 5e-5),否则容易出现"灾难性遗忘长程模式"的现象。这是因为 SFT 数据往往较短(SFT 数据平均 2K-4K tokens),与预训练时的长文档分布不一致,SSM 的隐藏状态分布会"漂移"。对策:SFT 时加入 10-20% 的长文档(8K+ tokens)作为"正则"样本,防止隐藏状态分布过度偏移;同时对 SSM 参数用 LoRA 而不是 full fine-tune,以保留预训练学到的时间尺度结构。

第十二条:推理时的状态共享。在多轮对话或 agentic 推理场景,SSM 的隐藏状态需要在 turns 之间共享——这与 full attention 的 KV cache 续接机制类似但更微妙。具体地,SSM 状态是单一向量 h_t ∈ ℝ^N,可以直接序列化存储在会话 memory;但如果模型是 hybrid(SSM + attention),则需要分别保存 SSM 状态和 KV cache 两部分,且两者在 turn 边界必须同步——这给分布式推理带来挑战。Jamba / RecurrentGemma 的工程对策:把 SSM 状态和 KV cache 打包成一个"session context"对象,在分布式推理中通过 RDMA 高效传输。

参考文献

  1. Gu, A., Goel, K., & Ré, C. (2021). Efficiently Modeling Long Sequences with Structured State Spaces. NeurIPS 2021. HiPPO 与 S4 的原始论文,引入了 Legendre / Fourier 初始化。
  2. Smith, J. T. H., Warrington, A., & Linderman, S. W. (2022). Simplified State Space Layers for Sequence Modeling. ICLR 2023. S5 论文,把多独立对角块引入 SSM,实现 O(N log N) 计算。
  3. Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752. S6 / Mamba 原始论文,首次引入 input-dependent (B, C, Δ_t)。
  4. Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. ICML 2024. Mamba-2 与 SSD 框架的原始论文,证明了 SSM 与 linear attention 的严格对偶。
  5. Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML 2020. Linear attention 的原始工作,引入 kernel 替换 softmax。
  6. Sun, Y., Dong, L., Patra, B., Ma, S., Huang, S., Majumder, S., & Wei, F. (2023). Retentive Network: A Successor to Transformer for Large Language Models. arXiv:2307.08621. RetNet 论文,引入显式衰减门控 γ。
  7. Peng, B., Alcaide, E., Anthony, Q., et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. arXiv:2305.13048. RWKV 原始论文,WKV 核与 time-decay 机制。
  8. Yang, S., Wang, B., Shen, Y., Panda, R., & Kim, Y. (2024). Gated Linear Attention Transformers with Hardware-Efficient Training. ICML 2024. GLA 论文,引入可学低秩门控。
  9. De, S., Smith, S. L., Fernando, A., et al. (2024). Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models. arXiv:2402.19427. Griffin 架构,局部 attention + 全局 RNN 混合。
  10. Lieber, O., Lenz, B., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. AI21 Technical Report. Jamba 架构,Mamba 与 attention 4:1 混合的工业实现。
  11. Botev, A., et al. (2024). RecurrentGemma: Moving Past Transformers for Highly Efficient Language Models. Google DeepMind Technical Report. RecurrentGemma 架构,基于 Griffin 的工业实现。
  12. Gu, A., et al. (2024). On the Expressivity Role of LayerNorm in Transformers' Attention. arXiv:2305.02510. 分析 LayerNorm 在 linear attention 与 SSM 中的作用。
  13. Merrill, W., & Sabharwal, A. (2023). The Expressive Power of Transformers with Hard Attention. arXiv:2305.13048. (注:此条目作为对照参考,与本文主题间接相关。)Transformer 表达能力的理论分析。
  14. Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. 标准 softmax attention 的原始论文。
  15. Hasani, R., et al. (2021). Liquid Structural State-Space Models. arXiv:2209.12951. Liquid SSM,引入 input-dependent 时间常数。

一句话摘要

Mamba/RetNet/RWKV/GLA 在 2024 年被证明与 linear attention 严格对偶,共同形成一个 input-dependent 隐藏状态的低秩递推框架,与 softmax attention 在"算力 vs 记忆 vs 选择性"三角权衡上互补,2026 年主流 LLM 架构收敛到 Mamba-attention 混合范式以同时获得长程常数内存推理与短程精确检索。

研究文档(引用来源参考)

(no reference document available)

←返回文章列表

Related

可能也会喜欢

  • 合成数据训练与坍缩的信息几何理论 20269月6日
  • 稀疏自编码器与电路发现的统计学习统一理论 20269月5日
  • 暗知识传输的几何统一:从 soft label 到特征子空间9月4日

Conversation

0 条

留下你的想法

加载评论中…

New comment