The emergence of clusters in self-attention dynamics

Borjan Geshkovski; Cyril Letrouit; Yury Polyanskiy; Philippe Rigollet

Source-linked AI summary

The emergence of clusters in self-attention dynamics

Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, Philippe Rigollet

arXiv:2305.05465v6May 9, 2023cs.LGmath.APstat.ML

TL;DR

The paper studies the geometric structure of Transformer representations with time-independent weights by modeling tokens as interacting particles. It proves clustering toward spectrum-dependent limiting objects, including a low-rank Boolean self-attention matrix in one dimension, while identifying open problems for multi-head dynamics.

Problem
The paper addresses limited understanding of the geometric structure of learned Transformer representations.
Method
The paper analyzes Transformer dynamics using interacting-particle, dynamical-systems, and partial-differential-equation techniques.
Results
Tokens cluster toward spectrum-dependent limiting objects; in one dimension, the self-attention matrix converges to a low-rank matrix with entries 0 and 1, revealing leaders.
Takeaways & Limitations
The results mathematically support the emergence of a small number of leaders and clarify how value-matrix spectra shape representation geometry.
Takeaways & Limitations
Clustering and self-attention convergence for dynamics with multiple heads and constant head-specific weight matrices remain open problems.

Abstract

from arXiv · show

Viewing Transformers as interacting particle systems, we describe the geometry of learned representations when the weights are not time dependent. We show that particles, representing tokens, tend to cluster toward particular limiting objects as time tends to infinity. Cluster locations are determined by the initial tokens, confirming context-awareness of representations learned by Transformers. Using techniques from dynamical systems and partial differential equations, we show that the type of limiting object that emerges depends on the spectrum of the value matrix. Additionally, in the one-dimensional case we prove that the self-attention matrix converges to a low-rank Boolean matrix. The combination of these results mathematically confirms the empirical observation made by Vaswani et al. [VSP'17] that leaders appear in a sequence of tokens when processed by Transformers.

1. Introduction

The paper models self-attention as an interacting particle system to characterize the geometry and asymptotic clustering of Transformer representations. With fixed weights, tokens form parameter-dependent limiting configurations, including polytopes, hyperplanes, mixed structures, and low-rank attention patterns.

1. Introduction: Fixed-weight Transformer dynamics treat tokens as interacting particles whose collective evolution can be analyzed asymptotically.The continuous-time formulation focuses on pure self-attention rather than multiple heads, feed-forward layers, or layer normalization.
1. Introduction: In one dimension with V > 0, the self-attention matrix converges to a low-rank Boolean matrix, revealing a small number of leaders.The extension beyond one dimension is suggested by numerical experiments and prior empirical work, not established by the stated theorem.
1. Introduction: For V = Id, appropriately time-rescaled tokens converge to a convex-polytope boundary and, for almost all initial sequences, to fewer vertices than tokens.For V = -Id, all tokens collapse to the origin.
1. Introduction: When V has a simple positive leading eigenvalue, tokens cluster toward at most three hyperplanes determined by its leading eigenvector.This setting is presented as closer to learned value matrices.
1. Introduction: When the leading eigenvalue has multiplicity, clustering combines convex-polytope vertices in some directions with a linear subspace in others.The paper also introduces a broader mixed setting when V acts like the identity on the leading eigenspace.
1. Introduction: The analysis establishes global existence and uniqueness while showing that clustering patterns vary substantially with the spectrum of V.The authors also report numerical evidence that conclusions extend to more compound architectures.

2. Asymptotic low-rankness of the self-attention matrix

The self-attention matrix becomes asymptotically low-rank and typically Boolean in one dimension, with its limiting structure concentrating attention on a small number of tokens.

In one dimension, P(t) converges exponentially fast to a matrix that is typically both Boolean and low-rank.
The set P consists of structured n × n matrices whose starred entries are non-negative and sum to 1, allowing the limiting attention pattern shown in Figure 2.
Theorem 2.1 shows that, for d = 1 with V > 0 and QK > 0, P(t) converges to a matrix P* in P as t tends to infinity.
For almost all distinct initial token sequences, P* has rank 1 or 2, with the unconstrained row equal to e1 or en.
The limiting matrix means that at most three tokens capture the attention of all tokens except possibly one, with the extreme tokens typically acting as leaders.

3. Clustering toward vertices of convex polytopes

When V = Id, a rescaled self-attention dynamics drives tokens toward the boundary of an initial-condition-determined convex polytope, typically concentrating them at its vertices.

Rescaled dynamics: The rescaling preserves self-attention coefficients and serves as a mathematically justified surrogate for layer normalization.Original tokens are recovered through x_i(t) = e^tV z_i(t).
Theorem 3.1: Tokens converge either to the origin or to points on the boundary of a convex polytope K when V = Id and Q^T K > 0.The polytope depends on the initial token sequence.
Vertex clustering: For almost all initial sequences, tokens converge to vertices of K, whose number is often substantially smaller than the sequence length.These vertices form the leaders attracting the remaining tokens.
Proof mechanism: The proof combines a shrinking convex hull with continued growth of non-boundary particles until they reach the limiting boundary.Time-rescaling prevents the convex hull from collapsing prematurely.
Scope and caveats: The conclusions may extend beyond Q^T K > 0, but for generic initial sequences the polytope is not explicitly predictable without running the full dynamics.Exceptional null-set sequences can converge to facet interiors rather than vertices.
Convergence behavior: The convergence theorem does not specify a rate, although numerics suggest that most tokens cluster after only a few layers.The meaning of “few layers” depends on the initial token magnitudes.

4. Clustering toward hyperplanes

For value matrices with a simple positive leading eigenvalue, tokens approach at most three parallel hyperplanes, while the spectrum of V governs attraction, repulsion, and higher-dimensional clustering structure.

Theorem 4.2: Under the good-triple conditions, every token approaches one of at most three parallel hyperplanes.The theorem requires a real, positive, simple dominant eigenvalue and positivity along its eigenspace.
Spectral mechanism: The leading eigendirection determines the attracting clusters because tokens with the largest component along it dominate the attention exponent.The hyperplanes are perpendicular to the leading eigenvector in the illustrated diagonalizable case.
Interpretation: The hyperplane result yields a linearly separable token representation and can be interpreted as clustering into at most three flats of dimension d − 1.The paper relates this behavior to K-flats clustering.
Higher-dimensional extension: The codimension conjecture predicts at most three parallel subspaces of codimension k when V has k eigenvalues with positive real part.This extension is supported by numerical experiments for more complex value matrices.
Mixed spectral effects: Positive-eigenvalue directions attract and cluster tokens, whereas negative-eigenvalue directions can generate repulsion and divergence.The illustrated three-dimensional case has two positive and one negative eigenvalue.

5. A mix of hyperplanes and polytopes

When the leading positive eigenvalue has multiplicity, the dynamics combine polytope-like clustering in its eigenspace with convergence toward a structured set extending across the complementary subspace.

Geometric synthesis: This result combines the hyperplane-clustering and convex-polytope regimes when the dominant eigenvalue is not simple.The paper identifies paranormal matrices, including normal matrices, as a relevant class.
Assumptions: The paranormal value matrix acts as λI on F while its restriction to G has spectral radius below λ, separating dominant and subdominant dynamics.These spectral conditions define the multiplicity setting.
Main theorem: For a good triple with multiplicity, tokens approach H = (∂K \ {0}) × G, where K is a bounded convex polytope in the leading eigenspace F.The complementary subspace G contains directions with strictly smaller spectral radius.

6. Well-posedness

The paper establishes global existence and uniqueness for the particle and mean-field dynamics, using continuity-equation methods and stability estimates to connect discrete tokens with measure-valued solutions.

Behavior without normalization: The dynamics can diverge in some directions without normalization, even though global existence and uniqueness still hold.Negative eigenvalues are associated with divergent token coordinates in the illustrated examples.
Particle dynamics: For every finite initial token sequence, the original dynamics has a unique solution defined for all times.The rescaled dynamics is likewise well-posed through the change of variables x_i(t) = e^tV z_i(t).
Proof strategy: The proof uses local Lipschitz bounds for the attention vector field, compact-support control, and a Wasserstein-based stability argument.A time-stepping construction freezes the vector field on short intervals before passing to the limit.
Mean-field well-posedness: The continuity equation admits a unique compactly supported solution for each compactly supported initial measure.The result also provides a stability estimate for solutions with bounded initial support.
Continuity equation: The empirical measure of the particles, μ_t = (1/n) Σ_j δ_{x_j(t)}, solves the corresponding continuity equation.This connects the particle system to its mean-field formulation.

7. Proof of Theorem 2.1

In one dimension, particles remain separated while unbounded particles escape exponentially and their attention rows converge rapidly to endpoint selectors. These properties imply a low-rank Boolean limit for the self-attention matrix, with at most one exceptional bounded particle.

The proof shows that particle distances are non-decreasing, so initially distinct particles never collide.
Unbounded particles converge to either −∞ or +∞, with asymptotic form x_i(t)=γ_i e^t+o(e^t).
Rows associated with particles tending to +∞ or −∞ converge doubly exponentially to the selectors δ_nj or δ_1j, respectively.
At most one particle can remain bounded; if an endpoint particle is bounded, its attention row converges to the corresponding endpoint selector.
Consequently, all but at most one rows converge to endpoint basis vectors, while the exceptional row may converge to a probability vector, yielding the stated low-rank Boolean-limit structure.

8. Proofs of Theorems 3.1 and 8.5

For V=Id, the rescaled token convex hull shrinks toward a convex polytope and particles converge to finitely many distinguished boundary points. For V=−Id, all particles converge to the origin and attention becomes uniform.

Each particle converges to a finite candidate set contained on the boundary of K or at the origin.
The convex hull of the rescaled particles is non-increasing and converges to a convex polytope K.
Particles eventually remain near one connected component of neighborhoods around the candidate points, preventing indefinite circulation between components.
For V=−Id, every particle converges to the origin, and the self-attention matrix converges to the matrix whose entries are all 1/n.

9. Proof of Theorem 4.2

When the value matrix has a simple positive leading eigenvalue separated from the remaining spectrum, particles converge in their leading-eigenvector coordinate toward at most three initial-condition-dependent levels. Other coordinates may still diverge along the limiting hyperplanes.

Monotonicity and boundedness of spectral coordinates provide uniform control when the corresponding eigenvalue is nonnegative.
Each particle’s leading-eigenvector coordinate converges, and its limit belongs to a set of at most three real numbers determined by the initial tokens.
Particles converge toward one of the candidate hyperplanes H_0, H_a, or H_b defined by those limiting leading coordinates.
Convergence of the leading coordinate does not prevent the full token norm from diverging along its limiting hyperplane, especially when V has negative eigenvalues.

10. Proof of Theorem 5.2

With a simple positive leading eigenvalue, the analysis combines polytope geometry with spectral projections: projected convex hulls converge, and particles approach finitely many distinguished points in the leading eigenspace.

Projecting onto the leading eigenspace yields a convex hull that is non-increasing and converges to a convex polytope K.
The candidate limit set is finite, lies on the boundary of K, and is defined using the projected geometry.
After accounting for lower-eigenspace terms, each particle remains within arbitrarily small distance of one candidate point in the leading eigenspace.

11. Numerical experiments

The experiments examine how sequence length, dimension, eigenvalue structure, and model weights affect the clustering dynamics predicted by the paper. They provide numerical evidence for higher-dimensional extensions, spectral dependence, and leader-related low-rank behavior.

Selected ALBERT heads satisfy the assumptions of Theorem 4.2, showing that the theorem’s spectral conditions occur in a pretrained model.The illustrated heads are 5 and 14.
Theorem 2.1’s one-dimensional conclusion appears to extend to higher dimensions, supporting the conjectured persistence of clustering when V>0.
The rank of P(t) appears independent of sequence length n=100, consistent with rank reflecting the number of leaders.This expands the same setup used in Figure 3.
In a high-dimensional ALBERT-sized setup, 65 positive-eigenvalue coordinates appear to converge to one of up to three real scalars, while remaining coordinates can oscillate or diverge.The divergent-coordinate behavior is shown on a logarithmic scale.
For complex eigenvalues with positive real part, particles appear to collapse to two points rather than merely approach two hyperplanes.The observed hyperplanes have codimension 3, matching the stated conjecture.
Complex eigenvalues with negative real part produce particle rotation within two-dimensional hyperplanes, as predicted by Theorem 4.2.

12. Outlook

The outlook identifies open extensions beyond the paper’s pure, single-head self-attention setting. Numerical examples suggest clustering can persist under altered query-key assumptions and appended nonlinear feed-forward dynamics, but the general theory remains unresolved.

12. Outlook: Clustering and convergence for multi-headed Transformers with constant head-specific weights remain open problems.Preliminary numerical investigations suggest that clustering phenomena may still occur.
12. Outlook: When the leading eigenvalue of V has negative real part, rescaled tokens diverge, whereas a complex leading eigenvalue is not expected to yield clustering.These spectral effects concern the value matrix and are illustrated numerically.
12. Outlook: The clustering patterns associated with Theorems 3.1 and 5.2 appear to persist even when Q^T K violates the positive-semidefinite assumption.The examples use random query-key matrices.
12.2. Beyond pure self-attention: adding a feed-forward layer.: With a feed-forward mechanism applied in parallel, exact clustering is difficult to anticipate because its weights act like an additional value matrix.The conclusions depend strongly on the identity-like structure of V.
12.2. Beyond pure self-attention: adding a feed-forward layer.: In the appended nonlinear-network example, clustering appears to depend on both the activation function and weight matrix, and the problem is left open.The setup uses ReLU or tanh followed by a weight matrix W.