Mamba 表达力理论

引言

选择性状态空间模型（Selective State Space Models, S6/Mamba）作为 Transformer 的有力竞争者，在语言建模和视觉任务中展现出卓越的性能。然而，其理论根基——即这类模型”能够表达什么”以及”如何高效学习”——长期缺乏系统的形式化分析。近年来，三项关键研究从不同维度填补了这一空白：¹²³

1. SSM 的表达力界限

1.1 线性与非线性变换的表达力鸿沟

传统线性状态空间模型（Linear SSM）本质上是线性时不变系统，其核心变换可表示为：

$$h_{t+1}=\mathbf{A}{h}_t+\mathbf{B}{x}_t,y_t=\mathbf{C}{h}_t+\mathbf{D}{x}_t$$

其中参数 $\mathbf{A},\mathbf{B},\mathbf{C},\mathbf{D}$ 是与输入无关的常数矩阵。这种线性特性带来两个根本限制：

1. 表达能力上限：线性 SSM 的输入-输出映射可表示为一个低秩矩阵，无法捕捉高阶交互特征
2. 与多项式的关系：线性 SSM 等价于双线性多项式的特殊形式，其函数类包含于 $O(n)$ 次多项式中（$n$ 为序列长度）

非线性的引入来自两个途径：

• 选择性机制（Selection Mechanism）：参数变为输入依赖的 $\mathbf{B}(x_t),\mathbf{C}(x_t)$
• 门控机制（Gating）：通过激活函数（如 SiLU）引入非线性

1.2 电路复杂性视角

从电路复杂性（Circuit Complexity）角度分析，poly$(n)$ 精度和常数深度的 Mamba 层属于 $\mathsf{DLOGTIME}$-uniform ${\mathsf{TC}}^0$ 复杂度类。⁴ 这意味着：

模型类型	复杂度类	表达能力
常数深度电路	${\mathsf{TC}}^0$	阈值计算，可处理多数函数
对数深度电路	${\mathsf{NC}}^1$	更复杂的树结构计算
多项式深度	$\mathsf{P}$	图灵完备

值得注意的是，Transformer 和 Mamba 在此理论框架下具有相同的计算能力上界，但实际表现差异来自常数因子和架构偏置。

2. 选择性机制的核心作用

2.1 输入依赖参数的理论优势

Mamba 的核心创新在于将 SSM 参数参数化为输入的函数：

$${\mathbf{B}}_t=\mathbf{B}(x_t),{\mathbf{C}}_t=\mathbf{C}(x_t),{\Delta }_t=\Delta (x_t)$$

这种设计带来了质变：

1. 时间变异性：系统从时不变变为时变，允许捕捉非平稳依赖
2. 选择性过滤：通过 ${\Delta }_t$ 控制信息保留/遗忘，实现类似注意力机制的软选择

理论上可证明：输入依赖的时间变异性提升了层的表达能力——从线性映射跃升至能够表示任意有限多项式函数的函数类。¹

2.2 动力学性质分析

Vo 等人 (2024) 对深度选择性 SSM 的 token 动力学进行了系统分析。²

在连续时间极限下，Mamba 的隐状态演化可建模为：

$$\frac{d\mathbf{h}(t)}{dt}=f(\mathbf{h}(t),x_t;\theta )$$

一维情形的主要结论：对于预训练的 Mamba 模型，token 演化只有两种渐近行为：

• 收敛场景：所有 token 收敛到零
• 发散场景：所有 token 发散到无穷大

具体发生哪种场景由模型参数决定。进一步分析表明：

收敛场景对性能有负面影响：所有 token 趋于相同的零状态，模型丧失区分能力。

发散场景下：不同 token 以不同速率发散到无穷大，这种不均匀性是关键——允许模型通过调整发散速率来编码 token 的重要性差异。

基于此，研究者提出两项改进：

1. 排除收敛场景的参数初始化策略
2. 基于重要性分数的 token 重排序

3. 1-WL 图灵测试与 GNN 表达力对比

3.1 WL 测试与图神经网络的联系

Weisfeiler-Lehman (1-WL) 同构测试是衡量图神经网络（GNN）表达力的标准基准。1-WL 通过迭代颜色细化判断两个图是否非同构，其表达能力上限等价于 UF（并查集）/ MPR（Mutually Recursive Powerset）协议。

关键联系在于：

• 注意力机制的 soft-max 聚合可看作 1-WL 的连续松弛
• 选择性 SSM 中的 ${\Delta }_t$ 调节可视为一种隐式消息传递

3.2 Mamba 与 Transformer 表达力的差异

尽管两者都属于 ${\mathsf{TC}}^0$，但实际表达力存在以下差异：

维度	Transformer	Mamba/S6
注意力矩阵	$O(n^2)$ 显式存储	$O(n)$ 隐式表示
依赖建模	全局注意力	选择性线性投影
非线性来源	softmax + FFN	选择性机制 + SiLU
序列顺序	位置编码	隐式 SSM 顺序

核心洞见：选择性机制使 S6 层等价于多元多项式的高阶项展开，而非仅线性组合。¹

4. 多项式表达力分析

4.1 多元多项式框架

Cohen-Karlik 等人 (2025) 提出了分析选择性 SSM 表达力的系统框架。¹

核心定理（简化表述）：

选择性 SSM 层（S6）能够表示任何次数为 $k$ 的多元多项式函数，其中 $k$ 与模型隐藏维度相关。

这一定理的证明思路：

1. 将 S6 层分解为多项式时间不变操作的组合
2. 展示选择性机制如何”打破”线性不变性
3. 通过构造性证明给出多项式到 SSM 参数的映射

4.2 与线性注意力的比较

关键结论：选择性 SSM 在表达力上严格超越线性注意力机制。

$$\text{Linear Attention}\subset \text{Selective SSM}\subset \text{Polynomial Functions}$$

• 线性注意力：等价于单次多项式（双线性形式）
• 选择性 SSM：可表示任意多项式，允许高阶特征交互
• 计算优势：在长序列上保持 $O(n)$ 复杂度

5. 泛化与学习动态理论

5.1 训练动态的形式化分析

Shandirasegaran 等人 (2026) 首次对 Mamba 的训练动态进行了严格的理论分析。³

研究设定：

• 简化的 Mamba 块：单层、单头选择性 SSM + 输入依赖门控 + 两层 MLP
• 训练方式：梯度下降（GD）
• 数据模型：包含类别相关和类别无关模式的 token 序列

5.2 主要理论结果

样本复杂度界限：

模型以高概率达到 $\epsilon$ 泛化误差所需的样本数为：

$$m=O\left(\frac{d}{{\lambda }^2{\epsilon }^2}\right)$$

其中：

• $d$：模型有效参数量
• $\lambda$：有效信号强度
• $\epsilon$：目标误差

收敛速率：梯度下降的收敛速率随信号强度 $\lambda$ 增加而改善。

5.3 门控的特征过滤机制

最重要的发现是门控向量的特征选择作用：

定理：在训练过程中，门的激活值与类别相关特征对齐，同时忽略不相关特征。

这与注意力机制的功能高度相似，但通过选择性循环而非显式注意力实现：

• 注意力机制：通过 soft-max 权重对特征加权求和
• 选择性 SSM：通过输入依赖的 $\Delta$ 和门控实现隐式筛选

这一发现为 Mamba 的高效学习提供了理论解释：模型能够自动聚焦于相关信息，过滤噪声。

6. 与 Transformer 的表达力对比

6.1 表达能力比较

特性	Transformer	Mamba
函数类	多项式（softmax 非线性）	多项式（选择性机制）
参数依赖	全注意力 $n\times n$	线性投影 $O(n)$
非线性类型	soft-max + MLP	SiLU + 选择性门控
长期依赖	显式注意力	隐状态压缩
计算复杂度	$O(n^2)$	$O(n)$

6.2 各有所长

Transformer 的优势：

• 完美的全局信息整合
• 可解释性强（注意力权重可视化）
• 更成熟的训练技术

Mamba 的优势：

• 长序列处理的计算效率
• 隐状态压缩带来的归纳偏置
• 特定任务上的长度外推能力

7. 经验验证

7.1 多项式表达力实验

Cohen-Karlik 等人在多项式逼近任务上验证了理论：

• 合成任务：学习 $k$ 次多项式函数
• 结果：S6 层在 $k\ge 2$ 时显著优于线性注意力，与理论预测一致

7.2 动力学分析实验

Vo 等人通过以下实验验证动力学理论：

1. 收敛/发散场景检测：在大规模语言模型中识别收敛行为
2. Token 重排序：基于重要性分数重排序可提升困惑度
3. 消融实验：移除收敛场景后模型性能提升

7.3 泛化理论实验

Shandirasegaran 等人在受控视觉任务上验证：

• 长度外推：训练序列长度与测试序列长度的泛化关系
• 特征对齐：门控激活与类别标签的相关性分析
• 噪声敏感性：门控对不同噪声水平的鲁棒性

8. 总结与开放问题

已确立的理论

1. 表达力上界：选择性 SSM 属于 ${\mathsf{TC}}^0$，与 Transformer 相同
2. 多项式表达：S6 层可表示任意多元多项式，严格超越线性注意力
3. 门控机制：实现类似注意力的特征选择，通过循环实现 $O(n)$ 复杂度
4. 学习动态：门控自动对齐相关特征，提供可证明的泛化保证

未解决的开放问题

问题	现状	挑战
深度堆叠的表达力	单层理论成熟	多层组合的复杂性
长度外推机制	经验观察	缺乏形式化理论
与其他架构的等价性	部分结果	完整分类困难
优化 landscape	初步分析	非凸优化困难

未来方向

1. 超越 ${\mathsf{TC}}^0$ 的架构设计：探索具有更强计算能力的 SSM 变体
2. 自适应计算时间：理论分析可变深度 SSM 的表达力
3. 与其他模态的融合：将理论扩展到视觉、语音等非文本领域

参考文献

Footnotes

1. Cohen-Karlik, E., Zimerman, I., Galanti, L., Atad, I., Globerson, A., & Wolf, L. (2025). On the Expressivity of Selective State-Space Layers: A Multivariate Polynomial Approach. arXiv:2502.02209. ↩ ↩² ↩³ ↩⁴

2. Vo, T., et al. (2024). Demystifying the Token Dynamics of Deep Selective State Space Models. arXiv:2410.03292. ↩ ↩²

3. Shandirasegaran, M., et al. (2026). A Theoretical Analysis of Mamba’s Training Dynamics: Filtering Relevant Features for Generalization in State Space Models. arXiv:2602.12499. ↩ ↩²

4. The Computational Limits of State-Space Models and Mamba via the Lens of Circuit Complexity. arXiv:2412.06148. ↩