Clustering in pure-attention hardmax transformers and its role in sentiment analysis

TL;DR

This work analyzes pure-attention hardmax transformers as discrete-time dynamical systems to explain how deep transformers develop context. It proves that inputs converge to a clustered equilibrium organized by a finite set of leaders that correspond to vertices of a limiting convex polytope, with nonleader tokens clustering near leaders via a hyperplane based geometry. The authors then build an interpretable sentiment analysis model that leverages leader words to filter context, demonstrating that clustering around leaders captures meaning and informs predictions. The results offer a rigorous mechanism for context formation in transformers and outline key open challenges for extending the theory to more general parameterizations and architectures.

Abstract

Transformers are extremely successful machine learning models whose mathematical properties remain poorly understood. Here, we rigorously characterize the behavior of transformers with hardmax self-attention and normalization sublayers as the number of layers tends to infinity. By viewing such transformers as discrete-time dynamical systems describing the evolution of points in a Euclidean space, and thanks to a geometric interpretation of the self-attention mechanism based on hyperplane separation, we show that the transformer inputs asymptotically converge to a clustered equilibrium determined by special points called leaders. We then leverage this theoretical understanding to solve sentiment analysis problems from language processing using a fully interpretable transformer model, which effectively captures `context' by clustering meaningless words around leader words carrying the most meaning. Finally, we outline remaining challenges to bridge the gap between the mathematical analysis of transformers and their real-life implementation.

Figures (10)

• Figure 1: Geometric interpretation of \ref{['eq:transformer_b']} for $i=1$ with (a) $A=I$ and (b) $A = \left(2111\right)$. In (a), tokens $z_{2}^{}$ and $z_{3}^{}$ have the largest orthogonal projection on the direction of $A z_{1}^{} = z_{1}^{}$, so $\mathcal{C}_{1}(Z^{}) = \{ 2,3 \}$. In (b), token $z_{4}^{}$ has the largest projection on the direction of $Az_{1}^{}$, so $\mathcal{C}_{1}(Z^{}) = \{ 4 \}$. In both cases, tokens attracting $z_1$ can only lie on the closed half-space $\mathcal{H}_1 = \{z:\;\langle A z_{1}^{}, z_{}^{} - z_{1}^{} \rangle \geq 0\}$ (blue shading).
• Figure 2: Simulations of \ref{['eq:transformer']} with $\alpha = 0.5$, $A = I$, and four different initial token values. In each panel, stars denote tokens $z_{i}^{}$ satisfying $\mathcal{C}_{i}(Z^{k})=\{i\}$ at layer $k\in N$, while circles denote all other tokens. Colors indicate which tokens are being followed. Tokens painted in two halves follow two tokens. Tokens whose interior and edge colors are different, instead, follow tokens of their interior color and are followed by tokens of their edge color. The shaded region in the last column is the closed convex hull of leaders.
• Figure 3: Schematic illustrations of a deep neural network with normalization and feed-forward sublayers (top), a full transformer with self-attention, normalization, and feed-forward sublayers (middle), and a pure-attention transformer with only self-attention and normalization sublayers (bottom). Each model takes a matrix $Z^0 \in \mathbb{R}^{n \times d}$ as its input and outputs a matrix $Z^{K} \in \mathbb{R}^{n \times d}$ after being processed by $K$ transformer layers. Residual connections are incorporated within each feed-forward and self-attention layer.
• Figure 4: Sketch of the $\delta$-neighborhood of the attracting set $\mathcal{S}$.
• Figure 5: Illustration of the geometric intuition behind \ref{['lem:verticesSelfMax']} for $\mathcal{S'} = \{ s_1, s_2\}$. If $\delta > 0$ is small enough, then for all $x\in B_\delta (s_0)$ the hyperplane with normal direction $x$ passing through $x$ (in green) separates $s_0$ from the neighborhoods $B_\delta (s_1)$ and $B_\delta (s_2)$.

Theorems & Definitions (34)

• Definition 1.1
• Theorem 1.1
• Remark 1.2
• Remark 1.3
• Lemma 3.1
• proof
• Lemma 3.2
• proof
• Lemma 3.3
• proof