揭示Mamba如何通过选择性循环筛选关键特征,实现高效泛化。
A Theoretical Analysis of Mamba's Training Dynamics: Filtering Relevant Features for Generalization in State Space Models
- 构建简化Mamba块模型,分析其学习动态与泛化能力。
- 证明在信号强、噪声低时,模型可快速收敛并保证泛化性能。
- 发现门控向量自动聚焦于相关特征,类似注意力但由选择性递归实现。
Mamba等选择性状态空间模型(SSMs)的近期成功重新激发了对非注意力架构序列建模的兴趣,但其理论基础仍不清晰。本文首次对一个简化但具代表性的Mamba模块进行分析:单层、单头输入依赖门控的选择性SSM,后接两层MLP,通过梯度下降(GD)训练。研究采用结构化数据模型,包含类相关和类无关模式,并在令牌级噪声下考察两类典型场景:多数投票与局部结构数据序列。我们证明模型能保证泛化,建立了非渐近的样本复杂度与收敛率边界,且随着有效信号增强、噪声减小而提升。进一步表明,门控向量会与类相关特征对齐,忽略无关特征,形式化了类似注意力的特征选择机制,但通过选择性递归实现。合成数据上的数值实验验证了理论结果。整体上,本研究为Mamba式选择性SSMs为何及何时高效学习提供了原理性洞见,构成对以Transformer为中心解释的理论补充。
原文摘要 · Abstract (English)
The recent empirical success of Mamba and other selective state space models (SSMs) has renewed interest in non-attention architectures for sequence modeling, yet their theoretical foundations remain underexplored. We present a first-step analysis of generalization and learning dynamics for a simplified but representative Mamba block: a single-layer, single-head selective SSM with input-dependent gating, followed by a two-layer MLP trained via gradient descent (GD). Our study adopts a structured data model with tokens that include both class-relevant and class-irrelevant patterns under token-level noise and examines two canonical regimes: majority-voting and locality-structured data sequences. We prove that the model achieves guaranteed generalization by establishing non-asymptotic sample complexity and convergence rate bounds, which improve as the effective signal increases and the noise decreases. Furthermore, we show that the gating vector aligns with class-relevant features while ignoring irrelevant ones, thereby formalizing a feature-selection role similar to attention but realized through selective recurrence. Numerical experiments on synthetic data justify our theoretical results. Overall, our results provide principled insight into when and why Mamba-style selective SSMs learn efficiently, offering a theoretical counterpoint to Transformer-centric explanations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。