A mathematical perspective on Transformers

Borjan Geshkovski; Cyril Letrouit; Yury Polyanskiy; Philippe Rigollet

Source-linked AI summary

A mathematical perspective on Transformers

Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, Philippe Rigollet

arXiv:2312.10794v5Dec 17, 2023cs.LGmath.APmath.DS

TL;DR

Transformers are central to modern language models, yet the mechanisms underlying self-attention remain largely uncharted. The paper models Transformers as interacting particle systems and analyzes their measure-valued dynamics through continuity equations and gradient flows. Its simplified models exhibit long-time clustering, with limiting distributions that can be point masses, while broader parameter settings and some repulsive dynamics remain open.

Problem
The paper addresses the limited mathematical understanding of how Transformers, especially self-attention, process data.
Method
The paper interprets Transformers as mean-field interacting particle systems and studies their continuity equations, interaction energies, and Wasserstein gradient-flow formulations.
Results
The analyzed dynamics indicate long-time token clustering, including point-mass limiting distributions in the simplified parameter settings studied.
Takeaways & Limitations
The framework connects Transformer dynamics with nonlinear transport, gradient flows, collective behavior, and spherical point configurations, providing mathematical perspectives and open problems.
Takeaways & Limitations
The proofs focus on simplified parameter matrices, while extending them to more general matrices and fully characterizing the repulsive case remain open.

Abstract

from arXiv · show

Transformers play a central role in the inner workings of large language models. We develop a mathematical framework for analyzing Transformers based on their interpretation as interacting particle systems, which reveals that clusters emerge in long time. Our study explores the underlying theory and offers new perspectives for mathematicians as well as computer scientists.

1. Outline

The paper develops a mathematical framework that treats Transformers as interacting particle systems and studies their long-time clustering behavior. It presents modeling tools, clustering results, and open directions connecting Transformer dynamics with established mathematical theories.

Framework: Transformers are modeled as flow maps on probability measures generated by mean-field interacting particle systems.Tokens follow vector fields depending on the empirical measure, whose evolution is governed by a continuity equation.
Framework: The framework connects Transformer dynamics to nonlinear transport equations, Wasserstein gradient flows, collective behavior models, and optimal point configurations on spheres.The authors aim to provide an accessible mathematical perspective while highlighting connections to established areas.
Modeling: The idealized model retains self-attention and layer-normalization while viewing successive layers as time discretizations of interacting-particle dynamics.Layer-normalization constrains particles to the unit sphere, while self-attention couples them through the empirical measure.
Clustering: Long-time analysis establishes clustering results across temperature regimes and in high dimensions, including exponential convergence when d ≥ n.The paper also characterizes inter-particle distance histograms and the time by which particles are nearly clustered.
Further questions: The paper proposes further questions involving d = 2, Kuramoto oscillators, spherical configurations, and parameter tuning in Transformer architectures.These directions are largely formulated as open questions supported by numerical observations.

Part 1. Modeling

The paper studies a simplified Transformer as a highly nonlinear mean-field interacting particle system. Its continuity equation defines a flow from initial to terminal particle distributions.

Model scope: The model includes self-attention and layer normalization but excludes additional feed-forward layers commonly used in practice.Despite this simplification, it remains a highly nonlinear mean-field interacting particle system.
Flow-map formulation: The interacting particle system induces, through a continuity equation, a flow map from initial to terminal distributions of particles.This measure-to-measure viewpoint provides the mathematical representation used throughout the paper.

2. Interacting particle system

The interacting-particle formulation represents Transformer tokens as vectors evolving through continuous-time dynamics with self-attention coupling and layer normalization. The paper focuses on simplified constant parameter matrices, while numerical evidence suggests qualitatively similar clustering beyond the identity case.

Continuous-time formulation: The continuous-time viewpoint treats discrete neural-network layers as time steps in a dynamical system.The authors use neural ODEs for exposition and state that the results can also be derived in discrete time.
Token dynamics: Transformers evolve sequences of n token vectors rather than a single input vector, allowing each token to be processed together with sentence or paragraph context.Tokens and particles are used interchangeably in the paper.
Layer normalization: Layer normalization is modeled by constraining tokens to the unit sphere, simplifying the practical ellipsoidal evolution produced by RMS normalization.In trained ALBERT XLarge v2, the normalization matrix is constant across layers with mean entry value 0.44 and standard deviation 0.008.
Self-attention: Self-attention is the nonlinear coupling mechanism, with attention weights describing how strongly each particle attends to every other particle.The matrices Q and K determine neighborhoods, while the attention matrix is stochastic with probability-vector rows.
Parameterization: The analysis uses constant parameter matrices, often setting Q, K, and V equal to the identity for a simplified scenario.The parameter β is interpreted as an inverse temperature, and Q and K need not be square.
Empirical scope: Numerical experiments show qualitatively similar clustering for Q = K = V = λId with λ > 0 and for generic random parameter matrices.The paper does not analyze all practical mechanisms, including feed-forward layers and multi-headed self-attention.
Empirical illustration: Figure 1 plots histograms of pairwise token inner products across layers in ALBERT XLarge v2, where increasing mass at 1 indicates progressively emerging clusters.Extending the depth from 24 to 48 layers further enhances the clustering pattern.

3. Measure to measure flow map

The paper models Transformers as flow maps between probability measures, generated by mean-field interacting particles on the sphere. Their empirical measure follows a continuity equation, and the dynamics admit energy-based and gradient-flow interpretations.

Modeling: Transformers are represented as flow maps on probability measures, with tokens evolving through a vector field depending on the empirical measure.Layer normalization places the particles on the sphere, while self-attention supplies their nonlinear coupling.
Continuity equation: The empirical measure evolves according to a continuity equation driven by the particle vector field.The equation is interpreted in the sense of distributions, with an initial probability measure on the sphere.
Well-posedness and mean field: The model has global weak, measure-valued solutions for arbitrary initial probability measures on the sphere.The paper also relates empirical-particle and arbitrary-measure formulations through a Dobrushin-type mean-field estimate.
Energy and gradient flow: An interaction energy increases along trajectories, has the uniform measure as its unique global minimizer, and has Dirac masses as global maximizers.The original dynamics can be viewed as a gradient flow after modifying the tangent-space metric, although asymmetry prevents this interpretation for more general parameters.
Energy and gradient flow: A normalized variant replacing the partition function by n exhibits essentially the same qualitative behavior as the original dynamics.The paper uses this proxy to obtain a Wasserstein gradient-flow formulation for qualitative analysis.

Part 2. Clustering

The paper’s clustering results indicate that Transformer tokens can converge to a point mass in long time for selected parameter regimes. Clustering is relevant because token distributions encode possible outputs, while numerical evidence also suggests long metastable multi-cluster phases.

Clustering: The theory identifies regimes where convergence to a single cluster as t → ∞ can be proven, with dimension and inverse temperature governing the behavior.When β is close to 1 relative to d and n, convergence can be slow and dynamic metastability is expected.
Interpretation: For tasks including sentiment analysis, masked language modeling, and summarization, clustering of the output measure represents a small number of possible token outcomes.The paper emphasizes that its point-mass limits apply to specific parameter choices and potentially very long time horizons.
Mechanism: Self-attention forms weighted averages that emphasize similar particles, producing leaders that attract other particles and promote clustering.The inner product between tokens is interpreted as a measure of semantic similarity.

4. A single cluster for small β

In the small-β regime, the paper proves that, for almost every initial configuration under stated conditions, the Transformer dynamics converge to a single cluster. The argument uses the β = 0 dynamics and perturbation estimates for small positive β.

4.1. The case β = 0: At β = 0, for Lebesgue almost every initial sequence, all particles converge to a common point x* on the sphere.This is established for d,n ≥ 2 and is interpreted as convergence toward consensus.
4.2. The case β ≪ 1: For small positive β, the dynamics remain close to the β = 0 system over times shorter than β^-1.The remainder term is O(β), so particles initially behave as in the β = 0 regime during this time window.
4.2. The case β ≪ 1: There exists a threshold β_m > 0 such that initial conditions entering a clustered region by time m still produce a single cluster for β ≤ β_m.The threshold decreases toward zero as m increases.
4.2. The case β ≪ 1: For fixed d,n ≥ 2, almost every initial configuration converges to one cluster whenever β lies below a dimension- and configuration-dependent threshold.The same conclusion is stated for both the self-attention and uniform self-attention dynamics.
4.2. The case β ≪ 1: When d = 2, the same single-cluster conclusion holds for β ≤ 1.This improves the small-β range in the two-dimensional case.

5. A single cluster for large β

For sufficiently large but finite β, the paper proves single-cluster convergence for almost every initial configuration. The proof relies on gradient-flow convergence and the exclusion of nongeneric critical points.

5. A single cluster for large β: For sufficiently large β, the conclusion of Theorem 4.3 holds for both the self-attention and uniform self-attention dynamics.The threshold has the form β ≥ C(d), with C depending only on the dimension.
5. A single cluster for large β: Real-analytic gradient-flow structure implies convergence from any initial condition to a critical point of the interaction energy.Almost-everywhere convergence to the single-cluster state follows by applying the center-stable manifold theorem to exclude nongeneric initial configurations.

6. The high-dimensional case

In dimensions d ≥ 3, the clustering conclusions extend across all temperatures, with exponential convergence when d ≥ n. High-dimensional symmetry also yields quantitative phase-transition approximations, while lower dimensions exhibit metastable multi-cluster states.

For d ≥ 3 and β ≥ 0, both (SA) and (USA) converge to a single cluster.
When d ≥ n and initial points are uniformly random, convergence to a common point occurs almost surely at an exponential rate.The result applies to both (SA) and (USA).
The equal-angle configuration is exponentially stable, and orthogonal initial configurations preserve a common pairwise angle throughout the dynamics.
For d < n, clustering to a single point still occurs with probability at least an explicit p_n,d ∈ (0,1).
As d increases, the curve Γ_∞,δ increasingly approximates the exact clustering phase transition in (t, β).The approximation is derived from high-dimensional dynamics and is illustrated numerically in Figure 3.
For small d, an intermediate region contains finite multi-cluster states that eventually relax to one cluster, indicating metastability.The behavior exhibits distinct transient and long-time scales.

Part 3. Further questions

The paper closes by proposing open problems on circle dynamics, BBGKY descriptions, broader parameter matrices, and mechanisms for tuning Transformer architectures.

The authors identify several future directions for refining the mathematical understanding of clustering and extending the model.
The two-dimensional case is connected to Kuramoto oscillators and singled out for further analysis.
BBGKY-based analysis is proposed as an alternative route for studying clustering.
Beyond Q = K = V = Id, the paper considers optimal sphere configurations, layer-normalization variants, singular equations, and diffusive regularization.

7. Dynamics on the circle

On the unit circle, the Transformer dynamics can be expressed through angular equations related to the Kuramoto model and as a gradient flow. The paper establishes partial convergence results and formulates an open question about the remaining regime.

7.1. Angular equations.: For d = 2, particles on S^1 are represented by angles, yielding angular equations for the (USA) and (SA) dynamics.
7.2. The Kuramoto model.: At β = 0, the angular dynamics reduce to a particular Kuramoto model.
7.2. The Kuramoto model.: The Kuramoto model exhibits a coupling-dependent transition from nonsynchronization to partial and then total synchronization.
7.1. Angular equations.: The circle dynamics can also be written as a gradient flow whose energy is maximized when all angles coincide.
7.1. Angular equations.: The authors ask whether every non-global critical point of E_β is a strict saddle point.
7.1. Angular equations.: Existing proofs answer the related problem for β ≤ 1 or β ≥ n²π², while the complementary regime remains open.

8. BBGKY hierarchy

The BBGKY approach studies reduced particle distributions to characterize synchronization, but the resulting equation is unclosed because it depends on higher-order correlations.

The r-particle distribution ρ_n^(r) tracks the joint law of r angles and is a marginal of the full n-particle distribution.
Rotational invariance makes the one-particle marginal uniform at every time and reduces the two-particle marginal to a function ψ(t, ·).
Proving synchronization amounts to showing that ψ(t, ·) converges to a Dirac mass at zero.
The BBGKY-derived equation for ψ is not closed because it depends on a three-point correlation function.
The paper proposes finding an ansatz for the higher-order correlation that closes the equation and proves convergence.
Deriving BBGKY hierarchies for d > 2 and for (SA) remains an open direction.

9. General matrices

This section extends the Transformer clustering framework to general parameter matrices, including repulsive interactions and time-scaled dynamics. It also identifies unresolved questions involving limit shapes, singular limits, noise, and broader matrix choices.

The repulsive case: In the repulsive case V = −Id, the uniform distribution minimizes the interaction energy over probability measures, whereas empirical measures can have many distinct global-minimizing configurations.For n-point empirical measures, the minimization problem connects to optimal point configurations on spheres, sphere packing, and coding theory.
The repulsive case: Global minima of Hβ among n points on the sphere are either sharp configurations or the vertices of a 600-cell.Sharp configurations are spherical designs with additional constraints on their distinct pairwise inner products.
The repulsive case: The complete long-time behavior in the repulsive case remains open because sharp configurations are not known for all regimes of n, d, and m.Known examples include regular polytopes and configurations related to the E8 and Leech lattices.
General matrices: For general (Q, K, V), rescaled particles can cluster toward well-characterized geometric objects, with outcomes depending on the spectral properties of V.The rescaling z_i(t) = e^-tV x_i(t) acts as a surrogate for layer normalization, while preserving the self-attention coefficients.
General matrices: When V = Id, particles generally cluster at vertices of a convex polytope because the convex hull shrinks along the dynamics.The polytope's vertices attract all particles as t → ∞, outside exceptional situations.
Singular dynamics: The large-β singular limiting dynamics admit Filippov solutions but not uniqueness, motivating a selection principle analogous to viscosity or entropy solutions.Example 9.3 exhibits multiple valid evolutions for one particle while preserving the limiting equations almost everywhere.

10. Approximation, control, training

The paper places Transformer expressivity within approximation and control theory while noting that training dynamics are outside its scope. Existing work establishes universal approximation for certain discrete-time and flow-map Transformer settings.

Approximation and control: Interpolation and universal approximation are related expressivity notions describing exact sample matching and functional approximation, respectively.The paper presents expressivity as the ability to reproduce maps by tuning neural-network parameters.
Approximation and control: Universal approximation has been established for discrete-time Transformers using a variant with translation parameters and an increasing number of layers.The section also cites related results and a review of the topic.
Approximation and control: For flow maps, interpolation and approximation reflect controllability properties, linking expressivity analysis to control-theoretic and optimal-control techniques.The paper cites controllability and universal-approximation results for flow-map settings.
Training: The paper does not cover the training dynamics of Transformers and refers readers to other work on this developing topic.This is stated as a scope boundary rather than as a result about training.

Appendix A. Proof of Theorem 4.1

The appendix proves single-cluster convergence by combining gradient-flow structure with analytic dynamical-systems results. Non-global critical points are strict saddles, so almost every initial condition converges to the synchronized global maximum.

Proof strategy: The dynamics are gradient ascent for an energy E0 on the compact manifold (S^(d−1))^n.This gradient-flow formulation is the starting point for the proof of Theorem 4.1.
Convergence: Every trajectory converges to a critical point of E0 because the gradient ascent is real-analytic on a compact real-analytic manifold.The convergence follows from the Łojasiewicz theorem.
Theorem 4.1: Therefore, for almost every initial configuration, the particles converge to a single synchronized point on the sphere.The result excludes only a Lebesgue-null set of initial configurations.
Generic convergence: Initial conditions whose gradient ascent converges to a strict saddle have volume zero.The center-stable manifold theorem gives the zero-volume basin result used in the argument.
Critical points: Every critical point that is not the global maximum, where all particles coincide, is a strict saddle point.Consequently, all local maxima are global.
Theorem 4.3: For sufficiently large β, the analogous argument yields single-cluster convergence whenever β ≳ (d−1)n^2.The proof uses a geometric restriction on critical points and the behavior of the function g.