Dynamical Properties of Tokens in Self-Attention and Effects of Positional Encoding

TL;DR

The paper analyzes how tokens move in Transformers by modeling the self-attention mechanism as a continuous-time dynamical system. It derives broad conditions, in terms of the symmetric parts of key matrices, that determine whether token distances contract (convergence) or expand (divergence), and shows how absolute and rotary positional encodings influence these regimes. Empirical results on language and vision tasks confirm that convergence can harm performance while divergence or clustering can be beneficial, leading to simple architectural refinements that promote divergence. The work provides a theoretical framework and practical guidance for designing Transformers with more robust dynamics and better empirical performance. Overall, it bridges dynamical systems theory with Transformer design, offering transferable insights for pre-trained models.

Abstract

This paper investigates the dynamical properties of tokens in pre-trained Transformer models and explores their application to improving Transformers. To this end, we analyze the dynamical system governing the continuous-time limit of the pre-trained model and characterize the asymptotic behavior of its solutions. Specifically, we characterize when tokens move closer to or farther from one another over time, depending on the model parameters. We provide sufficient conditions, based on these parameters, to identify scenarios where tokens either converge to zero or diverge to infinity. Unlike prior works, our conditions are broader in scope and more applicable to real-world models. Furthermore, we investigate how different forms of positional encoding -- specifically absolute and rotary -- affect these dynamical regimes. Empirical evidence reveals that the convergence scenario adversely impacts model performance. Motivated by these insights, we propose simple refinements to Transformer architectures that mitigate convergence behavior in models with absolute or rotary positional encoding. These findings support theoretical foundations and design principles for improving Transformer models.

Figures (11)

• Figure 1: Distances between tokens over time. Left: When $A_{\text{sym}} \succ 0$, the distances between tokens do not decrease. Right: Conversely, when $A_{\text{sym}} \prec 0$, the distances between tokens do not increase.
• Figure 2: Token trajectories under two configurations. Left: When $A_{\text{sym}}$ is negative definite and $W_{\text{sym}}$ is positive definite, all tokens converge to zero as time $t \to \infty$. Right: When $A = 2W$, tokens diverge to infinity over time, forming a few distinct groups that move in aligned directions.
• Figure 3: Token dynamic in pre-trained model when omitting $\operatorname{LayerNorm}$ and feed-forward network
• Figure 4: Tokens' trajectory in the later layers of a Transformer forms clusters with similar shapes in model (left) with $\operatorname{LayerNorm}$, no feedforward, and (right) with both $\operatorname{LayerNorm}$ and feedforward.
• Figure 5: By introducing the additional term $\Delta W_{li}$, the system’s behavior shifts from a convergent regime (left) to a divergent regime (right).

Theorems & Definitions (32)

• Remark 2.1: Scope and Generality
• Theorem 3.1: Distances between Tokens
• Remark 3.2: Distances between Tokens
• Proposition 3.3
• Theorem 3.4: Convergence Scenario
• Remark 3.5: Convergence Scenario
• Theorem 3.6: Divergence Scenario
• Remark 4.1: Absolute positional encoding has minimal impact
• Remark 4.2: Rotary positional encoding encourages token divergence
• Lemma B.1