A mathematical perspective on Transformers

Overview

• Transformers are modeled as flow maps on probability measures generated by mean-field interacting particles, whose long-time dynamics tend to cluster tokens.In a toy model, a metastable clustered phase precedes eventual collapse to a point mass, paralleling token uniformity, oversmoothing [CZC`22, RZZD23, GWDW23, WAWJ24, WAW`24, DBK24, DBK24, SWJS24], or rank-collapse [DCL21, FZH`22, NAB`22, JDB23, ZMZ`23, ZLL`23, NLL`24, BHK24, CNQG24].Source: §1 · p. 2
• The paper connects Transformer dynamics to nonlinear transport equations, Wasserstein gradient flows, collective behavior, and optimal configurations of points on spheres.These connections support an accessible mathematical framework and motivate open problems centered on long-time clustering.Source: §1 · p. 3

Part 1: Modeling

• The modeling part idealizes self-attention and layer-normalization as interacting-particle dynamics on the unit sphere, with the empirical measure governed by a continuity equation.A surrogate self-attention model is formulated as a Wasserstein gradient flow for an energy functional related to optimal point configurations on spheres.Source: §1 · p. 3

Part 2: Clustering

• The clustering part establishes long-time token clustering in extreme temperature regimes and, for d ≥ 3, at any temperature, including exponential convergence when d ≥ n.It further characterizes inter-particle distance histograms and the time by which particles are nearly clustered.Source: §1 · p. 3

Part 3: Further questions

• Further questions examine the two-dimensional case through Kuramoto oscillators, sphere optimization under model modifications, and parameter tuning in interacting particle systems.These directions are supported by numerical observations and target features needed in practical Transformer architectures.Source: §1 · p. 3

Part 1. Modeling

• The model includes self-attention and layer normalization but excludes the additional feed-forward layers commonly used in practice.Despite this simplification, it remains a highly nonlinear mean-field interacting particle system.Source: Part 1. Modeling · p. 4
• Through the continuity equation, the interacting particle system defines a flow map from initial to terminal particle distributions.Source: Part 1. Modeling · p. 4

2.1. Residual neural networks.

• Residual networks motivate interpreting the discrete layer index as continuous time, yielding neural ODEs while preserving an equivalent discrete-time analysis.Skip connections make the layer index naturally interpretable as time; the paper uses continuous time because dynamical systems are more naturally defined there.Source: §2.1 · pp. 4–5

2.2. The interacting particle system.

• Transformers are modeled as sequences of particles on the unit sphere, where self-attention is the nonlinear coupling and attention weights define each particle’s neighborhood.The unit-sphere assumption simplifies RMS normalization, while Query and Key matrices determine the stochastic attention matrix and its relative attention among particles.Source: §2.2 · pp. 5–6
• The simplified clustering behavior remains qualitatively similar when Q = K = V = λId with λ > 0 or when Q, K, and V are generic random matrices.This numerical evidence supports studying the identity-matrix toy model despite its simplification.Source: §2.2 · p. 6
• Clusters progressively emerge in trained ALBERT XLarge v2, with clustering enhanced when the number of layers increases from 24 to 48.The histogram of pairwise token inner products shows growing mass at 1 through the 24th hidden layer and further enhancement by 48 layers.Source: Figure 1 · p. 7

2.3. Toward the complete Transformer.

• The framework extends the simplified self-attention model toward practical Transformers with multiple attention heads and a feed-forward layer combined through Lie-Trotter splitting.The number of heads H can be viewed as an infinite-width parameter as H approaches infinity, paralleling wide two-layer neural networks.Source: §2.3.1 · p. 8, §2.3.2 · p. 8

2.3.2. Feed-forward layers.

• Feed-forward layers are critical in applications and underpin existing Transformer approximation results, but their analysis lies beyond the paper’s current methods.Practical implementations also use mechanisms whose mathematical analysis remains open.Source: §2.3 · p. 8, §2.3.2 · p. 9

3. Measure to measure flow map

• Transformers can be viewed as flow maps between input and output probability measures on S^{d−1}, with layer normalization inducing spherical projection.The framework interprets token evolution through the continuity equation for interacting particles.Source: §3 · p. 9

3.1. The continuity equation.

• The particle system evolves according to a continuity equation with global weak measure-valued solutions, while its mean-field approximation satisfies W1(µ_n(t), µ(t)) ≤ e^{O(1)|t|}W1(µ_n(0), µ(0)).The cited discussion notes possible improvements to the metric and exponential time dependence but does not address them further.Source: §3.1 · pp. 10–11

3.2. The interaction energy.

• The interaction energy E_β increases along the dynamics, has the uniform measure σ_d as its unique global minimizer, and has Dirac masses as global maximizers.This characterization ranges over all probability measures; restricting to discrete measures produces multiple global minima.Source: §3.2 · p. 11

3.3. A Wasserstein gradient flow proxy.

• A surrogate obtained by replacing each partition function Z_β,i(t) with n is a Wasserstein gradient flow for E_β and exhibits essentially the same qualitative behavior as the original dynamics.The original vector field is generally not a standard Wasserstein gradient because of its logarithmic, nonsymmetric structure [SABP22].Source: §3.3 · pp. 12–14

3.4.1. The case of particles.

• When QJK is symmetric and V = QJK, the particle dynamics can be written as a gradient flow of E_β under a newly defined metric on (S^{d−1})^n.This reweights the metric rather than changing the interaction energy.Source: §3.4.1 · p. 15, §3.4.2 · p. 16

3.4.2. The case of measures.

• For absolutely continuous measures, the same metric-reweighting idea is formulated through a modified formal Wasserstein metric tensor and a variational distance between measures.The distance is defined by minimizing an action over measure–velocity pairs satisfying the continuity equation and endpoint constraints.Source: §3.4.2 · p. 16
• Curves of measures use tangent vectors determined by the continuity equation (3.16), which underpins the associated geometric setup.Source: §3.4.2 · p. 16
• The variational problem generalizes the Benamou–Brenier formula and defines the weighted Wasserstein distance W2,Eβ(µ0, µ1).Source: §3.4.2 · p. 16
• Within this geometry, the Wasserstein gradient ∇Eβ is defined as an element of the tangent space Tµ(t)Pac(Sd−1).Source: §3.4.2 · p. 16
• The vector field driving (3.15) is shown to be the Wasserstein gradient of Eβ, following the framework of [CNWR24, Definition 5.9].Source: §3.4.2 · p. 16

Clustering

• Figure 2 identifies parameter regimes where convergence to a single cluster as t → ∞ can be proven.For sufficiently large d, the long-time asymptotics can be characterized more finely.Source: Figure 2 · p. 17
• Convergence slows when β ≈ 1 relative to d and n, with exponential decay constant λ = O(e^-β) when d ≥ n, suggesting dynamic metastability.The paper treats small β, large β, and high-dimensional settings in Sections 4–6, respectively.Source: Figure 2 · p. 17
• The limiting distribution is a point mass in the analyzed regimes, while other parameters numerically produce long metastable phases with a small number of clusters.The theoretical results apply to specific parameter matrices and possibly very long time horizons; the metastable behavior remains unexplained theoretically.Source: Part 2. Clustering · p. 17
• Similarity-weighted averaging makes particles select nearby particles, producing leaders that attract particles in their vicinity.In the token interpretation, inner-product amplitude measures semantic similarity and attention acts as inter-token attraction.Source: Part 2. Clustering · pp. 17–18

4. A single cluster for small β

• The analysis focuses on β=0 and then small positive β while holding d and n fixed, since β controls the clustering regime.The β→∞ case is deferred to Section 5, while β=∞ leaves all particles fixed.Source: §4 · p. 18, §4.1 · p. 18

4.1. The case β “ 0.

• Almost every initial configuration converges to one common point x* for all particles when β=0.This is convergence toward consensus and sharpens the known result that allowed one antipodal particle.Source: §4.1 · p. 18
• For sufficiently small positive β, the set of initial conditions yielding one cluster retains full measure, including for arbitrary matrices Q and K.The proof compares the β>0 dynamics with the β=0 flow over a time interval of order β^-1, then uses clustering stability.Source: §4.1 · pp. 18–20
• In the special case QJK=Id, almost every initialization still converges to one cluster under a stronger parameter condition, with β≤1 when d=2.The d=2 improvement is attributed to [CRMB24].Source: §4.1 · p. 20

• Theorem 5.1 extends the conclusion of Theorem 4.3 to both (SA) and (USA) for sufficiently large β, with a constant depending only on d.Source: §5 · p. 20
• Łojasiewicz’s theorem [Loj63] ensures that every gradient-flow trajectory converges to a critical point of E_β as t → ∞.Source: §5 · p. 20
• The center-stable manifold theorem implies that, for almost every initial configuration, trajectories avoid strict saddles and generically converge to local maxima, whose landscape is then analyzed to ensure they are global.Source: §5 · p. 20

6. The high-dimensional case

• For d ≥ 3 and β ≥ 0, the clustering conclusion of Theorem 4.3 holds for both (SA) and (USA).The high-dimensional argument exploits a continuum of directions to move all particles closer simultaneously.Source: §6 · p. 21

6.1. Clustering

• When d ≥ n, uniformly random initial configurations almost surely lie in an open hemisphere and converge exponentially to a single point; for d < n, this occurs with explicit probability p_n,d ∈ (0,1).The hemisphere condition is supplied by Lemma 6.4, while the d < n statement follows from Wendel’s theorem [Wen62].Source: §6.1 · pp. 21, 23, 24

6.2. More precise quantitative convergence

• For n pairwise orthogonal initial points (β ≥ 0, d,n ≥ 2), all distinct pairwise angles remain equal, and their cosine γβ(t) obeys separate scalar evolution equations for (SA) and (USA).Source: §6.2 · p. 24

6.3. Metastability and a phase transition

• Theorem 6.9’s approximation increasingly matches the phase-transition curve as d grows, while small d exhibits long-lived multi-cluster states that eventually collapse to one cluster.Figure 3 shows the improving approximation, and Figure 4 motivates the metastability question for β ≫ 1.Source: §6.3 · p. 26, Figure 3 · p. 27

6.3. Metastability and a phase transition.

• The subsection asks whether clustering and phase-diagram conclusions persist for significantly more general parameter matrices.Source: §6.3 · p. 27
• Figure 5 provides illustrations of the behavior for more general parameter matrices.Source: §6.3 · p. 27
• A stated open problem is whether Theorems 6.8–6.9 generalize to random matrices (pQ, K, V q).Source: §6.3 · p. 27

Section 7

• Section 7 studies the special case d = 2, linking its dynamics to a phase diagram and metastable clustering at large β.For β = 4, trajectories settle into 2 clusters, while for β = 9 they settle into 3 clusters for a duration of time.Source: Part 3. Further questions · p. 27, Figure 4 · p. 28

Section 8

• Section 8 presents the BBGKY hierarchy from statistical mechanics as an alternative approach to analyzing clustering.Source: Part 3. Further questions · p. 28

Section 9

• Section 9 extends the analysis beyond QJK = V = Id through several directions, including sphere configurations, layer-normalization-free results, singular equations, and diffusive regularization.It begins by relating V = −Id to optimal configurations on the sphere, then discusses related singular and diffusive models motivated by results without layer normalization.Source: Part 3. Further questions · p. 28

Section 10

• Section 10 briefly discusses approaches that focus directly on tuning the parameter matrices.Source: Part 3. Further questions · p. 28

7. Dynamics on the circle

• On S1, particles are parametrized by angles, making the d=2 dynamics a mathematically tractable model related to Kuramoto dynamics.The circle case is presented as independently interesting and deserving complete analysis.Source: §7 · p. 28, §7.1 · p. 28

7.1. Angular equations.

• The angular formulation yields a continuity equation for the empirical angle measure, with a corresponding field also obtained for dynamics (SA).The analysis focuses on (USA) and derives the angle equation by differentiating the circle-coordinate representation.Source: §7.1 · pp. 28–30

7.2. The Kuramoto model.

• For β=0, the angular dynamics reduce to the Kuramoto model, whose behavior changes from nonsynchronization to partial and then total synchronization as coupling increases.The cited review literature includes [Str00, ABV`05, HKPZ16], alongside further results [CCH`14, Chi15, FGVG16, DFGV18, HR20, TSS20, ABK`22, MP22].Source: §7.2 · p. 30
• The dynamics can be written as a gradient flow whose global maxima are synchronized states, while proving that all other critical points are strict saddles remains partly open.A positive answer is known for β≤1 or β≥n^2π^2; it would imply convergence to a common limit for almost every initial condition.Source: §7.2 · p. 31

8. BBGKY hierarchy

• The analysis studies r-particle distributions for particles initialized i.i.d. uniformly on S1, with rotational invariance fixing the one-particle marginal at 1/(2π).The r-particle law is the marginal of the full n-particle distribution, while the one-particle distribution remains uniform for all t ≥ 0.Source: §8 · pp. 31–32
• Clustering and synchronization reduce to proving that the two-particle correlation ψ(t, ·) converges to a Dirac mass at 0 as t → ∞.The equation for ψ is derived from the Liouville equation using the method underlying the BBGKY hierarchy [GSRT13].Source: §8 · p. 32
• Equation (8.1) is not closed because its function g(t, x) depends on the three-point correlation function, as is typical for BBGKY hierarchies.The section notes that closure commonly uses physical theory or experimental evidence, citing the molecular chaos hypothesis in deriving the Boltzmann equation.Source: §8 · p. 32
• The section poses the problem of finding a cluster-consistent ansatz for g(t, x) that closes equation (8.1) and proves convergence of ψ(t, ·) to a Dirac mass at 0.It also identifies deriving BBGKY hierarchies for d > 2 and for (SA) as directions for further investigation.Source: §8 · p. 32

9. General matrices

• The clustering phenomenon likely extends beyond Q = K = V = Id, but extending the proofs to general parameter matrices remains an open problem.The section also leaves open whether the results summarized in Table 1 extend to other (Q, K, V) and what limit shapes result.Source: §9 · p. 32, §9.2 · p. 34

9.1. The repulsive case.

• In the repulsive case, minimizing E_β over n-point empirical measures becomes an optimal spherical-configuration problem linked to sphere packing [CK07, CKM`22] and coding theory.Every global minimum is either a sharp configuration or the vertices of a 600-cell.Source: §9.1 · pp. 32–33

9.2. Pure self-attention.

• [GLPR24] shows that, after rescaling z_i(t) = e^(-tV)x_i(t), particles cluster toward geometric objects for various (Q, K, V), with outcomes depending on V's spectrum.For V = Id, particles generally approach vertices of a convex polytope; a simple positive dominant eigenvalue yields clustering on parallel hyperplanes.Source: §9.2 · pp. 33–34

9.3. Singular dynamics.

• The singular β → ∞ dynamics admit Filippov solutions but not uniqueness, as Example 9.3 demonstrates; a selection principle is posed to recover uniqueness and revisit clustering questions.The limiting equations can permit multiple valid evolutions for the same initial condition.Source: §9.3 · p. 35

9.4. Diffusive regularization.

• Adding diffusion raises the unresolved question of Transformers' long-time limit for noise strength σ > 0, across the models considered in the paper.Although related consensus-based optimization results exist, their assumptions and analyses do not apply directly because this setting lacks an analogous objective f(x).Source: §9.4 · pp. 35–36

10. Approximation, control, training

• Expressivity concerns a neural network’s ability to reproduce maps by tuning parameters, with interpolation and universal approximation as closely related notions.Interpolation exactly matches arbitrarily many samples, whereas approximation captures functional relationships in an appropriate topology.Source: §10 · p. 36
• Discrete-time Transformers have universal approximation results using translation parameters and infinitely many layers [YBR`19], with related work in [ADTK23] [JL23] and review [JLLW23].Source: §10 · p. 36
• For flow maps from R^d to R^d, interpolation and approximation reflect system controllability, motivating control-theoretic analyses of Transformer expressivity.Source: §10 · p. 36
• Control-theoretic transfers have produced controllability results [CLT20] [VR23] [CLLS23] and optimal-control insights [LCT18] [EGPZ20] [GZ22].Source: §10 · p. 36
• Universal approximation results for Transformers as measure-to-measure maps use control-theoretic tools and are reported in [AL24] [FdHP24].Source: §10 · p. 36
• Transformer training dynamics remain a major challenge outside this section’s coverage, with selected references provided in [TLTO23] [ACDS23] [DGTT24].The cited works are offered as entry points to a broader and rapidly growing literature.Source: §10 · p. 36

Appendix A. Proof of Theorem 4.1

• All non-global critical points of E0 are strict saddles, so gradient ascent converges to a critical point and avoids strict saddles for almost every initialization.The proof uses the Łojasiewicz theorem for convergence and a center-stable-manifold argument showing strict-saddle basins have zero volume.Source: Appendix A · pp. 37–38

Appendix B. Proof of Theorem 5.1

• For sufficiently large β, non-global critical points of Eβ are strict saddles, yielding convergence to a single cluster whenever β ≳ (d − 1)n^2.The argument first treats d = 2 by confining points to a half-circle, then extends the conclusion to arbitrary d ≥ 2 and to (SA) by metric invariance of the saddle property.Source: Appendix B · pp. 40–42

Appendix C. Proof of Theorem 4.3

• For β ≤ c n^-1, every critical point of Eβ except the synchronized one is a strict saddle.This follows by comparing Eβ and E0, whose nonsynchronized critical points have a uniformly positive Hessian direction, and using continuity of critical points and Hessian eigenvalues.Source: Appendix C · pp. 42–43

Appendix D. Proof of Theorem 6.9

• The proof of Theorem 6.9 focuses on (SA) and begins by establishing that its flow map is Lipschitz with respect to initial data.It compares two solutions of the Cauchy problem and uses the β-Lipschitz property of softmax in Euclidean norm.Source: Step 1 · pp. 43–44, Appendix D · p. 43

Step 1.

• Step 1 derives an integral bound for the maximum trajectory discrepancy, expressed through the initial discrepancy and the time integral of pairwise differences.The displayed estimates compare x_i(t) and y_i(t) coordinatewise before taking the maximum over indices.Source: Step 1 · p. 43
• Step 1 applies Grönwall’s inequality after bounding the interaction terms, yielding a uniform estimate on the distance between corresponding trajectories.The estimate is obtained by maximizing over particle indices and integrating the differential inequality.Source: Step 1 · p. 44
• Step 1 constructs pairwise orthogonal points close to the random i.i.d. initial points on the sphere.The construction begins with y1(0) = x1(0) and proceeds inductively using concentration of measure and a union bound.Source: Step 2 · p. 44

Step 2.

• Step 2 establishes the high-probability event needed for the proof, with probability at least p0 = 1 − 2n2d−1/64 when d ≥ n.The event combines the initial-data approximation with the required dynamical bound.Source: Step 2 · p. 44
• Step 2 uses the approximation and Theorem 6.8 to control trajectories and pairwise inner products under the event A.The argument applies the constructed orthogonal initial points to the solution of the Cauchy problem.Source: Step 2 · p. 45
• Step 2 proves that all trajectories converge to a common point x* as t → +∞ and that distinct-particle inner products become arbitrarily close to 1.The proof combines monotonicity of γβ, a lower bound on its growth, and differential inequalities for the pairwise correlations.Source: Step 2 · pp. 45–46