arXiv:2410.03292cs.LG2024-10ICLR被引 1

揭示Mamba模型中令牌动态机制,提出改进方案提升性能。

Demystifying the Token Dynamics of Deep Selective State Space Models

  • 构建Mamba连续时间动力学模型,分析令牌收敛或发散行为。
  • 发现收敛场景导致性能下降,发散场景中令牌更新不均等。
  • 提出排除收敛态与按重要性重排序令牌的优化方法。

选择性状态空间模型(SSM),如Mamba,因其在序列建模中的优异表现而备受关注。尽管其经验性能突出,但对深层选择性SSM的理论理解仍不充分,限制了其在高保真应用中的发展与采用。本文研究预训练Mamba模型中令牌的动力学特性,推导出模型连续时间极限下的动力系统,并刻画其解的渐近行为。在一维情况下,证明仅存在两种可能:所有令牌收敛至零,或全部发散至无穷。基于模型参数提供判断条件。实验验证收敛情形会负面影响模型性能;发散情形下,不同令牌以不同速率发散,导致训练中贡献不均。据此,提出两项改进:排除收敛情形,以及根据重要性得分重排令牌,旨在提升实际性能。实验结果证实了这些改进的有效性,为提升Mamba在真实场景中的效果提供了新思路。

原文摘要 · Abstract (English)

Selective state space models (SSM), such as Mamba, have gained prominence for their effectiveness in modeling sequential data. Despite their outstanding empirical performance, a comprehensive theoretical understanding of deep selective SSM remains elusive, hindering their further development and adoption for applications that need high fidelity. In this paper, we investigate the dynamical properties of tokens in a pre-trained Mamba model. In particular, we derive the dynamical system governing the continuous-time limit of the Mamba model and characterize the asymptotic behavior of its solutions. In the one-dimensional case, we prove that only one of the following two scenarios happens: either all tokens converge to zero, or all tokens diverge to infinity. We provide criteria based on model parameters to determine when each scenario occurs. For the convergent scenario, we empirically verify that this scenario negatively impacts the model's performance. For the divergent scenario, we prove that different tokens will diverge to infinity at different rates, thereby contributing unequally to the updates during model training. Based on these investigations, we propose two refinements for the model: excluding the convergent scenario and reordering tokens based on their importance scores, both aimed at improving practical performance. Our experimental results validate these refinements, offering insights into enhancing Mamba's effectiveness in real-world applications.

Mamba状态空间模型令牌动态模型优化

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