A multiscale analysis of mean-field transformers in the moderate interaction regime
Authors: Giuseppe Bruno
Federico Pasqualotto
Andrea Agazzi
Venue: NeurIPS 2025 oral ()
Keywords: mean-field limits
moderate interaction
mean-field transformers
self-attention models
clustering
multiscale
Primary area: deep_learning

Abstract: In this paper, we study the evolution of tokens through the depth of encoder-only transformer models at inference time by modeling them as a system of particles interacting in a mean-field way and studying the corresponding dynamics. More specifically, we consider this problem in the moderate interaction regime, where the number $N$ of tokens is large and the inverse temperature parameter $\beta$ of the model scales together with $N$. In this regime, the dynamics of the system displays a multiscale behavior: a fast phase, where the token empirical measure collapses on a low-dimensional space, an intermediate phase, where the measure further collapses into clusters, and a slow one, where such clusters sequentially merge into a single one. We provide a rigorous characterization of the limiting dynamics in each of these phases and prove convergence in the above mentioned limit, exemplifying our results with some simulations.

A multiscale analysis of mean-field transformers in the moderate interaction regime
Giuseppe Bruno
Federico Pasqualotto Department of Mathematics
Department of Mathematics and Statistics University of Bern
University California, San Diego
giuseppe. bruno@umibe.cr ch
fpasqualotto@hcca. edu
Andrea Agazzi
Department of Mathematics and Statistics
University of Bern andrea. Sadrea..gazzi @unibe. ch
Abstract
We study the evolution of tokens through the depth of encoder-only transformer models inference time by modeling them mean-field interacting particle system, and analyzing the corresponding dynamics. More specifically, we consider this problem in the moderate interaction regime, where the number N of tokens is large and the inverse temperature parameter ,3 of the model scales together with N. In this regime, the dynamics of the system displays a multiscale behavior: a fast phase, where the token empirical measure collapses on a low-dimensional subspace, intermediate phase, where the measure further collapses into clusters, and slow phase, where such clusters sequentially merge into single one. We characterize the limiting dynamics in each phase, exemplifying our results with some simulations.
1 Introduction
The transformer architecture [49], through its extensive use in Large Language Models, has played crucial role in the recent, unprecedented developments in machine learning and artificial intelligence. One of the key innovations at the heart of this architecture are self-attention modules [6], allowing to capture long-range dependencies in the data, e.g., in prompts with large number N of tokens. To further improve their performance, practitioners have implemented these models in different hyperparameter regimes, e.g., choosing the model's inverse temperature parameter / (a parameter that scales the query-key dot products in the self-attention layer) a function of N [37, 43]. However, the groundbreaking empirical success of these machine learning models remains largely unexplained from the theoretical perspective. In particular, precise mathematical description of the internal representations learned by transformers, and of how these representations behave different hyperparameter regimes, is still lacking.
promising approach to fill this gap was presented in the work [46], where the authors interpret tokens traveling through a deep stack of transformer layers particles evolving in time and interacting mean-field way. A subsequent line of work [24] has then observed that tokens in this model tend to organize into clusters, offering in simplified setting a compelling qualitative explanation of how transformer models build representation complex input data.
Despite its apparent simplicity, this interacting particle system exhibits remarkably rich dynamical behavior. Indeed, recent studies have identified distinct dynamical phases, characterized by quali-
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

tatively different clustering patterns, which depend on specific choices of parameters, timescales, and initial conditions. However, these results often rely on restrictive and sometimes unrealistic assumptions, and they typically capture only limited aspects of the collapse dynamics, providing partial views of the transformer's complex dynamical landscape that are difficult to reconcile into a consistent and global dynamical picture. Such a global characterization of the clustering phenomenon in a realistic parameter regime is arguably fundamental to understanding how internal representations form in deep models and how to operate such models in the optimal hyperparameter regimes.
Contributions In this work, we study the dynamics of the mean-fichl srean-field transformer model developed in [46, 25), constrained the d-dimensional sphere, in the limit of large context size, i.e., when the number N of input tokens is large. Motivated by recent scaling strategies in state-of-the-art LLMs such as SSMax [37] and YaRN [43] (used respectively in Llama 4 and Qwen 3), we consider the setting where the inverse temperature parameter 3 grows with N. In this regime, our contributions can be summarized follows:
1. We identify three distinct dynamical phases (respectively denoted the alignment, heat and pairing phases), corresponding to different scales of time as a function of the parameter 3. In each phase the model dynamics displays, asymptotically in B, qualitatively distinct behavior, characterized by a different limiting equation.
2. In the alignment phase, occurring on fast timescale of order O(1), we prove our main technical result: under general assumptions on the parameter matrices, the finite particle dynamics converges to a linear transport equation modeling the collapse of the token measure onto a low-dimensional manifold dictated by the spectral properties of such matrices. To the best of the authors' knowledge, this phase was not yet identified in the literature.
3. In the heat and pairing phases, occurring respectively on timescales of order ((3) and ( we identify, under stronger conditions on the parameter values, the limiting dy- namics as a forward or backward heat equation on the aligned manifold (leading, in the backward case, to metastable clustering) and a miiite-dmmessioall system of ODEs describing sequential cluster merging along geodesics.
Together, these phases reconcile various previously identified dynamical regimes as different timescales of a single unified dynamical picture. Furthermore, our multi-phase analysis allows to relax some of the restrictive assumptions imposed by previous works, extending their applicability to more realistic scenarios.
Related works The model studied in this paper was introduced in [46], where the authors also identify the heat equation as describing the dynamics of the partiele system in IRd in the large 3 regime. This limit emerges as a correction term in their analysis upon subtracting an appropriate leading order term from the prelimit equation. In this paper, by considering the dynamics on the sphere resulting from the inclusion of the layer normalization in our model we provide a justification for the spontaneous collapse of the system's 's state to subspace where this correction term becomes of leading order, dominating the dynamics of the model on a certain timescale.
In [25, 24, 29, 13, 44], the authors identified the clustering behavior occurring in this and closely related models as - oo for B, N fixed. Analogous convergence results, under different assumptions, are provided in [18, 34), while quantitative contraction rates for such convergence are given in [16]. These works, however, do not address the dynamically meta-stable phases characterized by partial clustering numerically highlighted in [25]. This intermediate phase is explored in [9] in the large N limit for tokens distributed uniformly at initialization and in [23]. where the authors study the formation of meta-stable clusters under the assumption that the system is initialized into well- separated configurations. Furthermore, in [23] the authors identify different dynamical timescales in the finite N case and characterize for the first time what we refer in this paper as the pairing phase in the large 3 limit. In all these cases, however, the results are limited to the setting where the model's 's key, query and value matrices, Q. K and V. were multiples of the identity. More recently, the work [1 1] analyzes the stability of fixed points of the same model based on the eigendecomposition of Q, K, V under the weaker assumption that parameters satisfy a modified Wasserstein gradient flow condition (Q" K = V D), but does not study the dynamical landscape connecting such fixed points. Finally, in [3, 4], the authors discuss clustering for hardmax transformers. Our work provides framework to combine the observations listed above in unique dynamical picture.

---

Our modeling approach shares concepiual roots with the neural ODEs literature [14, 21]. However, key distinction is that we consider . particles interacting through a mean-field PDE, as opposed to one in the previous references. This connects our work to the broader literature on mean-field models for neural networks [45, 36, 17, 2, 19), where timescale analysis has also been a subject of interest (see [7]. In contrast to these works, which typically focus on training dynamics, our study centers on the inference-time evolution of representations through network depth.
Finally, our research relates to the study of moderate scaling limits in interacting particle systems. For instance, Oelschläger [39] proved the convergence of certain systems to the porous medium equation with noise. These results were subsequently extended to cases without noise [41, 40], with different exponents [22], or employing different techniques and equations [12, 10, 42]. Another relevant line of work investigates the convergence of specific interacting particle systems to the heat equation, explored both numerically and theoretically [20, 33).
2 F Framework and notation
We consider the framework introduced [46, 25, 24), modeling the transformer architecture as a discrete-time dynamical system governing the evolution of N tokens (2,(O))==1,...)) through its layers via:
1 2E + 1]=1 =N (2((0)+ + (B,(()) Σ (Reeet
e= L
L- (1)
(0( =
where N :R - 1 denotes the normalization operator on the d dimensioonal unit sphere gd-1,
denotes the depth of the transformer architecture and (e) = Σy= ((Qzz.()) (e,))a,(()) is a normalization constant. The dynamics depends on the parameters with matrix values Qe, Kg, and Vi that represent the query, key, and value matrices at each layer, respectively.
In the spirit of neural ODEs [14], the authors then consider the infinite-depth limit of (1), leading to the following continuous-ting time model, describing the evolution of x,((): : (0, - Sd-1;
ix(t) (t)
(t) Σ
(SA)
Here and throughout, := y-<,, y) denotes the orthogonal projection of y onto the tangent space T,Sa-1, ( is the Euclidean inner product in Rd. and (t) N ΣØ= (0Q (,a,))))) is the time-dependent normalization factor. The parameter 3 > 0 is interpreted as the inverse temperature. Remark 2.1. The MLP would act as drift term the dynamics, whose consequence should still be investigated further. We expect different dynamical behavior depending on the relative scale of the MLP coefficients and the attention part. Although the framework allows for the inclusion of feedforward layers via a Lie-Trotter splitting scheme (see [26]), we choose to isolate exclusively the self-attention mechanism, both because our interest lies specifically its dynamics, and for the sake of clarity. For the same reason, we assume, as in the works cited above, the parameter matrices to be shared across layers: Q = Q K₁ = K, and Vi V.
further
As the positional information each token is encoded in its initial condition, the dynamics (SA) is invariant under permutations of the particles' indices. This symmetry allows to fully characterize the system's 's state through the particles' empirical measure p(t)= == A Σm" Sz.(0), where 8 denotes the Dirac measure centered x. The measure μ(t) evolves according to the continuity equation:
0. + div(x: u) 0 on R2 x sd-1
(2)
14|=0 = x0)
on 1
where the vector field x [u] gd-1 Tgd-1 1 is defined as [
x0(14] Ps
BiQeK du(v)).
(x) Jga-1
with (z, :- 3lQ (p This formulation extends the token dynamics a flow on the space P(Sd-1) of probability measures over the sphere gd-1, encompassing both empirical and absolutely continuous distributions.
3

---

3 Main results
As discussed in the introduction, in this paper consider the limit as 3 + of the dynamics (3). To present the dynamical scales arising in this limit, consider a formal Taylor expansion of the vector field [4] generated by sufficiently smooth measure j.:
S e v) PaV y ((dy)) ulx') J es v) odyy)
I e P P.Va I', Vp(z')) o(dy)
(x)x
(4)
plx') J e
where a denotes the Lebesgue measure on S and x' = KT Qz. For large 3. Laplace approximation suggests that (I) typically dominates (II) at initialization, giving rise to first, fast dynamical phase:
Alignment Phase: a timescale of (1), the limiting dynamics governed by linear transport equation (Eq. (5) below) and the token distribution rapidly collapses onto a lower- dimensional subspace determined by the spectral properties of the matrix VKTQ. Q.
After the dynamics collapses this low-dimensional subspace, identify some classes of parameters for which the leading-order contribution to the vector field, approximated by term (I), vanishes. In such scenarios, the dynamics becomes governed by term (II), which involves the gradient of the measure μ. This gives rise to a second, intermediate phase:
Heat Phase: operating timescale of ((3) (achieved by rescaling time t' = t/B), the dynamics within the previously identified subspace exhibits diffusive or anti-diffusive behav- ior. Depending on the model parameters (specifically the sign related to V V restricted to the subspace), this phase can lead to further concentration into distinct clusters (backward heat equation) to smoothing/spreading of the distribution (forward heat equation).
In the attractive case, we identify the limiting dynamies up until the formation of clusters. We expect the clusters to be invariant in this timescale, and to interact only on much longer ones.
Pairing Phase: on an exponentially long timescale in (e.g.. 0(e*9) for some O, where e depends on the distance between clusters), the clusters formed in the previous phase sequentially merge. Typically, the closest pair of clusters collapses first, governed by system of ODEs describing their interaction, eventually leading to a single clustered state.
We refer to Appendix E for a graphical representation of the three phases introduced above.
We outline the structure of the remainder of this Section. In in Section 3.1, we recall quantitative result connecting the large behavior of the ODE system with the behavior of the corresponding PDE in the relevant timescale. This will allow us to focus solely on the PDE analysis when we describe the three main dynamical phases in Sections 3.2, 3.3, and 3.4.
3.1 Large N convergence
To connect the timescales analysis above to the N-particle system of ODEs (SA), we consider the regime where N - and B = Bx -> slowly enough with respect to N. This is relevant for context scaling techniques and LLMs (see introduction). To proceed, we use the following lemma: Lemma 3.1. Assume that the initial tokens (0:((O)iE[M)) are sampled independently and identically distributed from a reference measure μág E P(gd-1). Let μ4 N.8 be the empirical measure for particles {a.(t) evolving via the ODEs (SA), and let μe be the solution to the continuity equation (2) with initial condition po- Fix time interval [0, Ta) where T3 is a B-dependent timescale. If 3 = 3N
depends N and diverges slowly enough N - oo, then:
Wi
N. 14 0 as +
uniformly on Ts].
Proof. This follows from the Dobrushin-type stability estimate: (I where Lg is a positive constant depending on the Lipschitz constant of the vector field X as discussed in [9]. The claimed convergence follows from Wi (1g) no) - 0 provided that La TAN grows sufficiently slowly with N such that the overall term tends to zero.
N.8 p < WnG (o) chat

---

Our goal is to understand the behavior of H4 N. in the joint limit N. SN + co. We denote the limiting distribution of iu B - by pa. Following an argument analogous that in [22], though in different setting, we can decompose the convergence problem as:
W1 N. .14 S (14 N.3 .14) + W1 (1)
In our regime, Lemma 3.1 guarantees that the first term vanishes as N - cc. Consequently, the analysis of the N N-particle system in this coupled limit reduces to studying the behavior of the solution ii to the continuity equation (2) as B + co. The PDE analysis in this limit will therefore be the focus of the following sections.
3.2 The Alignment Phase
To characterize the limiting dynamics in the first phase make the following assumptions:
Assumption 1. Q, K, V are invertible square matrices.
Assumption 2. The probability measure Mo on S is absolutely continuous with respect to the Lebesgue measure on gd-1, 1 Its density is bounded from above and below (mingeg*-11 po(()))))))))))))))))))())) > 0) and Lipschitz, continuous.
These technical assumptions, significantly milder than the ones made in most related works, are needed to guarantee that the terms appearing in the analysis of the limiting equation, e.g., the denominator in (5), are sufficiently well behaved. Under these conditions, we show that the limiting dynamics in this regime coincides with the formal Laplace approximation of term (I) in (4), i... the integrals in the definition of the vector fields can be replaced by the value of the integrand at the maximum point x' KT Q=//KTT
leading to the significantly simplified expression (5) below. Theorem 3.2. Lei Assumptions 1, 2 hold, then the solutions {uz} y of the continuity equation (2) converge in C(00 P(S to the solution nTdh of the partial differential equation:
O,u(z)) = didiv (u(x)P VK (m(0,z) =pu()), x EES
[K
(5)
Proof Sketch. To establish the result we must prove well-posedness of the family of equations leading to the desired limit and obtain sufficient regularity uniformly in 3 to ensure that the formal simplifications from (4)(I) to (5) are allowed. This is particularly important as the derivatives of the kernel tend infinity in the limit B - co. The core argument proceeds in three main steps. First, we establish the relative compactness of the family of trajectories [u²}as in the space C( T), P(,)(1))))) using a variant of Ascoli-Arzela theorem and the boundedness of the vector field X. The second, crucial step involves deriving uniform in 3 estimates on the regularity (i.c., Lipschitz bounds) of the vector field [=[ along the solution trajectories Srر. This is achieved by analyzing the concentration behavior of the kernel ((Qr as 3 + oo, leveraging properties related to the cumulants of the Von Mises-Fisher distribution, and employing continuation argument to propagate regularity over time. Finally, using the compactness and uniform regularity, we pass to the limit 3 + oc in the weak formulation of the continuity equation (2). The uniform estimates allow us to conclude that 8 is solution of (5), while uniqueness follows from [5]. The full proof is deferred to Appendix A.
consequence of Theorem 3.2 is that in the large-B limit, the tokens, to leading order, evolve independently of each other, driven primarily by the structure of the Q, K, and V matrices. In this regime, self-attention behaves like a composition of linear layers followed by layer normalization, with minimal influence from inter-token interactions.
Combining the above result with Lemma 3.1 we obtain the following convergence result:
Corollary 3.3. Under Assumptions 1, 2, for every > O we have Wi (pt N,8 ,182) - 0 as N + 0o. provided that 3N + oo slowly enough.
Having established that n is a solution of equation (5), we can investigate its long-time behavior. In particular, show below that the support of ju, is asymptotically flattened onto lower-dimensionall subspace determined by the spectral properties of the matrix V K Q.
5

---

Proposition 3.4. Let of be a probability measure on S absolutely continuous with respect to the Lebesegue measure, and let μr be the corresponding solution of (5). Then for every v w(Ho) (the w-limit set of (o) it holds:
supplv) C Emaa n gd-1
where Emaz is the generalized eigenspace associated to the eigenvalue of VVKQQ with largest real part.
Proof. The proof is provided in Appendix A.2, where we reduce the analysis to a linear system of ODEs in IE with matrix V V'' Q. identifying the corresponding asymptotics with the ones of (5). O
Remark 3.5. At a first glance, this result might appear inconsistent with those of [11], since in some
cases measures supported on Emax do not maximize the energy. However, this apparent discrepancy is consequence of the order of the limits being taken, with B + preceding t - oo in our case.
Proposition 3.4 demonstrates that the token representations rapidly collapse onto lower-dimensional subspace determined by the model's 's matrices. This can be interpreted as the initial phase of the inference process, where information is compressed into smaller, more relevant subspace. This phenomenon is consistent with the rank collapse observed, e.g., in [38, 27].
Remark 3.6. Apart from the collapse to Emax, one cannot in general conclude the existence of a limiting (stationary) dynamics for (5). Indeed, it is not difficult to construct examples where the particles continue rotate on the sphere indefinitely, e.g., when V is a rotation and Q" K = ld. Remark 3.7. Recent works have studied transformer models with stochastic perturbations [48]. where the token dynamics is influenced by random noise. In this setting, the convergence to the corresponding equation (5) (with an additional Laplacian term) is rypically easier establish due to the regularizing effect of the noise (see [39]).
3.3
The Heat Phase
Having established the rapid collapse onto the subspace Emar ngd-1, now investigate the slower cvolution within this subspace, assuming that the initial measure (Ag is supported in Emax n Sd as consequence of the previous analysis:
Assumption 3. The initial condition pig in the heat phase satisfies supp(po) S. Emmas n ngd-1.
Since the intersection Emax n gd-1 can be identified with a gdiml 1m we will, with a slight abuse of notation, continue to denote it by gd-1
sphere, specifically
To demonstrate that the heat equation, described in a different setting in [46], emerges an interme- diate dynamical phase due to the spherical geometry induced by LayerNorm, we assume:
Assumption 4. Q" KK |Emaz = I and V|Emar =\2I when restricted to Emar, with X A2 > 0.
Under this condition, Eman is an invariant subspace for Eq. (2) and, without loss of generality, we can suppose \1 A2 =
Remark 3.8. Assumption 4, for example, is satisfied under the global assumption Q K = S and
V +S, with S symmetric definite positive matrix. This is a fairly standard assumption in recent studies within this framework and it endows the model with an additional structure of gradient flow on P() with respect a modified metric (see [11, 251).
In this regime, the vector field [=(1]] vanishes on the support of as B + co, but its rescaled version admits the formal limit (see Corollary B.2):
3xxlu](x) + Vga 1p (z), as 3 + 00,
where 7 =1, depending the sign choice in the definition of V. This scaling of the vector field by B corresponds to a time rescaling dt = :ds,, explaining the phase duration of order O(3).
Proposition 3.9. Let Assumption 4 hold and let n E P(Sd-1) be the initial measure. Assume that there exist T 0, k positive integer, and pu E C] ( for all I E [0, T), such that ni solves the heat equation [0, T] x 1d-1.
with min ,²²(x) 0
{ @ose = (m(0)
(6)
Po:

---

where A denotes the Laplace-Beltrami Beltrami operator on S Then, for large B. j4 solves the mean-field PDE:
jōu = -div( Bxala)) + Ra,
(u(0) = Mo.
where the residual term satisfies Rg - O in C(10, T), (3(-11) as 3 + cc.
This proposition, whose proof is provided in Appendix B, characterizes the limiting dynamics within the lower-dimensional manifold, connecting the transformer model with heat flow on the sphere, thereby justifying the name of this phase. Remarkably, this connection holds without the need for correction terms, in contrast [46]. We now need distinguish between two different cases:
Forward diffusion. When y < 0 in equation (6), the dynamics corresponds to forward heat equation. In this setting, local existence and regularity for n²² (and in particular the assumptions of Prop. 3.9) automatically satisfied due the smoothing properties of forward diffusion, provided that fa0 E ( Notably, interacting particle systems of the specific form given by equation (SA) (under the assumption Q' K = = -V) have been studied in the literature and are known as viffusiion-velocityy methods; see, for instance, [33, 20, 8, 31, 32, 35].
Backward diffusion. When 7 > O, the dynamics corresponds backward heat equation. In this case, the regularity assumptions on 1 ensuring local existence and regularity are significantly more restrictive (e.g. requiring that 14g is in the Gevrey- space). Nonetheless, we construct explicit examples of solutions below. The backward heat equation is a prototyp- ical ill-posed problem, which explains why the statement of Proposition 3.9 is necessarily weaker than that of Theorem 3.2.
A family of initial conditions po that satisfies the assumptions of Proposition 3.9 is given by
M
a N
NM
where 0g 2 O. M aj 1, and Nm,, o²) denotes the heat kernel (the forward-in-time evolution under the heat semi-group exp(IA) of a Dirac delta, analogous to a Gaussian N (m, o²) in Euclidean space) centered at E gd-1 with concentration related to o². By linearity of A, the explicit solution to 0.p Ay is then given by:
M
11
Σ
a 03
For forward diffusion (7 < 0), this solution a smooth function for all t 2 O. while in the backward case (> > 0) this only holds for & E [0, mmin), where T'in min, of is the time at which the first Gaussian component collapses to Dirac delta 5my. more general class in which local existence and well-posedness hold in both the forward and backward directions is the set of positive, Gevrey-1/2. functions.
Motivated by the aggregation behavior observed in the finite-3 particle system, conjecture that the collapsed dm. remains invariant under the limiting dynamics, while other components continue evolving independently according to the backward heat equation until their respective collapse times. From the practical perspective, this observation suggests that the transformer's 's behavior in this regime form of regularized denoising (when 3 is finite) acting on the input. This
can be interpreted aligns with the clustering phenomena extensively studied in previous works on the model. The dynamies in this phase, governed by heat equation, drive the formation of distinct token clusters (via backward diffusion) or the smoothing of the token distribution (via forward diffusion). This behavior can be interpreted more defined semantic groups.
representation refinement stage, where tokens organized into
Remark 3.10. In [25], a simplified model, referred to the Unnormalized Self-Attention (USA) model, is proposed, where the normalization factor Zz, (xx) is replaced by a constant Zz, significuntly simplifying the mathematical analysis. By choosing Z3 = 3
1 fgg-1 (a.a) dy) (or equivalently by rescaling time), the limiting behavior of the model no longer yields the heat equation, but rather the porous medium equation: dip A(p²). Even in this case, the convergence of particles system to this nonlinear PDE has been extensively studied (see for example [22, 33, 39, 41] or [47] for gd-1),
7

---

3.4 Pairing Phase
The initial conditions we consider for the dynamics on longer timescales must be compatible with the steady states of the preceding phase. Motivated by the discussion at the end of the previous section, we therefore formulate the following assumption:
Assumption 5. The initial condition Po in the pairing phase can be written as po = Ei= azōs,
for an m E N. with zj E E aj > OVj E {l....m} and 2 aj =1 1.
Under this assumption, further supposing for the sake of clarity that a, 1/m for all j E (I....,,)) m we can interpret each cluster as a particle, and the dynamics of the system is given by the set of ODEs (SA). In the regime of large 3, clusters interact very weakly due their separation and the exponential tails (in 5) of the interaction kernel, resulting in exponentially long timescales for the nontrivial dynamics. Here, analogously to [3], interactions are dominated by the closest pair of clusters (t. 1). assumed unique, satisfying (a,, mz) max+4 (Zi, Tz)) at initialization. We note that this hardmax particle interaction, as well as the timescale where it arises in the large 3 limit, was introduced in [23, Section 6] in the case d 2. We present an analogous result here in arbitrary dimension, without claiming originality, to provide a complete dynamical picture across phases.
Proposition 3.11. The solutions xt(t) of the ODE system (SA), under Assumptions 4 and positive V, with the rescaled time dt~ 211 (((i))))))) ds, converge as 3 + oo to the solutions of the system:
(1) P (w)(W))
ifk = i. ifk = j. otherwise,
(1.(1)
U/ (0) (0)
on finite intervals T.). with T. such that (yi, U3) < 1 throughout the interval. for any 0.
In other words, all clusters remain stationary except for the closest pair, which collapses along the geodesic connecting them, in a time exponential in 33 Note that this result only holds up to an arbitrary moment before the first collapse. We refer to [23, Section 6] for detailed explanation of the challenges to bypass this limit and a proof of an analogous result until and beyond the collapse time in related but simplified model. The above proposition is proven for completeness in Appendix C.
This final, slow phase models the sequential merging of the closest token clusters. This can be interpreted as the construction of higher-order abstractions, where previously formed groups are hicrarchically combined to create more complex representations.
4 Numerical experiments
This section presents numerical simulations of the transformer model in Eq. (1), All experiments are conducted in dimension d 3 or or d 2 to facilitate visualization and are designed to validate our theoretical findings. The attention mechanism implemented using the official PyTorch wwl.prduct_aattettionn
function torch. nn. .n..nn.uu www.tuncciie ooal.. scaled_ _tl_d_pooouuctt.....m attention () and the experiments are performed on a single Nvidia H100. The code is available at [28].
First Phase Dynamics. Figure 1 illustrates the dynamics of the alignment phase, showing distinct behaviors based on the parameters choices for Q, K, V. For both scenarios presented in Figure 1, the initial state consists of N 1 tokens sampled independently and identically uniformly from the sphere 8
We set the inverse temperature parameter 3 = 30 and use a time step of dt = 10-2
Scenario la (collapse to 1D subspace). The matrix VKJ Q chosen such that it possesses a unique eigenvalue with maximal real part. As predicted by our theory, this configuration leads to the tokens collapsing onto a one-dimensional subspace (i.e. two antipodal points).
Scenario 1b (non-gradient flow dynamics and rotation). This example employs param- eter choice for Q, K, V that falls outside the gradient flow regime. Nevertheless, the tokens are observed to collapse toward two-dimensional subspace (a great circle), accompanied by collective rotation of the particles along this circle.
8

---

Both observed behaviors are consistent with the results in Theorem 3.2 and Proposition 3.4.
==
0
2
Q'k=
0
==
10 1
Q'K
-2
Figure 1: Simulations of two different scenarios (one per row, four timesteps) for the first phase.
Second Phase Dynamics. We support the conclusions of section 3.3 through two examples, comparing empirical dynamics with analytical solutions of backward and forward diffusion equations.
Scenario 2a (collapse to 2D subspace and backward diffusion). For the experiment in Figure 2 we set 10. The initial configuration comprises = 10 i.i.d tokens. Their elevation angle & is sampled uniformly on I- 2. 3), while their azimuthal angles, 0; [0, 2T), are distributed according to the mixture density g(0):
g(0) 0.2 (O; + 0.5 (O; oo) 0.3 (O 4
where Ni:: p4, oo) denotes the probability density function of a wrapped normal distribution on S with mean and standard deviation do = 0.2. The parameters Q, K, and V are chosen so that, after the first phase, the tokens collapse onto the zy-plane, with distribution g(0). The analytical solution to the backward heat equation with initial condition g(f) (computed as in Eq. (7)) is plotted red curve in Figure 2. The positions of the clusters agree with this solution, numerically confirming the predictions of Proposition 3.9.
Scenario 2b (forward diffusion comparison). In Figure 3, we compare the empirical token distribution with the analytical solution of the forward heat equation characterizing possible example the second phase of the dynamics. Specifically, we simulate the evolution of 5 x 10 tokens, initially sampled from superposition of three Gaussian densities on S! through the transformer model with parameters 3 50, 2, Q = K Id. V Id, and dt 10-3 The analytical solution of the forward heat equation (in red) closely matches the token distribution histogram (in blue) over time (i.e., depth). Note that, as expected, the forward diffusion process is significantly more stable numerically than the backward one.
==
0
Lo
0.1.
Q"K=
0
0.1
Analytical solution
Figure 2: Numerical simulation the backward scenario for the second phase on $

---

6:5
212
0.22
L 2.03
rabery Araystcalten Eo.54
0.0
Prace (nndion
Provee Jostoon))
Prace
Praton Upvdiance
Figure 3: Evolution of the tokens distribution in the forward scenario for the second phase in S!.
5 Conclusions
This work provides a mathematical analysis of token dynamics in mean-field transformer models within the moderate interaction regime, where interaction strength (3) scales with context size (N), motivated by scaling practices in modern LLMs. Our study reveals a fundamental multiscale structure governing the evolution of token representations through network depth in this setting. More specifically, we showed that, under this scaling, the system progresses through a sequence of three different dynamical phases characterized by qualitatively distinet dynamical behavior, operating on separated timescales. Through our analysis, we offer a unified dynamical picture describing how deep transformers might achieve progressive representation refinement.
This unified dynamical picture is, however, not yet fully rigorous. Indeed, while we establish convergence results for the alignment phase and provide partial justification for the intermediate and late phases under specific assumptions, a complete mathematical treatment of the full dynamics particularly of the backward heat regime beyond the first collapse and of the slow clustering interactions remains an open challenge due to significant technical difficulties in the analysis of the strongly unstable limiting equations. Additionally, characterizing phase transitions in self-attention for different 3, scalings is an interesting separate question, with progress made in [15] under assumptions on inter-token angles.
Furthermore, while the first dynamical phase has quite general assumptions on the parameter matrices, the following phases still require relatively limited assumptions (although less limited than in most previous works). Relaxing these assumptions further, in particular in the case of non-gres non-gradient dynamics, would constitute interesting, but also technically quite challenging, avenue of future research.
There are several important directions in which our work could be extended. Most notably, incor- porating the MIP, which could be interpreted as introducing a drift term in the dynamies, acting independently on each token without accounting for mutual interactions. Another natural extension involves studying the dynamics under more general parameter settings. For instance, during the heat phase, assume that Q" K S is symmetric positive definite, which holds, for example, when Q Q=A K. This "shared-QK" assumption is not novel and has been adopted in prior empirical work (e.g., in [30]). While different choices of these parameters (both MLP and Attention matrices) can have a dramatic effect on model behavior, with adversarial choices potentially leading to qualitatively different dynamics from the one predicted in this paper, believe results to be relevant first step towards understanding the development of representations in transformers, capturing some important qualitative features of these models shown in [25].
further direction of future research consists in providing sufficient conditions for the stability of the space Emaz emerging in the alignment phase under the prelimit model (i.e., for large but finite B), thereby justifying Assumption and, ultimately, connecting in a rigorous way the alignment and heat phases identified in this paper.
While the path from this theoretical analysis to direct application is not immediate, we believe our work opens several potential avenues for future investigation. The characterization of the alignment phase, for instance, offers a potential mechanism for interpreting how token representations evolve into learned subspaces. Finally, by focusing on the "moderate interaction regime", we hope our analysis provides a theoretical foundation for a more principled understanding of parameter scaling, particularly models are adapted for longer contexts.
10

---

Acknowledgements
AA and GB acknowledge partial support by the Institute of Mathematical Statisties and Actuarial Sciences at the University of Bern. FP acknowledges support by the American Mathematical Society through the Bergman fellowship. AA thanks B. Geshkovski, J. Mattingly, Polyanskiy, and P. Rigollet for interesting discussions. All authors thank the anonymous reviewers for comments that have improved the quality the paper. Calculations were performed on UBELIX, the HPC cluster at the University of Bern.
References
[1] Milton Abramowitz and Irene A Stegun. Handbook of mathematical functions with formulas, graphs, and mathematical tables, volume 55. US Government printing office, 1948.
[2] Andrea Agazzi and Jianfeng Lu. Global optimality of softmax policy gradient with single hidden layer neural networks in the mean-field regime. In International Conference on Learning Representations (ICLR 2021), 2021.
[3] Albert Alcalde, Giovanni Fantuzzi, and Enrique Zuazua. Clustering in pure- pure-attention hardmax transformers and its role in sentiment analysis. SIAM Journal on Mathematics of Data Science, 73)):1367 2025.
[4] Albert Alcalde, Borjan Geshkovski, and Domènec Ruiz-Balet. Attention's 's forward pass and wolfe- arXiv preprint arXiv:2508. 09628, 2025.
[5] Luigi Ambrosio, Luis Caffarelli, Michael G Crandall, Lawrence Evans, and Nicola Fusco. Transport equation and cauchy problem for non-smooth vector fields. Calculus of Variations and Nonlinear Purtial Differential Equations: With a historical overview by Elvira Mascolo, pages 1-41, 2008.
[6] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint (wt.:v::400..0 0093,0 2014.
[7] Raphaël Berthier, Andrea Montanari. and Kangjie Zhou. Learning time-scales in two-layers neural networks. Foundations of Computational Mathematics, pages 1-84, 2024.
[8] Yann Brenier. Geometric diffusions of 1-currents. In Annales de la Faculté des sciences de Toulouse: Mathématiques, volume 26, pages 831-846, 2017.
[9] Giuseppe Bruno, Federico Pasqualotto, and Andrea Agazzi. Emergence of meta-stable clustering in mean-field field transformer models. In International Conference on Learning Representations (ICLR 2025), 2025.
[10] Martin Burger and Antonio Esposito. Porous medium equation and cross-diffusion systems as limit of nonlocal interaction. Nonlinear Analysis, 235: 13347, 2023.
[11] Martin Burger, Samira Kabri, Yury Korolev, Tim Roith, and Lukas Weigand. Analysis of mean-field models arising from self-attention dynamies transformer architectures with layer normalization. Philosophical Transactions A, 383322299: 2025.
[12] José Antonio Carrillo, Katy Craig, and Francesco S Patacchini. blob method for diffusion. Calculus of Variations and Partial Differential Equations, 58:1-53, 2019.
[13] Valérie Castin, Pierre Ablin, José Antonio Carrillo, and Gabriel Peyré. A unified perspective on the dynamics of deep transformers. arXiv preprint arXiv:2501.18322 2025.
[14] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[15] Shi Chen, Zhengjiang Lin, Yury Polyanskiy, and Philippe Rigollet. Critical attention scaling in long-context transformers. arXiv preprint arXiv:2510.005555 2025.
[16] Shi Chen, Zhengjiang Lin, Yury Polyanskiy, and Philippe Rigollet. Quantitative clustering in mean-field transformer models. arXiv preprint 27X1 2025.
11

---

[17] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing systems, 2018.
[18] Christopher Criscitiello, Quentin Rebjock, Andrew D McRae, and Nicolas Boumal. Synchro- nization on circles and spheres with nonlinear interactions. arXiv preprint arXiv:2405.18273, 2024.
[19] Valentin De Bortoli, Alain Durmus, Xavier Fontaine, and Umut Simsekli. Quantitative propa- gation of chaos for sgd in wide neural networks. Advances in Neural Information Processing Systems, 33:278-222 288, 2020.
[20] Pierre Degond and Francisco-José Mustieles. A deterministic approximation of diffusion equations using particles. SIAM Journal on Scientific and Statistical Computing, (1 11(2):293-310, 1990.
[21] Weinan E. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 1(5):1-11, 2017.
[22] Alessio Figalli and Robert Philipowski. Convergence to the viscous porous medium equation and propagation of chaos. ALEA Lat. Am. J. Probab. Math. Stat, 4:185-203 2008.
[23] Borjan Geshkovski, Hugo Koubbi, Yury Polyanskiy, and Philippe Rigollet. Dynamic metasta- bility in the self-attention model. arXiv preprint arXiv:2410.06833 2024.
[24] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2024.
[25] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on transformers. Bulletin of the American Mathematical Society, 62(3):427-479, 2025.
[26] Borjan Geshkovski, Philippe Rigollet, and Domènec Ruiz-Balet. Measure-to-mmeauure interpola- tion using transformers. arXiv preprint arXiv:2411.0455 2024.
[27] Alessio Giorlandino and Sebastian Goldt. Two failure modes of deep transformers and how avoid them: unified theory of signal propagation at initialisation. arXiv preprint arXiv:2505.24333, 2025.
[28] GitHub-Repositor
http:////ggghhuubcoombmebbuuutee wtt.://rruuolom/mmmissss.......a orrar///ersstt.taeeeee.rrrrrrrrr transformers..
[29] Nikita Karagodin, Yury Polyanskiy, and Philippe Rigollet. Clustering in causal attention masking. Advances in Neural Information Processing Systems, 37:115652-115611 2024.
[30] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[31] Gilles Lacombe. Analyse d'une équation de vitesse de diffusion. Comptes Rendus de l'Académie des Sciences-Series I-Mathematics, 32965):383-386, 1999.
[32] Gilles Lacombe and Sylvie Mas-Gallic. Presentation and analysis of a diffusion-velocity method. In Esaim: Proceedings, volume 7, pages 225-233 233. EDP Sciences, 1999.
[33] Pierre-Louis Lions and Sylvie Mas-Gallic. Une méthode particulaire déterministe pour des équations diffusives linéaires. Comptes Rendus de I l'Académie des Sciences-Series I- Mathematics, 332(():369 2001.
[34] Johan Markdahl, Johan Thunberg, and Jorge Gonçalves. Almost global consensus on the n-sphere. IEEE Transactions Automatic Control, 6366):1664-1675, 2017.
[35] Sylvie Mas-Gallic, The diffusion velocity method: a deterministic way moving the nodes for solving diffusion equations. Transport Theory and Statistical Physics, 31(4-6)555-660, 2002.
12

---

[36] Song Mei, Andrea Montanari. and Mhan-Minhh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E765-- E7671, 2018.
[37] Ken M Nakanishi. Scalable-softmax is superior for attention. arXiv preprint arXiv:2501.13399 2025.
[38] Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27198-27211, 2022.
[39] Karl Oelschläger. A law of large numbers for moderately interacting diffusion processes. Zeitschrift für Wahrecheinlichkeeissheeree und verwandte Gebiete, 69(2):77--322, 1985.
[40] Karl Oelschläger. Large systems of interacting particles and the porous medium equation. Journal of differential equations, 88(2):294-346, 1990.
[41] Karl Oelschläger. A sequence of diffrr--iffeeentialll equations approximating viscous porous medium equation. Zeitschrift für Analysis und ihre Anwendungen, 20(1):55 2001.
[42] Thierry Paul and Emmanuel Trélat. Universal approximations of quasilinear pdes by finite distinguishable particle systems. arXiv preprint urXiv:2501.13387 2025.
[43] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
[44] Yury Polyanskiy, Philippe Rigollet, and Andrew Yao. Synchronization of mean-field models on the circle. arXiv preprint arXiv:2507.22857, 2025.
[45] Grant Rotskoff and Eric Vanden- Eijnden.. Trainability and accuracy artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 15(9):1889-1955, 2022.
[46] Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transform- with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics, pages 3515-3530. PMLR, 2022.
[47] Anna Shalova. Noisy gradient flows: with applications in machine learning. PhD Thesis, 2025.
[48] Anna Shalova and André Schlichting. Solutions of stationary mckean-vlasov equation high-dimensional sphere and other riemannian manifolds. arXiv preprint arXiv:2412. 1488133,, 2024.
[49] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in newral information processing systems, 30, 2017.
13

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Abstract and introduction summarize the claims made in the paper that are proved or discussed the following sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, see section conclusions.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: The assumptions are reported in the corresponding sections, while the proofs are in the supplementary materials
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
The details of the experiments are reported in the corresponding section.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. .., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
15

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
An anonymous github repository is provided.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /nips. wttpcccpppp//eeeeeeee..ec...llle gguidde/Codeeeiieeee cearbbiissiooooccyyy. for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [NA]
Justification:
The experiments don't 't have training and test.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: The kind of experiments does not need statistical significance, they solution ODEs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See the section about numerical experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : /ttp///eerppp.e//ee//eee/ee/eeee cc htcppplicc/thiccuuuieeeeeeo
Answer: [Yes]
Justification: Anonymity and all the other rules have been respected.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
is a theoretical paper.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
is a theoretical paper.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: It is theoretical paper.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
See numerical experiments section and the github link.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cooe/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
19

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy buttps://neurppe:: http///rpps.ccoooeeeeeeeeeeeeeee cc//oonfeeeecccs/222//III for what should or should be described.
20

---

A Proofs of the alignment phase
This section is divided into two parts. The first part contains the proof of Theorem 3.2, while the second one contains the proof of Proposition 3.4.
A.1 Moderate scaling limit
Consider the family {a²) J220 of solutions to the usual continuity equation (2):
Sape -div(n Setu(nxslee x3)x))), x0, = -o(x)) a E8
where Xe the vector field given by:
xalu((x)= Jsa-1 Bl2, P. yp(dy) By) p(dy)
SSu eB(z.
Remark A.1. For notational simplicity, we will refer to the matrix Q K in the main body B.
Remark A.2. In the following C will be constant depending only V, B, d and po- Its value may change line by line.
Under assumptions 1, i.e.:
A1: Q, B are invertible square matrices.
A2: The probability measure μo 5d-1 is absolutely continuous with respect to the Lebesgue measure. Its density bounded and Lipschitz continuous and its minimum satisfies min (z) 0.
one prove Theorem 3.2: μı converges weakly to poo in C(O.TT):P(5+-1)) where Moo is the unique solution the PDE:
0,p -diu((p P.VBTa [8
((0,2)) =mo(z), I E
3d-1
and the metric in
is given by:
1 f(z)) (aami)f((x)) (x))
d(H,v)= == sup p(u(, v.) tE(O,T)
sup aE(0.T] (EBL(SH-1)
sup
where (Sd-d is the set of all the Lipschitz continuous functions on gd-1 which are bounded together with their Lipsehitz constant by 1.
The idea of the proof follows five steps:
Relative compactness in C(OO, P(St-11))
Bounds cumulants of the Von Mises-Fisher distribution,
Prove a relationship between the derivatives of μ and the regularity of the vector field,
Apply continuation argument to show the uniform regularity of xoluø).
Use the regularity to pass to the limit in the PDE.
A.1.1 Relative compactness
Proposition A.3. The set {pstazo is relatively compact in C( 1], P(Sdd11))
Proof. By Prokhorov's theorem, since 1d-1 is compact, we can conclude that P(Sd-1) is weakly compact. Since p metricizes the weak topology, then also (P(Sd-1), )) is compact. To apply Ascoli-Arzelà theorem, just need the equicontinuity of the set { Jaso-
21

---

Given 0 5 2 T and 3>0 0:
J ((0)( - (d(r))
SEAL (g0-1
sup
[.V (VI()), (x) (n)))
sup SEBL(g4-1)
ga-1 Ix s -ss- s/-
This is sufficient to conclude the proof.
A.1.2 Bounds the vector field
The aim of the following paragraphs is to obtain some bounds D'x 0, 1,2. To fix the notation we define the probability measure 11 on Sd-1 as:
BBcz u(dy) J.B..B (..B.B(dy) (dy) = JSu- Buta pully)
Remark A.A. The measure v2 v....B is the Von Mises-Fisher distribution with mean direction and concentration parameter B) 3)B²aa z/. Some properties of this distribution are studied later.
1
Then the vector field X [u] can be written as:
Xøle)(x) V&V (E.g.eeMr) -
(M (B. P&V
(BBT x),
where denotes the composition with respect the parameter of the measure up U. Lemma A.S. The derivatives of the vector field xs(uu] are bounded by:
[x)lu[(x)| cc
[D)xo||[<)) s 1 3 E 1p
(r Byg.s. [x])
)1)
02
11)
[D(xolu[()) < c 1+ 1 p (r Beg.n (M1)
@2
]| 32 |L |Bgg (r B 83 (1]
where a constant depending only on V, B,d.
Proof. Let's compute the derivatives of x [x)]u](x)1 E [Y] S Cv, 1D - PaVI| (T ( [Y] (3B* | |PAVIID! (E (Buga.n(r))) |
Cv (1 + DD E3 Eegeo (x)))
[Dxxsl|((0) < DDPP DOI J. [Y] + 2.D. (E (B.g (B ++1P2
(E
< Cy + 2 (B,....(Y))
D E (.
Hence we need to compute the derivatives with respect to x of E E. 2.8 [Y] Thanks to the Faa Bruno formula:
Egg [Y] (3BTx) x).
DE e [ = Σ
(D) TD!TEyea Re (V]) BB
X DD(33T
TEIL.
PER
with II, the set of all the possible partitions of {1...., ni and the tensor product. The previ- ous expression can be bounded by C' Ei= 3'1 (D (YT) where Cp is constant depending on the matrix B.
22

---

Thus, the aim is to compute a bound for D (Y) This is related to the d- dimensional cumulants (tensors) of the distribution pg 1lIl,Td Indeed, we can write:
D BB (Baga.aa(1)
D [B stv
[ eletoaal p(dy) J elrtovi p(dy)
D
v=0
elowny cl p(dy) S fele (ldy)
I felt p(dy)
D
(8)
eeviieelaa u(dy) 2=0
[ Selvvby (dd [eirie) v/ll.10 (dy)
D
v=0
D
(log BoggEygnnneee (ongeagnn
It is well known that the first three cumulants correspond to the central moments:
Do E
(M))
E gu (Y],
x
D2 E (eggh.sa(Y))
E Big.d.dd (Y - E 1 (11)
@2
].
(9)
x
D
R (egga gggat (Y L (M)
63
The thesis then follows by replacing these equalities in the initial bounds (after renaming Cv and Cn).
Lemma A.6. Let a be the uniform measure on Sd-1, Then:
E E [1]
1,
82 c_
E gggarral (Y
a. [1]
3'
03 B B.g (Y- Ba M1)
s C
3
Proof. By (8) and Schwarz's theorem the three tensors are invariant by permutations of the indices and by definition of vg. B,id they are also invariant by rotations that fix x. Hence (see lemma D.2) they must have the form:
aa [Y] = a1æ,
E. Eyg (Y (Y- E gi [8]
1 == D (Bgga.a(Y))
=
02 2 x B21,
(10)
8x
E Bggaral (Y ( B Bugiaraa (81
1 D2 (B.ga.n(M1)
u3
x 2 33Sym(# @ Id).
8z
Where Sym(x 220ju 22 + We need to compute the coefficients a1 022 32 0g P3. Define A(3) Jga-1 (z. (I. v) (dd Similarly to what has been done in (8), can relate this to the cumulants of (I., Y) by noticing that:
( (x, v)o(dy) 0g (A(3)); = 0r (A(B + 1))- = 0r (Telo-iiewolta (dy)
t=1
= 0." ( [ (z, (lav)) o(dy)
[ v) Sella.e)v 3.1. 3. (d
=
f o (dy)
t=
t=
= 0,-1 (logE,g.m 1ole (e6
1=
23

---

This give us immediately the following identities:
A(3) E (eg, X),,,,
A((3) E. [( Y) A(8))),
(1I)
(3 E. [(xx A())))
Suppose without loss of generality that x E1. Then ai is given by:
a B u. (en Y) A(8).
(12)
The coefficients 02 and B2 can be obtained comparing the representations in (10) with the representa- tions in (9), and exploiting the relations in (11):
02 + 82 = D2 (E (B
[e] = Eyg. ((e1, -
Am
(d - 1) = Σ DJ (Bagaee(Y1) [en = Σ
E B (((e, X1)
(13)
Br
i>1
131
=1 =1 A, 8)) =1 -B Bigia.ia [(ee1 Y) - A)3 - A²
2
And the same can be done for Or and B3:
0g + 333 = D2 E (Egg...e(11) 8x
[1,, E.e,,, A(B))* AB))),
(d - (d-1)33- = Σ D2 (B (Bngnra(M))
[e,,er =
Σ
B Ba ((((,, (/e, p) A))
B
8x
U/
i>1
i>1
=-B, v
Y) AA]]
= B Bgga,, Y) 1)3] 2A 2A [((e1, Y) A)
-
To conclude it is sufficient to show that 1 - A2 =0( and A" -(7).. Now, using the identity:
Z = ga-
p(a.)) -a/3--d22 C. (3)2
we can explicitly compute A(3) as:
0gZs
(1-
1eal/22//222 1)8 d/2 14/2-111111 (3) 31- (1
(La/2(B) + + - 11)444//2-11(1)
A(B)
8r d/2 12/21(111
(14)
(B) d-
T8)2-1(8)
2,3
where we used the derivatives rules for the modified Bessel function:
[[(=) I--1(2) V/ (,);
[x)x = 10+1(2) - (z),
and its asymptotic behavior (in both cases see [1]).
24

---

In similar way we can also compute:
(3)2
(3/2(5)
(3)
18/2-1(8y
13)
(Ia/2-1(3)
1a/2-1(3)
4/2
B (5/2(5)
T3/2
[aapa(3) + d/2-1
Las2
1a/2-1(3)
1d/2-1(83
13/2-1(3)2
(15)
d/2
d/2 A(B)
A(3) A(B)²
(())
3
3
d-1
A(3) (d- 1)²
1 (B)
3
432
and
(3) = -2 1 N00 E
where we used the asymptotics in (14). This is sufficient to conclude the proof.
Lemma A.7. Given a strictly positive probability measure μ the following holds:
/ (Iy anl + CB- Ug.0.B.
B. (y) 0. (y)1 ك
min |u|
B where TB
Proof. Indeed:
Ba, p(y)
3 (v) JS (ee Bw) ply)
By) u(xB) + (B, (m(y) (((B))
u(5B) Ssa- e Bo + Jgc-1 Bu) (u(v) - u(zn))ddo()
1+ u(v)=u(zE)
(g..B (g)
m(nn)
1 ( (u(y) )dv? 3.B
((y) ((u)) m(aa] [u(()) (d
3.
1
u(#B)
1 x(=z) [M(lV) - u(EB)) o. (.B))
3. (v) (1 (R1 + R1) +2))
|| (Iy - JBl + CB-1/22)
S 3.
1
min [n]
where we used:
RRl S (u())= u(xB)
| yy - x/.
u(In)
min |u|
/ (u(v) - u(an))dd2 Ral S m(xE)
1
||
min |14| / [g - xn 3.B. S
C3-11/2
|P 80 min |
and the last inequality is consequence of vg,8.B Proposition A.8. The derivatives of the vector field x² [M]] satisfy: xlu] < C,
0.3| Bd xl.Jd and lemma D.I.
Vxe
|/ 1122
D'xIuSS
1+
min |u||
| g-1/2 +
/
D2xlusS
1+
min |a|| min [n]
25

---

Proof. Thanks to Lemma A.5 we just need to bound the cumulants. The first one is already done.
Second cumulant:
E Eggin. 1(Y - E B.g. (M)) 82 =
/ / (w)
- -2))uf.... (dy) (dy))oz.BB (dz) (uf (dzz) 21) @0 (y- (1)
R Bugin (Y Begaav (v1) 82 ]
+ R.
We have already shown in lemma A.6 that the first term is S C Now we need to bound the second term. R can be expanded by multi-linearis and using lemma A.7 the worst case is either of the form:
c Valloc 3-1/2
//
JJW (w-
@ (y- J 21)
22) (dy) 2. (dz1) (uz)) 8.B. (dzz)
min |a|
Vulloc 3-1/2 B (dy)
II / y 2z ((
2
< C min |U|
/illlo 3-1/2
(ʃ1b 41 12
C
Uo (2 (dy)
0. B. (d21)
min |a||
Vullos 3 1/2 U2 (dD)) / 1 21 222.0. (dz1)
C min [p||
< 4C
1/2 V B 3-1/2
/ /wv- -
xn 1 (dy) S 40
Vul 8 0/2 8-12--1 < min
Vulloo min |p| B3/2*
min [s]
where in the last line we used lemma D.I, or of the form:
Vullo
If f ( 21) s (v - 22)ly - 2alug 3.B. (dy) 0. (dz) ((2)) (2. B.B
C min |a|
(dzz) ((zz))
/ulloo
/ / |1 21l1y zelly - BB 2 B.B (dy) 3. (2z) (d2)) ( (d22)
C min |p|
p
(flv- - 2B 2 R. (dy)
1/2
y 2
1/2
min |م||
// loo
(fls 11 nB 2 x B. (dg)
1/2
2110-201'00010 + 12 |12 2B3 201""200000 (dz) 1/2
S 40
B(
min lea|
Vulla 3 1/23-11
/pll
min| [u]
min |u| g3/2
or of the form:
s pa] f (V 21) 8 (y - 22)121 - 20112 (dy)) (d21)) B. (d22)
//llo
min
/
{ / 11 21 11y 22l1a1 - =B 2 1.
(dy) 3. ( 3.
min |M4]
/ (/12) - 2B 2B82 2. (.
1/2 ( by - 22 - 22 22 (dy) 3. (d21
1/2
B. )g (dzz)
min |A|
| (flu 11 =B P Pug.B.1 (dg)
1/2
(fl0- 1110-28l" 1. (dy) + / 1a =B 8,. (dz))
1/2
40 min [u||
/ 3-1/2 < C min |||| /3/2"
|
<C min |م|
1/2 by lemma D.I. Hence These are the worst cases because every ly 2| produces an additional / /-1/2 we proved, thanks to lemma A.5, that:
V 3-1/2)
|Dbxlu|| s C
1
min [p]]
26

---

The bound for the third cumulant is similar to what we have done above:
[R g (v B Begie a [M] 83
(y (---z1 21) 2 (y (2-22) (y za)ut...B ( (d.)) (ddz))n...BB (2 (122))
(dzz) (dzz))
-lRug.n.n (Y- [r]) 03
+1
We have already shown that the first term is O( 1
ge ). Now we need to bound the second term. R can
be expanded again as in lemma A.7 and the worst case is either of the form:
/ulloo 8
C
8-1/2 1/2
(y 21 2
(y
22) @ (y 2.. (dy)v,, (
(2 ).2 (d2))) D. (d2z)
min p
Vullz /-1/2 1/2 min pr|
II / ly 2/ 3. (d2s)
3
0
1/ (dy)
==1
| g g-1/2 1/2 12 (dv)
(ʃ | 21 122000 D. (dzz) 3
C
min jal
/ g 1/2 3-1/2 / (2(d))
C min (
1/ - 11 13 (dz1)
| lo 3 1/2 / 11 - EB 3.....BBBBBB (dy) & C3-1/23-312 < <c || 1
$ 22
min [p|]
min |u| _
or of the form:
<0 min a
/
oc
[[/[v- (y- (2)2))))))))))))))0)))))) (y 23)18 2Bl 3. (dy) (d)))) (dZ)) 3. (d2))) (d223)
<0
[ / ll - 1z A1123 2al
- B 3 3.B (dy) B.B (dzl)= 3.B. (dz)) (dza)
R.
min [
||/p
oo
/ a1 - 23 22 3. (dz, (d2;) B. B.
1/4
(flu-zeluu2 22 (2
1/4
<C
min |e
/ //1w-2/* 21 R.8. (dz) B.))
3/4
(fw-
1/4
min [M|
V do (8 (3-1418-21112 C. 1 3-2.
<0
min [a]
min [x]]
Thus, replace the bounds the second and third cumulants that obtained above in the estimates of Lemma A.5 to conclude that:
DDx11u11 SC 1 IE o... (r
B.g (11
63 "] 11 B² IE 1Bgg..e (Y (By.g. _(r)) [11]
1|
V g-1/2
Vuello min |M|
SC 1+ min [p]
1/2 +
Lemma A.9. If solves the PDE:
Opp (m(0) μ4-
then:
. Bellulloo S (Dllo/1Exxit
27

---

de (min )²) 2 (min )) x(u]],
0el V S || Dix(u]l + alllloo) x
Proof. Let be point of maximum for (ual. Then se(=) = 0 and:
(xp)( -div(px(u)((a) + (Vpe(z:), -((x))dd((((x(l)(()))
And for min μ use the same argument.
Now, let x, be point of maximum for V sa then Vg 0, hence:
Oe pp(r))²
(Vu(s,), (dio( + (Va(a,),
(Vu(s,), V(V V(Vxx))))) (V)(x)), V(DDi)))))))
(Vu(z)), (((())))))(( (V,(z,), hu
2 V2(z,) + 110e0////00
Using that V = (0 and dividing by [u((z)) we get the thesis.
Lemma A.10. Consider again μr solution of the PDE:
-dio(uxslu)),
(M(0) Po.
Define:
C1 201110ll20
| aCT
C2 = 20
[+
min μr
for B large enough (depending just on pu and C)::
Melloo s 2(11pollo0)e201,
min pz 2 (min Lmminuoe" -20 20t
Vullo 22 (C) + IVHoloo) eCst.
Proof. The thesis true time 0. Let us assume that true on [0,4] Then 33 big enough
(where "big" depends only on C₁, C2, i.e. just μo, B, V, d) such that |ull jM] 3-
1/2 < 1 on [0,4].
Hence Dx(us] < 2C [0, f] thanks proposition A.8. By Gronwall applied the first two bounds in lemma A.9 we can conclude:
|| x (1 (x0)
min μr 2 (min mo)e 20t
For // we have, again by lemma A.9:
0 V & [/V x + Z |11111o0 12
pelloo
2 + (1lvolboc220)) c [1+1+
min [p]
22C1|oollo 2CT +C (2+
11lollso ,CT)
/Pellso
min μi
= = =C
28

---

where in the second row we used the assumption on [0. and proposition A.8. Hence by Gronwall:
& v (GL 2 + |Vpolloo ecte
This concludes the continuation argument and the proof.
Corollary A.11. For & large enough (depending on po, B, V, d) the vector fields {xoluta, p.a are jointly Lipschitz in 3 and E T.
Proof. This is a consequence of lemma A.10 and proposition A.8.
Corollary A.12. For every a E gd-1 and E [0, T):
P. P&VBL
xa x2)
as 8 - o0
Proof.
With the usual notations one can write:
Xalu(](x)) = P&V (Eog. (X1) = PaV (Bng...0Mii PPaV(Egga + R.
The reminder R bounded using lemma A.7 by:
| / 1V xBl TBBL 8.8.
R S C
B 1/2
min p
|
c
1/2
B 8-1/2 1/2
0(3-1/22,
min H
where the last line follows from lemma D.1 and lemma A.10. Hence, the proof can be concluded by noticing that:
VB BT
P&V (Baga.e = P&V B
= (BrgIBTaNN = = V2g
where we used the identities in equations 10, 12, 14.
We can finally pass to the limit in the PDE. Indeed consider a subsequence μæ of solutions to the PDE converging in C( 21, P() certain probability measure u If we define the vector field (x) = P
n² x |
then for every f E (Sd--))::
(f.P)) (s, (5 - I (VI.XxoM)) ds < (f,4
S: KVV - -xelu(uu) ds,
pi +
where we used that p po o/ and that the PDE in weak form for 4 is:
(, ( (5 1." (V (ulle) ds
Moreover, as 3 > oo:
10 -
thanks to the fact that / is Lipschitz and by the definition of convergence in C((O,T]) P(Sd-1)). For the second term:
L KVS, KVSxX X ds 1 x ds
. (V KV - ds.
The first part goes to 0 by dominated convergence (V f. Xoo x [[u]] are bounded, and xolus] - Xoc point-wise by lemma A.12). The second part goes to 0 by definition of the convergence Hoxc - μ4 and by equi-lipschitzianity of x [uy] (see corollary A.11).
The uniqueness standard, since μ is probability measure and the vector field smooth (see, for example, [5]).
29

---

A.2
a time-reparamccterizaionn of: Asymptotic behavior
This section studies the asymptotic behavior of the support of the solution to the partial differential equation:
P
= -
1
BT x0
(16)
and in particular, we prove Proposition 3.4. Lemma A.13. The ODE:
d r() |B3 x((a(
PaVB²z(t x(t)
(17)
dt
dt (() = PyVB"y('
where y(t) (t) and f(t)= = So BT ds.
Remark A.14. The reparameterization well defined since B invertible.
Proof. We have:
d dt v(t)) xz d BT x) |BTx]
(t) dt'
(1)-11"""((
Py(VB"')
This shows that the ODE a time-reparamettriizaioon of the ODE for yt).
Lemma A.15, If 2t) solves:
d
z(1) 2() VBT=2(1)
di
x(t) then y(t) [=(e)
Proof. We have:
d.
d x(1)
z'
1
di u()
=(=
dt [=(()]
12]
2 2/2
VB 2
(ZT)B22)
2)
VBT
(w. V B" 1)
|2|
12|3
2
y
Pg V BT y.
This concludes the proof.
Corollary A.16. For Lebesgue almost every Tg. the w-limit set w(xo) C Emnax-
Proof. This is a consequence of lemma A. 13, lemma A.15 and of the classical theory for linear ODEs. after reducing the Jordan canonical form of the matrix V
Remark A.17. This technical result parallels Lemma 3.1 I in [29], which was used to analyze the dynamics of the first token in causal attention.
And now can finally prove Proposition 3.4:
30

---

Proof. Denote D the flow of the ODE (17) and let 1) be test function with supp(o) C Ee ngd-1, Fix Moe w(Ho). Then there exists a divergent sequence of times (t such that Nta 7 Mo weakly. As a consequence:
ISnn (d))
lim J 21 #4 (o 1-000
lim k gd-1
((x) (dx)
lim k-000
go o( (z)))o(d)) 0,
where we used corollary A.16 and the dominated convergence theorem.
B Proofs of the heat phase
In this section, we prove Proposition 3.9, which characterizes the second phase using the heat equation on the sphere.
Lemma B.1. Given a measure μt E (
strictly positive, then the following holds:
(x)lm((x)= Vap(œ) //lll (1+ ||Hajell>>) O1=088--5(08
m(x) min p
min
with the gradient Va and Hessian He defined with respect to the standard Riemannian metric S
Proof.
(xolu((x) = 3 I ((a.))) Pay u(dy)
I Se0lz.0 p(dy)
I ealz.0) PaY m(dy) 1
m(x) I c 1 1 R₁
B [ Pay (p(I) (y Vp(x))) o(dy) M(z) I (
1
+ Rs
1 Ri
where, by lemma D.1:
I Jellus)
Pullo
/ |ww 3. (dy) = Villo
R&
0 0(8-112).
(xx) I o
min jr
min j
I ((ein) Pay)uldd ((e,)) Pay (u(()) + (y - x, vn(x))) (dy)
R2
(18)
u(x) I (( (ddy
H / I1 x (dy)-
|| 0(8-3/2).
minu
min μ
Hence:
(P Ps 3. + Ps (V--))000 x) 3 (dy) VzH(Z) ((I)
8xolu((>) 3
+ RR2
I Bi
and noticing that E a [Y] is parallel to (see proof of lemma A.6):
{xolu|(x) = 3 Ps (y - E. [1] 82) 1/ 8, (dy)
(U-Ezg
V=H(I)
+ Rs
1
M(x)
1 R
P.B P. [[Y - E R. (2 02 V [zM(Z)
3
+ RR2
M(T)
1 R
Palazz & x + 321) V=H(T)
1
=
+ R2
((x)
1 Ri
1 (( 2 V [z)(T)
1
3
+ R
d d-1
((z)
+ R1
31

---

where 0z,, 52 are defined in the proof of Lemma A.6) and we used equation 13. To conclude, it suffices replace equations 18 and the asymptotic estimates 14 and 15:
3
1 V=f(a)
V
0 3/23)
1
B p(x)
min
+
Haalll 0(3-1/2)
min M
Vau(())
V 0(9-1/2)
HaH 1oc -0-30033-10 112))
1(x))
min μ
min μ
Vap(x)
Vplle
(1+
| min
0 1/2).
u(x)
min
Corollary B.2. Given a family {pf) }B of probability measures on Sd-1, suppose that there exist C, C. > O such that [| Ioarg4 1) C and μ² cfoor every 0. Moreover, assume there exists u such that μ* + n in C (Sd-1) Then
Vett (x) 1 (x)
(x))
Va S
where V. is the gradient with respect to the standard Riemannian metric on @d-1
Proof of Proposition 3.9. Without loss of generality, set = -1, the other case is analogous. The residual term is given by:
R₂ = 3 div Ap div (N [xslull Vul
1 1
is sufficient to show that:
Vp P ((t+)
[xxlul]
0 as + 00.
Corollary B.2 guarantees convergence in (S(S1-1) thanks to the assumptions μ4, To improve this to higher regularity, we can use an interpolation argument through uniform bounds in Ck+2.
Define the kernel (t( Ka*
et
where Kg Iga-1 3(m.)) do Then,
xx]u[(x) = V(Wa * (x)
(Wo (I)
By the product rule, for every 0 S j S l: there exists a polynomial Pj such that:
D
(W
1))
Vp
Pb [ pllcs1,, |llloa,, min / -
WA *
where we used that min Wa * 2 min μ.. The only thing left is to notice that |[ plles < Callullas
though proving this on gd-1 requires some care.
Consider the case = (j > follows by induction) and fix e Tz(S*-1). Let A so(d) (a skew-symmetric matrix) satisfying Ar e, and define R(t) = etA In such way R(0)r = x and R' (0)a v. Then:
d
V (W * p)[1| = (W ()(((t)z)
dt
W
dt 84
d {
dt gg Wa((x, R(1)
d
[
(o(())=))
dt 33-1
/ V do
32

---

where we used the change of variable R(t)"y, and the invariance of the measure on the sphere.
Since [W&llz i. it follows that
V (W * ()llou(s--)) s Cllullca(ss11)-
Higher derivatives follow similarly, completing the proof.
C Proofs of the pairing phase
In this section we provide the proof of Proposition 3.1 1. Consider the interacting particle system on gd-1 described by the following ODEs corresponding to the case (Q" K V Td):
N
(t)
((a Pz. (2z))
(o(
where (z; Σin (2 and Pa. (xz) = 2j my);; is the projection on the hyper- plane orthogonal to 24. Suppose that there exists a unique pair (i. such that at initialization (ig, ,z)) = m (2;, TT) and denote IT))),,,()) = d,. Define also m, == max xzyli of j and {i.j} * {i.1)}.
Let a == arccos(mm) (rocos(do) > and consider the time rescaling given by the inverse of dr = (l1-d)) dl, that we will still denote by L. Then:
e.(1-d,)
a.(()
Σ
((aray)) Pa: (I,))
Za(a;)
As usual the constant C can change from line to line, but does not depend on a or
Lemma C.1. If 3 is such that Ce Ce-0(1 -a(1-oo((a/44 S a, then:
at i
d2 m 2 1 - cos
on [0, T'.
Proof. We proceed by a standard continuation argument. At we have do - mmo 2 1 - cos((a) 2
1 - cos(a/4)) Suppose the thesis holds on [0, (]. Then we have:
if # i and j #
dr arccos((22 x;))
(t),,(t),z;(1))
V1 (t) x (t))
((1
Manna (Ps, (Ix),
V1 )2 Zz(0x)
(11-))
3(0,,20) (aaizi((Ppp (xx), xi)
V1 x,)
Zz(z,)
e0 4e)
(an (Pa, (Tx);
P2
(x(
|
((1-)))
Pai
((0 (Pr, (Fk)
Zz(1g)
1
2 (mm
2 Ce Ce- -0(1-ccs((a/4)
i = i and j #
33

---

1
Or (occos((2i,,0)
2 (zz(()) xz(())
vI
(0t, xz}
e d+)
N
O(g (Pa. (2x),
e
((1--d))
2(0).20) (Paz ( 2x)
VI (2t, I,)²
(3((r))
N (ll-d))
((222))) (Pzz (1x))
Zz(x4)
|
k=1
(11-))
N
Pa.
( (Pan (Ps,
Zz(x,)
|Pa,5a
k=l
eB
Zz(wn)
(11-))
N
P2.
( (Paz
(z(
|P2.ag]
k#1
e0
Pa. 12 (l3 (Pan (xx);
Zz(
PP Ti
Pa. az Poggy
by Cauchy-Schwarz inequality, (Po.71)
S |a. hence:
6 2
e0 Zz(x,)
PP
Zz(2z)
312 de)
N
Poç
Morgoaa (Pag (3:k) |Pr,0r]
e
(3(
k#1 N
(1-di))
Pa, 3/0 2(0 (Pa, (Ik);
Zz(
Pa, xig
2 02 arccos (aocoo(((2),))) x3)) - (mm)
And in the last line we used that:
e
e | + (m)
Üj arccos x,))
|
Zz(
Zz(
In both cases the following holds:
a arccos( arccos((s;, 2 0 (1;, Ce- ((00Cccco)))/4)
hence:
arc - arccos(mp) 2 (rccos(d)) - -drccos(do) - -Ce-B(1-con(a/4)))
that implics
arccos(m:)) arccos(de) > arccos(mo) - arccos(do) Ce ((1o-o(o(( > a
2
We can conclude:
de - m4 2 - cos(a/2) > - cos(a/4).
This is sufficient to close the continuation argument.
34

---

Remark C.2. We used the fact that, if (rrcos(()) arccos(y) 2 a, then r 2 1 - cos(a).
Now, recall Proposition 3.11:
Proposition C.3. The solutions (t) of the ODE system (SA), under Assumptions 4 and positive V,
with the rescaled time dt (1 ((l1-(a==)) ds, converge 3 to the solutions of the system:
(y)
if i
(1-(1)
Pod (1h)
if = j: otherwise,
0
(0)
(0)
on finite intervals [O, T.l. with T. such that (yis My) < throughout the interval, for any
First we need the following lemma:
Lemma C.4. B is large enough then 6z = (V1.V)) V2) S 1 1-c 1-coom on [0, T) implies de = (11 S $1 c/22 on [0.
Proof. We proceed again using a continuation argument. The derivatives of the differences are bounded by:
(5t(() d(t)) = -vi (vt. (11))
N (
-2 wy
((a (Pa (Tx))
(1(
k=
((1-
3/0 0(02 (Pas (1)))
Zz(x2))
e²
2 2 2 $
Ps (xy)) + Ce 9(1 ((1-coo(aa/4)
e
+ Ede
e
Pu 2 2
Pa, (xgg
2 + Ce 961 ((1-cco(///4)
e
Edd
21))(1) d112 + 2
Bd. 1 + Ce
(01-coo((a)))
e
< 201))(() - de(() + 2e -B/2 + Ce - 601 -0(1-cos(a/4])
where we used 2 = = 1 22 and the previous lemma. The conclusion is again an application of Gronwall's lemma.
Proof of Proposition 3.11. Thanks to the previous lemma for 33 large enough, on [0, we have d: cz. Now we can proceed with the proof:
Consider the case le = i.
(t)(t) vc(1)] / (ig(() (g(s)(ds= f(' (19(1-d)) N
Blan Pe. Paonneas Pu (v)) ds
< |
Pa Pag( (r.) - Py. (vy) (y))dds |d + Ce -0 (3(1)cos((a/4)
(x));))
/ [0i Vil + |zj U ds + Ce 8(1 8(1-co((a/4)))
<
35

---

where we used the previous lemma and the fact that el propertyon T. The case + i.j similar. Hence:
2 1 - 2
thanks to the
(xz)
(x((t) (t)
Σ mx(() (x((t) Ce -C
The conclusion is then just an application of Gronwall's lemma.
D Useful lemmas
Lemma D.1. Let > 0, 3 + co. Then:
J (1 (x,, v)) elash dy C e
Proof. In the following Ca, is a constant that depends just on the dimension and on k and could change each line:
1omn (1 (x, e
((a,)) dy = Ca.
[ (1 - () +
Cd, / (2 2u) (2u-1 u)) 222 du
CB S(¹' (1 u) virー
= = 23u
dd-1 1,
Ca.Ee SM
where M is the Kummer's confluent hypergeometric function and its asymptotic behavior for B (see [1]) is given by:
IN CaRD 111
Lemma D.2. Suppose that T is a tensor such that:
T invariant under permutations of the indices.
T invariant under rotations that fix a unit vector
Then i is a 2-tensor, then it must be of the form:
T a(x @ x) + BId.
I a 3-tensor, then it must be of the form:
T a(x 2 & 3Sym(( &
Proof. For simplicity, suppose d 2 5. Without loss of generality with the case of the 2-tensor. For every + j. we can consider another index 1 [1,i,] and the rotation R such that Re; -e, (wlog * 1, otherwise use j), Re = - and elsewhere is the identity. Then:
assume Let's start
T(e es] Re, Re,] -T(e,,,))
that implies e(e,,e,) = 0. Ifi > 1, then there exists a rotation R such that Re; = e2, Rez
and the identity elsewhere:
T(ere Rei Rei] Tlea, e2)
This conclude the proof for the 2-tensor.
36

---

For the 3-tensor: consider i,J, k such that >> and (i & {j.k} or i = =-= k) Consider another index & {1,i,j,) k} and the rotation such that Re; Cis Re; - and identity elsewhere.
Then:
Tie Cy, 64] = = T Rej, Rex] -Tlei ej, ek);
that implies T(e,, eg, ex] 0. The only cases left given by = j and If = > and k = 1, then construct the rotation R such that Re, = and Re -e, to conclude. as above, that the tensor must be of the form a(zz0jk + zjōu + xxdig) + BZGZYZK.
and their permutations.
E Supplementary figures
This experiment uses the same settings as in Figure 2 in the backward regime, but with an initial distribution given by mixture of four wrapped Gaussians on the xy-plane. The red curve shows the interaction energy of the system over time on logarithmic timescale. We highlight the three distinct timescales and the corresponding behaviors discussed the paper.
3.0
2.5
Interaction Energy
10-1
10f
Time
20
10
Alignment
Paring Phase
Figure Evolution of the dynamics with Q' S definite positive.
37


Reviews:
[{"review_id": "02mavSKtyV", "reviewer": "Reviewer_aiFk", "summary": "This paper gives a mean-field analysis of encoder-only transformer models -- it shows that if the number of tokens $N$ is large, and the inverse temperature parameter $\\beta$ scales with $N$ slowly enough, the dynamics display a multi-scale behavior -- the token empirical measure collapses onto a low-dimensional subspace at first in the \"fast\" phase, then collapses into clusters during an intermediate phase, and finally collapses sequentially into a single cluster in the final \"slow\" phase. The authors rigorously show the limiting dynamics in each of these three phases, and validate them empirically.", "questions": "- What does this work tell us about the behavior of modern language models? Is there any immediate practical application?\n- Are there potential applications that may arise from further developing this area?\n- Why should the LLM community pay attention to mean-field analyses more broadly?\n- What are the technical challenges involved in extending the analysis to handle differing Q, K, V matrices in different layers?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- Easy to follow paper with a clear set of theoretical and empirical results\n- Gaining any theoretical understanding of transformers/deep learning has historically proved challenging, so it's impressive that such an analysis sheds light on some limiting cases\n- Compared to previous works which either fail to analyze the intermediate phase characterized by a partial clustering, or do analyze but only in the case where the Q, K and V matrices are multiples of identity, this work analyzes this phase for more general Q, K, V matrices. \n- It also handles clustering for hardmax transformers\n\nWeaknesses:\n- The analysis seems restricted to the case when all the Q, K, V matrices are the same for the different layers. Furthermore, there are additional assumptions on these matrices. It is not clear how reasonable these assumptions are\n- Overall, it's unclear what this analysis really tells us about modern language models and their behavior. No explanation/motivation is provided\n- The empirical results again don't seem to really shed any light on modern language models\n\nOverall, while I am skeptical of the practical applicability of the results, it does build upon previous works in the area in a meaningful way, by extending the analysis of the intermediate regime for general Q, K, V matrices. I am not an expert in the area of mean-field analysis.", "was_revised": false, "final_justification": "", "rebuttal": "- **[Q1, Q2, Q3]**: We thank the reviewer for the feedback and the opportunity to discuss the potential implications of our work. \n    We agree with the reviewer that some aspects of the original architecture have been simplified in our analysis, given its technically challenging nature, and that, as a consequence, it is not clear that our results apply directly to real world architectures. Still, since a complete theoretical analysis of transformer models remains an open problem, we see our contribution as a step in this direction, that we aim to improve in future works considering increasingly realistic assumptions.\n    While the path from this  theoretical analysis to direct application is not immediate, we recall that analogous\n    clustering behavior in deep transformer networks have been shown to emerge in real-world models [1, 2], and we  believe our findings open several possibilities for future investigation:\n    - **Model interpretability**: our analysis of the alignment phase presents a concrete mathematical mechanism for how token representations might evolve through the depth of transformer networks, predicting their collapse onto subspaces determined by the spectral properties of the $Q, K, V$ matrices, along the same line as what was pointed out by previous works [1, 2, 3]. We believe this finding could open an interesting avenue for model interpretability, allowing researchers to explore the semantic meaning of these learned subspaces.\n    - **A link to diffusion models**: the emergence of the backward heat equation in our analysis reveals an interesting mathematical link between the dynamics of transformers and denoising diffusion models. While the practical implications require further investigation, this connection bridges two of the most important architectures in modern machine learning and may inspire new theoretical insights or hybrid model designs.\n    - **Parameter scaling**: while carried out in an idealized regime, our analysis clearly identifies, from a mathematical point of view, the effect of choosing model hyperparameters at a certain scale on the representations developed by transformer architectures in the given regime. This information, combined and contrasted with analogous studies considering complementary scaling regimes, can inform practitioners on the effect of various choices of such hyperparameters (e.g., scale of inverse temperature, depth, residual connections) on model representations they aim to observe/avoid (analogously to the effect of various choices of parameter scales at initialization for wide feedforward neural networks, resulting in either the NTK or feature learning regimes during the training process [4, 5, 6, 7, 8]). \n\n- **[Q4]**: The main difficulties in extending our analysis to handle matrices that vary from layer to layer are analogous to those encountered when establishing the continuum limit of ResNets as Neural ODEs (as discussed, for instance, in [9]).\nAssuming the matrices $Q_t, K_t, V_t$ vary sufficiently smoothly with depth $t$, one could expect that the alignment dynamics would still be described by a linear transport equation, but with a time-dependent velocity field:\n$$\n\\partial_t \\mu = - \\mathrm{div}\\left(\\mu P_x \\frac{V_t K^T_t Q_t x}{|K_t^TQ_t x|}\\right).\n$$\nOn the other hand, when the matrices vary rapidly with depth, as is the case at i.i.d. initialization, in the infinite-depth limit considered in our paper we expect an averaging phenomenon to emerge, resulting in dynamics that can be approximated by a system with constant $Q$, $K$, and $V$ matrices.\nHowever, making these arguments rigorous would require a careful analysis to establish the uniform bounds on the discrete dynamics that are necessary for convergence to this continuous-time limit and is left for future research. \nIn any case, we wish to emphasize that the use of shared parameters across layers is a design choice adopted in several models, most notably in [10]. \n\n\nOverall, we have improved the conclusion by better highlighting both the limitations and the implications of our work, based on the points discussed above:\n- Line 334: \"While the path from this theoretical analysis to direct application is not immediate, we believe our work opens several potential avenues for future investigation. The characterization of the alignment phase, for instance, offers a potential mechanism for interpreting how token representations evolve into learned subspaces. Furthermore, the emergence of the backward heat equation suggests an interesting mathematical link to diffusion models, which may bridge two of the most important architectures in modern machine learning. Finally, by focusing on the \"moderate interaction regime\", we hope our analysis provides a theoretical foundation for a more principled understanding of parameter scaling, particularly as models are adapted for longer contexts.\"\n\nWe thank again the reviewer for their comments and are happy to continue discussing during the next phase of the review process.\n\n## Bibliography\n\n[1] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of\nclusters in self-attention dynamics. \n\n[2] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical\nperspective on transformers. \n\n[3] Martin Burger, Samira Kabri, Yury Korolev, Tim Roith, and Lukas Weigand. Analysis of\nmean-field models arising from self-attention dynamics in transformer architectures with layer\nnormalization. \n\n[4] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers\nneural networks.\n\n[5] Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-\nparameterized models using optimal transport. \n\n[6] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of\ntwo-layer neural networks. \n\n[7] Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks:\nAn interacting particle system approach. \n\n[8] Jacot, Arthur, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks.\n\n[9] Sander, Michael, Pierre Ablin, and Gabriel Peyré. Do residual neural networks discretize neural ordinary differential equations?.\n\n[10] Lan, Zhenzhong, et al. Albert: A lite bert for self-supervised learning of language representations."}, {"review_id": "OmY8fjl54z", "reviewer": "Reviewer_dHkY", "summary": "This paper analyzes the evolution of transformers inference time representations in in the mean field interaction regime, where the number of tokens ($N$) and softmax temperature ($\\beta$) scale together. The authors identify three distinct phases that emerge depending on the order of magnitude of the chosen timescales, as a function of $\\beta$, with the characterization of the \"alignment phase\" on a fast timescale being the main novelty. Additionally, the work serves to combine it, under simpler assumptions, with the intermediate \"heat phase\" and an exponentially long \"pairing phase,\" into a single dynamical picture. Finally, the work is supported by numerical simulations.", "questions": "- With respect to the heat phase, I understand that $\\gamma$ controls the forward vs backward diffusion. The analysis of the Pairing Phase proceeds under the assumption of a positive V, which corresponds to the clustering (backward diffusion) regime. Could the authors comment on the long-term behavior in the forward diffusion regime ($\\gamma$<0)? How does one determine which regime the model evolves through in the heat phase (in other words, why is gamma assumed to be positive)?\n- Akin to [1], is possible to have simulations with actual models to better justify the proposed system and its assumptions?\n\n\nReferences\n- [1] The emergence of clusters in self-attention dynamics. Geshkovski et al., 2024.", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "I will preface my review by saying that I am likely not the intended audience for this work, due to limited familiarity with related literature or proof techniques. Therefore, I structure my review largely from the point of view of relevance and accessibility for the general audience.\n\nStrengths:\n- The proposed contributions of the paper are listed clearly and appear novel, and the paper seems well placed within the context of related work.\n- The analysis appears comprehensive and presents interesting insights into the evolution of transformer representations\n\nWeaknesses:\n- My biggest qualm is that this paper often assumes a very high level of notational and technical familiarity, and pursuit of rigor sometimes leads to sacrificing the high level perspective. I list a few instances below:\n    - It is unclear what the inverse sampling temperature (which generally only relates to the logit layer softmax in popular context) or timescales are denoting in the introduction, until one reads through section 2.\n    - While the introduction to section 3 does a decent job of delineating the phases at a high level, there is no attempt to tie them back to any notion of practically observed behaviors in any of the exposition in section 3. It would significantly strengthen the paper if each subsection briefly motivated the analysis by connecting it to a hypothetical but grounded representation learning behavior/goal.\n    - For the latter criticism, the exposition around the technical result often focuses on technicalities and does not offer enough high level intuition on the result itself or why it would be important. For instance, the proof sketch of Theorem 3.2 immediately references concepts like the 'Ascoli-Arzelà theorem' and 'cumulants of the Von Mises-Fisher distribution'. While technically precise, this gives the non-specialist reader little intuition about why the system should behave this way. (also see remarks proceeding proposition 3.4, or discourse around proposition 3.9).\n- I understand that the model/assumptions would be considered general (or even weaker) in the context of related work. But from a general standpoint, the model still appear fairly unrealistic and it would be helpful to briefly discuss what is lost by these simplifications. For instance, what complexities are lost in Remark 2.1 or Assumption 4?\n\nOverall: Despite my critiques on the exposition, I want to stress that I'm positive about the work's technical contributions. My main suggestion is to make the paper more self-contained to better reach the general conference audience. For example, along with the suggestions above, perhaps adding some background on key concepts in the appendix could significantly improve accessibility and let non-familiar readers appreciate the novelty of this work. My score is based on these presentational concerns, and I am willing to champion this paper if these are addressed in the rebuttal.", "was_revised": true, "final_justification": "Following discussion with the authors and their efforts towards addressing my concerns about the presentation, as well as reading through other reviews and the discussion, I am confident in recommending this paper for acceptance, increasing my score 4->5", "rebuttal": "We thank the reviewer for the valuable suggestions aimed at improving the presentation of our work to make it more accessible to the general audience of the conference. \n- **Weaknesses**: To address their punctual remarks, we applied the following modifications to an updated version of the paper:\n    - Line 18: \"...choosing the model’s inverse temperature $\\beta$ (a parameter that scales the query-key dot products in the self-attention layer) as a function of N [31, 36].\"\n    - Line 168: \"Under these conditions, we show that the limiting dynamics in this regime coincides with the formal Laplace approximation of term (I) in (4), i.e., the integrals in the definition of the vector fields can be replaced by the value of the integrand at the maximum point $x' = K^T Q x/|K^T Q x|$, leading to the significantly simplified expression (5) below.\"\n    - Line 171: \"To establish the result we must prove well-posedness of the family of equations leading to the desired limit and obtain sufficient regularity uniformly in $\\beta$ to ensure that the formal simplifications from (4) (I) to (5) are allowed. This is particularly important as the derivatives of the kernel tend to infinity in the limit $\\beta \\to \\infty$.\"\n    - Related works: we added references to the literature on rank collapse and oversmoothing to connect the geometric phenomena described in the paper to more applied consequences.\n    - Remark 2.1: the MLP would act as a drift term in the dynamics, whose consequence should still be investigated further. We expect different dynamical behavior depending on the relative scale of the MLP coefficients and the attention part.\n    - Conclusions and limitations: we have expanded in the paper both on the potential practical implications and limitations of the paper, summarizing the points in the answer to Reviewer aiFK.\n    - Appendix: we included paragraphs preceding each technical result and lemma providing a high level explanation of the result, its role in our proof and referring the reader, when necessary, to complementary external literature on the topic.\n    - Regarding the timescales, they can be viewed as a conceptual framework illustrating how transformers progressively build understanding.  From this perspective, the process proceeds in three distinct parts: an initial, fast Alignment Phase establishes a global context by collapsing tokens into a shared feature space. This is followed by an intermediate Heat Phase, where these representations are refined into discrete semantic clusters. The final, slow Pairing Phase then constructs a conceptual hierarchy by merging these clusters to form higher-order abstractions.\nWhile we believe this interpretation is a suggestive lens through which to view our results, we are cautious about including such speculation in the main body of the paper. Our primary focus is the rigorous, provable mathematical treatment of the model's dynamics, and we wish to focus on rigorous contributions rather than providing heuristics that can't be substantiated by mathematical proofs.\n\nWe are happy to further exchange with the reviewer during the discussion phase if they believe there is further space for improvement. \n\n- **[Q1]**: In the forward regime, the heat phase proceeds in the opposite direction, spreading the particles far apart and reaching the unique steady state of the forward heat equation, which is the uniform measure on $E_{max}$. This is an equilibrium of the model, so there is no pairing phase in this case, and our theoretical analysis concludes there (justifying our assumption for $\\gamma > 0$ in the pairing phase). What one could expect in numerical simulations with a large but finite $\\beta$ is that the particles spread around $E_{max}$, and then because of finite-$\\beta$ effects or the possible instability of $E_{max}$ beyond this timescale, new phenomena that are not described by our analysis may occur (see also answer to Q1, Q2 of reviewer uk1K). \n- **[Q2]**: In [22] ([1] in the review), the authors provide numerical simulations with ALBERT.  While we stress that our model is designed to isolate the core self-attention mechanism and, as the reviewer rightly points out, does not yet incorporate other architectural features of models like ALBERT, \n we can still draw a parallel between the phenomena in ALBERT and our description of the dynamics by running a simulation where we increase the inverse temperature of the model tested in [22] to bring it in the regime studied in our paper. Since the clustering tendency is studied in [21], our simulation aims at establishing the presence of the alignment phase in the above model.\nWhile we cannot upload the results of the simulation because of this year's review policy, they show a rapid collapse of the effective dimension of the tokens in the first layers of the model. More specifically, we added to the appendix a graph depicting a sudden collapse of the effective dimension of the token empirical measure as a function of layer depth in the first layers of the model (the alignment phase) followed by a slower decrease in the heat phase."}, {"review_id": "LWMRjOfV0d", "reviewer": "Reviewer_7zc7", "summary": "The authors study signal propagation in Transformers, which are now well-accepted as interacting particle systems on the unit sphere (due to layer normalisation). They consider the continuous time perspective, in which time plays the role of layers, and study the appearance of multiple time scales of clustering, in which tokens go from context length equal to n at initial time (layer 0), to several clusters (k<n, say), to a single cluster in infinite time, over multiple time scales. This behavior has already been observed empirically in several previous works, and this paper is one of the first which completely elucidates the transition from n to k points, along with the infinite time collapse to a single point, in a particular scaling of the temperature parameter ß (a magnitude of the key-query matrices) and the context length.", "questions": "I have few questions and remarks for the authors which I’d like to see clarified.\n\n1. The pairing phase, in which ß goes to +\\infty, is very remeniscent of the results presented in Section 6 of the paper by Geshkovski, Koubbi, Polyanskiy and Rigollet. In fact, the time re-scaling is almost identical, and their result describes the full collision dynamics over all time (albeit, with a slightly modified model). The authors should discuss this more and provide more comparisons.\n\n2. An example after Proposition 3.9 of initial data is provided. Perhaps the authors can comment on the fact that any positive Gevrey-1/2 initial datum also yields the desired assumption (I think), because the equation satisfies a maximum principle. If I am mistaken, I welcome the authors to correct me.\n\n3. In the « heat phase », the authors use the scaling to get a limit to the heat equation, for which attention dynamics are a particle method of. The authors also very nicely comment on the fact that unnonmalized attention dynamics are a particle method of the porous medium equation. Do the results presented in this paper also transfer to unnormalised attention dynamics? I wonder if in the regime in which Q^T K \\neq I_d, the partition function is not just a scaling factor, but plays a more important role (some kind of homogenization?)\n\n4. Regarding the presentation: \n\n- Perhaps a simple figure in which the time variable is considered in the horizontal axis, and the different results/phases are superposed, in page 4, would further help the detection of the different phases.\n\n- There are some typos here and there: \\R\\geq0 in (2), ..", "limitations": "yes", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "This is one of the best mathematical papers I have seen in NeurIPS over the past 5 to 10 years. The mathematics are deep and different from almost anything used in the machine learning theory literature (deep insights from partial differential equations), and the model is neither trivial nor over-simplified. The results are truly excellent both from an applied and theoretical point of view. I would be surprised if this is not the best theoretical submission of the year.", "was_revised": true, "final_justification": "The authors have answered all of my questions appropriately. All issues are resolved.", "rebuttal": "We thank the reviewer for their supportive comments and their appreciation of our work. We address their questions below:\n\n- **[Q1]**: We fully agree with the reviewer that the dynamics of the pairing phase are closely related to the results of Geshkovski, Koubbi, Polyanskiy, and Rigollet [20].\nIndeed, the work in [20] offers a detailed and rigorous treatment of the long-timescale collision dynamics, and we have updated the manuscript with a more thorough discussion of the connection, clarifying the parallels in the analysis and directing interested readers to that reference for an in-depth study of this regime. While [20] was cited in our original submission, we now expand on the previously incomplete discussion comparing our results with those in [20]:\n    - Line 274: \"We note that\nthis hardmax particle interaction, as well as the timescale where it arises in the large $\\beta$ limit, was\nintroduced in [20, Section 6] in the case $d = 2$. We present an analogous result here in arbitrary\ndimension, without claiming originality, to provide a complete dynamical picture across phases.\"\n    - Line 279: \"Note that this result only holds up to an\narbitrary moment before the first collapse. We refer to [20, Section 6] for a detailed explanation of\nthe challenges to bypass this limit and a proof of an analogous result until and beyond the collapse time in a related but simplified model. The above proposition is proven for completeness in Appendix C.\"\n\n- **[Q2]**: The reviewer’s observation is indeed correct. Positive initial data in the Gevrey-1/2 class satisfies the required assumptions for both the forward and backward problems: globally-in-time in the forward case, due to the maximum principle, and locally-in-time in the backward direction, thanks to Hadamard-type stability results in the Gevrey-1/2 class (applicable to the backward heat equation). \nWe will explicitly mention this fact in the revised version:\n    - Line 250: “A more general class in which local existence and well-posedness hold in both the forward and backward directions is the set of positive, Gevrey-1/2 functions.”\n\n- **[Q3]**: While a similar multi-scale approach could be applied to the unnormalized self-attention model, leading to a different set of limiting dynamics as the reviewer rightfully predicts, we expect that this extension would entail a slightly different set of mathematical challenges:\n    - The alignment phase would no longer be described by a linear transport PDE. Instead, the limit would formally be a non-linear transport equation, roughly $\\partial_t \\mu = - \\text{div}(\\mu V[\\mu])$, where the dependence on $\\mu$ of the vector field $V\\[\\mu\\](x)=\\mu(\\frac{K^TQx}{|K^TQx|})P_xV(\\frac{K^TQx}{|K^TQx|})$ arises from the unnormalized attention. Establishing a collapse result (similar to Proposition 3.4) for this equation has proven to be more challenging than in the normalized model because of the non-local structure of the limiting PDE.\n    - The heat phase would lead to a porous medium equation (PME) with exponent 2. Interestingly, proving convergence to the PME is more straightforward than in our case, thanks to the Wasserstein gradient flow structure of the USA model proved in  Geshkovski, Letrouit, Polyanskiy, and Rigollet [21].\n\nFrom this comparison, one could conclude that in this regime, the normalization factor in self-attention acts as a \"linearization\" of the flow from a PDE perspective. It transforms the non-linear transport equation into a linear one and the PME into the heat equation. Because of this,  we indeed expect a more complex interaction structure to arise in the limiting model in the unnormalized case.\n\n- **[Q4]**:  We thank the reviewer for these helpful suggestions. A timeline figure illustrating the different dynamical phases would greatly improve the paper's readability.\nWe have plotted the evolution of the interaction energy $\\sum_{i,j} \\exp(\\beta \\langle x_i, x_j\\rangle)$ for an experiment with initial condition given by a superposition of four gaussian distributions and $Q^tK=V=diag(1,1,0.1)$, corresponding to the backward regime, as a function of time on a logarithmic scale, where the three phases identified in the paper appear as different rates of increase of such energy.  We also appreciate the reviewer for pointing out the typos. We will correct them throughout the manuscript.\n\nWe stand by for further discussion on the interesting questions asked by the reviewer in the next phase of the reviewing process."}, {"review_id": "SYN3KfjS4w", "reviewer": "Reviewer_uk1K", "summary": "This paper presents a theoretical analysis of the token dynamics in a continuous-depth, mean-field transformer model. The authors study the \"moderate interaction regime\", where the parameters $N, \\beta$ tend to infinity together. The central contribution is the identification of a multi-scale dynamical structure of the system, composed of three distinct phases:\n\n1. Alignment Phase ($O(1)$ timescale): A novel finding where the tokens rapidly align with a low-dimensional subspace ($E_{\\max}$) determined by the leading eigenspace of the matrix $V K^{\\top}Q$. The dynamics are shown to converge to a specific linear transport equation, that is significantly simpler than the original equation. The assumptions in this part are very weak, noticeably so compared to other results in the field. The statement itself is a groundbreaking result for the field.\n2. Heat Phase ($O(\\beta)$ timescale): On a slower timescale and under more restrictive assumptions, the dynamics evolution within $E_{\\max}$ are approximated by a forward or backward heat equation, leading to either smoothing or clustering of tokens.\n3. Pairing Phase ($O(e^{c \\beta}$ timescale): On an exponentially long timescale, pre-formed clusters are shown to merge sequentially, with the closest pair collapsing along a geodesic.\n\nAuthors provide rigorous proofs for stated results, and the statements combined together contextualize the dynamics into a unified and clear, albeit idealized, picture of token evolution. It is supported by some numerical experiments, that are aligned with provided picture.", "questions": "**1. Transition from the Alignment to the Heat Phase.**\n    The analysis of the Heat Phase assumes that the system's state is fully supported on $E_{\\max}$. However, the convergence is asymptotic, suggesting that at finite time, residual parts of the measure may exist outside $E_{\\max}$. Could you please elaborate if this impact the dynamics significantly?\n\n**2. Numerical experiments**\n    The pictures present justify the complete picture. Could you please check other variations of matrices Q, K, V to see if that is still true? For instance, take $Q^\\top K = -Id$ in Figure 1a, or otherwise misaligned $V$ and $Q^\\top K$ in Figure 2. I believe the initial phase is still intact, while others are not entered under such assumptions.\n\n **3. Conditions on $\\beta** \n    It is hard to track what is assumed of $\\beta$ exactly. For example, in referenced SSmax approach $\\beta \\sim \\log n$. Could you share if you expect the alignment phase to still hold for say larger values of $\\beta$ (as far as I can see, it is assumed to grow very slowly in the paper)?", "limitations": "Yes, authors addressed the limitations of the paper. Assumptions in their statements are clear and the proofs are well-written.", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n1. The paper's most important result is the rigorous identification and characterization of the initial \"alignment phase\", supported by a complete proof of the statement under rather weak assumptions (invertible Q, K, V). The statement is interesting both theoretically, as an important approximation that simplifies the ODE analysis, and practically, as a strong evidence that early layers of a transformer might perform dimension reduction of the system to an explicitly given eigenspace. This observation alone gives rise to many questions: to check what said space might be in real transformers, whether it is shared by different layers, and others. In theoretical research, it simplifies the analysis for the system with general parameters. This makes it an important result that deserves recognition to foster further research.\nThe proof is complex and technically impressive. Because of weak assumptions, it might be a more general phenomenon (that holds under weaker assumptions on N, beta), making it a compelling explanation for rapid dimensionality reduction in transformers.\n2. The alignment phase, together with Laplace approximation of the dynamics, is an original observation about the system with important consequences. It is clearly conveyed in the paper and can be clearly tracked, providing clarity of exposition.\n3. The overall paper layout provides interesting picture of the tokens geometry evolution on various time scales. Second and third phases are important for overall understanding of the dynamics. Their analysis is not completely novel, as it in many parts coincides with established results, but it provides clarity and completeness to the paper. Overall presentation is well-written and gives a clear picture about the system.\n4. The paper is of high quality and well-written. It presents novel and significant results, while being easy to read. Its clarity is supported by good organization and experimental pictures. \n\n**Weaknesses:**\n1. The phase transitions between Alignment, Heat, and Pairing phases are idealized. The unified narrative relies on simplifying assumptions at the transition between phases. Most notably, the analysis of the Heat phase assumes the system's state is already fully confined to the subspace $E_{\\max}$. This overlooks the potential impact of residual components outside $E_{\\max}$ after the alignment phase, which could alter the next slower dynamics.\n2. Sequentially presented stages Alignment -> Heat -> Pairing might lead readers to believe that this is a complete description of the dynamics. In particular, this would imply that the final configuration is present in $E_{\\max}$. This does not actually follow from the paper results. In addition, presented numerical simulations seem to agree with that conclusion, but they rely on some specific properties of $Q, K, V$. For example, taking $Q^\\top K = - Id$ in the Figure 1a experiment (with the same $V$ and other parameters) leads to a similar initial picture (alignment phase) and a different final token position (because alignment does not happen completely and tokens continue to be influenced by components outside $E_{\\max}$). In general, importance of $E_{\\max}$ is clear for the initial stage, but it does not automatically transfer to later stages.", "was_revised": true, "final_justification": "Great contribution! The paper rigorously proves a complex and interesting phenomena.", "rebuttal": "We thank the reviewer for their thorough review and for their positive assessment of our work. We discuss their questions below:\n\n- **[Q1]**: We agree with the reviewer about the importance of rigorously connecting the phases presented in the paper:\nwhile our analysis suggests that in certain stable regimes the full three-phase picture holds (for example in the backward regime, when $Q^T K = S = V$ for a positive definite matrix $S$, from our simulations and from the existence of an energy function and the fact that Dirac deltas on $E_{\\max}$ are stable equilibria [9]), and that the three phases are  achieved sequentially, a rigorous unification of the different dynamical phases for general parameter choices is indeed a challenging open problem, central to which is the stability of the subspace $E_{\\max}$.  Due to space constraints, we could only briefly touch upon this in Remark 3.5 and our conclusions, but we now expand in section 5 on the need to identify under which conditions the transition from alignment to heat phase observed in our experiments can be made rigorous, increasing the transparency of the paper. \n    -  Line 334: A further direction of future research consists in providing sufficient conditions for the stability of the space $E_{max}$ emerging in the alignment phase under the prelimit model (i.e., for large but finite $\\beta$), thereby justifying assumption 3 and, ultimately, connecting in a rigorous way the alignment and heat phases identified in this paper.\n- **[Q2]**: We thank the reviewer for suggesting an interesting choice of parameters ($Q^TK=-Id$ and same parameters of Figure 1a).\nWe have run additional simulations for this case. As the reviewer anticipated, the alignment phase is still achieved, with the tokens collapsing to the predicted y-axis. However, after this initial phase, finite-$\\beta$ effects, which depend on the order of the $\\beta, t \\to \\infty$ limits, cause the two clusters to merge along the z-axis. While the example provided by the reviewer does not satisfy Assumption 4, and therefore we make no claims in our paper on such a choice of parameters, we agree that, connecting to the previous question, exclusively based on the results of the paper we cannot rule out situations where the long-term dynamics of the system differs from our results because, e.g., of residual parts of the measure outside $E_{\\max}$ and the stability of this subspace.\nWe stress again that investigating this phenomenon across different parameter regimes is a crucial and exciting future research direction. We believe this is key to providing a more complete picture of the dynamics, including the unstable cases, and to extend the picture presented in our paper to more general and realistic choices of parameters.\nOther numerical experiments also highlight other types of finite-$\\beta$ effects. For instance, in the forward diffusion case, the alignment and heat phases are still visible, but their separation is less sharp. We also note that in the backward case, if the initial measure collapses to a uniform distribution on the $xy$-plane (an unstable equilibrium), the heat phase effectively disappears, and the system proceeds to the clustering behavior described in [7].\n\n- **[Q3]**: The exact maximum scaling rate of $\\beta$ with respect to $N$ depends on several factors, including the spectral properties of the matrices $Q, K, V$ and the initial condition. We opted not to provide a specific rate in the paper because the technique used in our proof, particularly in Section 3.1, although simplifying certain mathematical steps, yields to a rate that we believe to be suboptimal. A naive calculation based on our approach gives a very slow growth of order $O(\\log(\\log(N)))$ and numerical experiments would require an extremely large amount of GPU memory. However, we have strong reasons to believe there is significant room for improvement. For comparison, in the related problem of convergence of a model analogous to the noisy unnormalized self-attention to the porous medium equation in $\\mathbb{R}^d$, Oelschläger [32] established convergence for a much faster scaling of $\\beta = O(N^{\\alpha/d})$ for $\\alpha \\in (0,1)$. We conjecture that a more refined analysis could establish a similarly improved rate for our setting.\n\nWe again thank the reviewer for their positive assessment and their insightful comments. We are happy to continue discussing in the next phase of the reviewing process."}]

Author rebuttal:

Meta-review: My reading of the paper confirms the reviewers' high scores. This is an excellent paper on a topic of theory of transformers, where contributions are rarely clean and simple to understand. This one discovers a very simple and clean idea of deriving a high-beta versions of the Transformers PDE, which in many ways are more tractable for analysis than the original PDEs of Geshkovski et al. I thus recommend this paper for acceptance and congratulate the authors on this excellent result.
Decision: Accept (oral)