Emergence of meta-stable clustering in mean-field transformer models
Authors: Giuseppe Bruno
Federico Pasqualotto
Andrea Agazzi
Venue: ICLR 2025 Oral (2025)
Keywords: Mean-field limits
Transformers
Meta-stability
Clustering
Primary area: foundation or frontier models, including LLMs

Abstract: We model the evolution of tokens within a deep stack of Transformer layers as a continuous-time flow on the unit sphere, governed by a mean-field interacting particle system, building on the framework introduced in Geshkovski et al. (2023). Studying the corresponding mean-field Partial Differential Equation (PDE), which can be interpreted as a Wasserstein gradient flow, in this paper we provide a mathematical investigation of the long-term behavior of this system, with a particular focus on the emergence and persistence of meta-stable phases and clustering phenomena, key elements in applications like next-token prediction. More specifically, we perform a perturbative analysis of the mean-field PDE around the iid uniform initialization and prove that, in the limit of large number of tokens, the model remains close to a meta-stable manifold of solutions with a given structure (e.g., periodicity). Further, the structure characterizing the meta-stable manifold is explicitly identified, as a function of the inverse temperature parameter of the model, by the index maximizing a certain rescaling of Gegenbauer polynomials.

Published as a conference paper at ICLR 2025
EMERGENCE OF META-STABLE CLUSTERING IN MEAN- FIELD TRANSFORMER MODELS
Giuseppe Bruno Federico Pasqualotto Andrea Agazzi**
Department of Mathematics and Statistics, University of Bern 2 Department of Mathematics, University of California, San Diego
ABSTRACT
We model the evolution of tokens within deep stack of Transformer layers as a continuous-time flow on the unit sphere, governed by mean-field interacting par- ticle system, building on the framework introduced in (Geshkovski et al., 2023). Studying the corresponding mean-field Partial Differential Equation (PDE), which can be interpreted as a Wasserstein gradient flow, in this paper we provide a math- ematical investigation of the long-term behavior of this system, with a particu- lar focus on the emergence and persistence of meta-stable phases and clustering phenomena, key elements in applications like next-token prediction. More specif- ically, we perform perturbative analysis of the mean-field PDE around the iid uniform initialization and prove that, in the limit of large number of tokens, the model remains close meta-stable manifold of solutions with given structure (e.g.. periodicity). Further, the structure characterizing the meta-stable manifold is explicitly identified, function of the inverse temperature parameter of the model, by the index maximizing a certain rescaling of Gegenbauer polynomials.
INTRODUCTION
The transformers architecture (Vaswani, 2017), widely used in large language models, has played pivotal role in fueling the recent Al revolution. A key element of this architecture are self-attention modules (Bahdanau, 2014), allowing to capture long-range relationships within the data, e.g., in sen- tences where the number of tokens is large. However, despite the impact of transformers and the key role played by this mechanism, the theoretical understanding of these models still in an embryonic stage. In particular, a mathematical description of the representations developed by these models, and of the effect of different choices of hyperparameters on such representations is still lacking.
Recently, the work (Geshkovski et al., 2023) introduced a simplified model for the dynamics of tokens through the layers of the transformer architecture. This model consists of nonlinear Partial Differential Equation (PDE) modeling the evolution of the tokens interpreted as N interacting particles - through the layers of the network. As showcased by the authors in their work, this model displays clustering dynamics of the tokens through the layers, feature indicating the emergence of representations in the network (e.g., the emergence of concepts/consensus in next-word predictions).
Despite its simplicity, this model which, in its simplest form, results from the choice of specific value, key and query matrices in the self-attention mechanism displays a number of dynamical properties that have not been fully understood. In particular, while the mathematical description of the emergence of clusters was proven by studying the long-time (infinite depth) behavior of the token dynamics, which can be shown to generically converge single cluster as time (depth) goes to infinity, numerical simulations show that this phase of total collapse is preceded by several phases of partial clustering, which long-time time results fail characterize. Characterizing this phenomenon is crucial for understanding how representations develop in deep models: on the one hand, deep models are expected to produce more diverse and rich representations of the data, leading to improved generalization and performance. On the other hand, due to the finite depth of transformer models and the prohibitive computational and memory cost required to reach the single-cluster regime,
* Part of this work was carried out while at the University of Pisa. Contacts: {gi usappe. bruno, andrea agazzi iuunibe.ci bh, fpasqua lessquulottoqucdd.. ottc@ucsd.

---

Published as a conference paper at ICLR 2025
these intermediate clustering phases are likely play crucial role in the model's effectiveness, highlighting the need for a deeper understanding of their meta-stable behavior.
Contributions In this paper we study the clustering dynamics of the model developed in (Geshkovski et al., 2023) in the intermediate time range, i.e., between initialization and the emer- gence of the first set of clusters. To do so, we give a detailed mathematical description of the dynamics of the network in the large-context regime, i.c., when the number of tokens on which the transformer acts is large, characterizing the intermediate, meta-stable phase of the network before collapses to the single-cluster limit. More specifically we prove the following:
1. We rigorously prove the convergence, the number of tokens N' becomes large, to the dy- namics of the limiting law under the mean-field PDE. We further show that this convergence is non-uniform in fime.juusifyiin
the more refined metastability analysis carried out below.
2. We characterize the evolution of the N- N-particle measure close to initialization, identifying an initial phase where a certain number of clusters starts to form. We explicitly characterize the structure of the solution forming in this first phase e.g., its periodicity of the temperature parameter, the number of tokens and the embedding dimension d.
a function
3. We perform an in-depth analysis of the dynamics on longer time scales to prove that the periodicity developed in the first phase is maintained over time intervals of length O(In N). This proves the existence of meta-stable phase where the network learns richer repre- sentation of the data, before the total collapse phase predicted by (Geshkovski al., 2023).
We note that the above characterization has important practical implications, it clarifies the rela- tionship between hyperparameters such as the temperature parameter, the number of tokens or the embedding dimension and the emergence of representations in transformers by providing quantita- tive estimates on the "richness" of such representation and on the depth required to achieve it.
Related works This paper is closely related to the works (Sander et al., 2022; Geshkovski et al., 2023), where the model we study here was developed. In (Geshkovski et al., 2023), the authors highlight the problem of investigating the dynamic meta-stability of such models, problem that, to the best of our knowledge, remains open to this day. In this setting, the clustering phenomenon was studied in (Markdahl et al., 2017) in d >> 2 and in (Criscitiello et al., 2024) for d = 2. These works give sufficient conditions for the occurrence of the collapse of the dynamics for finite number N of tokens to a single cluster, which occurs with probability over the initial conditions of the system. These results, however, only hold in the -> oo limit, while our work provides a more detailed analysis, in the large N limit, capturing the meta-stable properties of the dynamics as it approaches slow manifolds of intermediate representations in the form of structured solutions before reaching the ultimate single-mode collapse. The dynamic meta-stability of mean-field transformers was also studied in the complementary setting of sepl-sepe configurations in (Geshkovski et al., 2024a).
The approach taken in this paper is closely connected to the one developed in (Chen et al., 2018; E, 2017), where the dynamics of the network's 's state through its depth are modeled by a differential equation. In our case, however, the state of the network can be broken down in N "particles" (as opposed to one in (Chen et al., 2018)) that interact in a mean-field way through a nonlinear PDE. More generally, this paper connects pioneered by the papers (Rotskoff & Vanden- 2022; Mei et al., 2018; Chizat & Bach, 2018), later extended more general settings and more quantitative estimates (e.g.. in (Agazzi & Lu, 2021; De Bortoli et al., 2020)). While those papers also model the state of the network as an empirical measure whose evolution is described by a mean-field PDE, in this series of works the dynamics describe the training of the model and not the evolution of the model's state through its depth. Furthermore, in these cases the propagation of chaos estimates are established on bounded time intervals, while we extend these estimates O(log N) time intervals.
vast literature on mean-field models for neural networks,
Another line of related works studies the fluctuations around the mean-field limit of nonlinear PDEs such the one discussed in this work. In a more general setting, (Carrillo et al., 2020) analyzed the connection between Fourier coefficients of a given potential and the stability of the homogeneous steady state for general McKean- -eeen--llsoov equations on the torus, (Lancellotti, 2009) characterized the fluctuations in the linear regime, and (Grenier, 2000) devised a general method for investigating the instability of the Fuler and Prandtl equations.
2

---

Published as a conference paper at ICLR 2025
Structure of the paper In Section 2 introduce the framework and some notation. In Section 3 we present the large-N convergence results. Section we discuss the long-time analysis of the dynamics and the corresponding meta-stability results. In Section 5 we present some numerical simulations and in Section 6 the conclusions.
FRAMEWORK AND NOTATION
We start by introducing the models under consideration in this paper. These were first derived in (Sander et al., 2022; Geshkovski et al., 2023; 2024b), by considering the transformer architecture discrete-time dynamical system describing the evolution of N tokens (t)()))=1,...... given by
22 (k 1) / (z. + Zu.i E/= e (lQzz(k) K2, Vzz(k)) k == O.... --1
(0)
where N _ gd denotes the normalization operator and Zg,; is normalization factor. The dynamics depend on (matrix) parameters Q, K. V (query, key. value) whose significance is inherited from the transformer architecture. For simplicity, (Geshkovski al., 2023) set Q = K = V ld and introduce the following continuous-time, simplified model problem describing the dynamics of (t): [0, oo) + gd-1 as a (layer-wise) limit of (1) in the spirit of (Chen et al., 2018);
(t) Pai(t)
(I Σ
(SA)
Here, Pay is the projection onto T given by Pay (x,, v)xr, (..)) denotes the Eu- clidean inner product in Rd. (t, is layer-wise normalization term defined by Zt, = Ek=re (fxx(t))ax)))) and 3 > positive constant, identified throughout as the inverse tem- perature. They also introduce a related model problem'
(t) = (()
1
e ((x10002 (0)
(USA)
N
modifying the normalization factor of the sum in (SA).
It is convenient to parametrize the dynamics (SA) and (USA) by "modding out" permutation invari- ance. This is classical and can be achieved by setting uut) == x Σm Sng(t), where we denote by
Sg the Dirac delta distribution at the point æ. This yields the continuity equation?:
O.ps + div(x())p)
R200 x Sd-1
(1:=0 (0)
on 1d-1
In the above formula, xl44 : Sd-1 + Tgd-1 is given by either XSA or XUSA, respectively defined as
/ aland) du(v)
XsAlu((x) = Pa
(SA-MF)
(x)
(fre Clown y
(us)lu((x) P.
(USA-MF)
where we used (z) = ((2.)) du(y).
Note that the dynamics are then interpreted as flow map between probability measures on Sd-1 In addition, as shown in (Geshkovski et al., 2023), (2) admits Wasserstein gradient flow structure.
Noting that (2) evolves both empirical and absolutely continuous measures, it reasonable to guess that the continuous dynamics (2) approximate the particle dynamics (USA) in the limit of an infinite number particles, at least in a certain timescale. the following section, we show propagation of chaos result, which puts this consideration rigorous footing and is prelude to our results dynamical metastability.
Which has the advantage of having gradient flow structure.
For these specific (non absoluully continuous with respect to Lebesgue) measures this equation needs to be interpreted with the weak formulation. Further, note that div here and throughout indicates the divergence intrinsic to gd-1 i.c. the covariant divergence induced by the standard metric @d-1
3

---

Published as a conference paper at ICLR 2025
3 PROPAGATION OF CHAOS RESULTS AND MEAN FIELD LIMIT
In this section, we establish the rigorous statements for the mean field limit of models (USA) and (SA). In what follows, we denote by Wi( the 1- 1-Wasserstein distance on P(S (the space of probability measures on gd-1).
Consider cloud of N points Su-1, 3(M) (N)
(N)
(s
which we think
of as initial conditions for either (SA) of (USA). Let the associated evolution be (() = (N) #1 (t), .x (t)). (t) and consider its empirical measure
(M)
MEON)( = 6 m0000
The following theorem proves that the convergence of a sequence of empirical measures E(M) to a limiting one μo (c.g.. if {zy} } are drawn iid from (4) is preserved by the flow of the PDE for any finite time f, i.e, that the N N-particles measure remains close to the dynamics of the particles' law as described by the PDE.
Theorem 3.1 (Mean field limit). Assume that there exists Ho P(S such that Wi (ME(N) (0), Po) + 0 as N + co. Let put) be the unique weak solution of the associated mean field dynamics (2) with initial condition p(0) = p4o. Then, for any fixed 2 0, as 0C,
W1 (1 (t), u(t)) +
Remark 3.2. Global existence of weak solutions to (2) in P(S is classical, and follows from a-priori estimates on the kernel x/14].
The proof of the above theorem is provided for completeness in Appendix A. Following classical references (Sznitman, 1991), we prove the desired estimate by propagating the error estimate at time t = 0 to any given positive time by using an exponential (in f) bound on the rate of separation of trajectories, also called Dobrushin's bound (Dobrushin, 1979), summarized in A.4.
To conclude this section, we highlight the fact that the estimate presented in the above theorem only holds on finite time intervals. This a result of the exponential degeneration in time of Dobrushin's 's bound, which is expected to hold in general (as it is a form of Cauchy stability). This degeneration is indeed inevitable for the system at hand (i.e., this system does not obey uniform time propagation of chaos) proven in counterexample in Appendix B.
As we shall see in the upcoming section, despite the above exponential degeneration, estimates on longer time intervals may still be established, providing insight on the qualitative behavior of the system beyond the finite-time time horizon. This is a necessary step to characterize the emergence of meta-stable phases, i.e., states existing for long in N time intervals.
4 DYNAMIC META-STABILITY RESULTS
In (Geshkovski et al., 2023), it is observed that the tokens (1) models (SA) and (USA) ex- hibit exponential convergence to single clustered configuration in certain regimes, time goes to infinity. In an intermediate time range, however, the authors observe the existence of long-timi time meta-stable states. Their simulations reveal two-phase dynamic: initial phase characterized by the formation of multiple clusters, followed by pairwise merging of these clusters, ultimately leading to single point mass distribution. In this section, we rigorously describe both the onset and the development of such dynamical metastability phenomenon.
4.1 META-STABILIH SETUP AND DISCUSSION
We restrict attention class of models which arise generalization of (USA):
0ip + div(n VWW * pd) = 0 for (t, Rzo x gd-1 for gd-1
11:=0== 1o
See Appendix A.1 the precise definition.
4

---

Nonlinear and collapsing phase. This phase consists of any finite-time interval after the quasi-linear phase. In this phase, the solution preserves the structure (e.g., periodicity) that emerged during the preceding phases. The exact evolution in this phase depends on Published as a conference paper at ICLR 2025
Here, the convolution operator is defined canonically as follows*:
* (()))))= = 1 ga--
(4)
where f : [-1,1] + IR, g: : gd-1 + R and døgs is the standard Lebesgue measure S In addition div and V are the intrinsic divergence and gradient in Sd-1, induced by the standard metric on gd-1
Note that, setting (1) == 8-1 exp(3g) we recover the case of the (mean-field) dynamics (USA- MF). However, the spirit of keeping our discussion general, we analyze the above equation for more general kernels W, making the following assumptions about its properties.
Assumption 1. W is a smooth (C function on [-1,1]-
Assumption 2. Let Ws be the k-th Gegenbauer coefficient of W (see (C.1.2) for the definition), and " == k(k+ d d-2)Wx, The sequence {7k { has a finite maximum "imax O. Moreover, this maximum is attained only at single value kmoz > O.
Remark 4.1. Note that Assumption is trivially satisfied in the transformers model. Furthermore, for d = 2 the coefficients Wy. are simply given by 1-11x(3, where 15 is the k-th order modified Bessel funcion of the first kind. Classical asymptotic estimates on Bessel functions (Abramowitz & Stegun, 1948, p. 360, 9.1.10) show that the maximum "imax is finite for all B and it is non-unique only on set of Lebesgue measure Furthermore, as 8 3- 00, kiax 23 /B.
Assumption 2 is related to the concept of H-stability introduced by (Ruelle, 1999) and, as diseussed in (Carrillo et al., 2020) for the torus, determines the instability of the uniform measure and the properties of phase transitions in the presence of noise. We will in particular show that Wkmax also controls the time scale and the type of emerging symmetries.
In what follows, motivated by the absence of preferential direction in embedding space, we con- sider as the initial condition the empirical measure associated with N tokens sampled independently and uniformly on gd-1, Although there is an extensive body of work on fluctuations of the mean- field limit ((ernandez & Méléard, 1997; Lancellotti, 2009)), these studies generally focus on finite time interval [0, 2J. which is insufficient for observing the meta-stable phase. Indeed, in our setting, the size of the perturbation scales ONN21 vanishing at any finite time T in the N - do limit. To observe the effect of fluctuations at initialization at macroscopic scales, a much more refined analysis is required. To carry out such analysis. we decompose the meta-stable phase in three parts: the linear phase, the quasi-linear phase, and the clustering phase, as depicted in Figure 1:
Linear phase. Starting from a very small neighborhood of the uniform measure, the pertur- bation coalesces towards the dominant mode determined by kmaz until it reaches roughly size -NN44)) We describe the dynamics in this phase in a precise fashion, providing quantitative estimates for the distance between the solution to the nonlinear PDE in (3) and its linearization around the uniform measure.
O(N-1/4), the perturbation remains close to the nonlinear evolution of the dominant mode sufficiently long so that its size exceeds a small threshold
Quasi-linear phase. After exiting the ball of radius
0 independent of the number of particles. Moreover, when exiting the quasi-linear phase, the solution is arbitrarily close (as N + oo) to the invariant manifold selected by kmaz. This manifold characterized by measures invariant under 2x rotations for d 2. For d 2 3, the description of the above manifold is kmaz provided in Section 4.4.
the specific form of the interaction potential W. In the case of Transformers, for W(g) = 3_ numerical simulations show that, after exiting the quasi-linear phase, the dominant mode collapses into kmaz clusters.
We recall the definition for integrable functions, however it is standard to extend this the convolution of smooth W with measure p..
5

---

Published as a conference paper at ICLR 2025
Linear Phase
Quasi-lineds Phase
5
Collapsing Phase
pa.
Pd 62
Poo
Kmaz 2
Kpax 8z
i
Figure 1: Schematic representation of our decomposition of the dynamics. Here, represents the "true" evolution of the system, while az represents the linearized evolution around the uniform measure /oo: fa denotes the evolution of the mean-field PDE with initial conditions fö the invariant manifold selected by kmaz, and fo. K is the approximation by Grenier's iterative scheme.
After the three phases described above, final phase occurs, which we refer to as the late-time phase. During this phase, shown (Markdahl al., 2017) and (Criscitiello al., 2024), the multiple clusters eventually converge to a single point, possibly going through several meta-stable phases with an intermediate number of clusters.
4.2 META-STABILITY. PRECISE STATEMENTS
To avoid technical complications and to improve clarity of exposition, we focus on the system (3) in dimension 2, identifying S' with [0, 2x). The analysis conducted here for the linear and quasi-linear phase can be immediately generalized the case d > 2 as we outline in Section 4.4.
We start by recalling our setup. We draw N tokens (0) at initialization, independently and uni- formly at random on gd-1, and consider the resulting (random) evolution under (USA). Let be the empirical measure associated to the tokens (x,(t) and let μo be at time = O. We define the perturbation Pt μe Mox. omitting its dependence on N for clarity, and the corresponding characteristic time for the linear phase by:
N-1/4
N
1/4
T
In
|PPll
Tmaz
where here and throughout ||.1| H-1 denotes the Sobolev -1 defined in (14). Moreover, 7, is well-defined and grows logarithmically in N, as consequence of the the Central Limit Theorem (see also Lemma E.3).
with p and negative exponent k
Theorem 4.2 (Linear phase). The measure ((T1) can be decomposed as:
((Ti) Moo + N-1/4 (po)
cos( (oo(kmaa00) + R.
(5)
||
where a remainder which satisfies RR|H=1 -111 in probability, and (po) denotes the Fourier coefficient of pn with index kimaz. Thus, N - the evolution at time T is kinaz
periodic and has rypical size N- 1/4 away from the uniform measure.
From Ti onwards, the dynamics are well-approximated by the mean-field PDE with initial data at Ti:
a - (f f() (0) = Prx + a cos 0),
(6)
where 1-4 (Pole pollx-1 Importantly, since the first equation in display (6) is invariant der rotation (i.c., translation in 0), will preserve the periodicity (if any) of its initial condition. Consequently, note that the function fo will be periooi for all times.

---

Published as a conference paper at ICLR 2025
Let 8 > 0 be fixed small parameter, independent of N. As the following theorem shows, the characteristic time to reach the end the quasi-linect phase is given by:
1 In(6/a).
T2
Tmax
Theorem 4.3 Douasi-linea phase). Let fo be the solution to the initial value problem (6) and small enough. Then, we have
+ +0,
probability as N - oo and
Wi Wi(u(Ti + (oc)
Remark 4,4. Recalling that f has periodicity kmaz. Theorem 4.3 states that the evolution in the mean field limit, when exiting the quasi-linear phase, converges to a tmarpeeii function, while being fixed umount away from the equilibrium measure. In other words, this result proves that when exiting the quasi-linear phase, p. (approximately) displays a very specific structure.
Finally, after the quasi-linear phase, establish the mean-field limit in the clustering phase. To do so we require the following assumption:
Assumption 3. Let ft == lima-:0 f (T2). and consider the initial value problem (6) with initial data fo imposed at time T2. Then, DC.
Kmas
W (F) Peluster) 0
where
Peluster
8
2m
Kmaz
Kmaa
Let us motivate Assumption Indeed, is relatively straightforward (using the gradient flow struc- ture) to show that, as time goes to infinity, f. (t) converges (in the weak sense) to sum of a number of delta masses located at several points (see Lemma E.5). Moreover, again by translation invariance in 0, the limit will be kmmz periodic. Thus, Assumption 3 is true if we only require the limit to be kmaz periodic superposition of a number delta masses.
To motivate why Meluster should have the specific form (8), let focus our attention on the dynamics of equation (6) on each interval of periodicity (for simplicity, consider the interval I : [- / /mmx)) <//mmx))... Following an analogous reasoning in (Markdahl et al., 2017) it can be shown that the evolution under the dynamics of (6) (restricted to via periodicity) every empirical measure will tend, time goes to infinity, to single delta function. Thus, Assump- tion 3 holds for almost every empirical measure, which leads us to believe that should hold for the particular f. under consideration.
almost
Under Assumption 3, we obtain the following statement.
Theorem 4.5 (Clustering phase). Let T3 > be a finite time. As tends to infinity, M(T, M(Tx-Ta+Ts) is approximated by f (T3) Wasserstein 1-distance, with an ervor vanishing in probability as N tends to infinity. In particular, provided that Assumption 3 holds, and picking the time T3 sufficiently large, as N goes to infinity, M(T.+Ta+Ts) is arbitrarily close (in the Wasserstein I-distance and in probability) to the measure Peluster defined in display (8).
Remark 4.6. Note that, without requiring assumption 3, analogous rigorous statement can be obtained in which the limiting measure is a kmax periodic superposition of delta masses. However, in this case number l 2 kmax of clusters, possibly of different intensity, may emerge.
4.3 DESCRIPTION OF THE PROOF
The technical machinery which enables us to achieve the proof of the above Theorems is contained in Appendix C. First, introduce key technical definitions and statements in Section C.I, and the linearization of our problem in Section C.2.
Linear phase. Section C.3 is devoted to the proof of Proposition C.10. the quantitative statement of Theorem 4.2. First, we improve, via Lemma C.7, the Lyapunov exponent on times smaller than T,, by studying explicitly the linearized operator around (lxo- This Lemma employs duality argument

---

Published as a conference paper at ICLR 2025
to avoid derivative loss, and allows to show quantitative improvement of Dobrushin's estimate, which in turn is enough to show propagation of chaos for times comparable to T1. These ingredients enough to conclude the proof of Proposition C.10.
Quasi- Minsi-linear phase. Section C.4, we construct a genuinely nonlinear solution which arises from the unstable mode cos(kmaz20 after time T1. The main technical issue in this regard is that we only have a poor estimate of the nonlinear growth rate C > of solutions to the mean-field equation (appearing in Dobrushin's estimate). In particular, direct approach would work if 2mmar > C (i.e., the nonlinear growth rate is slower than twice the growth rate of the unstable mode) for the relevant times Tj + T2, but this estimate does not seem be available. By employing Grenier's scheme (Grenier. 2000), we consider higher-order approximation to the dynamics, which allows us to relax the above condition to nkmmar > C, where n > the order of our approximation. This reasoning allows us to prove Proposition C.15 on the quasi-linear phase.
Collapse phase. At time T2, as N - oC, the evolution concentrates around the kiax periodic func- tion f. considered in assumption 3. We then use Dobrushin's 's estimates, together with the specific form of the infinite time limit provided by assumption 3 to conclude the statement of Theorem C.23. Note in particular that, after time T2, the time required for clustering to occur does not depend on N. Consequently, Dobrushin's 's estimates apply directly here.
4.4 SETA-STABLE PHASE IN HIGHER DIMENSIONS
In higher dimensions, a similar analysis of the meta-stable phase can carried out as in the lower- dimensional case (see D for details). The definitions and arguments from the previous sections naturally extend to the case d 2 2 due to the properties of spherical harmonics and Gegenbauer coefficients. The uniform measure remains an unstable equilibrium. The key difference is that the dominant mode f", which emerges in the linear phase as described in Theorem 4.2, is now superposition of spherical harmonics of degree kmax- solving:
0.ff- =- (f fo) (Ti) Mo + ELmner ajYkmannd
where Zimaz the multiplicity of spherical harmonics of degree kmar in dimension d.
This function belongs to the subspace Hkmas invariant for the continuity equation (3), defined as:
H4s = {u P(S (M, Yra) O vl not divisible by k),
that generalizes the invariance under this higher-dimensiona context as:
2x rotations. Henee, Theorem 4.3 can be reformulated in Amaz
Theorem 4.7 (Quasi-linear phase). Let far be the solution to the initial value problem (9) and 8 small enough. Then, we have
+ +
in probability as N 0o and
Wi W1 + T2), 8.
Again, Theorem 4.7 states that the evolution in the mean field limit, when exiting the quasi-linear phase, converges function in Hamerr while being fixed amount ó away from the equilibrium measure. This proves that the measure displays a very specific structure at the end of the quasi-linear phase. A natural extension of the analysis is characterize the stable equilibria of the dynamics when restricted to the subspace Hemaz. excluding the uniform measure. These equilibria must be finite unions of submanifolds of dimension at most d (see Lemma E.5). Furthermore, if the measure the empirical measure representing a set of points, being in Ht... (Bannai et al. (2015)) implies that these points form a spherical kmaz I-design. This condition indicates specific properties regarding the minimum number of clusters, their geometric structure, and the distribution of tokens, which have been extensively explored in the literature and an active area of research.
8

---

Published as a conference paper at ICLR 2025
Figure 2: On the left, the plots of " as function of for (top) and (bottom) are shown. The dashed blue line represents the value of yk generalized to real k. From these graphs, we observe that the corresponding values of kimax are 3 and 4, respectively. representing the predicted number of clusters in the large-N limit. In the center, numerical simulations depict particle trajectories for = 5 (top) and 9 = (bottom), with 10t particles whose initial conditions are sampled uniformly at random. On the right, the histograms of particle distributions at the end of the simulations displayed, showcasing the formation of 3 and clusters respectively.
5 NUMERICAL EXPERIMENTS
The first experiment aims to illustrate the relationship between the parameter 3 and the number of clusters formed: an immediate manifestation of the symmetry described Theorem 4.3 and the clustering tendency of the dynamics. Specifically, the left column of Figure 2 illustrates the values of the coefficients "/h for the Transformer model USA with 3 set to 5 and 7. Assumption 2 is satis- fied with kmaz values of 3 and 4, respectively. By Theorem 4.2 and Proposition C.20, these values correspond to the expected number of clusters observed during the meta-stable phase. To confirm this, we simulate in Figure 2 the trajectories of 10 particles on S! whose initial conditions are sam- pled uniformly at random. The angular ODEs system (10), equivalent Eq. (USA), is numerically solved using the Fuler method with time step di = x 10-4 Notice the different timescales of the two simulations, the meta-stable phase and the emergence of the predicted 3 and 4 clusters. The computation performed using Py Torch in double precision on a Nvidia Tesla GPU.
With the second experiment demonstrate the decomposi- tion of Eq. (5) and emergence periodicity in the token distribution. To further investigate the asymptotic limit on the number tokens, and given the computational constraints on simulating large number of particles, chose to examine the limiting case an absolutely continuous initial condition. Given that the results in Theorem 4.2 apply to general pertur- bations of uniform measures, we consider as initial condition, Mo. the uniform measure slightly perturbed by white noise. Figure illustrates the evolution of starting from Ag. The continuity equation (3), more precisely its angular counter- part (11(1)), is numerically solved using the Lax-Friedrichs method by discretizing the spatial domain into 10a grid points over the interval [0, 2T). with spatial step dæ 2n 10 and time step di = 0.05dr. The resulting dynamics consis- tently confirm theoretical expectations.
Hambor of STNNS (5)
Figure 3: Plot of the average time required for the empirical measure to exceed a fixed threshold. Note the logarithmic scale the a-axis.
In the third experiment, we demonstrate that the duration of the meta-stable phase is (ln(N)), predicted by Theorem 4.2 and Theorem 4.3. We fix 3 2 and run multiple simulations with the same setup the first experiment, with the number of particles ranging from N = 1000 to N 16000. The simulations terminated when the approximate total variation distance between

---

Published as a conference paper at ICLR 2025
Figure 4: Evolution of the initial condition given by the uniform measure perturbed by white noise (o 0.01) with 8 = Note that the emerging solution is 3-periodic as predicted by Theorem 4.3.
the uniform distribution and the token distribution (computed using 100 bins for the histogram) exceeds fixed threshoid. Figure 3 shows the average results from 20 simulations per particle number. We observe a clear logarithmic dependence in N of the exit time from the quasi-linear regime, aligning with the theoretical expectations. Code is available at this GitHub-Ree Repository.
CONCLUSIONS AND FUTURE WORKS
In this paper we rigorously investigate the dynamics of N tokens, interpreted as exchangeable par- ticles, through the depth of simple transformer model in the limit of a large N. In this limit, the dynamic meta-stability of the system can be captured by studying the emergence of symmetry at initialization through a careful analysis of the evolving fluctuations in the systems. In particular, we show that, after a time that depends logarithmically on N, the system exits from a ball containing the uniform measure displaying, approximately, specific structure (a certain periodicity in the case = 2) that depends explicitly on the inverse temperature parameter 33 and on the dimension d. Since the equation describing the dynamics of this model preserves such structure, this observation allows to capture the emergence of dynamically meta-stable solutions and to characterize them explicitly. In practice, this translates in detailed understanding of the representations produced by the network at finite but large depths, and how these representations vary with the depth as a function of N.
While 10t tokens might seem large number, many modern models can handle significantly larger context lengths. For instance, ChatGPT can process up to 128k tokens, and Gemini 1.5 Pro supports up to 2 million tokens. As the context lengths and the depth of large language models continue to expand, the large-N limit and the meta-stable effects discussed here will become increasingly relevant. Indeed, we demonstrated the existence parameters of Transformer model: number of tokens, temperature, embedding dimension, and model depth. These hyper-parameters cannot be adjusted independently of one another.
strong interconnection among the main hyper-
Admittedly, our work makes strong assumptions the structure of the model, e.g., absence of MLP layer and Q = K = = Id. The latter assumption can be immediately relaxed to Q² = AId, and quite possibly to QT assumptions are still far from realistic, numerical experiments in (Geshkovski et al., 2023) show that BERT displays qualitatively similar token clustering behavior to the one of our model. On the other hand, the choice of MIP layer will have a dramatic impact on the dynamical landscape (Cowsik al., 2024). The qualitative effect of this extra term on the dynamics depends heavily on the specific choice of MLP weights, but a linearized analysis similar to ours is expected to go through in neighborhood of the (new) unstable equilibria of the model. Despite these limitations, we believe isolating properties of the full model in simplified settings such as the one considered here is a crucial step toward building a comprehensive understanding of real-worke transformers. Combining this result with complementary ones on MIP and nontrivial Q, K, V is left for future research.
V. as the gradient flow structure is preserved in that case. While these
This work hints at a number of further interesting questions. First, an interesting avenue of future research consists of proving that Assumption 3 holds, i.e., that the solution to the PDE (2) as it exits the quasi-linear phase indeed converges to a certain number of delta measures. This result goes well beyond the ones presented in (Cohn & Kumar, 2007), as it requires characterising the solution of the PDE with a specific initial condition that may fall in the set of measure 0 in the previously established result. It would also be interesting investigate the effect of adding noise- motivated by considering randomly initialized linear layers between attention layers before training to the ODEs (SA) and (USA) and, consequence, on the emergence of symmetry and clustering in the resulting system.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Philippe Rigollet for discussions that motivated this work. AA and GB thank the Mathe- matics Department at the University of Pisa, where part of this work was carried out. AA acknowl- edges partial support of Dipartimento Eccellenza, UNIPI, the Future of Artificial Intelligence Re- search (FAIR) foundation (WP2), PRIN 2022 project ConStRAINeD (2022XRWY7W), and PRA Project APRISE.
REFERENCES
Milton Abramowitz and Irene A Stegun. Handbook of mathematical functions with formulas, graphs, and mathematical lables, volume 55. US Government printing office, 1948.
Andrea Agazzi and Jianfeng Lu. Global optimality of softmax policy gradient with single hidden layer neural networks the mean-field regime. International Conference on Learning Repre- sentations (ICLR 2021), 2021.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
Dzmitry Bahdanau. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv: 1r0v.04409.0473. 2014.
Eicchi Bannai, Takayuki Okuda, and Makoto Tagami. Spherical designs of harmonic index Journal of Approximation Theory, 195:1-18, 2015.
Jose Antonio Carrillo, Rishabh S Gvalani, Grigorios A Pavliotis, and Andre Schlichting. Long-time behaviour and phase transitions for the mckean- vlasov equation on the torus. Archive for Rational Mechanics and Analysis, (3551))63566999 2020.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances neural information processing systems, 31, 2018.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- tems, 31, 2018.
Henry Cohn and Abhinav Kumar. Universally optimal distribution of points on spheres. Journal of the American Mathematical Society, 20(1):99-148, 2007.
Aditya Cowsik, Tamra Nebabu, Xiao-Liang Qi, and Surya Ganguli. Geometric dynamics of signal propagation predict trainability of transformers. arXiv preprint arXiv:2403.02579, 2024.
Christopher Criscitiello, Quentin Rebjock, Andrew D McRae, and Nicolas Boumal, Synchronization on circles and spheres with nonlinear interactions. arXiv preprint wxXiv:2405.18275 2024.
Valentin De Bortoli, Alain Durmus, Xavier Fontaine, and Umut Simsekli. Quantitative propagation of chaos for sgd in wide neural networks. Advances in Neural Information Processing Systems, 33:278-288 288, 2020.
Roland L'vovich Dobrushin. Vlasov equations. Functional Analysis and Its Applications, 13(2): 115-123, 1979.
Weinan E. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 165):1-11, 2017.
Begona Fernandez and Sylvie Méléard. A hilbertian approach for fluctuations on the mckean-vlasov model. Stochastic processes and their applications, 71(1):33-53, 1997.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical per- spective on transformers. arXiv preprint arXiv:2312. 2023.
Borjan Geshkovski. Hugo Koubbi, Yury Polyanskiy, and Philippe Rigollet. Dynamic metastability in the self-artention model. arXiv preprint 10.00833,, 2024a.

---

Published as a conference paper at ICLR 2025
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clus-
ters in self self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2024b.
GitHub-Repository.
com/ cowphuuooooooo//oooooooooooooooo 6/
transformers metastability.
Emmanuel Grenier. On the nonlinear instability of euler and prandtl equations. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences, 53(9):1067-1099 2000.
Weimin Han, Kendall Atkinson, and Hao Zheng. Some integral identities for spherical harmonics in arbitrary dimension. Journal of Mathematical Chemistry, 50:1 10.1126-1135, 2012.
Daniel Han-Kwan and Toan T. Nguyen. Instabilities in the mean field limit. Journal Statistical
Physics, 162(6))1639-1653 February 2016. ISSN 1572-9613, doi: 10. 1007/110955-01--4455-6 URL http: doi org/ 1////1011000 100.11007011
wtw.11007ss119550 s11.///199955111 01 1455-6.
Emmanuel Hebey. Sobolev spaces on Riemannian manifolds, volume 1635. Springer Science & Business Media, 1996.
SG Krantz. Primer of Real Analytic Functions. Birkhääser Advanced Texts: Lehrbicche/Birkhäuserr Boston, Inc, 2002.
Basler
Carlo Lancellotti. On the fluctuations about the vlasov limit for n n-particle systems with mean-field interactions. Journal of Statistical Physics, 136(4))643-665, 2009.
Johan Markdahl, Johan Thunberg, and Jorge Gonçalves. Almost global consensus on the n-sphere. IEEE Transactions Automatic Control, 63(6):1664-100) 2017.
Song Mei, Andrea Montanari, and Phan Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115533):E7665-E777) 2018.
Victor M Panaretos and Yoav Zemel. An invitation to statistics in Wasserstein space. Springer Nature, 2020.
Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 75 (9): 198:-189-1935 2022.
David Ruelle. Smooth dynamics and new theoretical ideas in nonequilibrium statistical mechanics. Journal of Statistical Physics, 95:393 468, 1999.
Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transform- with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics, pp. 3515-33113 2515-3530. PMLR, 2022.
Robert T Seeley. Spherical harmonics. The American Mathematical Monthly, 135442) 1966.
Alain-Sol Sznitman. Topics in propagation of chaos. Ecole d' été de probabilités de Saint- Flont-Flour XIX-1989, 1464:165-251, 1991.
Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
12

---

Published as a conference paper at ICLR 2025
Appendix
A PROPAGATION OF CHAOS
A.1 PRELIMINARIES: WASSERSTEIN DISTANCE
Let (X, d) be compact metric space and let P(X) be the set of probability measures on X. The p- Wasserstein distance is defined as:
/ [((x; ((l.,))dy((z,))) 1/D
Wp(p, v) ==
min (Elll(u.))
is a well defined distance and in compact spaces satisfies for all p 2 1:
Wi(H,) v) < Wp(n, v) & diam(X) ==1 (i(,, レ)
Moreover, under the compactness assumption, for p to0% we have An - if and only if Wp(Jnn (i) -
Another common characterization of the distance W1 is a consequence of the duality theorem of Kantorovich and Rubinstein: when j and have bounded support,
(H, = suu (
- ())= X + R continuous Lip(f) $1).
For other properties and results related the Wasserstein distance, the reader refer (Ambrosio et al., 2008; Panaretos & Zemel, 2020).
A.2 ANGULAR EQUATION AND MEAN-FIELD PDE CONVERGENCE
In this section provide proof of the Dobrushin's 's estimate for the Transformers model. For clarity we restrict our attention to the case where d = 2 for the model USA. Nevertheless, the same arguments can be extended prove the estimate for d 2 3 and for the model SA.
An interacting particle system on S can equivalently described by associating the particles (t) S' with the corresponding angles (,(t) T. The angular representation of the system of ODEs for the model USA can be expressed as follow:
0,
cos(8.-0)) sin 0,).
(10)
Then the empirical measure of the angles, v(t) 1/2 E dø,(e), is solution of the continuity equation:
+ + (x(((v(t))((t) O.
(11)
where:
x()((0) = h₂(00 - w) de(w)
and (0) 8 @cm((0) With this notation and these definitions, Dobrushin's estimate can be
formulated follows:
Theorem A.1. Let pt, Vi be two solutions of Eq. (11), with initial conditions respectively po, Mo P(S)). Then VI. 2 0:
Wi(me, v4) 20 vo),
with C 1.
13

---

Published as a conference paper at ICLR 2025
We denote (z) the flow of system with initial condition E S' and vector field x[me]- Furthermore, given a generic measurable map f. we denote the forward of the measure j under the map as ## In particular Eq. (11) is equivalent to 2 T, #v
The proof of Theorem A.I is a consequence of the Lipschtzianity of the vector field x(a1] and the flow Te(pe), which are shown in the next two lemmas:
Lemma A.2. Let μ4, E P(S!), then:
lbx[=] x|u]lle~ S (olll.-W((un)
Proof. By definition of X:
/ 10310 w)] - du(w)] s
[|x|a] x[v]ll
sup
8E81
/
1z(00 - (dplw) de(w)) |17 10
w)
S | sup 0E31
Since kg(0-w) [ a
is 1-Lipschitz conclude by definition of W1.
Lemma A.3. Let μ E P(St), then:
Ixx S C, [Tilellesp S 50
where C' |1//lle..
Proof. For the first part:
(1x))](11) - xlx1(02)1 = / [ - w) - -hg(02 w)ldu(w) & WWllee~ (01 02).
For the second part:
[(i (12)) - T((i)(())) d dt
d
dt TB2 = =xxlxc - xlaa] (()) T,(B2))
The conclusion follows consequence of Gronwall's lemma.
Now we can return to the theorem:
Proof of Theurem A.I.
Wi(u,, = Wi (oo S < Wi (o + Wo
(12)
For the first term:
sup / do +vbu W1 W1 uu))
1 g
sup 1ollesp51
(6o dlus - -vol
=ect sup
(e ce Telps]) divo vol= s
51 g1
5 ce Ct
1 dduo vo] -WooW((Ho
sup
|| 21 S1
where the last follow from Lemma (A.3).
14

---

Published as a conference paper at ICLR 2025
For the second term:
Wi Wo Hvo))
lolll S1 s2 VboTelll - Φ Ti(Ml) dvo 1
sup
- g2
[Tel4] - Tz(u]l dep = 2(1).
Now,
dA(() dt
- /
d
Telm)
T:(ul dup dup
dt
[xlee] Tilml x(vi] Ti(uildes den s
- f [xled] o Tilml - xlua o Telvil duo duo + / [xlee Tely] x(vs] o Tipp]] des-
By Lemma (A.3) for the first term, and by the properties of T, for the second term (in particular He Ti # we get:
d'(t)
SC / [Tlp - Tilm] duo / Ixleel x(vi]l dv
dt
$ Ct((t) + [Ixlm] - xlville= S C(A(t) + (1(,,
where the last step follows by Lemma (A.2). As consequence of Gronwall's inequality it holds:
x(t) S S." oa (l Dr)dT
Going back to Equation (12) we get:
Wi < C vo) + +C f. Wi
And to conclude just need to apply another time Gronwall's lemma: Wi(me, 14) S 20ttWWWooo vo).
An analogous reasoning in the higher dimensional case, for the dynamics (SA-MF) and (USA-MF) yields the following more general result.
Theorem A.4 (Dobrushin's estimate in d dimensions). There exists constant C 0 depend- ing on 3 such that the following holds. Suppose that μr, Ut are two solutions of (2) in the class C(R 20, P(gd-1)), with initial conditions respectively μo, Vo E EP(( (Sd-1)). Then VI 0.
W1 (Pe, v1) 2Ct (HO, (HO,PO). Mo).
Moreover, C can be computed explicitly. In the case of (SA-MF),
C 1V (3t
where
Kg(r, y) ((r.0) Pa(y).
Finally, Theorem 3.1 is a direct consequence of Theorem A.4 and we omit the proof here.
B ON THE EXPONENTIAL IN TIME DEPENDENCE
In this section, we demonstrate through a straightforward counterexample that the exponential time- dependence in Dobrushin's 's estimate for the Transformers model cannot generally be improved. The time control is not only non-uniform but must also exhibit exponential growth.
Let's consider for simplicity the case 3 1 and N The ordinary differential equations system obtained from Eq. (SA) is:
i1
Pa,
1 (ele + (((1)22)
Z
32
Pay
(eleg,01) + g022 (2)
Zz.2
15

---

Published as a conference paper at ICLR 2025
It's easier to examine the angle version of the previous ODE system:
01 - c000 92) sin(O) 02)
Zp..
(22 Z0. cec0l(0=-00 01) sin(02 01).
Using that Zz,1 = Z3,2 = Z, then w = 02 - 02 satisfies:
ecos(w)
sin(w).
Since (() = 0 and (t) constant solutions, then for every wo E [0, x the corresponding trajectory ((t) is confined in [0, x Vt 2 0, and this implies sin(w(t)) 2 0. Moreover, + S pos(w) hence:
W 2 22 sin(w),
which implies, by an ODE comparison theorem:
(t) s 2 cot 2ot (c
(13)
where the right-hand side is the solution of 7 - 2/4 sin(?) and depends the initial condition.
Now consider the initial conditions 0f == and 0g = क - e, The corresponding empirical measure is:
1
Pe (0) 2 (ō, + Ôx-
We want a bound of this kind:
w (m),
S f(t) Ve V2 0. (O,
for some function f-
Notice that n4 (0) is constant solution of our system. The Wasserstein distance w.r.t. the geodesic distance on S² is given by:
1 - (000)))())0=0000------000-000 02(1) w(t)
W (1, = -0
2
hence the function f must satisfy:
(tt)
2 f(t
Ve > V2 0,
e
with initial condition w(0)
2c. Using the bounds in Eq. (13):
w(t)
cot - 2222))
6
with s.t. cot (co)-1(ee) 2e, i.. e tan(E). Hence it must hold:
1
f(t)
é 2 cot cot--l(tani) (tan(e) Ve 0.
Passing to the limit + 0, using first order Taylor expansion in 0 for cot (x) x get:
x,
f(t) 2 lim c- - - cot cot-tootaaa (tan(e))))))))
tan 21/2² 2t/c²
lim
This implies that cannot improve the exponential time dependence for the stability.
16

---

Published as a conference paper at ICLR 2025
C APPENDIX: METASTABILITY
C.1 PRELIMINARY DEFINITIONS
C.I.1 SOBOLEV SPACES OF NEGATIVE ORDER
Given domain & C Rd define for every E N and p E [1, +oo the corresponding Sobolev space of negative order as the dual:
W (9) = W
where p/ p/ 1) and (S2) denotes the closure of ( the usual Sobolev space (.). The corresponding dual norm for W- (S2) is given by:
|uulww (o) sup { (u, v) W (2), Iellww = 1}
(14)
For any and p E [L, +oo the space W-s (S)) is Banach space. Moreover, if < 400, W- (I)) is separable and reflexive.
These definitions can be naturally extended to compact Riemannian manifolds, while the non- compact case more subtle (the reader can find a reference (Hebey, 1996)).
Most of our discussion is set within the space Wf)))( also referred to H *(2). These spaces can be conveniently characterized using Fourier series:
H'(T') {u (T"))) (1l) co}.
where the norm on H*(TM) is defined as:
|
Wx (1 [8:13)*. Σ
kEZd
Additionally, note that the Dirac delta is an element of H (IR)) for every s < -d/2.
C.1.2 SPHERICAL HARMONICS AND GEGENBAUER POLYNOMIALS
spheres gd-1 in the following way. Consider P the space of homogeneous polynomials p(x) Rd C of degree / in d variables, and let A denote the subspace Pr of harmonic polynomials. Then the spherical harmonics of degree are an orthogonal basis for:
The classical spherical harmonics can be generalized to
H+ = ( + C exists E A such that Yı(x) p(x) for all a E 8d-1).
While explicit formulas can be given by induction, we will just need the following:
Theorem C.1. If A is the Laplace- Beltrram operator on Sd-1, then for all Yi E Hi, one has:
AY 1(1 + - 2jy
Furthermore, spherical harmonics have strict relationship with the convolution defined in Eq. (4), through the so called Gegenbauer polynomials P
These polynomials, also known as ultraspherical, are orthogonal on [-1, with respect to the prob-
ability measure with density vxT(a+ formula:
[(a+1) 12)-- An useful definition through Rodrigues
P (1) = (a + 2)
+ark++a+i a dth [0 - [1)ttt---1]]
d²
They generalize Legendre polynomials and Chebyshev polynomials, and satisfy the so called Funk- Hecke formula:
17

---

Published as a conference paper at ICLR 2025
Theorem C.2. Let f
[-1,1] - R such that t (1 (1- (1)) f(t) is integrable. Then the Funk-
Hecke formula holds:
[ ga f((x,)))
\,H((a),
with the constant da given by:
2T r( (( N [" P. (t)((t)
(1- ( dt,
where [()) is the gamma function and Y; is any spherical harmonics on gd-1 of degree I.
Other properties discussed in (Han et al., 2012; Seeley, 1966). Although the nomenclature may not be standardized the literature, we refer to ò as Gegenbauer coefficent f of degree I.
C.2 THE EMERGENCE OF SYMMETRY
For clarity, the following discussion is restricted to S' (refer to Section 4.4 for adaptations to higher dimensions). Specifically, notice that identifying each pair of points I, E S with the corre- sponding angle 8, y [0, 2T] and using the relation W((r, v)) W (Ccos(0 11100~00 r)). the definition of convolution provided in Eq. (4) aligns with the classical notion of convolution on $ Furthermore, thanks to this change of variable, the Gengenbauer coefficients We correspond to the Fourier cosine expansion of W cos. Thus, for simplicity of notation we will continue using the notation W to represent the composition W cos.
Remark C.3. In the case of Transformers, both Assumption I and Assumption 2 are satisfied since 2 W(0): = ge cos(9) = -E B I (B) cos(k0), where (3) are the modified Bessel functions of the first kind.
Consider Moo as the uniform measure on 81, and let the intial perturbation be 140 == 11o0 + Pa in P(S)). If b4 evolves according to the following PDE (consequence of Eq. (3)):
{ 10,pe = - (Me VW pr),
(15)
(u(0)= po,
then the perturbation Pt satisfies:
{ [..p: = Lu (pt) - Colp. pc), ((()) = Pu.
(16)
where Cu (p) -1:0 * p denotes the PDE's linearized operator around Plo.
Remark C.4. Since VW * Mo 0, the uniform measure is an equilibrium of the system.
The objective of this section is to demonstrate that the dynamics of the system near Poo are governed by the maximal eigenvalue of L
The core estimate is provided by the following lemma:
Lemma C.5. Let n 2 O. For any g (S') and Wm..o(S'), the following inequality holds:
(S0og. (0 n= Cx /C
Proof. For every k S n:
k
k) (g) ff fonn-saattg.a )))
(f*g)) 0 g)12 L2 =
Σ
0=
(00-1 gu+1 g, 0rg) 1fott1g0000
11 ||r 9llz= 2 + (f. 1 (0:0000)))
fllz f
Calllfllww..
18

---

Published as a conference paper at ICLR 2025
For sufficiently small time intervals, the evolution of μ4 be approximated by the solution of the lincarized operator. The deviation from this approximation will be estimated in the next paragraph.
C.3 THE LINEAR PHASE
Dobrushin's 's estimate applies globally on S and may be suboptimal in the vicinity of the equilibrium Moo- While in general is not possible to have a uniform or polynomial dependence with respect to time (see Appendix B), this estimate can be locally refined by considering the singular values of the linear operator Luze= H (S') + H (8'). Specifically, under Assumption (1) and Assumption (2), for every f E (S*) with k E Z (refer to Lemma E.2):
cl 2mmaz
2T
Remark C.6. Hereafter, 2x Th will be denoted simply as "
Utilizing the Lagrangian flow, it's possible to bound the growth of the perturbation in small neighborhood Poo using Tmez- In particular, the following lemma can be proved:
Lemma C.7. Let k 2 0. There exists e >> 0 such that, if [| e, then
Ti(c, =
In
(IPoll
Ymaax
dell-defined, and for every [o, (e, the following inequality holds:
|| S 331ppollss=e22matt
Remark C.8. The constant is independent of Po- Furthermore, this lemma implies that, within the time interval [o, Ti] the evolution of the measure μ4 remains within B (Hos, 3c).
Proof. Given test function 3 E (S)), let the be the solution of:
( (o.4/ Gu (1ha) (VW pa)dsrl. [0,
(17)
1/4 = V.
Hence, by construction:
0, (ps, v.)) = (0,p, v² + (p, 0a8))
(Erm (ps) Oo(p. VW pa), (ps, Fr (1b) + (VW xpa)0oo))
This implies:
Ipelll-
sup
(pe, (3
sup
(pt: 2/r)
11 m1
1elllg==
sup
(po. voo) s
(18)
11
sup 1
|
Next, we need an estimate for [/ullle The energy of the time-reversed Eq. (17) satisfies:
1 d
2 ds 1/65 112gs (Cam Vi--o V1-o) + (VW Pt- Pr Ve-e) S
Ymax 1 faxx C1 1 VW Pi-sol $
Ymar (// 11 + C1 VW |
< (mmaz + C C'Ilpt 11
where the second row is an application of Lemma C.5. Hence, by Gronwall's lemma: Tmar fo ll
|| 1r
19

---

Published as a conference paper at ICLR 2025
Substituting this bound into Eq. (18), we obtain for all 2 0:
| S || ke lip-ll
(19)
The conclusion follows from standard continuation argument: assume
llpall==
3 Po IH - 7 Then, by Eq. (19), in the interval [0, (i(e, it holds:
S Poll km?mar t-O C To 11olll <
I | = Tmant 3010-00 - Tmoo
emaax
- | a KE
- Tmar
Touns Ty(c.p))
S [|(0))ll==ee Tmartt
Amas.
If is such that 3Cc 1, then |ppll 311poll e in the interval [0, Ti (c, po)]- concludes the continuation argument.
This
Twux
Now we want to control the distance between and the solution p of the linearized PDE:
{ 0 Fux (p"))
(20)
P6 Po-
By considering loss of regularity derive the following estimate:
Lemma C.9. Ler l: 2 0. For all E [0, (e,((,,)))) the following holds:
ll p/'1 -- x-1 S Cllpoll?r-+ee 2imas e
Proof.
The difference p p² satisfies the PDE:
0.(p - p-) Fu (pi- v!)) -0)-
ColpvV * pc).
Thus,
d pll p 1|- S Tmaz Pl 11-11 + 1lee VW polln= S
dt
S Tmaz l| P: 15: || - + ///p|l==//VVM Pellwno S
I Imaz Pr pt: -| + Cl
Applying Lemma C.7, we obtain:
d llc Pt - S "mar [p: -----P-.ll------ + Cllpoll?r se 27mxx
dt
The conclusion is consequence of Gronwall's lemma (Lemma E.
can conclude this section with the following proposition, which describe the measure μ at the exit point from B (Joo, 3e). In particular, we prove that after time Ti, defined as:
(e)
In
(IPolle
(21)
Yimax
the perturbation of the uniform measure consists of a dominant mode characterized by kmaas along with remainder made up of two components: the error from discarding the other modes and the error resulting from linearizing the PDE.
Proposition C.10. Let l 2 O. There exists co > such that if [I H < < co. then:
(Ti(e.P) = Poo + (Pu)kmaz cos(kman + R(e, (|poll=T=))
[poll
where Tí (e, po) defined in Eq. (21). Furthermore, let 7 = m The then:
R(c, x) = Oh x-1 [Ipolln= ( e
+ è
|lollles
20

---

Published as a conference paper at ICLR 2025
Proof. Notice that the solution Eq. (20) can be explicitly expressed in terms of the Fourier coef-
ficients of the initial condition:
pf
(22)
KEZ
Let y ma 7h. By substituting Eq. (22) into Lemma C.9 and using the definition of Ti, we get:
Pn =ph, OH (1l -e
2mmaz Ti
= =(pu)kma emmurll cos(k k) +
Σ
(Pu) + +0 (llooll - 2mazll))
by thmmxx
(0o)kmus cos( (o +
(Do)
Ymas
Likz OH-A-1 (e2)
-Ipolla++
|IPolles-
te
(po)kmaz cos(kmaant) + OF--8-1
|10ollor
(Tooller
Tmaa
+
- [|Po|HH- -H
||
We turn to our discussion the subsequent "quasi-linear" phase.
C.4 THE QUASI-LINEAR PHASE
When the measure 14 exits the ball By-+ 3c), the estimates based the linearized equation are no longer sufficient to describe the system's 's evolution. To address this, we build nonlinear approximation following Grenier's iterative scheme (GGrenier, 2000; Han-Kwan & Nguyen, 2016)).
In the preceding paragraph, we isolated the dominant mode cos(kmaar0) 0). Let a >> 0 and consider the solution of the following nonlinear PDE:
f - 0g (faV f()
(23)
J Mox. + com(kmaz0))
We will demonstrate the existence of 8 0, independent of a, such that the evolution of fa within the ball Bн x (Poo, 5) is still controlled by "maax In particular, we prove in Proposition C.15 that after time Tz:
(8/0a))
(24)
V/max
the nonlinear evolution of the dominant mode moves outside neighborhood the uniform mea- sure, with the radius of this neighborhood being independent of a. Moreover, in Proposition C.17, we provide a bound on the distance between fo and any distribution close to it at time 0, valid over the time horizon [0, T2).
Let's s introduce the following definitions essential for Grenier's scheme:
(t, 0) == e emaaz (os(kman
K
fa (t, 1oo Σ a gj.
where the functions gr are defined recursively by:
f (0, Lp )gk = - - Og (gg VW 9k-))
(25)
(k((0) Ak 2
Notice that:
a. K + (o NWW fo,k) = k.. =
Σ
a' @''0oll, *g.)
(26)
IS jH2K+1
21

---

Published as a conference paper at ICLR 2025
Lemma C.11. Let K 0 and s 0. Then vk S K:
| S CKeetmort
Proof. The proof proceeds by induction, showing that ll ll H S Chethmmat where == K + By definition, we immediately get |192llæ= S Ceaat
For the inductive step, consider k > 0 and assume the hypothesis holds up to k 1. Then by Eq. (25):
d di 11 Iam S 11 Pxll 9all
100(9) VW gh-s)lla" & S
Tmaz 7mazl19k|1 Ha +
119, VW gk-jill +11 S
119k Tmaz
Hm- +
119511 611 VW * -e-illlw-kiaa kI S
- Ymaz llah Hm + 5=1
19 e V [e S
S TmazllgN Ilgh| Hm + C
9gs IH 11x S
5=1
Tmaz 1l Il Ha-N + Cke
Therefore, by Gronwall's 's lemma (Lemma E.1):
11ghl H S Che
To conclude, recall the definition of n as K 8 Remark C.12. This lemma implies that if amatt S. < 1, then:
I| Moollet S CCK] Σ lae tmxtt S (ae
and similarly:
r 1100 M 2 |1ggi|1 - Cx K
Tmaxtys 2 Taar
We now need to estimate the remainder Ra that appears in Eq. (26): Lemma C.13. Let K > 0, 2 0 and 6 < 1. If Imart < 8, then:
R K || Üx (mor)) K+1
Proof. By definition of Ra.K;
R.. || <
Σ
ISjICK
H+12K+1
Σ
ا [g VW g1l 5
15HKKK
j422K+1
< CK
Σ
-
ISHICK
j412K+1
22

---

Published as a conference paper at ICLR 2025
Finally, Lemma C.11 and the assumption imply:
|Ra I|= S CK
@ S CK (artmaat) K+1 ISAE j412
Σ
Now, we can compare the evolution of the dominant mode fo defined in Eq. (23) and Grenier's approximation fa,k, Let ra. := f Lemma C.14. Given 8 2 O, there exist constants Ko and 5 > 0 (independent of o)) such that, for every K N Ko for every t such that we aieata < Ö, the following bound holds:
I| I| S (((oetmaat)
where the constant C' depends only on 8, K and W.
Proof. Combining Eq. (23) and Eq. (26), ra solves:
arak 0g (fa,k VW ra. 0g ar VW pa.k) 0g (ra. VW x pa.
Ha.k
To study its energy:
1 d
| lr 11
(0+ (fa k VW * ))a paa.k FoH
2 dl
(0o (ra.k VW VWa
fa.k) (o.k tuh
(27)
(0o (raK VW * ra.) jark
+ (Ra.K pa.k) Ho.
The three terms can be estimated separately:
first term: -(0g (fa, VW k) rark) = ((0ofm.)) VW * aaak (0:K) H
(fo.k AW * park jack) S
<C1lfa.K ||aa||||111 1 park 1
second term (using Lemma C.5):
-(8g (ia, VW * fa.k) ,a,K) S Cal 1| alrlaaa
third term (again by Lemma C.5);
-(0 (ra. VW * n), ra,k) H S Cl 1
Combining these estimates, and utilizing Lemma C.12 and Lemma C.13:
d
Hot1 + (lo. ol || + a 1| di I | SC(])aaki I
SC(1) + Cx + || (| al I| +Ck t
The conclusion follows via the standard continuation argument.
Assume || S C (actmaz)) k and t 00 سيل In (1) If sufficiently small such that & < 1, and K sufficiently large such that 30 K Kmman then:
d dt
C(llooll1 Cgac + CK (met) ")l Ilist
+ Cx (at (ae Imaat) K+1 <
< 3. 1 1| + CK mma K+1
23

---

Published as a conference paper at ICLR 2025
By applying Lemma E.1, which requires the assumption 3C < K Ymaz obtain:
Ckak-1
|ra. I S
S &C &
Kymant
Ymmo
Since 8 < the solution remains within the desired bound and can be continued.
The results presented in this section lead to the following proposition, which provides a lower bound on growth rate of the dominant mode.
Proposition C.15. Let fa be the solution of Eg. (23). Given & 2 O, there exists d > 0 (independent of (x) such that for T2(ō, () defined as in Eq. (24), the following holds:
11 PoolH 2 C8.
(28)
Proof. By Lemma C.12 and Lemma C.14, we have:
|l Me
2 | ma - Mooll - 11 -pakklla 11
2 Crae Ck (ar Tomoet)
Cgae?must
To conclude, observe that the maximum t for which the assumption in Lemma C.14 holds is pre- cisely (2(6, a).
Remark C.16. Note that the same estimate holds for < O. Indeed, the only lower bound is needed for 91. which grows as emmuxx in every H space.
To conclude this section, we need to establish an estimate on the distance between given solution Me and the evolution of the dominant mode ft²
Proposition C.17. Consider a solution μa of Eq. (15) with initial condition po. and define h4 H4 - fe- If | H I|mo - föllH S Co, then there exists ó (independent of a)) such that:
|1ell S 3/1boll ke
for every E (2,6 a)].
Remark C.18. We will apply this proposition μi Moo + p(Ti).
Proof.
Using Eq. (15) and Eq. (23), the distribution h satisfies:
0,h- - ( VW & μ) + VW fa)
Og (WV fa) VV h) do(WW h)
--HoAAW h
-VVV - Moo)
-Oo((f- - VW hh)
-oVWWW h).
We can adapt the I Lagrangian flow argument used in Lemma C.7 to bound [[b]] Let E (SL), and define 1/2 as the solution of:
O.v..= = Mo Mo AW t/s
VW (f -Moo))00
-W [(S - -1xx)000111
-VW * h2100ve
8/ J.
To estimate the energy of 1/t, proceed follows:
24

---

Published as a conference paper at ICLR 2025
0s 1. Mox (AW * Vr- 4
2
+ (VVW (f - -4cc (000) 1/- x)H++ H
(VW |( -1oc )o0 s), 1he- H
VVV h100040-0: Ve-
We estimate each term separately:
first term (Lemma E.2):
Poo AWW tt--s 1/2 S Tmaz 1lver-1111
second term (Lemmaa C.5):
([VW (f 001
V/-5) 41 S011 -l VW * (fi - Woollwwr
<0 11 flft Mool112.
third term:
(VW+ 1(fa - Moc )0x =(fa Mox )0 VW v/ 127
+( 1100 )o 00 V t' 124 12
+(f 1o0
10 11 V 1h 12
SlI
1
11 m
VW -1-olle +
SCIl 100 |1> 1000=>/1/0-111
fourth term (Lemma C.5):
([VW h10o0t-s, wh-5) S -1
Summarizing, we obtain the following differential inequality:
01llv=-llam (Tmaz + CIIIE-= Moo Mooll 21 l1ee-s Ht )|
and by Gronwall's 's lemma:
| 1///2 e Jo ClI - w ads
(29)
Since by construction 0 (he, v/x) = O, that implies together with Eq. (29):
= sup =
(he, +) = sup
(he, the) =
sup (ho, vo) I | sup 1| -==1
|
< | [[boll=-w emant C1I - 12a hr -
The conclusion follows by continuation argument. Suppose [helll= 31 Maoolll--ee mma and t E [0, T2]. Then:
|/4 Ia <11((0)|l H be Ammxxttff D+ Cae Tmas +3] holl ll m
<|[b(O)]l 11- Tmnas Co ?muut
<1) - Timass Smaz
where the hypothesis - Ca was used in the second line. By choosing small enough, such that C8
< 1, we get the desired estimate.
Tmaz
25

---

Published as a conference paper at ICLR 2025
C.5 COMBINING THE TWO PHASES
In Section C.3, we analyzed the initial evolution of perturbation Moo + p of the uniform mea- sure by comparing it with the solution of the linearized PDE. After time 7i, the dominant mode emerges. As the perturbation exits the neighborhood By (Moo, 3e), a higher order approximation using Grenier's scheme is required track the nonlinear evolution this dominant mode, as de- scribed in Section C.4. We now consolidate these findings into the following proposition, which we specify to the case of empirical measures.
Let ffr) JNEN denote sequence of random measures in P(S'), (()),, where Mo is the empirical measure associated with i..d. particles sampled from the uniform measure on S Define the radius of the ball in which the linear phase occurs as follows:
(HD) 1oo )
1I/ Meo =
and consider the corresponding solutions pa and fir of Eq. (15) with initial conditions μú and Moo0 + x cos(kimaze respectively.
Remark C.19. Since (16 kmme O11 - Msal H- ON + 0 N oo. Proposition C.20. There exist a constant 6 and two sequences of times T and Ty (depending on v such that:
1 - Moo ll 2(81) > 6
uniformly in N, and
Ilu +T 11g (1 -
as + 3o in probubility.
Remark C.21. By Proposition C.15 and Remark C.12 we have both lower and upper bound for I - Moc Hl 2(8t), uniform in N.
Proof. For clarity, denote p = p. Moo. Applying Proposition C.10 with e = N- and noting that | Ia + 0 in probability (Lemma E.3), we obtain:
PTP
Mox + a cos cos(kimaz0) R.
where TN In
NNT IIr-1
and:
Tmar
R = =0
||poller [| +
Tmaz
Note that, by Lemma E.3, we have:
N ||
1|
(30)
probability. This implies that the hypothesis Proposition C.17 are asymptotically satisfied for the initial condition p Therefore, there exists constant 5 0 such that f.N. defined as above, satisfies:
I 42
Tmatt
for all t E NN, with TNN 1 In (2x) Thus,
1 fng. 11 5 N
|R || + +
in probability, as recalled in Eq. (30).
Finally, observe that || I Poo Il 2(81) > 8 is immediate consequence of Proposition C.I5.
Remark C.22. The convergence also holds in Wasserstein distance, as proved in Lemma E.4.
26

---

Published as a conference paper at ICLR 2025
C.6 THE CLUSTERING PHASE
In the previous steps, only an approximate understanding of the trajectory of the dominant mode, f was necessary. Now, however, focus on more specific case. Numerical experiments (see Figure 5) conducted on the Transformers model, W(0) - reveal that after initial phase where f surpasses fixed threshold, i... Wi(fe' (loo) > the system converges to kmaz equally spaced clusters. The time required for the transition from such threshold to the clustered state is independent a. This assumption summarized in Assumption 3.
cos(1),
Although this remains an assumption, we believe it holds true, particularly in the case of Transform- ers. This belief strongly supported by the fact that probability measures invariant under kmaz- fold rotations form an invariant manifold for the system's dynamics, and the limiting measure must be sum of Dirac deltas (see Lemma E.5).
0.5
0.4
WotetAAl::a0...0 0.00
0.3
Wi =
Me 0.001
0.-
Wi. 62) 0.001
0.1
0.0
0.0
0.1 0.2 0.3
0.4 0.5
Figure 5: Plots showing the evolution of the Wasserstein distance between μe and the uniform measure Moo (solid line) and between μlt and 53 == Σ x2// (dashed line). Note that the simulations with the initial condition fa for a = 0.01 and a = 0.001 almost completely overlap, differing only by time shift. This observation supports the validity of Assumption 3.
Under Assumption we can more precisely characterize the final result of the meta-stable phase: Theorem C.23. Let (ur New denote sequence of random measures in P(St), where Pō is the empirical measure associated with N i.i.d. particles sampled from the uniform measure on St Under assumptions 1,2 and we can conclude:
lim inf W, ( .88tmaz)) =0, No
in probability, where 8mmr
Σ3
8 2rc kmaa
modulo rotations.
kmax
Proof. By Assumption 3, for every > O there exists i >> 0, independent of N, such that:
W1(fa" 8kmmaz) C
2
where N and T. are defined as in Proposition C.20. Therefore, for all e > 0:
lim inf Wi (ny) 8kmazz <
lim
W1 8kmass) S
N 20 120
N
lim Wi(un +1,+i f1+)) + Wi (faa gl.mao 5
N
lim e W /a. ) 5'
N
where the classical Dobrushin's estimate (Theorem A.4) is used the last line, with finite con- stant L depending on the properties of W. By Proposition C.20 and Remark C.22, it follows that
eL'W, (WW, +2: f ) < & definitely in N. Hence, the conclusion follows by arbitrariness of n
27

---

Published as a conference paper at ICLR 2025
D EXTENSION TO HIGHER DIMENSIONS
In this section we discuss what happens the system is defined on Sd-1 by:
Oej
div()>V( * p)),
(m(0) μo:
(31)
where W -d-1 and W μ is the convolution defined in Eq. (4).
Remark D.1. The differential operators should be interpreted with respect to the undertying Rie- mannian manifold structure.
The uniform measure Mo remains unstable equilibrium of the dynamics. Furthermore, the lin- earized PDE for an initial perturbation Pox + Po is again:
{ O.p' 460 A(W p²)),
de = Po:
with A the Laplace-Beltrami Beltrami operator on the sphere s
As in the d = 2 case, we can study the linear operator p (f) := MM0 A(W f). Letting fn represent the spherical harmonics coefficients, we obtain:
(e (f)). M + 2)
where W., are the Gegenbauer coefficients. This result follows from the fact that spherical harmonics eigenfunctions of the Laplace-Beltran Beltrami operator, combined with the Funk-Hecke theorem for convolutions on spheres (see Section C.1.2). This provides the usual spectral bounds with a similar Tmoz for the linear operator, and an explicit solution p²:
Z
Pi
(Po)n 2n Ynnj.
)=
where 2 denotes the number of spherical harmonics Yn. on gd-1 of degree n.
It is important note that now have superposition of functions of the same degree Tmmaz all growing at the same rate. In particular, the steps outlined in the previous sections can be repeated, but with initial condition for fN given by:
zd
JNN
Poo
+
Σ
NN Y
j=0
Additionally, the calculations have to be performed in the space H with < -d/2, rather than in H-1 the proofs remain analogous.
In d = 2, the emerging symmetry was characterized by invariance under rotations by angles of 2n,, which was relatively straightforward describe. However, move to higher dimensions, Emon.
the nature of the symmetry becomes complex and challenging express.
Consider the subspace:
Hu :- {uE 11 (n, Yi,) VI not divisible by k).
This set is invariant for the dynamics of Eq. (31), indeed, since Hj is vector subspace of H it suffices show that 0,p E H4 for every μ Hy:
(0tt, Yim) = - (div(uV(WW ()), Yi,)
(W n), VYi.)
Pna Ynns
Wreneyy V n'
Σ
VYmm)
W Whe (Yns VYn. VYim)
Σ
W Wi-nfM-nn Unn (Yns VY-ny, VYL.),
28

---

Published as a conference paper at ICLR 2025
00.
-1
-055
05
y
0.5
05
Figure 6: Example of a curve in H3.
where we used Heeke theorem and triple product integral identity for vector spherical har- monics. If this derivative is non-zero, then at least one of the terms in the sum must be non-zero. This implies n = 0 mod and n = mod Thus, must be a multiple of k.
In conclusion:
lim inf Wi W14 He Ho.... (P((4-11) - O.
N -000 120
in probability.
AUXILIARY RESULTS
Lemma E.1. Let f [0, oo) R satisfy f'()) S af(t) + celt for all t (0, where 0 < < b and c 0. Then:
f(t) & f(
b-a ee
c
Proof. This result follows directly by Gronwall's inequality. Specifically:
S(1) < f(0) + e
C )ean) ae (t- eee + - +J. (S(O)
f(0) + c e + f(0)acut S( aas + aacat (' clo-a)o
b
b
5(0) +
c
+ +(0)eat - aa + c a
eat (cb(e -1) s
b
bb-a
< f(0)e +
br.
Lemma E.2. Les E Z. Under Assumptions and the linear operator C H*())) + (S')
defined as Cp (f) AW satisfies:
[ fll S
2T lfazzllfllar.
Proof. Using the Fourier definition of the spaces H² (
p f 1122=
Σ n 1-
f), 2=
neZ
= Σ n
(Moomil(fall) S nbo 2 Baauell.1 1f
nEZ
29

---

Published as a conference paper at ICLR 2025
Lemma E.3. Let } be sequence of random probabilities on St given by the empirical mea- sures of N i..d samples from the uniform measure. Then pD¹ = 11"h Poo satisfies:
11/41le6 -00,
Np Il (PG Jamax
|1P6 11 (N17741P6
Tmax (NN1/49) O.
with both of the limits in probability.
Proof.
First of all notice that, fixing an orthonormal basis {eLtR of L3 vB > and
> L. 2
- P [
-N 28
Nk 2# (pu
ex)
1 N
27
< N.-22 29
Nk
ea (2:)
N 1=1
N 29
Σ
0,
where we used Beppo Levi and the indipendence of samples. In particular by Chebychev's lemma this implies N - l|x -> in probability. Taking 1/4 and we get the first limit.
Before proving the second limit, observe that
VN(PM) VN ce)
Σ Ex((X,) - efxx) +N(0,1 1=1
VN
in distribution by the central limit theorem.
Now:
Ppoll
N-1/4
Yonax
+ 1/4 I|lPolll--
(po)Rmaz
||po|| 1-1
(po Amax
N N1/4 |1/41poll
N2/2|1Pol1
N2
Jmax
+
VN(Po)
VN ( (po)
N
NN// ))
Tman
||o0l
H
(po))
The terms between the parenthesis converge to zero in probability (indeed (1
Tmaas
-- and we can use the argument above). The term outside parenthesis instead con- Tmax verges in distribution 1/Z with Z standard normal (because of the continuous mapping theorem, since 1/2 is almost-surely continuous). Hence, by Slutsky's 's lemma, we get convergence to 0 in distribution and then in probability.
Lemma E.4. Given sequences of probability measures phy Vn in H-*(X), with X compact space, such that [| - v, (- + 0, then also Wi(Pn, Un) + 0 is true.
Proof. Notice that I|m,, Vn 1 O implies (44, Vn, f) - 0 for every E H*). Since (X)()) is dense (x) (it contains the smooth functions) and since μt, and v,, are both probability measures, then ph, Vn 0 in weak sense.
Let us fix a subsequence that satisfies limg W (Pmy, Pns) lup, W (Hmn V.). For each k consider On, E Tápi (X) such that f One (Hag Une) Wi (Wmng Ung) Without loss of generality,
30

---

Published as a conference paper at ICLR 2025
by adding constant (which does not affect the integral), assume that the functions ong vanish at the same point, hence they are uniformly bounded and equicontinuous. By Ascoli-Arzeli theorem we extract subsequence converging to some & Lipı (X). Hence, renaming the subsequence:
W1 Wi(uon Uns) = Φni Vnn2 S 21 @l +
/ Ødllume Une) -
This concludes the proof.
Lemma E.S. Let W [-1.1] + IR be an analytic function and suppose that W., for every n. Then the equilibria of Eq. (31) can only be the uniform measure or a finite union of submanifolds of dimension most d
Proof. Consider Hog to be equilibrium of Eq. (31). Then, for every / (² it must hold that:
0 (f, dive V(W peg))-
In particular, substituting ʃ = V(W Hag), we obtain:
f gd
V(W * [ea) = O.
This implies A = {i E - V(W (x) 0} is such that Peg (AC) = O, hence supp(Meg) A. Note that + |V(W (eg) 12 is still analytic function, hence A must be either $ or a finite union of submanifolds dimension at most 2 (by Lojasiewicz's structure theorem, (Krantz, 2002)).
(x)]l
To conclude the proof, we just need to study the case A = gd-1, The definition A implies that W Mey must be constant Sd-1. By the Funk- Hecke theorem, Vn > 0, Vj:
0 = (W Peq)a.j Wn (Meg)nss
and thanks to the hypothesis Wa * O, we obtain (Heg)nys = Vn + 0 and Vj. Thus, Heg is the uniform measure.
31


Reviews:
[{"review_id": "zt8eJHqeJY", "reviewer": "Reviewer_MWwW", "summary": "This paper investigates the evolution of tokens in simplified transformer models using a mean-field formulation. By analyzing the long-term behavior of interacting particle systems, the authors prove results on metastable clustering, demonstrating that the model remains close to a metastable manifold of solutions. The study identifies the existence of linear, quasi-linear, and clustering phases in the token evolution.", "questions": "- In the integral representation of (USA-MF) (lines 148-149), the coefficient $1/N$ seems redundant since $\\mu$ is an empirical distribution which contains the normalization term. Also, I’m not sure if (SA-MF) representation is also appropriate because of the same reason.\n\n- Taking sufficiently large $N$, $T_1$ becomes negative. Does this mean the non-existence of linear phase?\n\n- The convergence rate of W1 distance in Theorem 4.3 is not given in the main text. Can you clarify it?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors provide a comprehensive analysis of metastability that precisely describes the evolution of tokens. They successfully identify the existence of linear, quasi-linear, and clustering phases, contributing valuable insights into the behavior of simplified transformer models.", "weaknesses": "- The model considers all attention parameters $Q, K, V$ as fixed identity matrices. This significant simplification may limit the applicability of the theoretical results, as it does not capture the complexity of learned attention mechanisms.\n\n- Some notations in the paper are not properly defined, which can hinder readers' understanding of the results. For instance, the notation $\\hat{\\rho}\\_0$ related to Theorems 4.2 and 4.3 is not introduced. Additionally, $\\mu\\_{cluster}$ in Equation (8) seems to be defined on a one-dimensional space since $2\\pi j / k_{max}$ is a scalar, making it unclear how to measure the W1 distance from $f_*(t)$. Providing clearer definitions and explanations would enhance the paper's readability.\n\n- For sufficiently large $N$, the parameter $T_1$ becomes negative. This could restrict the applicability of the theorem to certain token counts. Clarification on how this affects the results and whether it imposes practical limitations would be beneficial.\n\n- Minor: At the first glance, I thought the propagation of chaos result (Theorem 3.1) is also one of the main contributions but is actually weak due to the exponential growth of the bound. Now I see that this result serves to emphasize the importance of the metastability analysis. Clarifying this role could prevent potential confusion for readers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "deIdJ0PqNg", "reviewer": "Reviewer_bdQd", "summary": "Transformers are performant but not theoretically well understood. In particular, it is not clear how to think about the dynamics associated with a large stack of transformer layers. Geshkovski et al. (2023) made a formal analogy between these dynamics and the dynamics of a mean-field interacting particle system, and were able to characterize various aspects of these dynamics, at least in a simplified theoretical setting. The authors of this work build upon this earlier work and characterize not just the long-time behavior, but associated metastable dynamics, which may be important in practice.", "questions": "1. Do slightly more realistic transformers behave in a way consistent with the predictions of this theory? Relatedly, if the theory assumes Q = K = V = Id, what is the right point of comparison for a language model?\n2. How does the theory change if certain assumptions (e.g., Q = K = V = Id, finite $N$) are modified or generalized?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors present a detailed mathematical analysis of a simplified model of transformers, and rigorously prove a variety of results. They are good about presenting their results at a high level in the main text, and leaving the fine technical details in SI. They closely connect their approach to a variety of recent work, principally the aforementioned work by Geshkovski et al.", "weaknesses": "The authors present their results clearly and do not appear to overstate them. My major concern is mostly that the authors' results concern kind of a toy setting somewhat far from the kinds of transformers people use in practice. One assumes Q = K = V = Id (line 123), one takes a limit where the number of layers go to infinity, the number of tokens $N$ is assumed to go to infinity, etc. This makes sense given that theoretical analysis of any kind is quite difficult, but makes it unclear to what extent the authors' results shed light on transformers 'in the wild'. It would be extremely helpful if the authors could provide some numerical experiments for slightly more realistic transformer models that show key predictions from their theory hold, at least approximately. \n\nAlso, the authors could provide additional commentary on what happens if one tries to generalize their results. What is expected if Q, K, and V are not all the identity? How does one expect finite $N$ corrections to behave, and what do they change about dynamics?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CUm6SdwADQ", "reviewer": "Reviewer_McN2", "summary": "This paper presents a rigorous investigation of meta-stable clustering phenomena of mean-field Transformer models. By modeling transformer layers as continuous-time flows and employing a mean-field PDE framework, the authors present a very interesting approach to understanding intermediate clustering phases through linearization and higher-order approximation. Numerical experiments support the theoretical findings. This work opens a venue for understanding the meta-stable dynamics of Transformer models.", "questions": "I only have a minor question. The authors make a key Assumption 3 regarding convergence to the clustering phase. While they discuss this assumption based on (Markdahl et al., 2017), relying on periodicity, I am curious whether it would still hold in the generic case--specifically, in the case where the initial condition in equation (6) contains multiple modes. Could the authors provide some intuition and discuss potential implications of their result in this case?", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. I appreciate the authors provide rigorous proofs that illustrate the convergence of the tokens to structured meta-stable manifolds, and identifies the influence of temperature, which is a key factor in understanding the mechanism of transformers. In particular, understanding the mean-field transformer model is beneficial for tasks involving long-context dependencies. \n\n2. Numerical experiments are presented to demonstrate that the theoretical predictions of the clustering dynamics align well with observed practical results. It is nice to see the number of clusters is influenced by the temperature, which I found surprising.\n\n3. Overall, the paper is well-written, and the results are sound. The decomposition of the transformer dynamics into linear, quasi-linear, and collapsing phases allows for a systematic analysis of the clustering phenomena. I'm also excited to see new PDE techniques (the Grenier's scheme) applied to the analysis of transformer dynamics. The findings are both highly relevant to the field and of significant interest to the community.", "weaknesses": "1. (Minor) In the contribution bullet point 3, the authors state that \"the periodicity developed in the first phase is maintained over exponentially long time intervals.\" The authors may want to clarify in what sense the time interval is \"exponential long\" and where in the paper this claim is made.\n\n2. (Minor) Page 10, Line 504, Should \"Figure 4\" be \"Figure 3\"?\n\n3. (Minor) Page 5, Line 233, should the coefficient $\\hat{W}_k$ be $\\tfrac{1}{\\beta}I_k(\\beta)$?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7LxlYif4TF", "reviewer": "Reviewer_WDSo", "summary": "The authors build on existing work which models the evolution of tokens within a Transformer model. The authors prove the existent of a meta-stable phrase in learning.\n\nNote: I fear I am not the right audience for this paper. I am not at all familiar with the literature, or its relevance in the ML community. I additionally understood little of the content of this paper.", "questions": "- Why does the  existence of a \"meta-stable phase\" imply that then the network learns a richer and more complex representation of the data? What is meant by richer and more complex?\n\n- Are there any clear takeaways from this paper that a practitioner can use to build intuition and how and why a transformer works?", "limitations": "", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper appears to have a high degree of rigor.\n- The goal of developing a rigorous theory of how the transformer works is important.", "weaknesses": "- It is not clear (to me) what the relevance of the conclusions made by the paper is.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LDuOBbMLRI", "reviewer": "Reviewer_DN47", "summary": "The authors consider the evolution of the token distribution in a limit of an infinitely deep transformer which yields a differential equation for the evolution of the tokens. This equation is further simplified by the assumption that the attention parameters are constant over layers, and all equal to the identity, and that there are no MLP subunits. An initial result is that they show that the limit of infinite particles commutes with the time evolution in the sense that there is a PDE that describes the evolution of an empirical distribution with $N$ large enough. From this starting point, they describe a departure from these solutions (initialized so that the limiting measure is the uniform measure)\n\n1) an initial phase where a periodic fluctuation of a certain frequency, $k$, grows \n2) a quasi-linear phase where the fluctuations are well-described by the behavior of this mode around the uniform background\n3) and finally a collapse phase where the evolution reaches close to a finite set of delta measures.", "questions": "1. What is the relationship between this paper and the recent Geshkovski et al. paper \"Dynamic metastability in the self-attention model\" [1]?\n2. Are there simple ways to extend this analysis as $N \\to \\infty$ in the discrete case? Is there a qualitative difference in this setting?\n3. In practice, should we think that transformers operate in phase 1, 2, or 3?\n4. How do we know that neglecting the remainder $R$ from equation 5 to 6 is acceptable?\n5. What is the interpretation of the $\\alpha \\to 0$ limit (i.e. should it be interpreted as $N\\to\\infty$ again)?\n\n**References:**\n\n1. Geshkovski, Borjan, et al. \"Dynamic metastability in the self-attention model.\" arXiv preprint arXiv:2410.06833 (2024).", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The strengths of this paper are its rigorous analysis of the problem, where they first extract and then rigorously prove that certain clustering dynamics arise directly from attention. This work makes it very clear what structures arise and how they evolve. This is a strong extension of prior work, which relies on less intuitive arguments, and so this paper finds a way to explain a complicated phenomenon more clearly than prior work.", "weaknesses": "The main weakness is the lack of connection to a practical setting. The problem considered is highly simplified, and there is little discussion concerning even qualitative arguments explaining what differs in the setting with the greatest interest -- deep transformers with MLP blocks and non-identity weight matrices. It is difficult to understand how to take these results and understand what they say about that setting. \n\nFor example:\n1. Prior work [1] suggests that the structure of the value, key, and query matrices is important and derives results in various (though simplified) cases\n2. The work cited as the foundation for this work [2] discusses the BBGKY hierarchy -- a way to understand the joint law of all the particles by breaking them down in terms of the $n^{\\text{th}}$ degree marginals.\n3. Other work [3] considers the behavior, including the MLP, and finds (though non-rigorously) that the MLP significantly affects dynamics.\n\nIt is of significant interest if connections to ideas like these could be explained in the paper. Of course, the limited scope of one paper may hinder detailed analysis in all possible settings, but it would make it clear how to interpret the results in context.\n\nIf the authors could discuss this issue satisfactorily, I would be happy to increase my score for this paper.\n\n\n**References:**\n\n1) Geshkovski, Borjan, et al. \"The emergence of clusters in self-attention dynamics.\" Advances in Neural Information Processing Systems 36 (2024).\n2) Geshkovski, Borjan, et al. \"A mathematical perspective on transformers.\" arXiv preprint arXiv:2312.10794 (2023).\n3) Cowsik, Aditya, et al. \"Geometric Dynamics of Signal Propagation Predict Trainability of Transformers.\" arXiv preprint arXiv:2403.02579 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)