Selective state models mix tokens like attention to find consensus
The Attention Within: Consensus Dynamics in Selective State Space Models
Machine LearningArtificial Intelligence
Summary
Attention mechanisms in transformers help different parts of the input work together by making them agree or 'cluster' into a common idea. This paper shows that selective state space models (SSMs), which work differently from transformers, also cause their input parts to come to a consensus through a similar process. The authors use math from dynamical systems to prove that this mixing leads to stable agreement patterns. They also find that some parts of the model control how much this consensus happens, keeping the parts from merging completely.
What this means in practice
- •For machine learning engineers: Design efficient sequence models by understanding how SSMs mix token information similarly to attention but with different control mechanisms.
- •For natural language processing developers: Improve inference speed of language models by exploiting SSM recurrence dynamics to control token consensus and representation mixing.
Tested on one dataset.
Authors
João Pedro Silvestre, Álvaro Rodríguez Abella, Paulo Tabuada
Abstract
Selective state space models (SSMs) have recently emerged as a compelling alternative to transformers, combining competitive performance with substantially improved inference efficiency. At each SSM layer, a sequence of hidden states are propagated by a recurrence, mixing information of different tokens. Despite using a different mechanism, this mixing plays a role analogous to attention in transformers. In fact, recent works have shown that the two architectures may be closer than they first appear, as this recurrence admits a formulation akin to linear attention. In transformers, attention is known to drive the tokens to cluster, i.e., to reach consensus, collapsing in the limit to a single direction. Thus, we ask: does the recurrence at the core of SSMs drive the tokens to consensus, as attention does in transformers? To answer this question, we take a dynamical systems perspective on SSMs, modeling the evolution of tokens across layers as an ordinary differential equation. By exploiting input-to-state stability arguments, we establish local exponential stability of the consensus equilibria and characterize their domain of attraction for time-varying weight matrices, a setting not addressed by previous results. We thereby show that the resemblance between SSMs and transformers does run deeper: the recurrence at the core of SSMs aggregates tokens just as attention does. Numerical experiments on a pretrained Mamba-2 model point to the output gate as the component that regulates the extent of this consensus, preventing the tokens from reaching it in full.