Obsidian/LLM/A mathematical perspective on Transformers.md at 271fa25ddc36e401d49be49befa587871ab047c9

admin/Obsidian

Fork 0

Files

T

coco 8b94d76799 a

2026-07-21 13:57:15 +08:00

263 KiB

Raw Blame History

title, source, author, published, created, description, tags

title

source

author

published

created

description

Abstract.

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, with a particular emphasis on long-time clustering behavior. Our study explores the underlying theory and offers new perspectives for mathematicians as well as computer scientists.

Key words and phrases:

Transformers, self-attention, interacting particle systems, clustering, gradient flows

1991 Mathematics Subject Classification:

Primary: 34D05, 34D06, 35Q83; Secondary: 52C17

1. Outline

The introduction of Transformers in 2017 by Vaswani et al. [VSP ⁺ 17] marked a significant milestone in the development of neural network architectures. Central to this contribution is self-attention, a novel mechanism which distinguishes Transformers from traditional architectures, and which plays a substantial role in their superior practical performance. In fact, this innovation has been a key catalyst for the progress of artificial intelligence in areas such as computer vision and natural language processing, notably with the emergence of large language models. As a result, understanding the mechanisms by which Transformers, and especially self-attention, process data is a crucial yet largely uncharted research area.

A common characteristic of deep neural networks (DNNs) is their compositional nature: data is processed sequentially, layer by layer, resulting in a discrete-time dynamical system (we refer the reader to the textbook [GBC16] for a general introduction). This perspective has been successfully employed to model residual neural networks —see Section 2.1 for more details—as continuous-time dynamical systems called neural ordinary differential equations (neural ODEs) [CRBD18, E17, HR17]. In this context, an input x(0)\in\mathbb{R}^{d} , say an image, is evolving according to a given time-varying velocity field as \dot{x}(t)=v_{t}(x(t)) over some time interval (0,T) . As such, a DNN can be seen as a flow map x(0)\mapsto x(T) from \mathbb{R}^{d} to \mathbb{R}^{d} . Even within the restricted class of velocity fields \{v_{t}\}_{t\geq 0} imposed by classical DNN architectures, such flow maps enjoy strong approximation properties as exemplified by a long line of work on these questions [LJ18, ZGUA20, LLS22, TG22, RBZ23, CLLS23].

Following [SABP22] and [VBC20], we observe that Transformers are in fact flow maps on \mathcal{P}(\mathbb{R}^{d}) , the space of probability measures over \mathbb{R}^{d} . To realize this flow map from measures to measures, Transformers evolve a mean-field interacting particle system. More specifically, every particle (called a token in this context) follows the flow of a vector field which depends on the empirical measure of all particles. In turn, the continuity equation governs the evolution of the empirical measure of particles, whose long-time behavior is of crucial interest. In this regard, our main observation is that particles tend to cluster under these dynamics. This phenomenon is of particular relevance in learning tasks such as next-token prediction, wherein one seeks to map a given input sequence (i.e., a sentence) of n tokens (i.e., words) onto a given next token. In this case, the output measure encodes the probability distribution of the next token, and its clustering indicates a small number of possible outcomes. Our results on a simplified but insightful toy model indicate that the limiting distribution is actually a point mass, leaving no room for diversity or randomness, which is at odds with practical observations. This apparent paradox is resolved by the existence of a long-time metastable state. As can be seen from Figures 3 and 5, the Transformer flow appears to possess two different time-scales: in a first phase, tokens quickly form a few clusters, while in a second (much slower) phase, through the process of pairwise merging of clusters, all tokens finally collapse to a single point. This appears to corroborate behavior observed empirically in trained Transformer models, which goes under the names token uniformity, over-smoothing [CZC ⁺ 22, RZZD23, GWDW23, WAWJ24, WAW ⁺ 24, DBK24, SWJS24], or rank-collapse [DCL21, FZH ⁺ 22, NAB ⁺ 22, JDB23, ZMZ ⁺ 23, ZLL ⁺ 23, NLL ⁺ 24, BHK24, CNQG24]; see also Figure 1.

The goal of this manuscript is twofold. On the one hand, we aim to provide a general and accessible framework to study Transformers from a mathematical perspective. In particular, the structure of these interacting particle systems enables concrete connections to established mathematical topics, such as nonlinear transport equations, Wasserstein gradient flows, collective behavior models, and optimal configurations of points on spheres, among others. On the other hand, we describe several promising research directions with a particular focus on the long-time clustering phenomenon. The main results we present are new, and we also provide what we believe are interesting open problems throughout the paper.

The rest of the paper is arranged in three parts.

Part: Modeling

We define an idealized model of the Transformer architecture that captures two of the main characteristics of transformers: self-attention and layer-normalization. Following a perspective put forward in classical architectures such as ResNets [CRBD18, E17, HR17], we view the successive layers of a neural network as time discretizations of a dynamical system of interacting particles. Layer-normalization effectively constrains particles to evolve on the unit sphere \mathbb{S}^{d-1} , whereas self-attention is the particular nonlinear coupling of the particles done through the empirical measure (Section 2). In turn, the empirical measure evolves according to the continuity partial differential equation (Section 3). Even after significant simplifications, this toy model retains macroscopic characteristics of trained Transformers, namely clustering. We also introduce a simpler surrogate model for self-attention which has the convenient property of being a Wasserstein gradient flow [AGS05] for an energy functional that is well-studied in the context of optimal configurations of points on the sphere and sheds a complementary light of the source of clustering.

Part: Clustering

In this part we recall existing and establish new mathematical results that indicate clustering of tokens in the long-time limit. (See Figure 2 for a summary.) Our first results (Theorem 4.3 in Section 4 and Theorem 5.1 in Section 5) are in extreme regimes of a temperature parameter \beta^{-1} that appears in the equation. We then move to the high-dimensional case in Section 6, where we begin by recalling Theorem 6.1 —a result of [MTG17], recently revisited in [CRMB24] —which entails long-time clustering at any temperature when d\geq 3 . We provide an exponential rate of convergence in Theorem 6.3 when d\geq n —here n denotes the number of particles—. We complement this result with an even more precise characterization of the rate of contraction of particles into a cluster. Namely, we describe the histogram of all inter-particle distances, and the time at which all particles are already nearly clustered (Theorem 6.9).

Part: Further questions

We propose potential avenues for future research, largely in the form of open questions substantiated by numerical observations. We first focus on the case d=2 (Section 7) and elicit a link to Kuramoto oscillators. We briefly show in Section 9.1 how a simple and natural modification of our model leads to non-trivial questions related to optimal configurations on the sphere. The remaining sections explore interacting particle systems that allow for parameter tuning of the Transformers architectures, a key feature of practical implementations.

Part I Modeling

We begin this part by presenting the mathematical model for a Transformer in Section 2. Throughout the paper we focus on a simplified version that includes the self-attention mechanism as well as layer normalization (Section 2.2), but excludes additional feed-forward layers commonly used in practice (Section 2.3). This nonetheless leads to a highly nonlinear mean-field interacting particle system. In turn, this system implements, via the continuity equation, a flow map from initial to terminal distributions of particles that we present in Section 3.

2. Interacting particle system

Before writing down the Transformer model, we first provide a brief preliminary discussion to clarify our methodological choice of treating the discrete layer indices in the model as a continuous time variable in Section 2.1, echoing previous work on ResNets. The specifics of the toy Transformer model are presented in Section 2.2, and a complete model is presented in Section 2.3.

2.1. Residual neural networks

One of the standard paradigms in machine learning is that of supervised learning, where one aims to approximate an unknown function f:\mathbb{R}^{d}\to\mathbb{R}^{m} , from data, \mathscr{D}=\{x^{(i)},f(x^{(i)})\}_{i\in[N]} say. This is typically done by choosing one among an arsenal of possible parametric models, whose parameters are then fit to the data by means of minimizing some user-specified cost. With the advent of graphical processing units (GPUs) in the realm of computer vision [KSH12], large neural networks have become computationally accessible, resulting in their popularity as one such parametric model.

Within the class of neural networks, residual neural networks (ResNets for short) have become a staple DNN architecture since their introduction in [HZRS16]. In their most basic form, ResNets approximate a function f at x\in\mathbb{R}^{d} through a sequence of affine transformations, a component-wise nonlinearity, and skip connections. Put in formulae,

(2.1)

{ x (k + 1) = x (k) + w (k) σ (a (k) x (k) + b (k)) for k ∈ { 0,…,L − 1 } x (0) = x.\begin{dcases}x(k+1)=x(k)+w(k)\sigma(a(k)x(k)+b(k))&\text{ for }k\in\{0,\ldots,L-1\}\\ x(0)=x\,.\end{dcases}

Here \sigma is a Lipschitz function applied component-wise to the input vector, while \theta(\cdot)=(w(\cdot),a(\cdot),b(\cdot))\in\mathbb{R}^{d\times d}\times\mathbb{R}^{d\times d}\times\mathbb{R}^{d} are trainable parameters. We say that (2.1) has L\geq 1 hidden layers (or L+1 layers, or is of depth L ). The output x_{i}(L) serves as a representation of the input x^{(i)} that is then fed into a last layer that corresponds to a classical learning task such as linear or logistic regression in order to predict the label f(x^{(i)}) . One can also devise generalizations of (2.1), for instance in which matrix-vector multiplications are replaced by discrete convolutions in order to reflect other common architectures such as convolutional neural networks [GBC16, Chapter 9]. The key element that all these models share is that they all have skip-connections, namely, the previous step x_{i}(k) appears explicitly in the iteration for the next one.

One upside of (2.1), which is the one of interest to our narrative, is that the layer index k can naturally be interpreted as a time variable, motivating the continuous-time analogue

(2.2)

{ x ˙ (t) = w (t) σ (a (t) x (t) + b (t)) for t ∈ (0,T) x (0) = x.\begin{dcases}\dot{x}(t)=w(t)\sigma(a(t)x(t)+b(t))&\text{ for }t\in(0,T)\\ x(0)=x.\end{dcases}

These are dubbed neural ordinary differential equations (neural ODEs). Since their introduction in [CRBD18, E17, HR17], neural ODEs have emerged as a flexible mathematical framework to implement and study ResNets. We use this abstraction here for convenience, as dynamical systems are more naturally defined in continuous time. Although this approach is helpful for exposition, we emphasize that all the results presented can also be derived in discrete time using the same techniques.

2.2. The interacting particle system

Unlike ResNets, which operate on a single input vector x(0)\in\mathbb{R}^{d} at a time, Transformers operate on a sequence of vectors of length n , namely, (x_{i}(0))_{i\in[n]}\in(\mathbb{R}^{d})^{n} . This perspective is rooted in natural language processing, where each vector represents a word, and the entire sequence a sentence or a paragraph. In particular, it allows to process words together with their context. A sequence element x_{i}(0)\in\mathbb{R}^{d} is called a token, and the entire sequence (x_{i}(0))_{i\in[n]} a prompt. We use the words “token” and “particle” interchangeably.

All practical implementations of Transformers make use of layer normalization [BKH16], most commonly in the form of root mean square (RMS) normalization [ZS19]. (See lines 105–116 in [Mis24] for instance.) RMS normalization takes the sequence of tokens output after each layer, divides each token by its Euclidean norm ¹ ¹ 1 The original form instead consisted in an entry-wise standardization of every token, namely subtracting the mean of all tokens, then dividing by the standard deviation. (plus a small parameter to avoid a possible division by zero), and multiplies the result by a trained diagonal matrix. This process aims to ensure that tokens do not diverge, thus avoiding rounding errors and overflow. The result is an evolution on a time-varying axis-aligned ellipsoid. To simplify the presentation and obtain insight and precise results, we assume that the trained diagonal matrix is equal to the idenitity, so that we work on the unit sphere \mathbb{S}^{d-1} throughout. This simplification is justified empirically in the trained ALBERT XLarge v2 model described in Figure 1, wherein this diagonal matrix is constant over all layers, with entries of mean value equal to 0.44 and standard deviation equal to 0.008 . Furthermore, current text embeddings provided by OpenAI, namely text-embedding-3-small and text-embedding-3-large, return norm-one embedding vectors. While we can only speculate as to the actual implementation of these models, this is an indication that layer normalization could be as simple as the one used in our toy model.

A Transformer is then a flow map on (\mathbb{S}^{d-1})^{n} : the input sequence (x_{i}(0))_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} is an initial condition which is evolved through the dynamics

(2.3)

x ˙ i (t) = 𝗣 x i (t) ⟂ (1 Z β,i (t) ∑ j = 1 n e β ⟨ Q (t) x i (t),K (t) x j (t) ⟩ V (t) x j (t)) \dot{x}_{i}(t)=\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}\left(\frac{1}{Z_{\beta,i}(t)}\sum_{j=1}^{n}e^{\beta\langle Q(t)x_{i}(t),K(t)x_{j}(t)\rangle}V(t)x_{j}(t)\right)

for all i\in[n] and t\geq 0 . (We refer the reader to (2.8) and Section 2.3 for the full model.) Here and henceforth

	`\bm{\mathsf{P}}^{\perp}_{x}y=y-\langle x,y\rangle x`

denotes the projection of y\in\mathbb{R}^{d} onto \mathrm{T}_{x}\mathbb{S}^{d-1} . The partition function Z_{\beta,i}(t)>0 reads

(2.4)

Z β,i (t) = ∑ k = 1 n e β ⟨ Q (t) x i (t),K (t) x k (t) ⟩.Z_{\beta,i}(t)=\sum_{k=1}^{n}e^{\beta\langle Q(t)x_{i}(t),K(t)x_{k}(t)\rangle}.

where (Q(\cdot),K(\cdot),V(\cdot)) (standing for Query, Key, and Value) are parameter matrices learned from data, and \beta>0 a fixed number intrinsic to the model ² ² 2 In practical implementations the inner products are multiplied by d^{-\frac{1}{2}} , which along with the typical magnitude of Q,K leads to the appearance of \beta ., which, can be seen as an inverse temperature using terminology from statistical physics. Note that Q(\cdot),K(\cdot) need not be square.

The interacting particle system (2.3)–(2.4), a simplified version of which was first written down in [LLH ⁺ 20, DGCC21, SABP22], importantly contains the true novelty that Transformers carry with regard to other models: the self-attention mechanism

(2.5)

A i j (t):= e β ⟨ Q (t) x i (t),K (t) x j (t) ⟩ Z β,i (t),(i,j) ∈ [n] 2,A_{ij}(t):=\frac{e^{\beta\langle Q(t)x_{i}(t),K(t)x_{j}(t)\rangle}}{Z_{\beta,i}(t)}\,,\hskip 28.45274pt(i,j)\in[n]^{2},

which is the nonlinear coupling mechanism in the interacting particle system. The n\times n stochastic matrix A(t) (rows are probability vectors) is called the self-attention matrix. The wording attention stems from the fact that A_{ij}(t) captures the attention given by particle i to particle j relatively to all particles \ell\in[n] . In particular, a particle pays attention to its neighbors where neighborhoods are dictated by the matrices Q(t) and K(t) in (2.5).

It has been observed numerically that the probability vectors (A_{ij}(\cdot))_{j\in[n]} (for i\in[n] ) in a trained self-attention matrix exhibit behavior related to the syntactic and semantic structure of sentences in natural language processing tasks (see [VSP ⁺ 17, Figures 3-5]). To illustrate our conclusions as pedagogically as possible, throughout the paper we focus on a simplified scenario wherein the parameter matrices (Q,K,V) are constant, and even all equal to the identity unless stated otherwise, resulting in the dynamics

(SA)

x ˙ i (t) = 𝗣 x i (t) ⟂ (1 Z β,i (t) ∑ j = 1 n e β ⟨ x i (t),x j (t) ⟩ x j (t)) \dot{x}_{i}(t)=\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}\left(\frac{1}{Z_{\beta,i}(t)}\sum_{j=1}^{n}e^{\beta\langle x_{i}(t),x_{j}(t)\rangle}x_{j}(t)\right)

for i\in[n] and t\geq 0 and, as before

(2.6)

Z β,i (t) = ∑ k = 1 n e β ⟨ x i (t),x k (t) ⟩.Z_{\beta,i}(t)=\sum_{k=1}^{n}e^{\beta\langle x_{i}(t),x_{k}(t)\rangle}.

The appearance of clusters in Transformers is actually corroborated by numerical experiments with trained models (see Figure 1). While we focus on a much simplified model, numerical evidence shows that the clustering phenomenon looks qualitatively the same in the cases Q=K=V=\lambda I_{d} , \lambda>0 , and generic random (Q,K,V) (see Figures 3 and 5 for instance). We refer the interested reader directly to Sections 4, 5, 6; here, we continue the presentation on the modeling of different mechanisms appearing in the Transformer architecture.

Figure 1. Histogram of { ⟨ x i ( t ), j ⟩ } ∈ [ n ] 2 ≠ \{\langle x_{i}(t),x_{j}(t)\rangle\}_{(i,j)\in[n]^{2},i\neq j} at different layers in the context of the trained ALBERT XLarge v2 model ( LCG + 20 and https://huggingface.co/albert-xlarge-v2 4 contains all the mechanisms described in this text, namely, is a system of the form ( 2.8 ) (or rather the discretization thereof) with 12 or 24 layers. The sequence length is of the order of 512 1024, and the tokens evolve in ℝ 4096 \mathbb{R}^{4096}. The dynamics are therefore high-dimensional, lending weight to assumptions made later on (Section 6 )., which has constant parameter matrices. Here we randomly selected a single prompt, which in this context is a paragraph from a random Wikipedia entry, and then generate the histogram of the pairwise inner products. We see the progressive emergence of clusters all the way to the th (and last) hidden layer (top), as evidenced by the growing mass at 1. If the number of layers is increased, up to 48 say, the clustering is further enhanced (bottom).

Remark 2.1 (Collective behavior).

The dynamics (SA) have a strong resemblance to the vast literature on nonlinear systems arising in the modeling of collective behavior. In addition to the connection to the classical Kuramoto model describing synchronization of oscillators [Kur75, ABV ⁺ 05] (made evident in Section 7.2), Transformers are perhaps most similar to the Krause model [Kra00]

	`\dot{x}_{i}(t)=\sum_{j=1}^{n}a_{ij}(x_{j}(t)-x_{i}(t)),\hskip 28.45274pta_{ij}=\frac{\phi(\\\|x_{i}-x_{j}\\\|^{2})}{\sum_{k=1}^{n}\phi(\\\|x_{i}-x_{k}\\\|^{2})}.`

which is non-symmetric in general ( a_{ij}\neq a_{ji} ), much like (2.3). When \phi is compactly supported, it has been shown in [JM14] that the particles x_{i}(t) assemble in several clusters as t\rightarrow+\infty . Other related models include those of Vicsek [VCBJ ⁺ 95], Hegselmann-Krause [HK02] and Cucker-Smale [CS07]. All these models exhibit a clustering behavior under various assumptions (see [MT14, Tad23] and the references therein). Yet, none of the opinion dynamics models discussed above contain parameters appearing within the nonlinearity as in (SA), whilst set on the sphere.

Remark 2.2 (Permutation equivariance).

A function f:(\mathbb{S}^{d-1})^{n}\rightarrow(\mathbb{S}^{d-1})^{n} is permutation equivariant if f(\pi X)=\pi(f_{1}(X),\ldots,f_{n}(X)) for any X\in(\mathbb{R}^{d})^{n} and for any permutation \pi\in\mathbf{S}_{n} of n elements. Otherwise put, if we permute the input X , then the output f(X) is permuted in the same way. Given t>0 , the Transformer (SA), mapping (x_{i}(0))_{i\in[n]}\mapsto(x_{i}(t))_{i\in[n]} , is permutation-equivariant on (\mathbb{S}^{d-1})^{n} .

2.3. Toward the complete Transformer

There are a couple of additional mechanisms used in practical implementations that we do not explicitly address or use in this study. The mathematical analysis of these mechanisms remains open.

2.3.1. Multi-headed attention

Practical implementations spread out the computation of the self-attention mechanism at every t through a sequence of heads, leading to the so-called multi-headed self attention. This consists in considering the following modification to (SA):

(2.7)

x ˙ i (t) = 𝗣 x i (t) ⟂ (∑ h = 1 H ∑ j = 1 n e β ⟨ Q h (t) x i (t),K h (t) x j (t) ⟩ Z β,i,h (t) V h (t) x j (t)) \dot{x}_{i}(t)=\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}\left(\sum_{h=1}^{H}\sum_{j=1}^{n}\frac{e^{\beta\langle Q_{h}(t)x_{i}(t),K_{h}(t)x_{j}(t)\rangle}}{Z_{\beta,i,h}(t)}V_{h}(t)x_{j}(t)\right)

where Z_{\beta,i,h}(t) is defined as in (2.4) for the matrices Q_{h}(t) and K_{h}(t) . The integer H\geq 1 is called the number of heads ⁵ ⁵ 5 In practical implementations, H is a divisor of d , and the query and key matrices Q_{h}(t) and K_{h}(t) are \frac{d}{H}\times d rectangular. This allows for further parallelization of computations and increased expressiveness. For mathematical purposes, we focus on working with arbitrary integers H , and square weight matrices Q_{h} and K_{h} ..

The introduction of multiple heads also allows for drawing some interesting parallels with the literature on feed-forward neural networks, such as ResNets (2.1). Considerable effort has been expended to understand 2 -layer neural networks with width tending to +\infty ; more precisely, consider (2.1) with L=1 , w\in\mathbb{R}^{d\times\ell} , a\in\mathbb{R}^{\ell\times d} , and \ell\to+\infty . The infinite-width limit for Transformers is in fact very natural, as it is realized by stacking an arbitrary large number of heads: H\to+\infty . Hence, the same questions as for 1 -hidden layer neural networks may be asked: for instance the question of universal approximation, in the vein of [Cyb89, Bar93].

2.3.2. Feed-forward layers

The complete Transformer dynamics combines all of the above mechanisms with a feed-forward layer; in the discrete-time context, this is actually done by using a Lie-Trotter splitting scheme for

(2.8)

x ˙ i (t) = 𝗣 x i (t) ⟂ (∑ h = 1 H ∑ j = 1 n e β ⟨ Q h (t) x i (t),K h (t) x j (t) ⟩ Z β,i,h (t) V h (t) x j (t) + w (t) σ (a (t) x i (t) + b (t))),\dot{x}_{i}(t)=\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}\left(\sum_{h=1}^{H}\sum_{j=1}^{n}\frac{e^{\beta\langle Q_{h}(t)x_{i}(t),K_{h}(t)x_{j}(t)\rangle}}{Z_{\beta,i,h}(t)}V_{h}(t)x_{j}(t)+w(t)\sigma(a(t)x_{i}(t)+b(t))\right),

where w(t),a(t),b(t) and \sigma are all as in (2.2). The interested reader is referred to [LLH ⁺ 20, PH22] for all the details ⁶ ⁶ 6 and lines 123–130 in [Ope24] for some relevant source code.. The feed-forward layers (convolutional layers can alternatively be considered) are of critical importance in applications and drive the existing results on approximation properties of Transformers [YBR ⁺ 19]. Nevertheless, the analysis of this model is beyond the scope of our current methods.

3. Measure to measure flow map

An important aspect of Transformers is that they are not hard-wired to take into account the order of the input sequence, contrary to other architectures used for natural language processing such as recurrent neural networks. In these applications, each token x_{i}(0)\in\mathbb{R}^{d} contains not only a word embedding w_{i}\in\mathbb{R}^{d} , but also an additional positional encoding (we postpone a discussion to Remark 3.2) which allows tokens to also carry their position in the input sequence. Therefore, an input sequence is perfectly encoded as a set of tokens \{x_{1}(0),\ldots,x_{n}(0)\} , or equivalently as the empirical measure of its constituent tokens \frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}(0)} . Recall that the output of a Transformer is also a probability measure, namely \frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}(t)} , albeit one that captures the likelihood of the next token. As a result, one can view Transformers as flow maps between probability measures ⁷ ⁷ 7 See [DBPC19, VBC20, ZB21] for further related work on neural networks acting on probability measures. on \mathbb{S}^{d-1} . To describe this flow map, we appeal to the continuity equation, which governs precisely the evolution of the empirical measure of particles subject to dynamics. This perspective is already present in [SABP22], the only modification here being that we add the projection on the sphere arising from layer normalization.

After introducing the continuity equation in Section 3.1, we show that a particular interaction energy functional, which is maximized at any point mass, increases along solutions thereof in Section 3.2. Motivated by this monotonicity property, in Section 3.3 we propose an illustrative modified model which has the nice property of being a Wasserstein gradient flow for this energy. Finally, in Section 3.4, we demonstrate that the original equation presented in Section 3.1 is itself a gradient flow for the same energy, upon changing the metric underlying the definition of the gradient.

3.1. The continuity equation

The vector field driving the evolution of a single particle in (SA) clearly depends on all n particles. In fact, one can equivalently rewrite the dynamics as

(3.1)

x ˙ i (t) = 𝒳 [μ (t)] (x i (t)) \dot{x}_{i}(t)=\mathcal{X}[\mu(t)](x_{i}(t))

for all i\in[n] and t\geq 0 , where

	`\mu(t,\cdot)=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}(t)}(\cdot)`

is the empirical measure, while the vector field \mathcal{X}[\mu]:\mathbb{S}^{d-1}\to\mathrm{T}\mathbb{S}^{d-1} reads

(3.2)

𝒳 [μ] (x) = 𝗣 x ⟂ (1 Z β,μ (x) ∫ e β ⟨ x,y ⟩ y d μ (y)) \mathcal{X}[\mu](x)=\bm{\mathsf{P}}^{\perp}_{x}\left(\frac{1}{Z_{\beta,\mu}(x)}\int e^{\beta\langle x,y\rangle}y\,\mathrm{d}\mu(y)\right)

with

(3.3)

Z β,μ (x) = ∫ e β ⟨ x,y ⟩ d μ (y).Z_{\beta,\mu}(x)=\int e^{\beta\langle x,y\rangle}\,\mathrm{d}\mu(y).

In other words, (SA) is a mean-field interacting particle system. The evolution of \mu(t) is governed by the continuity equation ⁸ ⁸ 8 Unless stated otherwise, \nabla and \mathrm{div} henceforth stand for the spherical gradient and divergence respectively, and all integrals are taken over \mathbb{S}^{d-1} .

(3.4)

{ ∂ t μ + div (𝒳 [μ] μ) = 0 on ℝ ≥ 0 × 𝕊 d − 1 μ | t = 0 = μ (0) on 𝕊 d − 1 \begin{cases}\partial_{t}\mu+\mathrm{div}(\mathcal{X}[\mu]\mu)=0&\text{ on }\mathbb{R}_{\geq 0}\times\mathbb{S}^{d-1}\\ \mu_{|t=0}=\mu(0)&\text{ on }\mathbb{S}^{d-1}\end{cases}

satisfied in the sense of distributions.

Remark 3.1 (Well-posedness).

Global existence of weak, measure-valued solutions to (3.4) for arbitrary initial conditions \mu(0)\in\mathcal{P}(\mathbb{S}^{d-1}) follows by arguing exactly as in [GLPR24, Lemma A.3]. Here and henceforth, \mathcal{P}(\mathbb{S}^{d-1}) stands for the set of Borel probability measures on \mathbb{S}^{d-1} .

Remark 3.2 (Positional encoding).

For the sake of completeness, in this brief segue we discuss a few ways to perform positional encoding. The original one, proposed in [VSP ⁺ 17], proceeds as follows. Consider a sequence (w_{i})_{i\in[n]}\in(\mathbb{R}^{d})^{n} of word embeddings. Then the positional encoding p_{i}\in\mathbb{R}^{d} of the i -th word embedding is defined as (p_{i})_{2k}=\sin(\frac{i}{M^{2k/d}}) and (p_{i})_{2k+1}=\cos(\frac{i}{M^{2k/d}}) for k\in[d/2-1] , and M>0 is a user-defined scalar equal to 10^{4} in [VSP ⁺ 17]. The i -th token is then defined as the addition: x_{i}(0)=w_{i}+p_{i} . Subsequent works simply use either a random ⁹ ⁹ 9 This rationale supports the assumption that initial tokens are drawn at random, which we make use of later on. positional encoding (i.e., p_{i} is just some random vector) or a trained transformation. The addition can also be replaced with a concatenation x_{i}(0)=[w_{i};p_{i}] . (See [LWLQ22, XZ23] for details.)

Remark 3.3 (Mean field limit).

Although the analysis in this paper is focused on the flow of the empirical measure, one can also consider (3.4) for arbitrary initial probability measures \mu(0)\in\mathcal{P}(\mathbb{S}^{d-1}) . Both views can be linked through a mean-field limit-type result, which can be shown by making use of the Lipschitz nature of the vector field \mathcal{X}[\mu] . The argument is classical and dates back at least to the work of Dobrushin [Dob79]. Consider an initial empirical measure \mu_{n}(0)=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}(0)} , and suppose that the points x_{i}(0) are such that \lim_{n\to+\infty}W_{1}(\mu_{n}(0),\mu(0))=0 for some probability measure \mu(0)\in\mathcal{P}(\mathbb{S}^{d-1}) . (Here W_{1} denotes 1-Wasserstein distance—see [Vil09] for definitions—.) Consider the solutions \mu_{n}(t) and \mu(t) to (3.4) with initial data \mu_{n}(0) and \mu(0) respectively. Dobrushin’s argument is then centered around the estimate

	`W_{1}(\mu_{n}(t),\mu(t))\leq e^{O(1)\\|t\\|}W_{1}(\mu_{n}(0),\mu(0))`

for any t\in\mathbb{R} , which in the case of (3.4) can be shown without much difficulty (see [Vil01, Chapter 4, Section 1] or [Gol16, Section 1.4.2]). This elementary mean-field limit result has a couple of caveats. First, the time-dependence is exponential. Second, if one assumes that the points x_{i}(0) are sampled i.i.d. according to \mu_{0} , then W_{1}(\mu_{n}(0),\mu(0)) converges to zero at rate n^{-\frac{1}{d-1}} [Dud69, BLG14], which deteriorates quickly when d grows. Dimension-free convergence has been established in some cases, for instance by replacing the Wasserstein distance with a more careful choice of metric as in [HHL23, Lac23] or more generally in [SFG ⁺ 12]. Similarly, the exponential time-dependence might also be improved, as recent works in the context of flows governed by Riesz/Coulomb singular kernels, with diffusion, can attest [RS23, GLBM24] (see [LLF23] for a result in the smooth kernel case). We do not address this question in further detail here. For more references on this well-established topic, the reader is referred to [Vil01, Gol16, Ser20] and the references therein.

3.2. The interaction energy

One can naturally ask whether the evolution in (3.4) admits some quantities which are monotonic when evaluated along the flow. As it turns out, the interaction energy

(3.5)

𝖤 β [μ] = 1 2 β ∬ e β ⟨ x,x ′ ⟩ d μ (x) d μ (x ′) \mathsf{E}_{\beta}[\mu]=\frac{1}{2\beta}\iint e^{\beta\langle x,x^{\prime}\rangle}\,\mathrm{d}\mu(x)\,\mathrm{d}\mu(x^{\prime})

is one such quantity. Indeed,

	d d t 𝖤 β [μ (t)] \displaystyle\frac{\,\mathrm{d}}{\,\mathrm{d}t}\mathsf{E}_{\beta}[\mu(t)]	= ∬ β − 1 e β ⟨ x,x ′ ⟩ d ∂ t μ (t,x) d μ (t,x ′) \displaystyle=\iint\beta^{-1}e^{\beta\langle x,x^{\prime}\rangle}\,\mathrm{d}\partial_{t}\mu(t,x)\,\mathrm{d}\mu(t,x^{\prime})
		= ∫ 𝒳 [μ (t)] (x) ⋅ ∫ ∇ (β − 1 e β ⟨ x,x ′ ⟩) ⁡ d μ (t,x ′) d μ (t,x) \displaystyle=\int\mathcal{X}[\mu(t)](x)\cdot\int\nabla\left(\beta^{-1}e^{\beta\langle x,x^{\prime}\rangle}\right)\,\mathrm{d}\mu(t,x^{\prime})\,\mathrm{d}\mu(t,x)
(3.6)			= ∫ ‖ 𝒳 [μ (t)] (x) ‖ 2 Z β,μ (t) (x) d μ (t,x) \displaystyle=\int\Big{\\|}\mathcal{X}[\mu(t)](x)\Big{\\|}^{2}Z_{\beta,\mu(t)}(x)\,\mathrm{d}\mu(t,x)

for any t\geq 0 by using integration by parts. Recalling the definition of Z_{\beta,\mu}(x) in (3.3), we see that e^{-\beta}\leq Z_{\beta,\mu}(x)\leq e^{\beta} for all x\in\mathbb{S}^{d-1} . The identity (3.2) therefore indicates that \mathsf{E}_{\beta} increases along trajectories of (3.4). (Similarly, should V=-I_{d} , the energy \mathsf{E}_{\beta} would decrease along trajectories.) This begs the question of characterizing the global minima and maxima of \mathsf{E}_{\beta} , which is the goal of the following result.

Proposition 3.4.

Let \beta>0 and d\geq 2 . The unique global minimizer of \mathsf{E}_{\beta} over \mathcal{P}(\mathbb{S}^{d-1}) is the uniform measure ¹⁰ ¹⁰ 10 That is, the Lebesgue measure on \mathbb{S}^{d-1} , normalized to be a probability measure. \sigma_{d} . Any global maximizer of \mathsf{E}_{\beta} over \mathcal{P}(\mathbb{S}^{d-1}) is a Dirac mass \delta_{x^{*}} centered at some point x^{*}\in\mathbb{S}^{d-1} .

This result lends credence to our nomenclature of the case V=I_{d} as attractive, and V=-I_{d} as repulsive. The reader should be wary however that in this result we are minimizing or maximizing \mathsf{E}_{\beta} among all probability measures on \mathbb{S}^{d-1} . Should one focus solely on discrete measures, many global minima appear—these are discussed in Section 9.1 —. This is one point where the particle dynamics and the mean-field flow deviate. We now provide a brief proof of Proposition 3.4 (see [Tan17] for a different approach).

Proof of.

The fact that any global maximizer is a Dirac mass is easy to see. We proceed with proving the rest of the statement. Let f(t)=e^{\beta t} . The interaction energy then reads

	`\mathsf{E}_{\beta}[\mu]=\frac{1}{2}\iint f(\langle x,x^{\prime}\rangle)\,\mathrm{d}\mu(x)\,\mathrm{d}\mu(x^{\prime}).`

The proof relies on an ultraspherical (or Gegenbauer) polynomial expansion of f(t) :

	`f(t)=\sum_{k=0}^{+\infty}\widehat{f}(k;\lambda)\frac{k+\lambda}{\lambda}C_{k}^{\lambda}(t)`

for t\in[-1,1] , where \lambda=\frac{d-2}{2} , C_{k}^{\lambda} are Gegenbauer polynomials, and

	`\widehat{f}(k;\lambda)=\frac{\Gamma(\lambda+1)}{\Gamma(\lambda+\frac{1}{2})\Gamma(\frac{1}{2})}\frac{1}{C_{k}^{\lambda}(1)}\int_{-1}^{1}f(t)C^{\lambda}_{k}(t)(1-t^{2})^{\lambda-\frac{1}{2}}\,\mathrm{d}t`

where C_{k}^{\lambda}(1)>0 (see [DX13, Section 1.2]). According to [BD19, Proposition 2.2], a necessary and sufficient condition for Proposition ˜ 3.4 to hold is to ensure that \widehat{f}(k;\lambda)>0 for all k\geq 1 . To show this, we use the Rodrigues formula [Sze39, 4.1.72]

	`C_{k}^{\lambda}(t)=\frac{(-1)^{k}2^{k}}{k!}\frac{\Gamma(k+\lambda)\Gamma(k+2\lambda)}{\Gamma(\lambda)\Gamma(2k+2\lambda)}(1-t^{2})^{-(\lambda-\frac{1}{2})}\left(\frac{\,\mathrm{d}}{\,\mathrm{d}t}\right)^{k}(1-t^{2})^{k+\lambda-\frac{1}{2}},`

and the fact that C_{k}^{\lambda}(-t)=(-1)^{k}C_{k}^{\lambda}(t) for t\in[-1,1] , which in combination with integration by parts yield

	`\int_{-1}^{1}t^{\ell}C_{k}^{\lambda}(t)(1-t^{2})^{\lambda-\frac{1}{2}}\,\mathrm{d}t\begin{cases}>0&\text{ if }\ell\geq k\text{ and }\ell-k\text{ is even}\\ =0&\text{ otherwise }.\end{cases}`

We conclude by using the power series expansion of f . ∎

3.3. A Wasserstein gradient flow proxy

In view of (3.2), one could hope to see the continuity equation (3.4) as the Wasserstein gradient flow of \mathsf{E}_{\beta} , or possibly some other functional (see the seminal papers [Ott01, JKO98], and [AGS05, Vil09] for a complete treatment). The long-time asymptotics of the PDE can then be analyzed by studying convexity properties of the underlying functional, by analogy with gradient flows in the Euclidean case.

For (3.4) to be the Wasserstein gradient flow of \mathsf{E}_{\beta} , the vector field \mathcal{X}[\mu] defined in (3.2) ought to be the gradient of the first variation \delta\mathsf{E}_{\beta} of \mathsf{E}_{\beta} . However, notice that \mathcal{X}[\mu] is a logarithmic derivative:

(3.7)

𝒳 [μ] (x) = ∇ log ∫ β − 1 e β ⟨ x,y ⟩ d μ (y).\mathcal{X}[\mu](x)=\nabla\log\int\beta^{-1}e^{\beta\langle x,y\rangle}\,\mathrm{d}\mu(y).

(This observation goes beyond Q=K=I_{d} and V=\pm I_{d} insofar as Q^{\top}K=K^{\top}Q=\pm V ; see [SABP22, Assumption 1]. Because of the lack of symmetry, it has been shown in [SABP22] that (3.7) is not the gradient of the first variation of a functional.

To overcome this limitation on \mathbb{R}^{d} , thus without layer normalization, [SABP22] propose two ways to "symmetrize" (3.4) that both lead to a Wasserstein gradient flow; see [SABP22, Proposition 2]. We focus here on the simplest one which consists in removing the logarithm in (3.7), or equivalently to removing the denominator in (3.2). This is one point where working on the unit sphere is useful: otherwise, the equation on \mathbb{R}^{d} without layer normalization (as considered in [SABP22]) is ill-posed for general choices of matrices V , due to the fact that the magnitude of the vector field \mathcal{X}[\mu] grows exponentially with the size of the support of \mu . On the contrary, on \mathbb{S}^{d-1} the resulting equation is perfectly well-posed.

Remark 3.5 (Doubly stochastic kernel).

Considering the Transformer dynamics on \mathbb{R}^{d} , thus without layer normalization, the authors in [SABP22] propose an alternative symmetric model: they replace the self-attention (stochastic) matrix by a doubly stochastic one, generated from the Sinkhorn iteration. This leads to a Wasserstein gradient flow, whereby the resulting attention mechanism is implicitly expressed as a limit of Sinkhorn iterations. Understanding the emergence of clusters for this model is an interesting but possibly challenging question.

In view of the above discussion, we are inclined to propose the surrogate model

(USA)

x ˙ i (t) = 𝗣 x i (t) ⟂ (1 n ∑ j = 1 n e β ⟨ x i (t),x j (t) ⟩ x j (t)),\dot{x}_{i}(t)=\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}\left(\frac{1}{n}\sum_{j=1}^{n}e^{\beta\langle x_{i}(t),x_{j}(t)\rangle}x_{j}(t)\right),

which is obtained by replacing the partition function Z_{\beta,i}(t) by n . As a matter of fact, (USA) presents a remarkably similar qualitative behavior—all of the results we show in this paper are essentially the same for both dynamics—.

The continuity equation corresponding to (USA), namely

(3.8)

{ ∂ t μ (t,x) + div (𝗣 x ⟂ (∫ e β ⟨ x,x ′ ⟩ x ′ d μ (t,x ′)) μ (t,x)) = 0 μ | t = 0 = μ 0 \begin{dcases}\partial_{t}\mu(t,x)+\mathrm{div}\left(\bm{\mathsf{P}}^{\perp}_{x}\left(\int e^{\beta\langle x,x^{\prime}\rangle}x^{\prime}\,\mathrm{d}\mu(t,x^{\prime})\right)\mu(t,x)\right)=0&\\ \mu_{|t=0}=\mu_{0}\end{dcases}

for (t,x)\in\mathbb{R}_{\geq 0}\times\mathbb{S}^{d-1} , can now be seen as a Wasserstein gradient flow for the interaction energy \mathsf{E}_{\beta} defined in (3.5).

Lemma 3.6.

Consider the interaction energy \mathsf{E}_{\beta}:\mathcal{P}(\mathbb{S}^{d-1})\to\mathbb{R}_{\geq 0} defined in (3.5). Then the vector field

	`\mathcal{X}[\mu](x)=\bm{\mathsf{P}}^{\perp}_{x}\left(\int e^{\beta\langle x,x^{\prime}\rangle}x^{\prime}\,\mathrm{d}\mu(x^{\prime})\right)`

satisfies

(3.9)

𝒳 [μ] (x) = ∇ δ 𝖤 β [μ] (x) \mathcal{X}[\mu](x)=\nabla\delta\mathsf{E}_{\beta}[\mu](x)

for any \mu\in\mathcal{P}(\mathbb{S}^{d-1}) and x\in\mathbb{S}^{d-1} , where \delta\mathsf{E}_{\beta}[\mu] denotes the first variation of \mathsf{E}_{\beta} .

We omit the proof which follows from standard Otto calculus [Ott01], [Vil09, Chapter 15], [CNWR24, Chapter 5]. We can actually write (3.9) more succinctly by recalling the definition of the convolution of two functions on \mathbb{S}^{d-1} [DX13, Chapter 2]: for any g\in L^{1}(\mathbb{S}^{d-1}) and f:[-1,1]\to\mathbb{R} such that t\mapsto(1-t^{2})^{\frac{d-3}{2}}f(t) is integrable,

	`(f\ast g)(x)=\int f(\langle x,y\rangle)g(y)\,\mathrm{d}\sigma_{d}(y).`

This definition has a natural extension to the convolution of a function f (with the above integrability) and a measure \mu\in\mathcal{P}(\mathbb{S}^{d-1}) . We can hence rewrite

	`\mathsf{E}_{\beta}[\mu]=\frac{1}{2}\int(\mathsf{G}_{\beta}\ast\mu)(x)\,\mathrm{d}\mu(x)`

where [-1,1]\ni\mathsf{G}_{\beta}(t)=\beta^{-1}e^{\beta t} , and so

	`\mathcal{X}[\mu](x)=\nabla(\mathsf{G}_{\beta}\ast\mu)(x).`

Thus, (3.8) takes the equivalent form

(3.10)

{ ∂ t μ (t,x) + div (∇ (𝖦 β ∗ μ (t,⋅)) ⁡ (x) μ (t,x)) = 0 for (t,x) ∈ ℝ ≥ 0 × 𝕊 d − 1 μ | t = 0 = μ 0 for x ∈ 𝕊 d − 1.\begin{dcases}\partial_{t}\mu(t,x)+\mathrm{div}\Big{(}\nabla\big{(}\mathsf{G}_{\beta}\ast\mu(t,\cdot)\big{)}(x)\mu(t,x)\Big{)}=0&\text{ for }(t,x)\in\mathbb{R}_{\geq 0}\times\mathbb{S}^{d-1}\\ \mu_{|t=0}=\mu_{0}&\text{ for }x\in\mathbb{S}^{d-1}.\end{dcases}

The considerations above lead us to the following Lyapunov identity.

Lemma 3.7.

The solution \mu\in C^{0}(\mathbb{R}_{\geq 0};\mathcal{P}(\mathbb{S}^{d-1})) to (3.8) satisfies

	`\frac{\,\mathrm{d}}{\,\mathrm{d}t}\mathsf{E}_{\beta}[\mu(t)]=\int\left\\\|\nabla\Big{(}\mathsf{G}_{\beta}\ast\mu(t,\cdot)\Big{)}(x)\right\\\|^{2}\,\mathrm{d}\mu(t,x)`

for t\geq 0 .

Interestingly, (3.10) is an aggregation equation, versions of which have been studied in great depth in the literature. For instance, clustering in the spirit of an asymptotic collapse to a single Dirac measure located at the center of mass of the initial density \mu(0,\cdot) has been shown for aggregation equations with singular kernels in [BCM08, BLR11, CDF ⁺ 11], motivated by the Patlak-Keller-Segel model of chemotaxis. Here, one caveat (and subsequently, novelty) is that (3.10) is set on \mathbb{S}^{d-1} which makes the analysis developed in these references difficult to adapt or replicate.

Remark 3.8 (Particle version).

Let us briefly sketch the particle version of the Wasserstein gradient flow (3.8). When \mu(t)=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}(t)} , the interaction energy (3.5) takes the form

	`\mathsf{E}_{\beta}(X)=\frac{1}{2\beta n^{2}}\sum_{i=1}^{n}\sum_{j=1}^{n}e^{\beta\langle x_{i},x_{j}\rangle}`

where X=(x_{1},\ldots,x_{n})\in(\mathbb{S}^{d-1})^{n} . Denoting by \nabla_{X} the gradient associated to the standard Riemannian metric on (\mathbb{S}^{d-1})^{n} , we get the dynamics

(3.11)

X ˙ (t) = n ∇ X 𝖤 β (X (t)).\dot{X}(t)=n\nabla_{X}\mathsf{E}_{\beta}(X(t)).

Indeed, the gradient on (\mathbb{S}^{d-1})^{n} is simply \nabla=(\partial_{1},\ldots,\partial_{n}) where \partial_{i} is the gradient in \mathbb{S}^{d-1} acting on the i -th copy in (\mathbb{S}^{d-1})^{n} . Therefore

	`{\partial_{i}\mathsf{E}_{\beta}}(X(t))=\frac{1}{\beta n^{2}}\sum_{j=1}^{n}{\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}}\left(e^{\beta\langle x_{i}(t),x_{j}(t)\rangle}\beta x_{j}(t)\right)=\frac{1}{n}\dot{x}_{i}(t)`

which yields (3.11).

Note that (SA) also corresponds to a gradient flow of the same interaction energy albeit with respect to a Riemannian metric on the sphere different from the standard one (for n=2 the two are conformally equivalent). We provide more detail in the following section.

3.4. () is a gradient flow for a modified metric

We will now briefly demonstrate that for a particular choice of parameters (Q,K,V) , the true dynamics (SA) can be seen as a gradient flow for \mathsf{E}_{\beta} upon a modification of the metric on the tangent space of (\mathbb{S}^{d-1})^{n} . This will facilitate qualitative analysis later on by using standard tools from dynamical systems. This insight can then be extrapolated to the corresponding continuity equation (3.4) as well, as seen in Section 3.4.2.

3.4.1. The case of particles

We suppose that

	`Q^{\top}K\text{ is symmetric},\hskip 28.45274ptV=Q^{\top}K.`

We define a new metric on (\mathbb{S}^{d-1})^{n} as follows. Let X=(x_{1},\ldots,x_{n})\in(\mathbb{S}^{d-1})^{n} . Consider the inner product on \mathrm{T}_{X}(\mathbb{S}^{d-1})^{n} given by

(3.12)

⟨ (a 1,…,a n),(b 1,…,b n) ⟩ X = ∑ i = 1 n Z β,i (X) ⟨ a i,b i ⟩,\left\langle(a_{1},\ldots,a_{n}),(b_{1},\ldots,b_{n})\right\rangle_{X}=\sum_{i=1}^{n}Z_{\beta,i}(X)\langle a_{i},b_{i}\rangle\,,

where a_{i},b_{i}\in\mathrm{T}_{x_{i}}\mathbb{S}^{d-1} , and

	`Z_{\beta,i}(X)=\sum_{j=1}^{n}e^{\beta\langle Vx_{i},x_{j}\rangle}.`

Set

	`\mathsf{E}_{\beta}(X)=\frac{1}{2\beta}\sum_{i=1}^{n}\sum_{j=1}^{n}e^{\beta\langle Vx_{i},x_{j}\rangle}.`

We now show that the dynamics (2.3) can be equivalently written as

	`\dot{X}(t)=\nabla\mathsf{E}_{\beta}(X(t)),`

where the gradient \nabla is computed with respect to the metric (3.12) on (\mathbb{S}^{d-1})^{n} . To this end, we ought to show that for all vector fields Y on (\mathbb{S}^{d-1})^{n} and for all X\in(\mathbb{S}^{d-1})^{n} ,

(3.13)

d d t | t = 0 𝖤 β (Φ Y t (X)) = ⟨ Y (X),B (X) ⟩ X \frac{\,\mathrm{d}}{\,\mathrm{d}t}\Big{|}_{t=0}\mathsf{E}_{\beta}\left(\Phi^{t}_{Y}(X)\right)=\langle Y(X),B(X)\rangle_{X}

holds, where \Phi^{t}_{Y} is the flow associated to the vector field Y , whereas B=(B_{1},\ldots,B_{n}) with

	`B_{i}=\bm{\mathsf{P}}^{\perp}_{x_{i}}\left(\frac{1}{Z_{\beta,i}(X)}\sum_{j=1}^{n}e^{\beta\langle Vx_{i},x_{j}\rangle}Vx_{j}\right)\in\mathrm{T}_{x_{i}}\mathbb{S}^{d-1}.`

By linearity, it is sufficient to show (3.13) for vector fields Y of the form

	`Y(X)=(Ax_{1},0,\ldots,0)\in\mathrm{T}_{X}(\mathbb{S}^{d-1})^{n}\,`

where A is an arbitrary non-zero skew-symmetric matrix. Clearly

(3.14)

Φ Y t (X) = (e t A x 1,x 2,…,x n).\Phi_{Y}^{t}(X)=(e^{tA}x_{1},x_{2},\ldots,x_{n}).

One first computes

	`\frac{\,\mathrm{d}}{\,\mathrm{d}t}\Big{\\|}_{t=0}\mathsf{E}_{\beta}\left(\Phi^{t}_{Y}(X)\right)=\sum_{j=1}^{n}e^{\beta\langle Vx_{1},x_{j}\rangle}\langle Ax_{1},Vx_{j}\rangle.`

Now observe that \langle Ax_{1},y\rangle=\langle Ax_{1},z\rangle for all skew-symmetric matrices A if and only if x_{1}(y^{\top}-z^{\top}) is a symmetric matrix. Since \bm{\mathsf{P}}^{\perp}_{x_{1}}=I_{d}-x_{1}x_{1}^{\top} , we see that

	`\sum_{j=1}^{n}e^{\beta\langle Vx_{1},x_{j}\rangle}\langle Ax_{1},Vx_{j}\rangle=\langle Y(X),B(X)\rangle_{X},`

as desired.

3.4.2. The case of measures

The above insight, which consists in reweighing the metric with respect to which the gradient is being taken, can be formally adapted to Wasserstein setting for (3.4)—which we recall, is not a gradient flow for the standard Wasserstein gradient of \mathsf{E}_{\beta} —. But note that (3.4) writes as

(3.15)

∂ t μ (t) + div (∇ δ 𝖤 β [μ (t)] δ 𝖤 β [μ (t)] μ (t)) = 0,\partial_{t}\mu(t)+\mathrm{div}\left(\frac{\nabla\delta\mathsf{E}_{\beta}[\mu(t)]}{\delta\mathsf{E}_{\beta}[\mu(t)]}\mu(t)\right)=0,

with \delta\mathsf{E}_{\beta}[\mu](x)=\int e^{\beta\langle x,y\rangle}\,\mathrm{d}\mu(y) . To avoid technicalities we henceforth only focus on the absolutely continuous case. For a fixed \mu\in\mathcal{P}_{\mathrm{ac}}(\mathbb{S}^{d-1}) , as in the well-known formal Riemannian reinterpretation of the Wasserstein space using Otto calculus [Ott01], [CNWR24, Chapter 5], we consider

	`\mathrm{T}_{\mu}\mathcal{P}_{\mathrm{ac}}(\mathbb{S}^{d-1})=\overline{\left\{\nabla\psi\colon\psi\in C^{\infty}(\mathbb{S}^{d-1})\right\}}^{L^{2}(\mu)},`

which, rather than endowing with the standard formal metric tensor given by the \dot{H}^{1} –inner product

	`\langle\nabla\psi_{1},\nabla\psi_{2}\rangle_{\mu}:=\int\langle\nabla\psi_{1}(x),\nabla\psi_{2}(x)\rangle\,\mathrm{d}\mu(x),`

we instead endow with

	`\langle\nabla\psi_{1},\nabla\psi_{2}\rangle_{\mu,\mathsf{E}_{\beta}}:=\int\langle\nabla\psi_{1}(x),\nabla\psi_{2}(x)\rangle\,\delta\mathsf{E}_{\beta}[\mu](x)\,\mathrm{d}\mu(x).`

Continuing the Riemannian geometry analogy, through this metric tensor we can define a distance between \mu_{0},\mu_{1}\in\mathcal{P}_{\mathrm{ac}}(\mathbb{S}^{d-1}) by solving the variational problem

	`\displaystyle\inf_{(\mu(t),v(t))_{t\in[0,1]}}\Bigg{\{}\int_{0}^{1}\\\|v(t)\\\|^{2}_{\mu(t),\mathsf{E}_{\beta}}\,\mathrm{d}t`	`\displaystyle\colon(\mu,v)\text{ satisfy }\eqref{eq: CE},\mu(0)=\mu_{0},\mu(1)=\mu_{1}\Bigg{\}},`

where

(3.16)

∂ t μ (t,x) + div (v (t,x) μ (t,x)) = 0 on [0,1] × 𝕊 d − 1.\partial_{t}\mu(t,x)+\mathrm{div}(v(t,x)\mu(t,x))=0\hskip 28.45274pt\text{ on }[0,1]\times\mathbb{S}^{d-1}.

This variational problem is a generalization of the celebrated Benamou-Brenier formula [BB00], the value of which we dub W_{2,\mathsf{E}_{\beta}}^{2}(\mu_{0},\mu_{1}) : it is a weighed Wasserstein distance. For a curve of measures (\mu(t))_{t\geq 0} with tangent vectors (v(t))_{t\geq 0} (meaning they solve (3.16)), the Wasserstein gradient \nabla\mkern-10.0mu\nabla_{\mathsf{E}_{\beta}}\mathsf{E}_{\beta} induced by this geometric setup is then the element of \mathrm{T}_{\mu(t)}\mathcal{P}_{\mathrm{ac}}(\mathbb{S}^{d-1}) such that

	`\partial_{t}\mathsf{E}_{\beta}[\mu(t)]=\langle\nabla\mkern-10.0mu\nabla_{\mathsf{E}_{\beta}}\mathsf{E}_{\beta}[\mu(t)],v(t)\rangle_{\mu(t),\mathsf{E}_{\beta}}.`

We can now demonstrate that the vector field driving (3.15) is the Wasserstein gradient of \mathsf{E}_{\beta} corresponding to this geometric setup. Indeed as in [CNWR24, Definition 5.9] we first have

	`\partial_{t}\mathsf{E}_{\beta}[\mu(t)]=\int\delta\mathsf{E}_{\beta}[\mu(t)]\,\mathrm{d}\partial_{t}\mu(t).`

We then find

	`\int\delta\mathsf{E}_{\beta}[\mu(t)]\,\mathrm{d}\partial_{t}\mu(t)=\int\langle\nabla\delta\mathsf{E}_{\beta}[\mu(t)],v(t)\rangle\,\mathrm{d}\mu(t)=\left\langle\frac{\nabla\delta\mathsf{E}_{\beta}[\mu(t)]}{\delta\mathsf{E}_{\beta}[\mu(t)]},v(t)\right\rangle_{\mu(t),\mathsf{E}_{\beta}},`

as desired. The literature studying weighed Wasserstein distances such as the one above is rather scarce, but a relevant reference is [Li21].

Part II Clustering

2 2 3 3 n n d ∗ d^{*} d d 1 1\hskip 4.26773pt n 2 π 2 \frac{n^{2}}{\pi^{2}}\hskip 4.26773pt β \beta fast ← \leftarrow → \rightarrow slow no result when d = 2 d=2

Figure 2. Green zones indicate regimes where convergence to a single cluster as t\to+\infty can be proven. Here n\geq 2 is fixed. When d is larger than specific thresholds, the long-time asymptotics can be chiseled out in finer detail. Convergence is slow when \beta\gg 1 (relative to the size of d,n ), as even the exponential decay constant when d\geq n is of the form \lambda=O(e^{-\beta}) and thus degenerates. One rather expects dynamic metastability. Section 4 addresses the case where \beta is small and Section 5 where it’s large, whereas Section 6 covers the high-dimensional case at arbitrary \beta .

As alluded to in the introductory discussion, clustering is of particular relevance in tasks such as sentiment analysis, masked language modeling, summarization, and so on. Therein, the output measure encodes the probability distribution of the missing tokens for instance, and its clustering indicates a small number of possible outcomes. In Sections 4, 5, and 6, we show several results (mostly summarized in Figure 2) which indicate that the limiting distribution is a point mass. While it may appear that this leaves no room for diversity or randomness, which is at odds with practical observations, these results hold for the specific choice of parameter matrices, and apply in possibly very long-time horizons. Numerical experiments indicate a more subtle picture for different parameters—for instance, there is an appearance of a long metastable phase during which the particles coalesce in a small number of clusters, which seems consistent with behavior in trained models (Figure 1 in Section 6.3)—. We are not able to theoretically explain this behavior as of now.

Ultimately, the appearance of clusters is somewhat natural since the Transformer dynamics are a weighted average of all particles, with the weights being hard-wired to perform a fast selection of particles most similar to the i -th particle being queried. This causes the emergence of leaders which attract all particles in their vicinity. In the natural language processing interpretation, where particles represent tokens, this further elucidates the wording attention as the mechanism of inter-token attraction, and the amplitude of the inner product between tokens can be seen as a measure of their semantic similarity.

4. A single cluster for small β\beta

As seen in Figure 2, the dimension d and inverse temperature \beta appear to play a key role in the clustering results we obtain. In this section and in Section 5, we begin by focusing on extreme choices of \beta whilst d,n are fixed. We first focus on the case \beta=0 (Section 4.1), before moving to the case \beta\ll 1 by a perturbation argument (Section 4.2). We cover the case where \beta is sufficiently large, but finite, in Section 5. The case \beta=+\infty is of little interest since all particles are fixed by the evolution.

4.1. The case β=0\beta=0

For \beta=0 , both (SA) and (USA) read as

(4.1)

x ˙ i (t) = 𝗣 x i (t) ⟂ (1 n ∑ j = 1 n x j (t)),t ⩾ 0.\dot{x}_{i}(t)=\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}\left(\frac{1}{n}\sum_{j=1}^{n}x_{j}(t)\right),\hskip 28.45274ptt\geqslant 0.

The following result shows that generically over the initial points, a single cluster emerges. It complements a known convergence result ([FL19, Theorem 2]) for (4.1). In [FL19, Theorem 2], the authors show convergence to an antipodal configuration, in the sense that n-1 particles converge to some x^{*}\in\mathbb{S}^{d-1} , with the last particle converging to -x^{*} . Moreover, once convergence is shown to hold, it holds with an exponential rate. Mimicking the proof strategy of [BCM15, Theorem 2.2] and [HKR18, Theorem 3.2], we sharpen this result by showing that the appearance of an antipodal particle is non-generic over the choice of initial conditions.

Theorem 4.1.

Let d,n\geq 2 . For Lebesgue almost any initial sequence (x_{i}(0))_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} , there exists some point x^{*}\in\mathbb{S}^{d-1} such that the unique solution (x_{i}(\cdot))_{i\in[n]}\in C^{0}(\mathbb{R}_{\geq 0};(\mathbb{S}^{d-1})^{n}) to the corresponding Cauchy problem for (4.1) satisfies

	`\lim_{t\to+\infty}x_{i}(t)=x^{*}`

for any i\in[n] .

This is also referred to as convergence toward consensus in collective behavior models. We refer the interested reader to Appendix A for the proof, which relies on a gradient flow reinterpretation of (4.1), much like the proofs of several subsequent theorems. We provide some comments in Section 5.

4.2. The case β≪1\beta\ll 1

Theorem 4.1 has some implications for small but positive \beta , something which is already seen in Figure ˜ 3 and Figure ˜ 4. This is essentially due to the fact that, formally,

	`\dot{x}_{i}(t)=\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}\left(\frac{1}{n}\sum_{j=1}^{n}x_{j}(t)\right)+O(\beta)`

for \beta\ll 1 . So, during a time \ll\beta^{-1} , the particles do not feel the influence of the remainder O(\beta) and behave as in the regime \beta=0 . This motivates

Theorem 4.2.

Fix d,n\geq 2 . For \beta\geq 0 , let \mathscr{S}_{\beta}\subset(\mathbb{S}^{d-1})^{n} be the subset consisting of all initial sequences for which the associated solution to the Cauchy problem for (SA) (or (USA)) converges to one cluster as t\rightarrow+\infty . Then

	`\lim_{\beta\to 0}\mathbb{P}(\mathscr{S}_{\beta})=1.`

More generally, if Q and K are arbitrary d\times d matrices, then the same result also holds for the Cauchy problem for (2.3) with V=I_{d} (or the natural analogue of (USA) with these parameters).

Proof.

We focus on the dynamics (SA), but the proof is in fact identical in the case of (USA).

For \alpha\in[0,1) , we say that a set formed from n points z_{1},\ldots,z_{n}\in(\mathbb{S}^{d-1})^{n} is \alpha –clustered if for any i,j\in[n] , \langle z_{i},z_{j}\rangle>\alpha holds. Observe that if \{z_{1},\ldots,z_{n}\} is \alpha –clustered for some \alpha\geq 0 , then the solution to the Cauchy problem for (SA) (for arbitrary \beta\geq 0 ) with this sequence as initial condition converges to a single cluster, since w=z_{1} satisfies the assumption in Lemma ˜ 6.4.

Now, for any integer m\geq 1 , we denote by \mathscr{S}_{0}^{m}\subset\mathscr{S}_{0} the set of initial sequences x_{1}(0),\ldots,x_{n}(0) in (\mathbb{S}^{d-1})^{n} for which the solution (x_{i}^{0}(\cdot))_{i\in[n]} to the associated Cauchy problem for (4.1) is \frac{3}{4} –clustered at time t=m , namely

(4.2)

⟨ x i 0 (m),x j 0 (m) ⟩ > 3 4 \langle x_{i}^{0}(m),x_{j}^{0}(m)\rangle>\frac{3}{4}

holds for all i,j\in[n] . We see that \mathscr{S}_{0}^{m} is an open set for any integer m\geq 1 . Moreover, \mathscr{S}_{0}^{m}\subset\mathscr{S}_{0}^{m+1} according to the proof of Lemma ˜ 6.4, and \bigcup_{m=1}^{+\infty}\mathscr{S}_{0}^{m}=\mathscr{S}_{0} . This implies that

(4.3)

lim m → + ∞ ℙ (𝒮 0 m) = 1.\lim_{m\to+\infty}\mathbb{P}(\mathscr{S}_{0}^{m})=1.

We now show that the solution to (SA) is near that of (4.1), starting from the same initial condition, when \beta is small. Using the Duhamel formula, we find

	`\displaystyle x_{i}^{\beta}(t)-x_{i}^{0}(t)`	`\displaystyle=\int_{0}^{t}\sum_{j=1}^{n}\left(\frac{e^{\beta\langle Qx_{i}^{\beta}(s),Kx_{j}^{\beta}(s)\rangle}}{\sum_{k=1}^{n}e^{\beta\langle Qx_{i}^{\beta}(s),Kx_{k}^{\beta}(s)\rangle}}\right)\bm{\mathsf{P}}^{\perp}_{x_{i}^{\beta}(s)}(x_{j}^{\beta}(s))\,\mathrm{d}s`
		`\displaystyle\quad-\int_{0}^{t}\frac{1}{n}\sum_{j=1}^{n}\bm{\mathsf{P}}^{\perp}_{x_{i}^{0}(s)}(x_{j}^{0}(s))\,\mathrm{d}s`
		`\displaystyle=\int_{0}^{t}\sum_{j=1}^{n}\left(\frac{1}{n}+O\left(\frac{\beta}{n}\right)\right)\bm{\mathsf{P}}^{\perp}_{x_{i}^{\beta}(s)}(x_{j}^{\beta}(s))\,\mathrm{d}s`
		`\displaystyle\quad-\int_{0}^{t}\frac{1}{n}\sum_{j=1}^{n}\bm{\mathsf{P}}^{\perp}_{x_{i}^{0}(s)}(x_{j}^{0}(s))\,\mathrm{d}s,`

where we used that all particles lie on \mathbb{S}^{d-1} for all times. Employing Grönwall, we deduce

(4.4)

‖ x i β (t) − x i 0 (t) ‖ ≤ O (β) e 3 t \displaystyle\left\|x_{i}^{\beta}(t)-x_{i}^{0}(t)\right\|\leq O(\beta)e^{3t}

for all t\geq 0 , \beta\geq 0 and i\in[n] . Due to (4.4), there exists some \beta_{m}>0 such that for any \beta\in[0,\beta_{m}] ,

(4.5)

‖ x i β (m) − x i 0 (m) ‖ ≤ 1 8.\left\|x_{i}^{\beta}(m)-x_{i}^{0}(m)\right\|\leq\frac{1}{8}.

For this to hold, we clearly need \beta_{m}\to 0 as m\to+\infty . Combining (4.2) and (4.5), we gather that for any initial condition in \mathscr{S}_{0}^{m} , the solution (x_{i}^{\beta}(\cdot))_{i\in[n]} to the corresponding Cauchy problem for (SA) is \frac{1}{2} –clustered at time t=m , namely satisfies

	`\langle x_{i}^{\beta}(m),x_{j}^{\beta}(m)\rangle>\frac{1}{2}`

for all i,j\in[n] and \beta\in[0,\beta_{m}] . Thus \mathscr{S}_{0}^{m}\subset\mathscr{S}_{\beta} for any \beta\in[0,\beta_{m}] by virtue of Lemma ˜ 6.4, which together with (4.3) concludes the proof. ∎

In the specific case where Q^{\top}K=I_{d} , we can in fact significantly sharpen Theorem 4.2 by relying on the gradient flow structure evoked in Section 3.4. Namely, we can show the following.

Theorem 4.3.

Fix d,n\geq 2 . There exists a numerical constant C>0 such that whenever

	`\beta\leq Cn^{-1}`

the following holds. For Lebesgue almost any (x_{i}(0))_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} , there exists x^{*}\in\mathbb{S}^{d-1} such that the solution (x_{i}(\cdot))_{i\in[n]}\in C^{0}(\mathbb{R}_{\geq 0};(\mathbb{S}^{d-1})^{n}) to the corresponding Cauchy problem for (SA) (resp. for (USA)) satisfies

	`\lim_{t\to+\infty}x_{i}(t)=x^{*}`

for all i\in[n] .

Moreover, when d=2 one can take \beta\leq 1 and the same conclusion holds.

The proof can be found in Appendix C. As for Theorem 4.1, we briefly discuss the general strategy in Section 5 just below. The improvement to \beta\leq 1 when d=2 is due to the recent paper [CRMB24].

5. A single cluster for large β\beta

The chain of thought leading to the proof of Theorem 4.3 also leads to

Theorem 5.1.

Fix d,n\geq 2 . There exists a constant C=C(d)>0 depending only on d such that whenever

	`\beta\geq Cn^{2},`

the conclusion of Theorem 4.3 holds for both (SA) and (USA).

We refer the interested reader to Appendix B for the proof, which, much like those of Theorem 4.1 and 4.3 relies on the gradient flow interpretation of the dynamics. We give a general outline thereof. Since all the intervening functions and metrics are real-analytic, the celebrated Łojasiewicz theorem [Loj63] implies that for any initial condition, the gradient flow converges to some critical point of \mathsf{E}_{\beta} as t\to+\infty . The genericity over the initial configurations seen in the statements comes from the center-stable manifold theorem (Lemma A.1). The former ensures that for almost every initial configuration, the gradient flow does not converge to a strict saddle point of \mathsf{E}_{\beta} (namely critical points where the Hessian has at least one positive eigenvalue). Whence, generically over the initial configurations, the gradient flow converges to a local maximum. One is then left analyzing the landscape of the underlying energy, with the goal of ensuring that all local maxima are necessarily global.

6. The high-dimensional case

We now elucidate the role that the dimension d plays in clustering results. To start, it turns out that the restrictions on \beta provided by Theorems 4.3 and 5.1 are specific to the two-dimensional case. The following result is shown in [MTG17, CRMB24].

Theorem 6.1 ([, ]).

Fix n\geq 2 , d\geq 3 and \beta\geq 0 . Then the conclusion of Theorem 4.3 holds for both (SA) and (USA).

In Section 6.1 we show that the convergence of Theorem 6.1 is exponentially fast when d\geq n (although, with a decay constant that is exponentially small in \beta ) and in Section 6.2 we describe the full dynamics when n is fixed and d\to+\infty . We discuss some numerical experiments and posit questions on the intermediate behavior of the dynamics when \beta\gg 1 in Section 6.3.

We were made aware of Theorem 6.1 after the first version of this manuscript. The downside of our original proof, which caused us to miss the full range of \beta , is that we focused on d=2 , where the natural perturbations to a critical point involve selecting which particles are moving and which are standing still. However, in d>2 , one can move all particles in the same direction (as is done in the proof in [CRMB24]). Using that there is a continuum of directions and only finitely many points, one can find some direction that brings all of them closer.

Remark 6.2 (Invariant measures).

In the theory of dynamical systems, often the first question of interest is to find smooth invariant measures. It is clear that whenever conclusions of Theorem 4.3 hold (in particular, for all \beta,n when d\geq 3 ), neither (SA) nor (USA) may possess a smooth invariant measure.

6.1. Clustering at an exponential rate when d≥nd\geq n

One can ask whether for almost every initial configuration, the convergence provided by all of the results above holds with some rate. The answer is affirmative—and the rate is in fact exponential—when the initial configuration lies in an open hemisphere.

Theorem 6.3.

Let n\geq 1 and \beta>0 . Suppose d\geqslant n . Consider the unique solution (x_{i}(\cdot))_{i\in[n]}\in C^{0}(\mathbb{R}_{\geq 0};(\mathbb{S}^{d-1})^{n}) to the Cauchy problem for (SA) or (USA), corresponding to an initial sequence of points (x_{i}(0))_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} distributed uniformly at random. Then almost surely there exists x^{*}\in\mathbb{S}^{d-1} and constants C,\lambda>0 such that

(6.1)

‖ x i (t) − x ∗ ‖ ≤ C e − λ t \|x_{i}(t)-x^{*}\|\leq Ce^{-\lambda t}

holds for all i\in[n] and t\geq 0 .

In fact, let Q and K be arbitrary d\times d matrices. Then the same result also holds for the solution to the corresponding Cauchy problem for (2.3) with V=I_{d} (or the natural analogue of (USA) with these parameters).

When d\geq n and the points (x_{i}(0))_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} are distributed uniformly at random, with probability one there exists ¹¹ ¹¹ 11 This weak version of Wendel’s theorem (Theorem 6.7) is easy to see directly. w\in\mathbb{S}^{d-1} such that \langle w,x_{i}(0)\rangle>0 for any i\in[n] . In other words, all of the initial points lie in an open hemisphere almost surely. The proof of Theorem ˜ 6.3 thus follows as a direct corollary of the following result, which holds for any n\geq 1 and d\geq 2 :

Lemma 6.4 (Cone collapse).

Let \beta>0 and let (x_{i}(0))_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} be such that there exists w\in\mathbb{S}^{d-1} for which \langle x_{i}(0),w\rangle>0 for any i\in[n] . Consider the unique solution (x_{i}(\cdot))_{i\in[n]}\in C^{0}(\mathbb{R}_{\geq 0};(\mathbb{S}^{d-1})^{n}) to the corresponding Cauchy problem for (SA) or (USA). Then there exists x^{*}\in\mathbb{S}^{d-1} and constants C,\lambda>0 such that

	`\\\|x_{i}(t)-x^{*}\\\|\leq Ce^{-\lambda t}`

holds for all i\in[n] and t\geq 0 .

Remark 6.5.

Lemma 6.4 implies that \left\{(\bar{x}_{i})_{i\in[n]}\in(\mathbb{S}^{d-1})^{n}\colon\bar{x}_{1}=\ldots=\bar{x}_{n}\right\} is Lyapunov asymptotically stable as a set. In fact, it is exponentially stable.

Lemma 6.4 is reminiscent of results on interacting particle systems on the sphere (see [CLP15, Theorem 3.7] for instance), and the literature on synchronization for the Kuramoto model on the circle ([ABK ⁺ 22, Lemma 2.8], [HR20, Theorem 3.1] and Section 7.2). We often make use of the following elementary lemma.

Lemma 6.6.

Let f:\mathbb{R}_{\geq 0}\to\mathbb{R} be a differentiable function such that

	`\int_{0}^{+\infty}\\|f(t)\\|\,\mathrm{d}t+\sup_{t\in\mathbb{R}_{\geq 0}}\left\\|\dot{f}(t)\right\\|<+\infty.`

Then \lim_{t\to+\infty}f(t)=0 .

The proof of Lemma ˜ 6.4 is an adaptation of [CLP15, Theorem 1]. We present it here for completeness.

Proof of.

We focus on the case (USA), and set

	`a_{ij}(t):=n^{-1}e^{\beta\langle x_{i}(t),x_{j}(t)\rangle}>0.`

The proof for (SA) is identical, and one only needs to change the coefficients a_{ij}(t) by Z_{\beta,i}(t)^{-1}e^{\beta\langle x_{i}(t),x_{j}(t)\rangle} throughout. Also note that since we only make use of the positivity of the coefficients a_{i,j}(t) throughout the proof, all arguments are readily generalizable to the case of arbitrary d\times d matrices Q and K appearing in the inner products.

Step 1. Clustering

For t\geq 0 , consider

	`i(t)\in\operatorname*{arg\,min}_{i\in[n]}\langle x_{i}(t),w\rangle.`

Fix t_{0}\geq 0 . We have

	`\displaystyle\left(\frac{\,\mathrm{d}}{\,\mathrm{d}t}\langle x_{i(t_{0})}(\cdot),w\rangle\right)\Bigm{\\|}_{t=t_{0}}`
	`\displaystyle=\sum_{j=1}^{n}a_{i(t_{0})j}(t_{0})\Big{(}\langle x_{j}(t_{0}),w\rangle-\langle x_{i(t_{0})}(t_{0}),x_{j}(t_{0})\rangle\langle x_{i(t_{0})}(t_{0}),w\rangle\Big{)}\geq 0.`

This implies that all points remain within the same open hemisphere at all times and the map

	`t\mapsto r(t):=\min_{i\in[n]}\langle x_{i}(t),w\rangle`

is non-decreasing on \mathbb{R}_{\geq 0} . It is also bounded from above by 1 . We may thus define r_{\infty}:=\lim_{t\rightarrow+\infty}r(t) . Note that r_{\infty}\geq r(0)>0 by assumption. By compactness, there exist a sequence of times \{t_{k}\}_{k=1}^{+\infty} with t_{k}\rightarrow+\infty , and some (\overline{x}_{i})_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} such that \lim_{k\rightarrow+\infty}x_{i}(t_{k})=\overline{x}_{i} for all i\in[n] . Using the definition of r(t) , we also find that

	`\langle\overline{x}_{j},w\rangle\geq r_{\infty}`

for all j\in[n] , and by continuity, there exists i\in[n] such that \langle\overline{x}_{i},w\rangle=r_{\infty} . Then

(6.2)

lim k → + ∞ ⟨ x ˙ i (t k),w ⟩ = ∑ j = 1 n a ¯ i j (⟨ w,x ¯ j ⟩ − ⟨ x ¯ i,x ¯ j ⟩ ⟨ x ¯ i,w ⟩) ≥ r ∞ ∑ j = 1 n a ¯ i j (1 − ⟨ x ¯ i,x ¯ j ⟩),\displaystyle\lim_{k\rightarrow+\infty}\langle\dot{x}_{i}(t_{k}),w\rangle=\sum_{\begin{subarray}{c}j=1\end{subarray}}^{n}\overline{a}_{ij}(\langle w,\overline{x}_{j}\rangle-\langle\overline{x}_{i},\overline{x}_{j}\rangle\langle\overline{x}_{i},w\rangle)\geq r_{\infty}\sum_{\begin{subarray}{c}j=1\end{subarray}}^{n}\overline{a}_{ij}(1-\langle\overline{x}_{i},\overline{x}_{j}\rangle),

where we set \overline{a}_{ij}:=e^{\beta\langle\overline{x}_{i},\overline{x}_{j}\rangle}>0 . Notice that

	`\lim_{k\to+\infty}\int_{t_{k}}^{+\infty}\langle\dot{x}_{i}(s),w\rangle\,\mathrm{d}s=r_{\infty}-\lim_{k\to+\infty}\langle x_{i}(t_{k}),w\rangle=0,`

and by using the equation (USA) we also find that |\langle\ddot{x}_{i}(t),w\rangle|=O(e^{2\beta}) for any t\geq 0 . Therefore by Lemma 6.6, the left-hand side of (6.2) is equal to 0 , and consequently the right-hand side term as well. This implies that \overline{x}_{1}=\ldots=\overline{x}_{n}:=x^{*} . Repeating the argument by replacing w with x^{*} , we see that the extraction of a sequence \{t_{k}\}_{k=1}^{+\infty} as above is not necessary, and therefore

(6.3)

lim t → + ∞ x i (t) = x ∗ \lim_{t\to+\infty}x_{i}(t)=x^{*}

for all i\in[n] .

Step 2. Exponential rate.

We now improve upon (6.3). Set

	`\alpha(t):=\min_{i\in[n]}\langle x_{i}(t),x^{*}\rangle.`

From (6.3) we gather that there exists some t_{0}>0 such that \alpha(t)\geq\frac{1}{2} for all t\geq t_{0} . Also, in view of what precedes we know that x^{*} lies in the convex cone generated by the points x_{1}(t),\ldots,x_{n}(t) for any t>0 . Thus, there exists some \eta\in(0,1] such that \eta x^{*} is a convex combination of the points x_{1}(t),\ldots,x_{n}(t) , which implies that

(6.4)

x ∗ = ∑ k = 1 n θ k (t) x k (t), for some ∑ k = 1 n θ k (t) ≥ 1,θ k (t) ≥ 0 ∀ k ∈ [n].x^{*}=\sum_{k=1}^{n}\theta_{k}(t)x_{k}(t),\hskip 14.22636pt\text{ for some }\hskip 14.22636pt\sum_{k=1}^{n}\theta_{k}(t)\geq 1,\quad\theta_{k}(t)\geq 0\quad\forall k\in[n].

For any t , we denote by i(t) an element of \operatorname*{arg\,min}(\langle x_{i}(t),x^{*}\rangle) for which \langle\dot{x}_{i}(t),x^{*}\rangle is smallest. It follows from a Taylor expansion of \langle x_{i}(t+h),x^{*}\rangle for h>0 and i\in[n] that

	`\dot{\alpha}(t)=\langle\dot{x}_{i(t)}(t),x^{*}\rangle.`

Therefore

(6.5)

α ˙ (t) = ⟨ x ˙ i (t) (t),x ∗ ⟩ ≥ ∑ j = 1 n a i (t) j (t) (1 − ⟨ x i (t) (t),x j (t) ⟩) α (t) \dot{\alpha}(t)=\langle\dot{x}_{i(t)}(t),x^{*}\rangle\geq\sum_{j=1}^{n}a_{i(t)j}(t)(1-\langle x_{i(t)}(t),x_{j}(t)\rangle)\alpha(t)

On another hand,

(6.6)

min j ∈ [n] ⁡ ⟨ x i (t) (t),x j (t) ⟩ ≤ ∑ k = 1 n θ k (t) ⟨ x i (t) (t),x k (t) ⟩ = ⟨ x i (t) (t),x ∗ ⟩ = α (t).\min_{j\in[n]}\langle x_{i(t)}(t),x_{j}(t)\rangle\leq\sum_{k=1}^{n}\theta_{k}(t)\langle x_{i(t)}(t),x_{k}(t)\rangle=\langle x_{i(t)}(t),x^{*}\rangle=\alpha(t).

Plugging (6.6) into (6.5) and using a_{ij}(t)\geq n^{-1}e^{-2\beta} we get

(6.7)

α ˙ (t) ≥ 1 2 n e 2 β (1 − α (t)) \dot{\alpha}(t)\geq\frac{1}{2ne^{2\beta}}(1-\alpha(t))

for t\geq t_{0} . Applying the Grönwall inequality we get

(6.8)

1 − α (t) ≤ 1 2 e − 1 2 n e β (t − t 0) 1-\alpha(t)\leq\frac{1}{2}e^{-\frac{1}{2ne^{\beta}}(t-t_{0})}

for all t\geq t_{0} . The conclusion follows. ∎

In the case d<n , we can still apply Wendel’s theorem (recalled below) together with Lemma 6.4 to obtain clustering to a single point with probability at least p_{n,d} for some explicit p_{n,d}\in(0,1) .

Theorem 6.7 (Wendel, []).

Let d,n\geq 1 be such that d\leq n . Let x_{1},\ldots,x_{n} be n i.i.d. uniformly distributed points on \mathbb{S}^{d-1} . The probability that these points all lie in the same hemisphere is:

	`\mathbb{P}\Big{(}\exists w\in\mathbb{S}^{d-1}\colon\langle x_{i},w\rangle>0\quad\mathrm{for}\,\mathrm{all}\quad i\in[n]\Big{)}=2^{-(n-1)}\sum_{k=0}^{d-1}\binom{n-1}{k}.`

6.2. More precise quantitative convergence

When n is fixed and d\rightarrow+\infty , in addition to showing the formation of a cluster as in Theorem 6.3, it is possible to quantitatively describe the entire evolution of the particles with high probability. To motivate this, on the one hand we note that since the dynamics evolve on \mathbb{S}^{d-1} , inner products are representative of the distance between points, and clustering occurs if \langle x_{i}(t),x_{j}(t)\rangle\rightarrow 1 for any (i,j)\in[n]^{2} as t\rightarrow+\infty . On the other hand, if d\gg n , n points in a generic initial sequence are almost orthogonal by concentration of measure [Ver18, Chapter 3], and we are thus able to compare their evolution with that of an initial sequence of truly orthogonal ones.

We begin by describing the case of exactly orthogonal initial particles, which is particularly simple as the dynamics are described by a single parameter.

Theorem 6.8.

Let \beta\geq 0 , d,n\geq 2 be arbitrary. Consider an initial sequence (x_{i}(0))_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} of n pairwise orthogonal points: \langle x_{i}(0),x_{j}(0)\rangle=0 for i\neq j , and let (x_{i}(\cdot))_{i\in[n]}\in C^{0}(\mathbb{R}_{\geqslant 0};(\mathbb{S}^{d-1})^{n}) denote the unique solution to the corresponding Cauchy problem for (SA) (resp. for (USA)). Then the angle \angle(x_{i}(t),x_{j}(t)) is the same for all distinct i,j\in[n] :

	`\angle(x_{i}(t),x_{j}(t))=\theta_{\beta}(t)`

for t\geqslant 0 and some \theta_{\beta}\in C^{0}(\mathbb{R}_{\geq 0};\mathbb{T}) . Furthermore, for (SA), \gamma_{\beta}(t):=\cos(\theta_{\beta}(t)) satisfies

(6.9)

{ γ ˙ β (t) = 2 e β γ β (t) (1 − γ β (t)) ((n − 1) γ β (t) + 1) e β + (n − 1) e β γ β (t) for t ≥ 0 γ β (0) = 0,\displaystyle\begin{dcases}\dot{\gamma}_{\beta}(t)=\frac{2e^{\beta\gamma_{\beta}(t)}(1-\gamma_{\beta}(t))((n-1)\gamma_{\beta}(t)+1)}{e^{\beta}+(n-1)e^{\beta\gamma_{\beta}(t)}}&\text{ for }t\geq 0\\ \gamma_{\beta}(0)=0\,,\end{dcases}

and for (USA), we have

(6.10)

{ γ ˙ β (t) = 2 n e β γ β (t) (1 − γ β (t)) ((n − 1) γ β (t) + 1) for t ≥ 0 γ β (0) = 0.\displaystyle\begin{dcases}\dot{\gamma}_{\beta}(t)=\frac{2}{n}e^{\beta\gamma_{\beta}(t)}(1-\gamma_{\beta}(t))((n-1)\gamma_{\beta}(t)+1)&\text{ for }t\geq 0\\ \gamma_{\beta}(0)=0\,.\end{dcases}

Here and henceforth, \mathbb{T}=\mathbb{R}/2\pi\mathbb{Z} denotes the one-dimensional torus. We provide a brief proof of Theorem ˜ 6.8 just below. The following result then shows that when d\gg n , t\mapsto\gamma_{\beta}(t) is a valid approximation for t\mapsto\langle x_{i}(t),x_{j}(t)\rangle for any distinct i,j\in[n] .

Theorem 6.9.

Fix \beta\geqslant 0 and n\geq 2 . Then there exists some d^{*}(n,\beta)\geq n such that for all d\geq d^{*}(n,\beta) , the following holds. Consider a sequence (x_{i}(0))_{i\in[n]} of n i.i.d. uniformly distributed points on \mathbb{S}^{d-1} , and let (x_{i}(\cdot))_{i\in[n]}\in C^{0}(\mathbb{R}_{\geq 0};(\mathbb{S}^{d-1})^{n}) denote the unique solution to the corresponding Cauchy problem for (SA). Then there exist C=C(n,\beta)>0 and \lambda=\lambda(n,\beta)>0 , such that with probability at least 1-2n^{2}d^{-1/64} ,

(6.11)

| ⟨ x i (t),x j (t) ⟩ − γ β (t) | ⩽ min ⁡ { 2 ⋅ c (β) n t log ⁡ d d,C e − λ t } \Big{|}\langle x_{i}(t),x_{j}(t)\rangle-\gamma_{\beta}(t)\Big{|}\leqslant\min\left\{2\cdot c(\beta)^{nt}\sqrt{\frac{\log d}{d}},Ce^{-\lambda t}\right\}

holds for any i\neq j and t\geq 0 , where c(\beta)=e^{10\max\{1,\beta\}} , and \gamma_{\beta} is the unique solution to (6.9).

Since the proof is rather lengthy, we defer it to Appendix D. It relies on combining the stability of the flow with respect to the initial data (entailed by the Lipschitz nature of the vector field) with concentration of measure. An analogous statement also holds for (USA), and more details can be found in Remark D.1, whereas the explicit values of C and \lambda can be found in (D.15). The upper bound in (6.11) is of interest in regimes where d and/or t are sufficiently large as the error in (6.11) is trivially bounded by 2 .

Proof of.

We split the proof in two parts. We focus on proving the result for the dynamics (SA), since the same arguments readily apply to the dynamics (USA).

Part 1. The angle θβ(t)\theta_{\beta}(t)

We first show there exists \theta\in C^{0}(\mathbb{R}_{\geq 0};\mathbb{T}) such that \theta(t)=\angle(x_{i}(t),x_{j}(t)) for any distinct (i,j)\in[n]^{2} and t\geq 0 . Since the initial tokens are orthogonal (and thus d\geq n ), we may consider an orthonormal basis (e_{1},\ldots,e_{d}) of \mathbb{R}^{d} such that x_{i}(0)=e_{i} for i\in[n] . Let \pi:[d]\rightarrow[d] be a permutation. By decomposing any x\in\mathbb{S}^{d-1} in this basis, we define P_{\pi}:\mathbb{S}^{d-1}\rightarrow\mathbb{S}^{d-1} as

	`P_{\pi}\left(\sum_{i=1}^{n}a_{i}e_{i}\right)=\sum_{i=1}^{n}a_{i}e_{\pi(i)}.`

Setting y_{i}(t)=P_{\pi}(x_{i}(t)) for i\in[n] , we see that y_{i}(t) solves (SA) with initial condition y_{i}(0)=P_{\pi}(x_{i}(0)) . But (x_{\pi(1)}(t),\ldots,x_{\pi(n)}(t)) is a solution of (SA) by permutation equivariance, and it has the same initial condition since P_{\pi}(x_{i}(0))=x_{\pi(i)}(0) . Consequently, we deduce that P_{\pi}(x_{i}(t))=x_{\pi(i)}(t) for any t\geq 0 and any i\in[d] . Hence

	`\langle x_{i}(t),x_{j}(t)\rangle=\langle P_{\pi}(x_{i}(t)),P_{\pi}(x_{j}(t))\rangle=\langle x_{\pi(i)}(t),x_{\pi(j)}(t)\rangle`

which concludes the proof.

Part 2. The curve γβ(t)\gamma_{\beta}(t)

By virtue of the orthogonality assumption we have \gamma_{\beta}(0)=\cos(\theta_{\beta}(0))=0 . To prove that \gamma_{\beta}(t) satisfies (6.9) for the case of (SA), recall that

	`\displaystyle\bm{\mathsf{P}}^{\perp}_{x_{i}(t)}(x_{j}(t))=x_{j}(t)-\langle x_{i}(t),x_{j}(t)\rangle\,x_{i}(t).`

Then for k\neq i ,

	`\displaystyle\dot{\gamma}_{\beta}(t)`	`\displaystyle=2\langle\dot{x}_{i}(t),x_{k}(t)\rangle`
		`\displaystyle=2\sum_{\begin{subarray}{c}j=1\end{subarray}}^{n}\left(\frac{e^{\beta\langle x_{i}(t),x_{j}(t)\rangle}}{\sum_{\ell=1}^{n}e^{\beta\langle x_{i}(t),x_{\ell}(t)\rangle}}\right)\left(\langle x_{j}(t),x_{k}(t)\rangle-\langle x_{i}(t),x_{j}(t)\rangle\langle x_{i}(t),x_{k}(t)\rangle\right).`

Since the denominator in the above expression is equal to (n-1)e^{\beta\gamma_{\beta}(t)}+e^{\beta} , we end up with

	`\displaystyle\dot{\gamma}_{\beta}(t)`	`\displaystyle=\frac{2e^{\beta\gamma_{\beta}(t)}}{(n-1)e^{\beta\gamma_{\beta}(t)}+e^{\beta}}\sum_{\begin{subarray}{c}j=1\end{subarray}}^{n}\Big{(}\langle x_{j}(t),x_{k}(t)\rangle-\langle x_{i}(t),x_{j}(t)\rangle\langle x_{i}(t),x_{k}(t)\rangle\Big{)}`
		`\displaystyle=\frac{2e^{\beta\gamma_{\beta}(t)}}{(n-1)e^{\beta\gamma_{\beta}(t)}+e^{\beta}}(1-\gamma_{\beta}(t)^{2}+(n-2)(\gamma_{\beta}(t)-\gamma_{\beta}(t)^{2})),`

as desired. ∎

6.3. Metastability and a phase transition

An interesting byproduct of Theorem ˜ 6.8 and Theorem ˜ 6.9 is the fact that they provide an accurate approximation of the exact phase transition curve delimiting the clustering and non-clustering regimes, in terms of t and \beta . To be more precise, given an initial sequence (x_{i}(0))_{i\in[n]}\in(\mathbb{S}^{d-1})^{n} of random points distributed independently according to the uniform distribution on \mathbb{S}^{d-1} , and for any fixed 0<\delta\ll 1 , we define the phase transition curve as the boundary

	`\Gamma_{d,\delta}=\partial\left\{t,\beta\geq 0\colon t=\operatorname*{arg\,inf}_{s\geq 0}\Big{(}\mathbb{P}(\langle x_{1}(s),x_{2}(s)\rangle\geq 1-\delta)=1-2n^{2}d^{-\frac{1}{64}}\Big{)}\right\}`

where (x_{i}(\cdot))_{i\in[n]} denotes the solution to the corresponding Cauchy problem for (SA). (Here the choice of the first two particles instead of a random distinct pair is justified due to permutation equivariance.) Theorem ˜ 6.9 then gives the intuition that over compact subsets of (\mathbb{R}_{\geq 0})^{2} , \Gamma_{d,\delta} should be well-approximated by

(6.12)

Γ ∞,δ = { t,β ≥ 0:γ β (t) = 1 − δ }.\Gamma_{\infty,\delta}=\Big{\{}t,\beta\geq 0\colon\gamma_{\beta}(t)=1-\delta\Big{\}}.

(a) d = 2 d=2

This is clearly seen in Figure 3, along with the fact that the resolution of this approximation increases with d\to+\infty .

Figure 3 appears to contain more information than what we may gather from Theorem 6.3, Theorem 6.8 and Theorem 6.9. In particular, for small d , we see the appearance of a zone (white/light blue in Figure 3) of parameters (t,\beta) for which the probability of particles being clustered is positive, but not close to one. A careful inspection of this region reveals that points are grouped in a finite number of clusters; see Figure 4. The presence of such a zone indicates the emergence of a long-time metastable state where points are clustered into several groups but eventually relax to a single cluster in long-time. This two-time-scale phenomenon is illustrated in Figure 4 and prompts us to formulate the following question.

Problem 1.

Do the dynamics enter a transient metastable state, in the sense that for \beta\gg 1 , all particles stay in the vicinity of m<n clusters for long periods of time, before they all collapse to the final cluster \{x^{*}\} ?

There have been important steps towards a systematic theory of metastability for gradient flows, with applications to nonlinear parabolic equations—typically reaction-diffusion equations such as the Allen-Cahn or Cahn-Hilliard equations [OR07, KO02] —. While these tools to not readily apply to the current setup, they form an important starting point to answer this question.

Figure 4. We zoom in on the phase diagram (Figure 3 ) for the dynamics on the circle: d = 2 d=2. For β 4, 9 \beta=4,9, we also display a trajectory of ( SA ) for a randomly drawn initial condition at times t 2.5 18 30 t=2.5,18,30. We see that the particles settle at clusters when \beta=4 (bottom right) and \beta=9 (top right), for a duration of time. This reflects our metastability claim for large \beta

Finally, one may naturally ask whether the clustering and phase diagram conclusions persist when the parameter matrices (Q,K,V) are significantly more general: some illustrations ¹² ¹² 12 See github.com/borjanG/2023-transformers-rotf for additional figures which indicate that this phenomenon appears to hold in even more generality. are given in Figure 5.

Problem 2.

Can the conclusions of Theorem ˜ 6.8 – Theorem ˜ 6.9 be generalized to the case of random matrices (Q,K,V) ?

(a) Q, K V Q,K,V in real Ginibre ensemble

Part III Further questions

We conclude this manuscript by discussing several avenues of research that can lead to a finer understanding of the clustering phenomenon and generalizations of our results, and which, we believe, are of independent mathematical interest. Specifically,

In Section 7, we zero in on the special case d=2 , where we make a link with the celebrated Kuramoto model when \beta=0 ;
In Section 8, we discuss an alternative approach for analyzing clustering, based on the so-called BBGKY hierarchy from statistical mechanics;
In Section 9, we foray beyond the case Q^{\top}K=V=I_{d} in a few directions. We start with Section 9.1 by relating the case V=-I_{d} to optimal configurations on the sphere. We then discuss existing results in the absence of layer normalization in Section 9.2, which motivate the study of a related singular equation (Section 9.3) and a diffusive regularization (Section 9.4);
Finally, in Section 10, we briefly discuss various ways focusing on the tuning of the parameter matrices themselves.

7. Dynamics on the circle

We study the dynamics (SA) and (USA) in the special case d=2 , namely on the unit circle \mathbb{S}^{1}\subset\mathbb{R}^{2} . This model, parametrized by angles and related to the celebrated Kuramoto model, is of independent interest and deserves a complete mathematical analysis.

7.1. Angular equations

On the circle \mathbb{S}^{1} , all particles x_{i}(t)\in\mathbb{S}^{1} are of course completely characterized by the angle \theta_{i}(t)\in\mathbb{T} : x_{i}(t)=\cos(\theta_{i}(t))e_{1}+\sin(\theta_{i}(t))e_{2} where e_{1}=(1,0) and e_{2}=(0,1)\in\mathbb{R}^{2} . We focus on the dynamics (USA) for simplicity. For any i\in[n] and t\geq 0 , we may derive the equation satisfied by \theta_{i}(t) from \cos(\theta_{i}(t))=\langle x_{i}(t),e_{1}\rangle : differentiating in t and plugging into (USA) we obtain

	`\displaystyle\dot{\theta}_{i}(t)=-\frac{n^{-1}}{\sin(\theta_{i}(t))}\left(\sum_{j=1}^{n}e^{\beta\langle x_{i}(t),x_{j}(t)\rangle}\Big{[}\langle x_{j}(t),e_{1}\rangle-\langle x_{i}(t),x_{j}(t)\rangle\langle x_{i}(t),e_{1}\rangle\Big{]}\right)`

where we used the definition of the projection (if \theta_{i}(t)=0 for some t , we differentiate the equality \sin(\theta_{i}(t))=\langle x_{i}(t),e_{2}\rangle instead, which also leads to (7.1) in the end). Observing that

	`\langle x_{i}(t),x_{j}(t)\rangle=\cos(\theta_{i}(t)-\theta_{j}(t)),`

we find

	`\displaystyle\dot{\theta}_{i}(t)`	`\displaystyle=-\frac{n^{-1}}{\sin(\theta_{i}(t))}\left(\sum_{j=1}^{n}e^{\beta\cos(\theta_{i}(t)-\theta_{j}(t))}\Big{[}\cos(\theta_{j}(t))-\cos(\theta_{i}(t)-\theta_{j}(t))\cos(\theta_{i}(t))\Big{]}\right).`

Using elementary trigonometry, we conclude that

(7.1)

θ ˙ i (t) = − 1 n ∑ j = 1 n e β cos ⁡ (θ i (t) − θ j (t)) sin ⁡ (θ i (t) − θ j (t)).\dot{\theta}_{i}(t)=-\frac{1}{n}\sum_{j=1}^{n}e^{\beta\cos(\theta_{i}(t)-\theta_{j}(t))}\sin(\theta_{i}(t)-\theta_{j}(t)).

The case \beta=0 is exactly the Kuramoto model recalled in Section 7.2. Suppose for the time being that \beta>0 . Defining the function h_{\beta}:\mathbb{T}\to\mathbb{R}_{\geq 0} as

	`h_{\beta}(\theta)=e^{\beta\cos(\theta)},`

we have effectively deduced that the empirical measure of the angles, \nu(t)=\frac{1}{n}\sum_{j=1}^{n}\delta_{\theta_{j}(t)} , which is a measure on the torus \mathbb{T} , is a solution to the continuity equation

	`\partial_{t}\nu(t)+\partial_{\theta}(\mathcal{X}[\nu(t)]\nu(t))=0,\hskip 28.45274pt\text{ on }\mathbb{R}_{\geq 0}\times\mathbb{T},`

where

	`\mathcal{X}[\nu](\theta)=\frac{1}{\beta}\Big{(}h_{\beta}^{\prime}\ast\nu\Big{)}(\theta).`

When the particles x_{i}(t) follow (SA), one readily checks that the same continuity equation is satisfied but rather with the field

	`\mathcal{X}[\nu](\theta)=\frac{1}{\beta}\left(\frac{h_{\beta}^{\prime}\ast\nu}{h_{\beta}\ast\nu}\right)(\theta).`

7.2. The Kuramoto model

As mentioned above, when \beta=0 , (7.1) is a particular case of the Kuramoto model [Kur75]:

(7.2)

θ ˙ i (t) = ω i + K n ∑ j = 1 n sin ⁡ (θ j (t) − θ i (t)),\dot{\theta}_{i}(t)=\omega_{i}+\frac{K}{n}\sum_{j=1}^{n}\sin(\theta_{j}(t)-\theta_{i}(t)),

where K>0 is a prescribed coupling constant, and \omega_{i}\in\mathbb{T} are the intrinsic natural frequencies of the oscillators \theta_{i}(t) . It is known that for sufficiently small coupling strength K , the oscillators \theta_{i}(t) in the Kuramoto model (7.2) do not synchronize in long-time. It is also known that when K exceeds some critical threshold value, a phase transition occurs, leading to the synchronization of a fraction of the oscillators. If K is chosen very large, there is total synchronization of the oscillators in long-time. For more on the mathematical aspects of the Kuramoto model, we refer the reader to the review papers [Str00, ABV ⁺ 05, HKPZ16] (see also [CCH ⁺ 14, Chi15, FGVG16, DFGV18, HR20, TSS20, ABK ⁺ 22, MP22] for a non-exhaustive list of other recent mathematical results on the subject).

When all the frequencies \omega_{i} are equal to some given frequency, \omega\in\mathbb{R} say, after a change of variable of the form \theta_{i}(t)\leftarrow\theta_{i}(t)-\omega t , the dynamics in (7.2) become the gradient flow

	`\dot{\theta}(t)=n\nabla\mathsf{F}(\theta)`

where the energy \mathsf{F}:\mathbb{T}^{n}\to\mathbb{R}_{\geq 0} reads

(7.3)

𝖥 (θ) = K 2 n 2 ∑ i = 1 n ∑ j = 1 n cos ⁡ (θ i − θ j).\mathsf{F}(\theta)=\frac{K}{2n^{2}}\sum_{i=1}^{n}\sum_{j=1}^{n}\cos(\theta_{i}-\theta_{j}).

The oscillators can be viewed as attempting to maximize this energy. The energy \mathsf{F} is maximized when all the oscillators are synchronized, that is, \theta_{i}=\theta^{*} for some \theta^{*}\in\mathbb{T} and for all i\in[n] . As the dynamics follow a gradient system, the equilibrium states are the critical points of the energy, namely those satisfying \nabla\mathsf{F}(\theta)=0 . The local maxima of \mathsf{F} correspond to equilibrium states \theta that are physically achievable, since small perturbations thereof return the system back to \theta .

Some authors consider a variant of the Kuramoto model where the oscillators are interacting according to the edges of a graph. In other words, the coefficients A_{ij} of the graph’s adjacency matrix are inserted in the sum in (7.3) as weights, and the dynamics are then the corresponding gradient flow. A recent line of work culminating with [ABK ⁺ 22] has established that synchronization occurs with high probability for Erdős–Rényi graphs with parameter p , for every p right above the connectivity threshold.

Coming back to our dynamics (7.1), we notice that it can also be written as a gradient flow on \mathbb{T}^{n} :

	`\dot{\theta}(t)=n\nabla\mathsf{E}_{\beta}(\theta(t)),`

for the interaction energy \mathsf{E}_{\beta}:\mathbb{T}^{n}\to\mathbb{R}_{\geq 0} defined as

(7.4)

𝖤 β (θ) = 1 2 β n 2 ∑ i = 1 n ∑ j = 1 n e β cos ⁡ (θ i − θ j),\mathsf{E}_{\beta}(\theta)=\frac{1}{2\beta n^{2}}\sum_{i=1}^{n}\sum_{j=1}^{n}e^{\beta\cos(\theta_{i}-\theta_{j})},

which is maximized when \theta_{i}=\theta^{*} for some \theta^{*}\in\mathbb{T} and for all i\in[n] . In the spirit of [LXB19], we suggest the following open problem—we recall that a critical point is called a strict saddle point of \mathsf{E}_{\beta} if the Hessian of \mathsf{E}_{\beta} at these points has at least one positive eigenvalue.

Problem 3.

With the exception of the global maxima, are all critical points of \mathsf{E}_{\beta} strict saddle points?

The proofs of Theorems 4.3 and 5.1 already yield a positive answer to Problem 5 in the regimes \beta\leq 1 or \beta\geq\frac{n^{2}}{\pi^{2}} . The complementary regime remains open. By classical arguments, recalled in Appendix A, a positive answer to Problem 3 would imply that for all initial conditions except a set of measure zero, all \theta_{i}(t) converge under the dynamics (7.1) to a common limit as t\rightarrow+\infty .

Extensions of the Kuramoto model of the form

(7.5)

θ ˙ i (t) = ω i + K n ∑ j = 1 n h (θ j (t) − θ i (t)),\dot{\theta}_{i}(t)=\omega_{i}+\frac{K}{n}\sum_{j=1}^{n}h(\theta_{j}(t)-\theta_{i}(t)),

for a general non-linearity h:\mathbb{T}\to\mathbb{R} , which contains both (7.2) and our model (7.1) as particular cases, have already been studied in the physics literature. For instance, we refer the reader to [Dai92] (see also [ABV ⁺ 05, page 158]), where many heuristics are proposed to address the behavior of solutions to these dynamics. We are not aware of mathematical results for (7.1) besides Theorem 5.1. We nevertheless have some hope that handling the dynamics (7.1) is easier than dealing with (7.5) for a general h ; for instance, we have

	`h_{\beta}(\theta)=e^{\beta\cos(\theta)}=\sum_{k\in\mathbb{Z}}I_{k}(\beta)e^{\mathrm{i}k\theta}`

where I_{k}(\beta) are the modified Bessel function of the first kind, whose properties have been extensively studied.

8. BBGKY hierarchy

For the sake of simplicity, we again focus on the dynamics on the circle \mathbb{S}^{1} , where recall that all particles are parametrized by angles (which we also refer to as particles). To carve out an even more complete understanding of the clustering phenomenon, it is natural to consider initial particles sampled i.i.d. from the uniform distribution on \mathbb{S}^{1} and to study the time-evolution of the r -particle distribution \rho^{(r)}_{n}(t,\theta_{1},\ldots,\theta_{r}) , defined as the joint law of the particles \theta_{1}(t),\ldots,\theta_{r}(t) . Otherwise put, it is the r -point marginal of the joint distribution \rho^{(n)}(t,\cdot)\in\mathcal{P}(\mathbb{T}^{n}) of all n particles. Note that because of rotational invariance, \rho^{(1)}(t,\cdot) is just the uniform distribution equal to \frac{1}{2\pi} for all t\geq 0 . For r=2 , again by rotational invariance, there exists some \psi(t,\cdot):\mathbb{T}\rightarrow\mathbb{R}_{\geq 0} such that

	`\rho^{(2)}(t,\theta_{1},\theta_{2})=\frac{1}{2\pi}\psi(t,\theta_{2}-\theta_{1}).`

Proving the clustering/synchronization of all \theta_{i}(t) in long-time amounts to proving that \psi(t,\cdot) converges to a Dirac mass centered at 0 as t\rightarrow+\infty . Using the fact that \rho^{(n)}(t,\cdot) solves the Liouville equation, by following the method used to derive the BBGKY ¹³ ¹³ 13 Bogoliubov–Born–Green–Kirkwood–Yvon. hierarchy [GSRT13, Gol16], it is possible to show that \psi(t,\cdot) satisfies

(8.1)

{ ∂ t ψ (t,x) + ∂ x (v (t,x) ψ (t,x)) = 0 in ℝ ≥ 0 × 𝕋 ψ (0,x) = (2 π) − 1 in 𝕋,\begin{dcases}\partial_{t}\psi(t,x)+\partial_{x}(v(t,x)\psi(t,x))=0&\text{ in }\mathbb{R}_{\geq 0}\times\mathbb{T}\\ \psi(0,x)=(2\pi)^{-1}&\text{ in }\mathbb{T},\end{dcases}

where

	`v(t,x)=\frac{2}{\beta n}h_{\beta}^{\prime}(x)-\frac{2(n-2)}{\beta n}g(t,x),`

and

	`g(t,x)=\mathbb{E}\Big{[}-h_{\beta}^{\prime}(\theta_{3}(t))\,\Big{\\|}\,\theta_{1}(t)=0,\ \theta_{2}(t)=x\Big{]}.`

Note that the equation (8.1) is not closed since g(t,x) depends on the 3 -point correlation function. This is typical in the BBGKY hierarchy, whereupon physical theory and experimental evidence is typically used to devise an ansatz for closing the system. For instance, the Boltzmann equation is derived from the BBGKY hierarchy by assuming the molecular chaos hypothesis (Stosszahlansatz) at the level of r=2 . We suggest to close (8.1) in a way that reflects the formation of clusters:

Problem 4.

Devise a realistic ansatz for g(t,x) which allows to close equation (8.1), and allows to prove the convergence of \psi(t,\cdot) to a Dirac mass centered at 0 as t\to+\infty .

The derivation of a BBGKY hierarchy when d>2 , as well as for (SA), are also problems which we believe merit further investigation.

9. General matrices

Figure 5 hints at the likelihood of the clustering phenomenon being significantly more general than just the case Q=K=V=I_{d} . However, extending our proofs to more general parameter matrices does not appear to be straightforward and is an open problem. Here we discuss some particular cases (without excluding other approaches).

9.1. The repulsive case

As seen from Lemma 3.7, in the repulsive case V=-I_{d} the interaction energy \mathsf{E}_{\beta} decreases along trajectories. Recall that the unique global minimum of \mathsf{E}_{\beta} over \mathcal{P}(\mathbb{S}^{d-1}) is the uniform distribution (Proposition 3.4). In contrast, we explain in this section that many different configurations of n points may yield global minima for \mathsf{E}_{\beta} when minimized over empirical measures with n atoms.

We thus focus on minimizing \mathsf{E}_{\beta} over the set \mathcal{P}_{n}(\mathbb{S}^{d-1}) of empirical measures, namely sums of n Dirac masses. Rewriting \mathsf{E}_{\beta} as

	`\mathsf{E}_{\beta}[\mu]=\frac{e^{\beta}}{2\beta}\iint e^{-\frac{\beta}{2}\\\|x-x^{\prime}\\\|^{2}}\,\mathrm{d}\mu(x)\,\mathrm{d}\mu(x^{\prime}),`

it turns out that minimizing \mathsf{E}_{\beta} over \mathcal{P}_{n}(\mathbb{S}^{d-1}) is precisely the problem of finding optimal configurations of points on \mathbb{S}^{d-1} , which has direct links to the sphere packing problem [CK07, CKM ⁺ 22] and coding theory [DGS91]. For \mu\in\mathcal{P}_{n}(\mathbb{S}^{d-1}) , we can equivalently rewrite \mathsf{E}_{\beta} in terms of the set of support points \mathscr{C}\subset\mathbb{S}^{d-1} , \#\mathscr{C}=n :

	`\mathsf{E}_{\beta}[\mu]=\mathsf{H}_{\beta}[\mathscr{C}]=\frac{e^{\beta}}{2n^{2}\beta}\sum_{\begin{subarray}{c}x,x^{\prime}\in\mathscr{C}\end{subarray}}e^{-\frac{\beta}{2}\\\|x-x^{\prime}\\\|^{2}}.`

In [CK07], Cohn and Kumar characterize the global minima \mathscr{C} of \mathsf{H}_{\beta} . To state their result, we need the following definition.

Definition 9.1.

Let n\geq 2 . A set of points \mathscr{C}=\{x_{1},\ldots,x_{n}\}\subset\mathbb{S}^{d-1} is called a spherical t -design if

	`\int p(x)\,\mathrm{d}\sigma_{d}(x)=\frac{1}{n}\sum_{i=1}^{n}p(x_{i})`

for all polynomials p of d variables, of total degree at most t . The set of points \mathscr{C} is called a sharp configuration if there are m distinct inner products between pairwise distinct points in \mathscr{C} , for some m>1 , and if it is a spherical (2m-1) -design.

The following result is a special case of [CK07, Theorem 1.2].

Theorem 9.2 ([]).

Let n\geq 2 . Any global minimum of \mathsf{H}_{\beta} among \mathscr{C}\subset\mathbb{S}^{d-1} , \#\mathscr{C}=n is either a sharp configuration, or the vertices of a 600 -cell ¹⁴ ¹⁴ 14 A 600 -cell is a particular 4 -dimensional convex polytope with n=120 vertices..

The set of sharp configurations is not known for all regimes of n,d or m (the largest m such that the configuration is a spherical m -design). A list of known examples is provided in [CK07, Table 1]: it consists of vertices of full-dimensional polytopes (specifically, regular polytopes whose faces are simplices), or particular derivations of the E_{8} root lattice in \mathbb{R}^{8} and the Leech lattice in \mathbb{R}^{24} . We refer the reader to [CK07] and the illustrative experimental paper [BBC ⁺ 09] for further detail. A complete picture of the long-time behavior of Transformers in the repulsive case remains open.

9.2. Pure self-attention

An alternative avenue for conducting such an analysis which has shown to be particularly fruitful consists in removing the projector \bm{\mathsf{P}}^{\perp}_{x} , leading to

(9.1)

x ˙ i (t) = 1 Z β,i (t) ∑ j = 1 n e β ⟨ Q x i (t),K x j (t) ⟩ V x j (t) \dot{x}_{i}(t)=\frac{1}{Z_{\beta,i}(t)}\sum_{j=1}^{n}e^{\beta\langle Qx_{i}(t),Kx_{j}(t)\rangle}Vx_{j}(t)

for all i\in[n] and t\in\mathbb{R}_{\geq 0} . In fact, in [GLPR24] we analyze precisely these dynamics, and show different clustering results depending on the spectral properties of the matrix V . We briefly summarize our findings in what follows.

9.2.1. A review of []

For most choices of value matrices V , without rescaling time, most particles diverge to \pm\infty and no particular pattern emerges. To make a very rough analogy, (9.1) "looks like" \dot{x}_{i}(t)=Vx_{i}(t) (which amounts to having P_{ij}(t)=\delta_{ij} instead of (2.5)), whose solutions are given by x_{i}(t)=e^{tV}x_{i}(0) . To discern the formation of clusters, we introduce the rescaling ¹⁵ ¹⁵ 15 The rescaling (9.2) should be seen as a surrogate for layer normalization.

(9.2)

z i (t) = e − t V x i (t),z_{i}(t)=e^{-tV}x_{i}(t),

which are solutions to

(9.3)

z ˙ i (t) = 1 Z β,i (t) ∑ j = 1 n e β ⟨ Q e t V z i (t),K e t V z j (t) ⟩ V (z j (t) − z i (t)) \dot{z}_{i}(t)=\frac{1}{Z_{\beta,i}(t)}\sum_{j=1}^{n}e^{\beta\left\langle Qe^{tV}z_{i}(t),Ke^{tV}z_{j}(t)\right\rangle}V(z_{j}(t)-z_{i}(t))\hskip 28.45274pt

for i\in[n] and t\geq 0 , where

	`Z_{\beta,i}(t)=\sum_{k=1}^{n}e^{\beta\langle Qe^{tV}z_{i}(t),Ke^{tV}z_{k}(t)\rangle},`

whereas the initial condition remains the same, namely x_{i}(0)=z_{i}(0) . It is crucial to notice that the coefficients A_{ij}(t) (see (2.5)) of the self-attention matrix for the rescaled particles z_{i}(t) are the same as those for the original particles x_{i}(t) . The weight A_{ij}(t) indicates the strength of the attraction of z_{i}(t) by z_{j}(t) . In [GLPR24] we show that the rescaled particles z_{i}(t) cluster toward well-characterized geometric objects as t\rightarrow+\infty for various choices of matrices (Q,K,V) . Our results are summarized in Table 1 below, whose first two lines are discussed thereafter.

`V`	`K` and `Q`	Limit geometry	Result in [GLPR24]
`V=I_{d}`	`Q^{\top}K\succ 0`	vertices of convex polytope	Theorem 3.1
`\lambda_{1}(V)>0` , simple	`\langle Q\varphi_{1},K\varphi_{1}\rangle>0`	union of `3` parallel hyperplanes	Theorem 4.1
`V` paranormal	`Q^{\top}K\succ 0`	polytope `\times` subspaces	Theorem 5.1
`V=-I_{d}`	`Q^{\top}K=I_{d}`	single cluster at origin ^∗	Theorem C.5

Table 1. Summary of the clustering results of [GLPR24]. ^∗ All results except for the case V=-I_{d} hold for the time-scaled dynamics (9.3).

When V=I_{d} , outside from exceptional situations, all particles cluster to vertices of some convex polytope. Indeed, since the velocity \dot{z}_{i}(t) is a convex combination of the attractions z_{j}(t)-z_{i}(t) , the convex hull \mathcal{K}(t) of the z_{i}(t) shrinks and thus converges to some convex polytope. The vertices of the latter attract all particles as t\rightarrow+\infty . When the eigenvalue with largest real part of V , denoted by \lambda_{1}(V) , is simple and positive, the rescaled particles z_{i}(t) cluster on hyperplanes which are parallel to the direct sum of the eigenspaces of the remaining eigenvalues. Roughly speaking, the coordinates of the points z_{i}(t) along the eigenvector of V corresponding to \lambda_{1}(V) quickly dominate the matrix coefficients P_{ij}(t) in (9.3) due to the factors e^{tV}z_{j}(t) . For more results and insights regarding clustering on \mathbb{R}^{d} , we refer the reader to [GLPR24]. We nonetheless leave the reader with the following general question:

Problem 5.

Is it possible to extend the clustering results of Table 1 to other cases of (Q,K,V) ? What are the resulting limit shapes?

9.3. Singular dynamics

We mention another intriguing question, whose answer would allow for a transparent geometric understanding of clustering for (9.3). Let (Q,K,V) be given d\times d matrices. For \beta>0 , we consider the system of coupled ODEs

(9.4)

z ˙ i (t) = 1 Z β,i (t) ∑ j = 1 n e β ⟨ Q z i (t),K z j (t) ⟩ V (z j (t) − z i (t)),\dot{z}_{i}(t)=\frac{1}{Z_{\beta,i}(t)}\sum_{j=1}^{n}e^{\beta\langle Qz_{i}(t),Kz_{j}(t)\rangle}V(z_{j}(t)-z_{i}(t)),

where once again

	`Z_{\beta,i}(t)=\sum_{k=1}^{n}e^{\beta\langle Qz_{i}(t),Kz_{k}(t)\rangle}.`

For any T>0 , and any fixed initial condition (z_{i}(0))_{i\in[n]}\in(\mathbb{R}^{d})^{n} , as \beta\rightarrow+\infty , we expect that the solution to (9.4) converges uniformly on [0,T] to a solution of

(9.5)

z ˙ i (t) = 1 | C i (t) | ∑ j ∈ C i (t) V (z j (t) − z i (t)) \dot{z}_{i}(t)=\frac{1}{|C_{i}(t)|}\sum_{j\in C_{i}(t)}V(z_{j}(t)-z_{i}(t))

where

(9.6)

C i (t) = { j ∈ [n]:⟨ Q z i (t),K z j (t) ⟩ ≥ ⟨ Q z i (t),K z k (t) ⟩ for all k ∈ [n] }.C_{i}(t)=\Big{\{}j\in[n]\colon\langle Qz_{i}(t),Kz_{j}(t)\rangle\geq\langle Qz_{i}(t),Kz_{k}(t)\rangle\quad\text{ for all }k\in[n]\Big{\}}.

However, defining a notion of solution to (9.5)–(9.6) is not straightforward, as illustrated by the following example.

Example 9.3.

Suppose d=2 , n=3 . Let Q=K=V=I_{d} and z_{1}(0)=(1,1) , z_{2}(0)=(-1,1) , z_{3}(0)=(0,0) . Consider the evolution of these particles through (9.5)–(9.6). The points z_{1}(t) and z_{2}(t) do not move, because it is easily seen that C_{i}(t)=\{i\} for i\in\{1,2\} . On the other hand, the point z_{3}(t) can be chosen to solve either of three equations: \dot{z}_{3}(t)=z_{1}(t)-z_{3}(t) , or \dot{z}_{3}(t)=z_{2}(t)-z_{3}(t) , or even \dot{z}_{3}(t)=\frac{1}{2}(z_{1}(t)+z_{2}(t))-z_{3}(t) . In any of these cases, both (9.5) and (9.6) remain satisfied for almost every t\geq 0 .

It is possible to prove the existence of solutions to (9.5)–(9.6) defined in the sense of Filippov ¹⁶ ¹⁶ 16 We thank Enrique Zuazua for this suggestion.: for this, we can either use a time-discretization of (9.5)–(9.6), or use a convergence argument for solutions to (9.4) as \beta\rightarrow+\infty . Uniqueness however does not hold, as illustrated by Example 9.3. This naturally leads us to the following question:

Problem 6.

Is it possible to establish a selection principle (similar to viscosity or entropy solutions) for solutions to (9.5)–(9.6) which allows to restore uniqueness? In the affirmative, is it possible to revisit the clustering results of [GLPR24] and Problem 5 in the setting of (9.5)–(9.6)?

9.4. Diffusive regularization

We believe that (9.5)–(9.6) is also an original model for collective behavior. There are some similarities in spirit with methods arising in consensus based optimization (CBO for short), [PTTM17, CJLZ21]. With CBO methods, one wishes to minimize a smooth and bounded, but otherwise arbitrary function f:\mathbb{R}^{d}\to\mathbb{R} by making use of the Laplace method

	`\lim_{\beta\to+\infty}\left(-\frac{1}{\beta}\log\int_{\mathbb{R}^{d}}e^{-\beta f(x)}\,\mathrm{d}\rho(x)\right)=\inf_{x\in\text{supp}(\rho)}f(x),`

which holds for any fixed \rho\in\mathcal{P}_{\text{ac}}(\mathbb{R}^{d}) . This is accomplished by considering a McKean-Vlasov particle system of the form

	`\,\mathrm{d}x_{i}(t)=-\lambda(x_{i}(t)-v_{f})H^{\epsilon}(f(x_{i}(t))-f(v[\mu_{n}(t)]))\,\mathrm{d}t+\sqrt{2}\sigma\\|x_{i}(t)-v[\mu_{n}(t)]\\|\,\mathrm{d}W_{i}(t)`

for fixed \beta>0 , with drift parameter \lambda>0 and noise parameter \sigma\geq 0 ; H^{\epsilon} is a particular smoothed Heaviside function, and \mu_{n}(t) is the empirical measure of the particles. The point v[\mu]\in\mathbb{R}^{d} is a weighted average of the particles:

	`v[\mu]=\frac{1}{Z_{\beta,\mu}}\int_{\mathbb{R}^{d}}e^{-\beta f(x)}x\,\mathrm{d}\mu(x)`

where Z_{\beta,\mu}=\int_{\mathbb{R}^{d}}e^{-\beta f(x)}\,\mathrm{d}\mu(x) . Morally speaking, particles which are near a minimum of f have a larger weight. The drift term is a gradient relaxation (for a quadratic potential) towards the current weighted average position of the batch of particles. The diffusion term is an exploration term whose strength is proportional to the distance of the particle from the current weighted average. Results of convergence to a global minimizer do exist, under various smallness assumptions on the initial distribution of the particles, and assumptions on the relative size of the coefficients. They rely on the analysis of the associated Fokker-Planck equation, see [CJLZ21, CD22], and also [FHPS21] for the analog on \mathbb{S}^{d-1} . We point out that similarities are mainly in spirit—these results and analysis are inapplicable to our setting because there is no analog for f(x) —. Nonetheless, they do raise the following interesting question:

Problem 7.

What can be said about the long-time limit of Transformers with a noise/diffusion term of strength \sigma>0 ?

The question is of interest for any of the Transformers models presented in what precedes.

10. Approximation, control, training

Understanding the expressivity, namely the ability of a neural network to reproduce any map in a given class (by tuning its parameters), is essential. Two closely related notions reflect the expressivity of neural networks: interpolation —the property of exactly matching arbitrarily many input and target samples—and (universal) approximation —the property of approximating input-target functional relationships in an appropriate topology—. We refer the reader to [CLLS23] for a primer on the relationship between these two notions in the contex of deep neural networks.

For discrete-time Transformers, universal approximation has been shown to hold in [YBR ⁺ 19], making use of a variant of the architecture with translation parameters and letting the number of layers go to infinity; see also [ADTK23, JL23] and the review [JLLW23].

In the context of flow maps (from \mathbb{R}^{d} to \mathbb{R}^{d} ), it is now well understood that interpolation and approximation reflect the controllability properties of the system. The transfer of control theoretical techniques to the understanding of expressivity has borne fruit, both in terms of controllability results [AS22, CLT20, TG22, LLS22, RBZ23, VR23, CLLS23] and optimal control insights [LCT18, EGPZ20, GZ22, EYG23]. We refer the reader to [AL24, AG24, FdHP24] for the first universal approximation results for Transformers, viewed as measure-to-measure maps, using control theoretic tools.

Besides approximation, understanding the training dynamics of Transformers is another major challenge which we haven’t covered herein. As it is impossible to reference all works on this burgeoning topic, we refer the interested reader to [TLTO23, ACDS23, DGTT24] and references therein.

Acknowledgments

We thank Pierre Ablin, Sébastien Bubeck, Gabriel Peyré, Matthew Rosenzweig, Sylvia Serfaty, Kimi Sun, and Rui Sun for discussions. We thank Nicolas Boumal for referring us to [MTG17, CRMB24] and for clarifying comments.

B.G. acknowledges financial support from the French government managed by the National Agency for Research under the France 2030 program, with the reference ”ANR-23-PEIA-0004”. The work of Y.P. was supported in part by the MIT-IBM Watson AI Lab and by the National Science Foundation under Grant No CCF-2131115. P.R. was supported by NSF grants DMS-2022448, CCF-2106377, and a gift from Apple.

263 KiB Raw Blame History Unescape Escape

Abstract.

Key words and phrases:

1991 Mathematics Subject Classification:

1. Outline

Part: Modeling

Part: Clustering

Part: Further questions

Part I Modeling

2. Interacting particle system

2.1. Residual neural networks

2.2. The interacting particle system

Remark 2.1 (Collective behavior).

Remark 2.2 (Permutation equivariance).

2.3. Toward the complete Transformer

2.3.1. Multi-headed attention

2.3.2. Feed-forward layers

3. Measure to measure flow map

3.1. The continuity equation

Remark 3.1 (Well-posedness).

Remark 3.2 (Positional encoding).

Remark 3.3 (Mean field limit).

3.2. The interaction energy

Proposition 3.4.

Proof of.

3.3. A Wasserstein gradient flow proxy

Remark 3.5 (Doubly stochastic kernel).

Lemma 3.6.

Lemma 3.7.

Remark 3.8 (Particle version).

3.4. () is a gradient flow for a modified metric

3.4.1. The case of particles

3.4.2. The case of measures

Part II Clustering

4. A single cluster for small β\beta

4.1. The case β=0\beta=0

Theorem 4.1.

4.2. The case β≪1\beta\ll 1

Theorem 4.2.

Proof.

Theorem 4.3.

5. A single cluster for large β\beta

Theorem 5.1.

6. The high-dimensional case

Theorem 6.1 ([, ]).

Remark 6.2 (Invariant measures).

6.1. Clustering at an exponential rate when d≥nd\geq n

Theorem 6.3.

Lemma 6.4 (Cone collapse).

Remark 6.5.

Lemma 6.6.

Proof of.

Step 1. Clustering

Step 2. Exponential rate.

Theorem 6.7 (Wendel, []).

6.2. More precise quantitative convergence

Theorem 6.8.

Theorem 6.9.

Proof of.

Part 1. The angle θβ(t)\theta_{\beta}(t)

Part 2. The curve γβ(t)\gamma_{\beta}(t)

6.3. Metastability and a phase transition

Problem 1.

Problem 2.

Part III Further questions

7. Dynamics on the circle

7.1. Angular equations

7.2. The Kuramoto model

Problem 3.

8. BBGKY hierarchy

Problem 4.

9. General matrices

9.1. The repulsive case

Definition 9.1.

Theorem 9.2 ([]).

9.2. Pure self-attention

9.2.1. A review of []

Problem 5.

9.3. Singular dynamics

Example 9.3.

263 KiB

Raw Blame History