arXiv:2509.19633cs.LGcs.AI2025-09NeurIPS被引 8

解决Mamba模型长文本失效问题,通过调整矩阵谱实现稳定长序列推理。

Mamba Modulation: On the Length Generalization of Mamba

  • 通过分析状态转移矩阵谱特性,揭示Mamba长文本敏感根源
  • 对预训练Mamba的矩阵谱进行缩放,显著提升长上下文性能
  • 适合研究长序列建模与状态空间模型优化的开发者

Transformer中注意力机制的二次复杂度催生了具有次二次扩展性的替代架构,如状态空间模型。其中,Mamba已成为领先架构,在多项语言建模任务中达到顶尖水平。然而,当输入长度超过预训练时的范围时,Mamba性能急剧下降,表现出对上下文长度延伸的强烈敏感性。我们通过深入分析发现,这一限制源于状态空间动态在分布外情况下的异常行为,特别是状态转移矩阵\mathbf{A}的参数化方式。不同于将此归因于离散时间步累积衰减项\exp(-\sum_{t=1}^NΔ_t)消失的观点,我们建立状态在输入长度趋于无穷时的收敛行为与矩阵\mathbf{A}谱之间的联系,为该作用提供了理论基础。为此,我们提出一种方法:对预训练的Mamba模型应用谱缩放,通过有选择地调节每层\mathbf{A}矩阵的谱,实现鲁棒的长上下文泛化。实验表明,该方法在单纯调节Δ_t失败的场景下仍能显著提升性能,验证了我们的见解,并为具有结构化转移矩阵的状态空间模型提供了更优的长度泛化路径。

原文摘要 · Abstract (English)

The quadratic complexity of the attention mechanism in Transformer models has motivated the development of alternative architectures with sub-quadratic scaling, such as state-space models. Among these, Mamba has emerged as a leading architecture, achieving state-of-the-art results across a range of language modeling tasks. However, Mamba's performance significantly deteriorates when applied to contexts longer than those seen during pre-training, revealing a sharp sensitivity to context length extension. Through detailed analysis, we attribute this limitation to the out-of-distribution behaviour of its state-space dynamics, particularly within the parameterization of the state transition matrix $\mathbf{A}$. Unlike recent works which attribute this sensitivity to the vanished accumulation of discretization time steps, $\exp(-\sum_{t=1}^NΔ_t)$, we establish a connection between state convergence behavior as the input length approaches infinity and the spectrum of the transition matrix $\mathbf{A}$, offering a well-founded explanation of its role in length extension. Next, to overcome this challenge, we propose an approach that applies spectrum scaling to pre-trained Mamba models to enable robust long-context generalization by selectively modulating the spectrum of $\mathbf{A}$ matrices in each layer. We show that this can significantly improve performance in settings where simply modulating $Δ_t$ fails, validating our insights and providing avenues for better length generalization of state-space models with structured transition matrices.

Mamba长序列状态空间谱调控

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。