Back to AI Research

AI Research

First-token attraction in Mamba models grows with depth

Key Takeaways

  • A Mamba preprint proves first-token attraction under specified conditions and reports growing concentration across layers of Falcon-Mamba-7B.
  • Researchers model how the first token can pull later token representations toward its direction and test the prediction on Falcon-Mamba-7B.
  • The first token can acquire an outsized role as a deep Mamba model processes a sequence.
  • A November 5 arXiv preprint from Trinh T.
  • Nguyen and colleagues examines that effect by tracking how representations change with model depth.

The first token can acquire an outsized role as a deep Mamba model processes a sequence. A November 5 arXiv preprint from Trinh T. Nguyen and colleagues examines that effect by tracking how representations change with model depth. The paper pairs a mathematical analysis with visual measurements from Falcon-Mamba-7B.

How the authors model token trajectories

The authors describe a continuous-time limit of Mamba in which token representations move on a sphere. In that view, depth becomes a way to track directions over time, and the question becomes whether those directions tend to converge. The first token has a special place in the causal sequence because later tokens can interact with earlier ones.
The analysis differs from studies of attention mechanisms because Mamba’s interactions are matrix-valued. The coefficients lack the positivity guarantees that would make an attention-style proof transfer directly. The authors use a high-dimensional invariant box and an exit-time argument to recover enough control over the signs to study convergence.

What first-token attraction predicts

Under the conditions examined in the proof, the direction approached by the first token attracts later token trajectories. The authors characterize possible limiting directions and derive exponential rates of convergence for the analyzed dynamics. The paper connects this clustering to an attention-sink effect in state-space models, where much of the sequence’s interaction concentrates on an early position.

What the Falcon-Mamba experiment measures

The team visualizes attention concentration at different depths in Falcon-Mamba-7B. It reports that concentration on the first token increases and becomes dominant in later layers. The experiments check the predicted pattern in that model, while the precise convergence statements rely on the paper’s assumptions.
For researchers studying how sequence representations change, the useful comparison is across layers. Measuring only the final layer would conceal when the concentration emerges.

Comments