Mamba 训练动力学的理论分析:为状态空间模型中的特征筛选与泛化提供理论依据
机器学习
2026-02-16 v1
摘要
Mamba 等选择性状态空间模型(SSMs)的近期实证成功, 重燃了对非注意力序列建模架构的兴趣, 但其理论基础仍未充分探讨. 我们对 Mamba 块的训练动力学与泛化进行首次分析, 聚焦于代表性的简化模型: 单层、单头选择性 SSM, 其输入依赖门控, 随后接两个层的 MLP, 通过梯度下降(GD)训练. 本研究采用具备 class-relevant 与 class-irrelevant 模式的结构化数据模型, 这些模式在 token 级噪声下呈现, 并考察两种典型情形: majority-voting 与 locality-structured 数据序列. 我们证明该模型通过建立非渐近样本复杂度与收敛速率界, 实现保证泛化, 其中有效信号的增加与噪声的降低会改善这些界限. 此外, 我们展示了门控向量与 class-relevant 特征对齐, 同时忽略不相关特征, 从而形式化了类似注意力机制但通过选择性递归实现的特征选择作用. 数值实验支持了我们的理论结果. 总体而言, 我们的结果为何时为何 Mamba 式选择性 SSM 能高效学习提供原则性见解, 为 Transformer 主导的解释提供了理论反证.
引用
@article{arxiv.2602.12499,
title = {A Theoretical Analysis of Mamba's Training Dynamics: Filtering Relevant Features for Generalization in State Space Models},
author = {Mugunthan Shandirasegaran and Hongkang Li and Songyang Zhang and Meng Wang and Shuai Zhang},
journal= {arXiv preprint arXiv:2602.12499},
year = {2026}
}