Demystifying the Token Dynamics of Deep Selective State Space Models

TL;DR

This paper provides a theoretical analysis of deep selective state space models like Mamba, revealing their dynamical behaviors and proposing refinements to improve their practical performance.

Contribution

It derives the dynamical system of Mamba, characterizes token behavior, and introduces model refinements based on theoretical insights.

Findings

Tokens either converge to zero or diverge to infinity.

Divergent tokens contribute unequally during training.

Refinements improve Mamba's real-world effectiveness.

Abstract

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.