The emergence of clusters in self-attention dynamics
Authors: Borjan Geshkovski
Cyril Letrouit
Yury Polyanskiy
Philippe Rigollet
Venue: NeurIPS 2023 poster (2023)
Keywords: Transformers
Self-Attention
Clustering
Interacting Particle Systems
Continuous Time
TL;DR: We mathematically prove that trained Transformers cluster in long time.

Abstract: Viewing Transformers as interacting particle systems, we describe the geometry of learned representations when the weights are not time-dependent. We show that particles, representing tokens, tend to cluster toward particular limiting objects as time tends to infinity. Using techniques from dynamical systems and partial differential equations, we show that type of limiting object that emerges depends on the spectrum of the value matrix. Additionally, in the one-dimensional case we prove that the self-attention matrix converges to a low-rank Boolean matrix. The combination of these results mathematically confirms the empirical observation made by Vaswani et al. [ VSP`17 ] that leaders appear in a sequence of tokens when processed by Transformers.

The emergence of clusters in self-attention dynamics
Borjan Geshkovski MIT jorjannmit..
Cyril Letrouit MIT letrouit@mit.c. edu
Yury Polyanskiy MIT yp@mit. edu
Philippe Rigollet MIT rigollet@mit.ed
Abstract
Viewing Transformers as interacting particle systems, we describe the geometry learned representations when the weights are not time-dependent. We show that particles, representing tokens, tend to cluster toward particular limiting objects time tends to infinity. Cluster locations are determined by the initial tokens, confirming
of representations learned by Transformers. Using techniques from dynamical systems and partial differential equations, that the type of limiting object that emerges depends on the spectrum of the value matrix. Additionally, in the one-dimensional attention matrix converges to a low-rank Boolean matrix. The combination these results mathematically confirms the empirical observation made by Vaswani et al. [VSP* 17] that leaders appear in a sequence of tokens when processed by Transformers.
show
prove that the self-
Introduction
The introduction of Transformers in 2017 [VSP* marked a turning point in the AI revolution, powering breakthroughs in natural language modeling and computer vision. With remarkable empiri- cal success, Transformers enable large language models to compute very powerful representations using the self-attention mechanism. Yet, little is known about the geometric structure of these representations. As the size of these models grows at astonishing rate, the need to understand their inner workings becoming a pressing scientific challenge. this work, we make a first step in this direction by describing the geometry of learned representations.
To provide transparent presentation of our findings, we take leaf the literature on continuous- time dynamics such as neural ordinary differential equations (ODEs) [CRBD18, Wei17, HR17]. By viewing layers as a time variable, this formalism has emerged flexible mathematical framework to implement and study ResNets [HZRS1 6a] as particular discrete-ting versions of parametrized dynamics of the form
+(t) = fo(x)()),
t [0,
Here f is the trained parameter of neural network and fo is characterized by the precise architecture of the ResNet'. In turn, an input (c.g.. an image) r(0) e Rd is mapped to its representation xx (T").
Unlike neural ODEs and ResNets, the representation map of Transformers is not solely function of an individual input x0) IRd but rather of a set/sequence (O)..... Ix (0) of n > 1 dimensional tokens. These tokens then evolve in time by interacting with each other per the self- attention mechanism. Namely, following [SABP22], we view tokens as particles, and the transformer dynamics as interacting particle system of the form
ix(()) Σ (o(V)r;(1),
te [0, +0),
classical choice is = (W, A, Raxed x Rdddd Rd and fo(x) = Wo(Ax + b) where a an elementwise nonlinearity such as the ReLU ([HZRS16b]).
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

for any E [n2], where Pi)(t) are the entries of a n x stochastic matrix P(t), given by
<Q,,100,2,,0))
Pb(()= Ee e(Q2,(t)
(4,5) [m]
Here the matrices Q (Query), K (Key), and V (Value) are learned from data. Note that Q, K need not be square. The
matrix P(t) is called self- self-attention matrix. The wording attention stems precisely from the fact that Pij (t)(t) captures the attention given by token i to token j relatively to all tokens & E [n]. The matrices Q and K in (2) warp the geometry of the input tokens, so that a trained attention matrix contains weights which indicate semantic relations between words. Such conclusions have been drawn in the context of language processing tasks in [VSP+ 17, Figures 3-5].
Our goal is to showcase the fact that self-attention, which itself is the core novelty of Transformers, entails clustering effect. To that end, we focus on the pure self-attention dynamics described in (1). In particular, we do not model variations such as multiple heads, feed-forward layers, and layer normalization that are typically adjoined self-attention dynamics of (I). However, on this last point, we note that our theoretical findings indicate that without any normalization, the dynamics (1) can diverge in some (or even all) directions over time. We leave these additional questions for future research; see Section 6.
1.1 Organization of the paper and summary of contributions
The goal of this paper is to characterize clustered represen- tations of a trained Transformer by studying the asymptotic behavior of a sequence of tokens (t). Zr (t) as they evolve through the layers of a transformer architecture using the dynamies (1). In this setup, a Transformer is completely described by the weight matrices (Q, K, V) obtained during training. Note that we assume that these three matrices are time- independent. While this assumption is motivated by mathemat- ical convenience, it is worth noting that such weight-sharing scenarios are in fact used in practice- see, e.g., ALBERT (LCG [0]]=0000000-0000 220]as as they drastically reduce the number of parameters of a network.
Figure 1: For V = I3 tokens clus- ter toward the vertices of convex polytope (Theorem 3.1).
With parameters (Q, K, V) fixed, tokens are subject to collective dynamics that we call transformer dynamics. While these dynamics are reminiscent of existing models for opinion dynamics and flocking, they present they own mathematical challenges requiring ad-hoc tools to study their asymptotic behavior.
The main conclusion of our analysis is that the set of tokens {in (I), (r,...,,,,,,))))) (()). appropriately rescaled, tends clustered configuration as ( - oc. Our theoretical findings justify the empirical observation made in [VSP+ 17] that leaders appear in sequence of tokens when processed by Transformers. We now list our main contributions.
(i) As a warm-up to the geometric characterization of the limits of sequences of tokens, we show in Section 2 that when d and V > 0, the self att-nttent matrix P(t) converges low-rank matrix with entries 0 and as t - +00 thus revealing the emergence of a small number of leaders that drive the transformer dynamics. The restriction d = 1 follows from technical considerations, and some pathological phenomena may occur in higher dimensions (see Remark 5). But numerical experiments (as well past empirical work) indicate that the result may extend higher dimensions for almost all initial sequences of tokens.
(ii) In Section 3
we first focus on the case V = La natural canonical choice that enables to establish some of the main tools of the paper. We introduce a time re-scaling reminiscent of the layer normalization heuristics to alleviate the possible divergence of tokens. We show that along this scale the tokens converge to the boundary of a convex polytope. For almost all initial sequences they even converge to the vertices of the polytope, the number of which is significantly smaller than This elucidates the clustering phenomenon. (See Fig. 1.) When V IIa all tokens following the dynamics (1) collapse to 0.
(ii) We build on these results and Section 4 consider the case wherein V is only assumed to have simple and positive leading eigenvalue. This setting is much closer to reality and corresponds to

---

actual learned matrices V (see Figure 10). We show that along the particular timescale, tokens cluster toward one of most three hyperplanes which are determined by the corresponding eigenvector.
(iv) In Section 5 we complete the results of Sections and 4 by addressing the case where the leading eigenvalue has multiplicity. This results in clustering toward the vertices of convex polytope in some directions, and linear subspace in the others.
(v) We also prove the global existence and uniqueness of solutions of all dynamics considered in this work (including the mean field limit), We refer the reader to Appendix A for more details.
We also observed numerically that our conclusions extend to more compound architectures (see Conjecture 4.2, Section 6 and Appendix F).
Value
Key and Query
Limit geometry
Reference
QTK> >0
vertices of convex polytope union of 3 parallel hyperplanes
Theorem 3.1
A (V) > 0, simple
(Qpi, K 0 Q 0'>0
Theorem 4.1
paranormal
polytope x subspaces
Theorem 5.1
=
KQ
single cluster at origin*
Theorem C.5
Table I:
Summary of the clustering results of this work. *All results except for the case V -Id hold for the time-scaled dynamics (4).
Remark (Discrete time). While we focus on the idealized setting of self-attention dynamics in continuous-tine this is solely done for convenience and all of our methods are straighfforwardly applicable to the discrete-ting setting. (See also Remark 4.) The discrete-tien time analog of (1) with time-step At > 0 (equal to 1 in practice) is simply the forward Euler iteration (Qs,(ksS),Ka 2;((K + 1) (ka.tt + At Σ EE_ (eRe((()),KK((aa))) Vzz(k
e (3)
Notation. We denote by (, and || the Euclidean dot product and norm respectively, and we use the shorthand [n] {1, (1...., n). For any matrix M E Ra we order its eigenvalues (repeated according to multiplicityy) by decreasing order of modulus: (MT) > 3 [Ma(M).. We denote by ||AM]] up the operator norm of the matrix M, equal to the largest singular value of M. Given
set S Rdd we define the distance point Rdd to S dist(æ, S) == inf ses || s||. and by (Snv(S) the convex hull of S.
Related work
Our study and results build several different lines of work, and we draw parallels in what follows.
Analysis of attention based models. Given the widespread use of Transformers in natural language processing, there has been surge of interest understanding the function and significance of atten- tion layers within these models. In [YBR* YBR+ the authors show that when treated as discrete-tin time systems with additional dense layers and multiple heads appended to the core attention mechanism, Transformers exhibit the universal approximation property. In [LLH 20], the authors present, to the best of our knowledge, the first interacting particle systems perspective on Transformers. They then leverage the similarities between Transformers (with an additional feed-forward layer compared to (1)) and convection -dffftsionn equations to slightly improve the performance of Transformers by employing Strang-Marchui. Marchuk splitting scheme for time discretization. In [SABP22], the authors interpret system (1) as the characteristics of a continuity equation. Drawing on the similarities between (1) and Sinkhorn iterations, they propose novel architecture dubbed Sinkformer, which possesses the desirable property of being a Wasserstein gradient flow.
Quadratic complexity of Transformers. The major computational challenge of Transformers is their high computational complexity, particularly when processing long sequences. Transformers require quadratic time and space complexity process sequences, because each self-attention layer contains n² products of the form (Qu,, Kxz> (for i, [m]). The empirical observation that the self- attention matrix P is close to low rank matrix see [LWLQ22, Section 4.4] for references cited

---

as the inspiration behind Linformers [WLK* 20] and the fine-tuning algorithm LoRA [HysW* 22]. For both approaches, the low-rank structure is imposed rather than extracted from P itself. Other methods called sparse attention and block attention have been proposed to reduce the quadratic complexity- see [WI.K Section 2.2] for references. In the spirit of these works, a foreshadowing of the clustering mechanism was invoked in [VKF20], where queries are clustered into groups, again in view of reducing the quadratic complexity of self-attention. We point out that [DCL21] previously demonstrated that without skip connections, the dynamics trivializes and all tokens quickly lump together into a single tight cluster. Our work, in contrast, shows that in the presence of skip connections rich cluster structure emerges.
Compared to the usual BERT, ALBERT [LCG * 20] uses seaameter-shaiingg across layers, meaning that the weight matrices Q, K, V in (1)-(2) do not depend on time, as in the present paper. This does not reduce the theoretical (n complexity of the original Transformer, but, quoting [I.CG+20] "significantly reduce[s] the number of parameters for BERT without seriously hurting performance, An ALBERT configuration similar to BERT-large has 18x fewer parameters and can be trained about 1.7x faster. The parameter reduction techniques also act as a form of regularization that stabilizes the training and helps with generalization".
thus improving
Neural collapse. Our results and conclusions bear a resemblance to some geometric aspects of neural collapse for classification tasks [PHD20]. A key geometric aspect of neural collapse is the observation that, during the training of deep neural networks, the representation of different classes in the later layers of the network tends to form a tight cluster around the vertices of a simplex. The emergence of simplex structure in the representation space provides insights into how the neural network organizes and separates the different classes.
Clustering in interacting particle systems. The transformer dynamics (1) have a strong connection to the vast literature on nonlinear systems arising in the modeling of opinion dynamics and flocking phenomena. In addition to the classical Kuramoto model describing synchronization/Cluterri oscillators [Kur75, OBV+ the model which most similar (1) is the Krause model [Kra00]
of
(1x;; x2,13)
(t( Σ (t)
aij =
2k=1 00 Ix, 12
which is non-symmetric in general (aij # aj;), much like (I). When é is compactly supported, has been shown in [JM14] that the particles (t) assemble in several clusters as - + Other models of opinion dynamics and flocking have been proposed and studied, among which the Vicsek model [VCBJ* +95], the Hegselmann- Krause model [HK02] and the Cucker-Smail model [CS07]. These models may also exhibit clustering behavior under various assumptions (see [MT14, CHH 16, HKPZ19] and the references therein). The transformer dynamics are also closely related to the dynamics employed in mean-shift clustering [Che95], and this work indirectly sheds some light on its theoretical properties.
The analysis of transformer dynamics presents unique mathematical challenges that cannot be addressed using the tools developed for these more primitive models. In particular, our work demonstrates how different choices for the parameters lead to remarkably diverse clustering patterss. Much more remains to be discovered and this work is first attempt a rigorous mathematical analysis of these synthetic dynamics.
2 Asymptotic low-rankness of the self-attention matrix
As mentioned in Section 1.1, numerical experiments in [WLK+t2) show that the self-attention matrix P. defined in (2), has an almost low-rank structure. This observation has then been leveraged to reduce the quadratic complexity in the sequence length which is inherent to Transformers, resulting in nen-nggl decrease in the cost of training.
As warm-up deriving complete geometric representations of the dynamics, our first result shows, in the simple 1d case that P(t) indeed converges exponentially fast toward a matrix which is typically both Boolean and low-rank (see Fig. 3). Although there are clear obstructions to rigorous extension of this result to higher dimensions (Remark 5), numerical experiments appear to show that this result holds in greater generality, for almost all initial sequences (Appendix F).

---

To set this up, we introduce the set S of n x matrices having the form illustrated in Fig. 2, where the asterisks denote arbi- trary non-negative real numbers which add up to 1. The row of asterisks may actually be any row between the first and the last one.
0
0
P
...
Theorem 2.1 (Self-attention matrix converges to a low-rank Boolean matrix). Let = Suppose that the scalars (Q, V) satisfy V 0 and QK 0. For any initial se- quence of pairwise distinct tokens (O)....... In there exists some p* e S such that the self-attention matrix P(t) defined in (2) converges to P* as +00.
,n (0)) E R",
Figure 2:
Elements in P, where Po, E IR=x are some permutation matrices, and asterisks denote arbi- trary non-negative reals which add up to 1.
The proof may be found in Appendix The rate of convergence toward pa is in fact doubly exponential in / for coefficients outside the row of asterisks Fig. 2. The proof the theorem also reveals that for almost all initial sequences of pairwise distinct tokens, p* is actually of rank 1 or 2, i.e., the row of asterisks is equal to either ej = O.. 0) or En = 00,
The interpretation of Theorem 2.1 is that in the 1d case, at most three tokens capture the attention of all tokens except most one. Typically, these leading tokens are those carrying the largest amount of information. This is also illustrated in Fig. 4. Since the tokens z; here evolve on I, the right-most and left-most ones (which typically tend toward +:0) capture the attention of all the others.
0.0, rank- 11
ronk-
5.0, runl-
mike
Figure 3: An illustration of the asymptotics of P(t) entailed by Theorem 2.1 for 40 tokens, with Q K and V = 1. (See Appendix F for details on computing.) Increasing n has no effect on this behavior of P(t) see Fig. 11.
0.0
2.0
f 9.0
Figure 4: The clouds {K xt) hef2o) (green) and (Q0z (t) (e(2o)) (purple) for d = 2 where pairwise points of clouds are connected by line of width equal to (t)(t). Here V > 0 and Q > are random matrices and K Iz. The creation of clusters is reflected by the rank & 2 structure of the self-attention matrix P()). This interaction echoes findings illustrated in the original paper [VSP+17] for instance, Figures 3-5 therein.
3 Clustering toward vertices of convex polytopes
In the rest of the paper, we seek taxonomize various clustering results for the solutions to (4) when t +00, depending the sign and the multiplicity of the cigenvalues of V. We begin by focusing
5

---

on what may appear to be the most natural case V = Ia. as is also done in [SABP22]. In fact, we demonstrate (theoretically and numerically) later on, clustering is generic phenomenon which holds under much less restrictive assumptions.
The transformer dynamics considered in (1) does not contain layer normalization mechanism typically encountered in practice [VSP*17] 17]. In absence of such a device, tokens may diverge to infinity as in Theorem 1. In fact, the norm of the tokens xx(t) typically diverges exponentially toward +00 for any d: this is expected, by analogy with the non-trivial solutions to ((() = v(t).
To remedy this situation, take inspiration from the solution y(t) etV y(0) to j(() Vy(t).
Namely, for any E [n]] we consider the rescaled tokens 2x())= == e iV xy (t), which solve
EQ Ke s,(0)
(t)
Σ J=1 E Ee_re x,(0, =x()))
((2((t) (t))
for t e +00.
(4)
The initial condition remains the same: J.
20 for any i [n]] More importantly, the coefficients of the self-attention matrix for the rescaled tokens (t) are the same as those for the original tokens z, (t). Whence, the conclusion of Theorem 2.1 also applies to the dynamics (4). We
see this rescaling of tokens as mathematically justified surrogate for the layer normalization.
The appearance of the exponential factor within the self-attention kernel facilitates the analysis of (4) compared to (1), and it is in fact instrumental in the proofs of all results that follow. Each result on the rescaled tokens 2, (t) then gives information on the dynamics of the original tokens 2,(t) by virtue of the relation (t) etV (t).
We now able to state the main result of this section the case V Ia- The following theorem shows that the tokens z; (t) evolving per dynamics (4) converge to the boundary of a convex polytope as +00. We present here simplified but weaker version of our result for convenience, and refer the reader to Theorem C.I in the appendix for a complete statement.
Theorem 3.1 (Convergence to points on the boundary of a convex polytope). Suppose V = Ia and Q K > O. Then, for any initial sequence of tokens (2,(0) [e(n] Rdd, there exists convex polytope K c Idd such that for any i E [ni] (t) converges either to 0 or to some point on ôK
+00.
The convex polytope K is completely determined by the initial sequence of tokens, and Q K (refer to Claim 1). Numerical experiments (e.g. Fig. 5) also lead us to claim that for almost all initial sequences of tokens, one should expect convergence of 2t) (i [n]] toward some vertex of K.
(Furthermore, the number of vertices of K. is often found to be significantly smaller than n.) It may however happen that for initial sequences taken in some null set (not seen when tokens are drawn at random) some tokens converge to other points of the boundary ôK, namely in the interior of facets. On the other hand, for generic choices of initial sequences, we do not see a way to predict K. explicitly besides running the full dynamics.
0.0
t 20
5
5
Figure 5: A toy example illustrating Theorem 3.1 with n = 40 tokens in RR3 Here Q K Ig. The tokens converge to one of the vertices (leaders) of the limiting convex polytope.
Recall that the points (t) = e' (t) when = Theorem 2.1, this result also shows the emergence of a set of leaders (given by the vertices of K.)
La follow the original dynamics (1), Akin to
Note that the case V = Los may appear equally natural. For such choice of V. we show in Appendix C.2 that the dynamics converge to a single cluster located at the origin. Multiplicative constants preserving the sign, ie.. V = Icla, c trivially yield the same conclusions.

---

attracting all tokens as t grows. It has been experimentally observed (first in [VSP* 17I) that in trained Transformers, tokens focus their attention on local leaders in a way that seems reproduce the syntactic and semantic structure of sentences.
The proof of Theorem 3.1 is postponed to Appendix C, and amounts couple of effects entailed by the dynamics. First of all, the convex hull of the particles is shrinking over time (Proposition C.2). This is due to the fact that the distance of the particle nearest to any half-space (not containing the particles) increases with time. On the other hand, the convex hull ought not collapse since particles which have not concentrated near the boundary of the limiting polytope will continue increase in magnitude until they themselves reach this boundary (Step 2 in the proof). This occurs due to the üimerrescaling
Remark 2. Assuming QTK. K > 0 does not seem be essential for our conclusions: instead,
guides the direction of the proof. To emphasize the broader validity of our conclusion beyond this specific assumption, we conducted additional experiments (refer to Section F.6) which suggest that Theorem 3.1 (as well as Theorems 4.1 and 5.1 stated below) holds in more generality.
Remark (Rate convergence). Although Theorem 3.1 (as well as Theorems 4. and 5.1 stated below) does not specify rate of convergence toward θk, expect (and observe through numerics) that convergence happens very quickly-affe few layers. most tokens are already clustered. What "few layers" " means here necessarily depends the typical modulus of the initial tokens, since the dynamics (1) is not invariant under multiplication of all initial conditions by a fixed real number.
Remark 4 (Discrete time). As alluded to Remark 1, all our results extend to the discrete-time Transformers (3). Indeed, just as in the continuous- time case, there is a natural rescaled dynamics, which is the discrete analogue of (4): if we set Ia Vtt, and assume that R is invertible (which is the case for sufficiently small At), then 2, (kt R. (x(ktt) == 2 satisfies
(QRA 1K) 1RR
[k+1] 2
[k] 2,
At
R R
[K] - )
keN.
J=1
The proofs of Theorems 2, C.5, 4.1, and 5.1 carry through with straightforward modifications.
Let us provide some comments on the proof of Theorem 3.1 the discrete-time setting, for the sake of completeness. First of all, Proposition C.2 holds intuitively because for all integers E [n]] and > 1,
[k+1]
M
At
[k]
E conv
jE(m)) (
1+ At
We then define the candidate set of limit points in (37), and Claim I holds without any change in the statement or in the proof. Then, just as in Steps and 3 in the proof of C.I, we can first show that if 2 is not already near some point in the candidate limit set, it will keep moving toward the boundary of the convex polytope. Finully, we can prove that tokens cannot circulate indefinitely between different points on the boundary. The combination of these arguments would establish the convergence of each token toward some point in the set given by (37).
Clustering toward hyperplanes
While being natural example consider, value matrices found empirically much more general than V = Ia. which we considered in the previous section. We now turn our attention to significantly more general setting of value matrices, which we formalize as follows.
Definition 1. We call (Q, K. V) a good triple if the two following conditions are satisfied:
the eigenvalue of with largest modulus is real, positive, and simple; namely,
>²(V) 122(V)) > ... [aa(V))-
. <Qui, Kyp) > 0 for any 6p1 E TR lying on the line ker(V (v)
The second condition simply states that the quadratic form (Q:, K-) is positive definite along the eigenspace associated to the leading eigenvalue of V. Note also that if all entries of V are positive,
7

---

the first condition is automatically satisfied by virtue of the Berron-Frobenius theorem. In fact, this assumption is generic. On the one hand, it is satisfied by some pre- pre-trained value matrices for ALBERT (Figure 10). On the other hand, numerical experiments indicate that a constant fraction (about 14%) of matrices from the real Ginibre ensemble dimension d 128-this proportion known to vanish as d - s0, albeit very slowly [RS14].
Our clustering result in the setting of good triples can be summarized as follows: the coordinate (z) To1
22 of any token (t) along the eigenspace spanned by (pi converges, as t - + toward one among possibly 3 real scalars. Consequently, all the tokens =x(t) converge toward among most three parallel hyperplanes; see Fig. 6 for an illustration.
Theorem 4.1 (Convergence toward & 3 hyperplanes). Assume that (Q, K, V) is good triple in the sense of Definition 1. Then, for any initial sequence of tokens (z((0)) = Rd, there exist at most
three parallel hyperplanes in Re such that for any i E [mx], the distance of the solution 2+( to (4) to one of these hyperplanes converges to 0 as - +00.
0.0
1.0
5.0
15.0
10-
10-
10-
10
0-
0.
-10-
-10-
-10-
-10-
10
0 10
10
0
10
Figure 6: Illustrating Theorem 4.1 with n = 40 tokens in R.2. Here Q = K = Iz. V is random symmetric matrix with eigenvalues {1.35, -0.07, and 1 = (0.76, 0.65). The components of the tokens in the direction of (orange arrow) cluster over time. (See Figures 13-14 for examples in RR.) We also observe that tokens typically cluster toward only two hyperplanes- third one (passing through the origin) may appear for non-generic initial sequences. The hyperplanes are perpendicular to 101 since V is diagonalizable.
The proof may be found in Appendix D. The important role played by A (V) in the dynamics may be seen in (4): the component of 21(1) along şi determines the size of eV (t) in the exponent appearing in (4). The tokens 2, (!) attracting other tokens (() are those for which this component
along 1 largest in modulus. This attraction process forms the clusters. These leaders, as in all our
results, have been empirically observed to be the ones carrying the largest amount of information in the sentence (see Supplementary material in [VSP* 17J).
Furthermore, Theorem 4.1 can also be interpreted more classical machine learning terms. On the one hand, it can be sequence are clustered, based on their intrinsic similarity, to at most 3 "flats" of dimension 1. On the other hand, ensures that for good triple (Q, K, V), (4) generates linearly separable representation of tokens.
instance of K-flats clustering [BM00, Vidl1]-pont in the input
Beyond single direction?
Numerical experiments (c.g.. Fig. 7) indicate that similar phenomenon emerges for more complex V. We formulate following conjecture which a natural generalization of Theorem 4.1. Conjecture 4.2 (Codimension conjecture). Let S be the number of eigenvalues V with positive real part. Then there exist at most three parallel Euclidean subspaces of Id of codimension k such that for any i e [n], the distance of ze(l) to one of these subspaces converges to 0 as l. + 400.
5 A mix of hyperplanes and polytopes
We now turn our attention to an even more general version of Theorem 4.1, which does not require the leading eigenvalue of V to be simple. The resulting theorem can be viewed combination of
8

---

0.0
Positive limits for clustered coordinates
10
-12
t- 10.0
1 15.0
Negative coordinates
-007
-00
100
100
-00
0
035
(a) Conjecture 4.2: low-dimensional case.
(b) Conjecture 4.2: high-dimensional case.
Figure 7: (a) n 40, d and Q - K I3 with V a random matrix with eigenvalues (1.96, 0.22, 0.25]. The k positive eigenvalues of V generate attraction between the tokens and even convergence in the corresponding eigenspaces-this explains the codimension & statement. The negative eigenvalue generates a repulsive effect between the tokens, and we see divergence along two lines (note the different scales between the four figures). (b) n 256, d = 128, with (Q,K, V) fixed random matrices and V symmetric. For each coordinate j corresponding eigenvalue, the variance of the set {4] (t)): i [n]] (shaded area) tends to O with t, while the mean (solid lines) converges one among two real scalars: one positive (top figure), one negative (bottom) figure. Coordinates corresponding to negative eigenvalues diverge (Fig. 15).
a positive
Theorem 4.1 and Theorem 3.1. Specifically, we assume that V behaves as the identity when acting on the eigenspace of the leading eigenvalue. This property is automatically satisfied if V normal-so that its eigenvectors form an orthonormal basis so we call such a V paranormal.
Definition 2. We call (Q, K, V) a good triple with multiplicity if the following conditions hold:
(i) Q K is positive definite: QTK- K 0:
(ii) V is paranormal: there exist linear subspaces F. 3 c Rd which are invariant under V. and such that 3 @ 3 = Rd, Vis Ald X > 0, and p P(Vi)) < X, where p(.) denotes the spectral radius (the maximal modulus of eigenvalues).
example of such is used for Fig. 8. We may now state our main result in the setting of good triples with multiplicity. The proof may be found in Appendix E.
Theorem 5.1 (Clustering for A1 with multiplicity). Suppose that (Q, K) good triple with multiplicity in the sense of Definition 2. Then, for any initial sequence (0)))))een)) Rd, there exists bounded convex polytope K. c. F such that setting H (ŌK u {0}) x for any [n]. we have (1s)(2;(1), H) - O as - +00.
Outlook
Several important directions regarding the mathematical theory of Transformers remain unex- plored. An important extension of our work would amount to studying multi-headed Transformers. borrowing the notation from Remark 4, they amount to:
Qnx! Ky
3
[[+1]
[b] I
EZ Quat (4) Kaze(k))
Vi
ke N.

---

6.0
t= 16.0
-0
Figure 8: Illustrating Theorem 5.1 with 40 tokens in R As before, Q= la and we take V = (iag(111, ). A convex polytope ΣC emerges before time 5, toward which two coordinates of the tokens cluster, and persists throughout the evolution, while the tokens diverge along the coordinate corresponding to the eigenvalue (note the different scales between the four figures).
For each E [H] (corresponding to a different head), the weight matrices Qb, KA, Va are constant. Proofs regarding clustering convergence of the self-attention matrix for such dynamics is open problem. Preliminary numerical investigations seem to indicate that interesting clustering phenomena also occur in this context. A characterization or properties of optimal weights by invoking the optimal control correspondence in the spirit of [Wei17] 7] is also an interesting avenue for future research.
Acknowledgments and Disclosure of Funding
We thank Pierre Ablin, Léonard Boussioux, Enric Boix-Adsera, Gabriel Peyré, Yair Shenfeld and Emmanuel Trélat for helpful discussions. C.L. was supported by the Simons Foundation Grant 601948, DJ. P.R. is supported by NSF grants IIS-1838071, RMS-2022448, and CCT-2106377. Y.P. is supported in part by the MIT-IBM Watson AI Lab.
References
[ABV+0S] 05] Juan A Acebrón, Luis L Bonilla, Conrad J Pérez Vicente, Félix Ritort, and Renato Spigler. The Kuramoto model: A simple paradigm for synchronization phenomena. Reviews of modern physics, 77(1):137, 2005.
[BM00] Paul S Bradley and Olvi L Mangasarian. K- K-plane clustering. Journal of Global optimization, 16:23-32 32, 2000.
[Che95] Yizong Cheng. Mean shift, mode seeking, and clustering. JEEE Transactions on Pattern Analysis and Machine Intelligence, 17(8):990-799, 1995.
[CHH+16) 16] Junghee Cho, Seung-Yea Yeal Ha, Feimin Huang, Chunyin Jin, and Dongnam Ko. Emer- gence of bi-cluster flocking for the Cucker-Smale model. Mathematical Models and Methods in Applied Sciences, 26(06): 191-1218. 2016.
[CRBD18] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018.
[CS07] Felipe Cucker and Steve Smale. Emergent behavior in flocks. IEEE Transactions on Automatic Control, 8265).852-86,, 2007.
[DCL21] Yihe Dong, Dean-Baptist Cordonnier, and Andreas Loukas. Attention is all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793-2803. PMLR, 2021.
[Dob79] Roland L'vovich Dobrushin. Vlasov equations. Prilozheniya, 13(2):88553 1979.
Funktsional nyi Analiz, i ego
[Gol13] François Golse. Mean field kinetic equations. Course of Polytechnique, 2013.
10

---

[HJ12] Roger A Horn and Charles R Johnson. Matrix analysis. Cambridge University Press, 2012.
[HK02] Rainer Hegselmann and Ulrich Krause. Opinion dynamics and bounded confidence: models, analysis and simulation. Journal of Artifical Societies and Social Simulation (JASSS), 5(3), 2002.
[HKPZ19] Seung-Yeal Ha, Jeongho Kim, Jinyeong Park, and Xiongtao Zhang. Complete cluster predictability of the Cucker-Smale flocking model on the real line. Archive for Rational Mechanics and Analysis, 231:319-365, 2019.
[HR17] Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks, Inverse problems, 34(1), 2017.
[HysW+22] 22] Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference Learning Representations, 2022.
[6ZRS16a]
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
[HZRS16b] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pages 630-645. Springer, 2016.
[JM14] Pierre-Emmanuel Jabin and Sebastien Motsch. Clustering and asymptotic behavior in opinion formation. Journal of Differential Equations, 257(11):4165 4187, 2014.
[Kra00] UIrich Krause. A discrete nonlinear and non-autonomous model of consensus. In Communications in Difference Equations: Proceedings of the Fourth International Conference on Difference Equations, page 227. CRC Press, 2000.
[Krr55] Yoshiki Kuramoto. Self-entrainment of a population of coupled non-linear oscillators. In International Symposium on Mathematical Problems Theoretical Physics: January 23-29, 1975, Kyoto University, Kyoto/Japan, pages 420-422. Springer, 1975.
[LCG+20] 20] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. ALBERT: A Lite BERT for Self-supervised Learning of Language Representations. International Conference on Learning Representations, 2020.
[LLH 20] Yiping Lu, Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie-Yan Liu. Understanding and improving transformer from multi muulti-articlle dynamic system point of view. In International Conference on Learning Representations, 2020.
[LWLQ22] Tianyang Lin, Yuxin Wang. Xiangyang Liu, and Xipeng Qiu. A survey of transformers. Al Open, 3:111-132, 2022.
[MT14] Sebastien Motsch and Eitan Tadmor. Heterophilious dynamics enhances consensus. SIAM Review, 56(4))577-621, 2014.
[PHD20] Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117440):4655-24466, 2020.
[PR13] Benedetto Piccoli and Francesco Rossi. Transport equation with nonlocal velocity in Wasserstein spaces: convergence of numerical schemes. Acta Applicandae Mathemati- cae, 124:73-105, 2013.
[PR16] Benedetto Piccoli and Francesco Rossi. On properties of the generalized Wasserstein distance. Archive for Rational Mechanics and Analysis, 222: 2016.
11

---

[PRT15] Benedetto Piccoli, Francesco Rossi, and Emmanuel Trélat. Control to flocking of the kinetic Cucker Srakee model. SIAM Journal Mathematical Analysis, 47(6))4685- 4719, 2015.
[RS14] Brian Rider and Christopher D. Sinclair. Extremal laws for the real Ginibre ensemble. The Annals of Applied Probability, 24(4):1621 - 1651, 2014.
[SABP22] Michael Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transformers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics, pages 3515-3530. PMLR, 2022.
[VCBJ*95] 95] Tamás Viesek, András Czirók, Eshel Ben-Jacob, Inon Cohen, and Ofer Shochet. Novel type of phase transition system of self-driven particles. Physical Review Letters, 72(6): 1995.
[Vid11] René Vidal. Subspace clustering. IEEE Signal Processing Magazine, 28(2):52-68, 2011.
[VKF20] Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. Fast transformers with clustered attention. Advances in Neural Information Processing Systems, 33:21665- 21674, 2020.
[VSP*17] 17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[Wei17] E Weinan. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 1(5):1-11, 2017.
[WHL19] Weinan, Jiequn Han, and Qianxiao Li. A mean-field optimal control formulation of deep learning. Research in Mathematical Sciences, 6(1): 10, 2019.
[WLK* + 20] Sinong Wang, Belinda Z Li. Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint N0W:22006.04768, 2020.
[YBR+20] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiy Kumar. Are transformers universal approximators of sequence-to-sequemee functions? In International Conference on Learning Representations, 2020.
12


Reviews:
[{"review_id": "PSObPxRzAK", "reviewer": "Reviewer_Qc4C", "summary": "This paper studies the asymptotic behavior of a sequence of tokens processed by infinitely deep self-attention only Transformers, viewed as interacting particle systems (1).\n\nThe authors first study the one dimension case and show that the self-attention matrix converges to a low-rank boolean matrix. \n\nThey then focus on particular choices of the value matrix $V$ to obtain clustering in higher dimension. Specifically, when $V = I_d$, the paper shows that a time-rescaled version of the tokens converge to the boundary of a convex polytope. \n\nThe paper then focuses on more realistic choices of $V$. When it has a single leading eigenvalue, the clustering happens toward one of at most $3$ hyperplanes. When the leading eigenvalue has multiplicity, the limit geometry is a convex polytope in some directions and a linear subspace in the others. \n\nAll the theoretical results are numerically illustrated. \n", "questions": "1. What is the main obstacle for considering the time-dependent dynamics where $Q$, $K$ and $V$ depend on time, when $d = 1$ for instance ? Are there some assumptions that can be made on the time dependency to obtain similar results ?\n\n2. Regarding Weakness 1. For instance, in practical implementation, I believe $Q^TK$ is low-rank. Maybe it should also be clarified that by $Q^TK > 0$  you also mean symmetric ? In practical implementation, is $Q^TK$ close to a symmetric matrix ? \n\n3. In Th 4.1, you do not assume anymore that $Q^TK$ is symmetric ? Have you verified whether the second condition in Definition 1. holds on some pre-trained triple (Q, K, V) on ALBERT ? \n\n", "limitations": "Yes. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is very clear and well written.\n2. The overall contribution is strong, as it is the first paper provably showing the emergence of clusters in self-attention only Transformers. \n3. More specifically, each theorem, (well summarized in Table. 1) is an interesting result of clustering in interacting particles systems. \n4. Most assumptions on (Q, K, V) are clearly discussed.  \n5. The figures illustrating the theorems are well presented and insightful. ", "weaknesses": "1. Regarding the assumption in Th. 3.1 and 5.1 that $Q^TK > 0$ is positive definite: I think that the strength of this assumption should be emphasized. See Questions. \n\n2. An intuitive explanation (as for instance the one made after th 4.1 (l. 257 to 262)) would be welcome after th 3.1. Especially, how does the convex polytope $K$ depend on $Q^TK$ ?\n\nTypo: \nl. 242, I believe a $Q$ is missing in the quadratic form. \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback, and respond to each point individually below.\n\n**Weaknesses.**\n\n1. We recognize that this assumption is substantial. However, it does not appear to be necessary for our conclusions; rather, it serves to direct the proof. To reinforce the broader validity of our conclusion beyond just this specific assumption, we have carried out additional experiments (Figures 1, 2 in the PDF), suggesting that our clustering results are more universal.\n\n2. Theorem 3.1 amounts to a couple of effects entailed by the dynamics. First of all, the convex hull of the particles is shrinking over time (Proposition C.2). This is due to the fact that the distance of the particle nearest to any half-space (not containing the particles) is decreasing. On the other hand, the convex hull ought not collapse since particles which have not concentrated near the boundary of the limiting polytope will continue to increase in magnitude until they themselves reach this boundary (Step 2 in the proof). The latter is due to the time-rescaling and the \"probabilistic\" structure of the self-attention mechanism. (A version of this comment will be added to the camera-ready version.) As for the convex polytope $\\mathcal{K}$--it depends both on the initial sequence of particles and on $Q^\\top K$. Unfortunately, we do not believe there is a way to predict $\\mathcal{K}$ explicitly besides running the full dynamics.\n\nWe also thank the Reviewer for spotting this rather important typo.\n\n--- \n**Questions.**\n\n1. We believe that if we assume that $Q(t)^\\top K(t)$ is bounded from below and above by positive multiples of the identity uniformly in $t$, the same conclusions as in Theorems 2.1, 3.1, 4.1 and 5.1 hold. This requires adaptations in the proofs. However, the conclusions of the theorems do not hold in general if $V$ depends on time.\n\n2. Indeed $Q^\\top K\\succ 0$ means in particular that $Q^\\top K$ is symmetric (the partial order within the set of symmetric matrices). We agree with the Reviewer that this assumption is at odds with the low-rankness in practical implementations, but, as alluded to in the answer to Weakness 1, we do not believe that it is essential to the result. However, at this moment, we can't identify a clear method to eliminate it from our proof.\n It does not appear evident that the pre-trained ALBERT matrices are approximately symmetric.\n \n3.  Indeed, no symmetry or positive-definiteness of $Q^\\top K$ is required. Since we only look at the behavior of the particles along the direction given by $\\varphi_1$, it suffices to assume that the induced quadratic form is positive only along said direction. With regard to the pre-trained triple: the first point is satisfied by $V_5$ and $V_{14}$ (indicated in the Supplementary material), and the second point is too satisfied by $(Q_5, K_5)$ and $(Q_{14}, K_{14})$ (the inner products evaluated at the eigenvector of norm 1 equal 1.3060 and 0.6719 respectively). In other words, the triples $(Q_h, K_h, V_h)$ corresponding to heads $h=5$ and $h=14$ in ALBERT satisfy the assumptions of Definition 1. This comment will be added to the Supplementary Material of the camera-ready version."}, {"review_id": "CQOKt2OOvJ", "reviewer": "Reviewer_U2EX", "summary": "In this work, the authors develop a theoretical analysis of self-attention mechanism. In particular, the authors study the setting of a trained transformer  and their goal is to characterize the output of a deep transformer with multiple layers of self-attention. For simplicity authors focus on weight sharing and also do not use MLPs and multiple heads. The key results include:\ni) the convergence guarantees of self-attention matrix to a low rank matrix. \nii) the characterization of convergence of initial inputs to final outputs as a function of query, key and value matrix, \n\nIn their characterization, the authors prove a clustering structure. For instance when value matrix is identity, then the authors show that the depending on the initial condition the output of dynamics of self attention converges to one of the vertices of a polytope. In another setting, the authors show that the output converges to a point that lies on one of three hyperplanes. All of these results point to a clustering phenomenon at the output of the transformer. This phenomenon has strong parallels to neural collapse studied in standard supervised learning setting. ", "questions": "Please see the weakness section. ", "limitations": "The authors point to several open problems and directions. They can be interpreted as gaps that can still be filled. A more extensive discussion on limtiations can be useful. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is a very well written paper. It studies a very important problem of understanding the core of transformers, i.e., self attention. \n\nOriginality: The paper has quite an original idea and approach to solve the problem. \n\nQuality: This is a high quality paper that can inspire important developments on the path of demistifying transformers. \n\nClarity: The paper is a joy to read. \n\nSignificance: I believe that the paper takes an important step in the right direction and can inspire the theory community in machine learning to furthering our understanding of transformers. ", "weaknesses": "The paper does not have significant weaknesses. I do have some questions and concerns that I would want to ask. \n\n1. I see that the convergence analysis and result that self attention converges to a low rank matrix is given for the case of d=1. What about the more general setting? This seems to be true empirically but what can we say in theory. \n\n2. I see that the authors state that discrete time system's analysis is straightforward. Intuitively, it feels not so straightforward and should rely on assumptions on length of step size. Can the authors elaborate? \n\n3. The authors shed some light on the ALBERT's pretrained matrices satisfying the condition in the theorem in some cases. A more extensive analysis of this can be quite useful to understand if these assumptions are indeed necessary to achieve such a clustering structure in practice or they are a matter of theoretical convenience. \n\n4. I would appreciate if authors could add some more intuition on what is geometric meaning of Definition 1. \n\n5. While I understand that authors have focused only on self attention purely. This makes a lot of sense from theory point of view. From an empricial point of view, I believe there is value in doing numerical experiments to see how sensitive are the clustering illustrations to addition of an MLP layer on top.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback, and respond to each point individually below.\n\n**Weaknesses.**\n\n1. A result similar to Theorem 2.1 certainly holds for $d\\geq 2$ and $V=I_d$, but its statement and proof would need the exclusion of numerous pathological and non-generic initial configurations of tokens. The proof is anticipated to be highly technically challenging (and admittedly tedious), and we have chosen to postpone it to future work.\n\n2. We agree with the referee, we will give more details on the discrete-time setting in the camera-ready version. Note that we impose that $I_d+V\\Delta t$ is invertible, and this is implicitly is an assumption on the step size, which holds for instance for sufficiently small $\\Delta t$. Let us give here some details on the proof of Theorem 3.1 in the discrete-time setting, following the intermediate results of Appendix C.1. First, Proposition C.2 (convex hull shrinkage) holds intuitively because for any $i$ and $k$, $z_i^{[k+1]}=\\frac{1}{1+\\Delta t}(z_i^{[k]}+\\Delta t\\sum_{j} P_{ij}^{[k]}z_j^{[k]})$ belongs to the convex hull of the $z_j^{[k]}$. Then we define the candidate set of limit points as in (35), and Claim 1 holds without any change in the statement and in the proof. Then, as in Steps 2 and 3 in Section C.1.2, we first prove that if $z_i^{[k]}$ is not already close to one point of the candidate set, it will keep moving toward the boundary, and finally we prove that tokens cannot circulate indefinitely between different points on the boundary. This proves convergence of each token toward a point of the set given by (35).\n\n3. We believe that some assumptions are purely made out of theoretical convenience. \nAs indicated to Referee Qc4c in Weakness 1, we have carried out additional experiments that supplement our claims beyond assumptions made for the theory. \nFor instance, the assumption $Q^\\top K\\succ 0$ in Theorems 3.1 and 5.1 does not appear necessary in experiments.\n\n4. The first condition in Definition 1 means that $V$ has one dominant eigenvalue: the evolution, which shares some similarities with the ODE $\\dot{x}=Vx$, is quicker in the eigenspace $\\mathrm{span}(\\varphi_1)$ (suppose $V$ diagonalizable for simplicity) of the dominant eigenvalue. Therefore, after some time (or equivalently, after a few layers), the tokens have the corresponding coordinate much larger than the other coordinates, and this is one of the key ingredients in our analysis.\nThe second condition means that $Q^\\top K$ is positive definite but solely on the subspace $\\mathrm{span}(\\varphi_1)$. This condition is useful to guarantee that the components of $x_i$, $x_j$ along $\\mathrm{span}(\\varphi_1)$ represent the principal contribution to the weight $\\exp(\\langle Qx_i,Kx_j\\rangle)$.\n\n5. This can be done in two ways. The first way consists in first running the pure self-attention dynamics up to time $T$ (or equivalently, for $O(T)$ layers), with then applying a pure MLP to the concatenated vector of clustered features at time $T$. This amounts to seeing the MLP as a map from $\\mathbb{R}^{nd}$ to $\\mathbb{R}^{nd}$, which can be studied independently by existing theory.  The second way consists in alternating the self-attention and the MLP layers. In this case, clustering in the sense of Theorems 3.1 and Theorems 5.1 should not be expected since the weights of the MLP play the role of a value matrix $V$, and the conclusions of these theorems strongly depend on the identity-like structure. In Figure 3 of the PDF we illustrate an extension of Theorem 4.1 in the presence of such an MLP layer. "}, {"review_id": "arBzW8TVHn", "reviewer": "Reviewer_8epq", "summary": "This paper analyses the self-attention mechanism in (trained) Transformers under the lens of dynamical systems. The authors focus on a bare-bone self-attention architecture without the bells and whistles of standard Transformers (e.g. multi-head attention, layer norm) and assume time-independent weights, i.e. shared weights across \"layers\".\n\nThey start their analysis by studying simple self-attention architectures, including 1d settings, and progressively move to more complex and realistic settings, including value matrices with simple and positive leading-eigenvalue.\nAcross all the scenarios studied, the authors observe that in the limit, tokens cluster towards a few objects, such as vertices of a polytope or hyperplanes. Thus, the authors mathematically confirm that a small number of leaders drive the transformer dynamics.", "questions": "1. In eq. 4, the authors introduce an exponential factor, which is \"instrumental in the proofs of all results that follow\". Would the theorems hold without the use of this factor? What are the practical implications of this additional factor?\n\n2. Furthermore, according to the authors, this operation is a \"mathematically justified surrogate for the layer normalization\". Do the authors expect this rescaling operation to be useful in an actual self-attention implementation, either replacing layer norm or used along with it?\n\n3. In Definition 2, the authors define what constitutes a good triple $(Q, K, V)$. How realistic is this assumption, and how robust is the theorem to the assumption holding approximately?", "limitations": "The authors should focus more on the limitations of the assumptions used for the theoretical results introduced in this work. It is unclear whether the theorems hold in practice, or whether they are based on very strong and unrealistic assumptions. See also weaknesses and limitations.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This work focuses on the mathematical under-pinnings of self-attention and transformer architectures, a research direction that remains unexplored, yet is much needed, given the importance of transformer architectures in modern deep learning. Such a theoretical foundation is original and significant for the community.\n\nThe authors showcase the clustering behaviour of self-attention in progressively realistic scenarios. The main points and conclusions of this work are clear.", "weaknesses": "Overall, it is unclear whether the assumptions used in to prove theorems (e.g. good triples, $Q^T K \\succ 0$) are realistic for standard Transformers. \nFurthermore, the analysis is performed assuming only time-independent weights. While that design choice has been used in practice (e.g. ALBERTA), it is not the standard design choice for transformers. It is unclear whether the theorems introduced in this work hold for time-dependent weights.\n\nThe paper is very math-heavy for a general machine learning audience, and as such it is often hard to follow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback, and respond to each point individually below.\n\n**Weaknesses.**\n\nWe have attempted to paint a more complete picture on the necessity of some assumptions made to facilitate the development of this theory through numerical experiments (Figures 1, 2, 3 of the PDF) and additional comments (e.g. the answer to Question 1 by Reviewer Qc4c).  \n\n---\n\n**Questions.**\n\n1. The main effect, namely the emergence of leaders (visible in the self-attention matrix) remains true without the rescaling factor, see Theorem 2.1. Without the rescaling factor however, there is no clustering: this rescaling factor, in practice, has an effect of normalization on tokens, without which tokens typically diverge to $\\pm\\infty$.\n\n2.  We indeed believe that the exponential rescaling can be used in actual implementations, as a surrogate to layer norm. (There is a slight difference however in the sense that layer norm amounts to having particles evolving on the unit sphere of $\\mathbb{R}^d$, hence bounded in all directions.) It could potentially be used in conjunction with layer norm, but it is not totally clear to us what this would entail.\n\n3. We recognize that point (i), namely $Q^\\top K\\succ 0$, appears stringent. This assumption is made since the proof follows the lines of that of Theorem 3.1 (which is conceived on elements using this condition). \nHowever, experiments indicate that (i) is not necessary for the conclusion of the theorem to hold.  \nPoint (ii) is a generalization of the first bullet point in Definition 1 (the latter is satisfied by some pre-trained value matrices of individual heads in ALBERT, see Figure 10 in Supplementary Material). \nAssumptions on $V$ are generally more sensitive to perturbations than assumptions on $(Q, K)$. Should the eigenvalue $\\lambda$ with largest real part be negative, all rescaled particles will diverge to infinity. Should $\\lambda$ be complex, we do not expect any clustering phenomenon. Yet none of the conclusions seem to change if $Q^\\top K$ is taken arbitrary. "}, {"review_id": "0GUYh6ZGwG", "reviewer": "Reviewer_P5zn", "summary": "For $Q,K,V$ fixed and different structures of $V$, analyze the distribution of $x(t)$ as $t\\to\\infty$, where tokens are seen as particles and the self-attention mechanism is seen as particle interaction, i.e. as a McKean-Vlasov SDE. The conclusion is as $t\\to\\infty$, i.e. going through the layers, $x(t)$ converges to a clustered configuration, and the attention matrix $P$ becomes low-rank, i.e. tokens depend on few tokens.", "questions": "1. Please comment on the effect of different embeddings. (since tokens converge to the boundary and vertices of a convex polytope, and the polytope depends on the initial distribution of tokens and the embedding).\n\n2. Please give some intuition on what the vertices mean. (are vertices more descriptive since everything inside is a convex combination of the vertices? how is the final distribution on the vertices related to the initial distribution? e.g. is the mean unchanged? will it speed up training/ execution if a neural network was designed to give the convex polytope and the weights on the vertices instead of repeated iteration of self-attention?)\n\n3. Please comment on the downstream effects of \"linearly separable representation of tokens (line 266-267)\". Perhaps $x(t)$ as $t\\to\\infty$ can be used as new embeddings.\n\n4. Read the proof for -V in the appendix. Would it also help to think of clustering dynamics in terms of potential and thus flipping the sign of $V$ flips the potential landscape?\n\n5. Please comment on the shrinking convex hull/ polytope (since it is finite) (Appendix C.1.1). What about infinite number of tokens?\n\n6. If the tokens cluster, why would depending on one token (low rank attention matrix $P$) within a cluster be more likely than depending on another token within the cluster? Or is it because $P$ is low rank and thus tokens cluster? Or a combination of both effects?\n\n7. Please comment on the effects of skip and non-skip connections [DCL21]. ", "limitations": "Authors addressed the limitation of fixed $Q,K,V$, simple structures for $V$. Authors mentioned multi-head attention as future work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Following [LLH+20] and [SABP22], the authors interpret self-attention and transformers through the lens of interacting particle systems/ McKean-Vlasov SDEs. This is perhaps the first paper to show via the above interpretation the emergence of clusters/ leaders in transformers.", "weaknesses": "1. Please state within the main paper the relationship between the clustered configuration and the initial configuration.\n\n2. The result in [SABP22] seems to be the opposite of clustering. Please let me know if I misunderstood. Otherwise please comment on the clustering/ non-clustering effects and use-cases for downstream applications. \n\n3. Please remark on why \"transformer dynamics present unique mathematical challenges that cannot be addressed using the tools developed for these more primitive models. (lines 147-148).\" \n\n4. Please strengthen the contribution with layer normalization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback, and respond to each point individually below.\n\n**Weaknesses.**\n\n1. We do not believe there necessarily is a relation, as there is no easy way to predict the clustered configuration from the initial one beyond simply running the dynamics. Moreover, the number of clusters is often not equal to the number of vertices of the polytope given by the convex hull of the initial sequence.\n\n2.  In our opinion the results in [SABP22] are of a slightly different nature. The authors show, by virtue of an additional bandwidth hyperparameter $\\varepsilon$ in the exponential, that the continuity equation for the transformer dynamics (replacing self-attention by the Sinkform kernel) can be seen as an approximation for the heat equation ($V=I_d$) when the bandwidth goes to $0$. In our case, the bandwith is fixed and equal to $1$. We believe that in practical applications, the self-attention mechanism serves a dimension-reduction purpose. Our clustering theorems can be seen as an indicator of this thesis.\n\n3. The related models which we mention (Vicsek, Cucker-Smale, and so on) have significantly more basic interactions--often involving symmetric interactions and a radial dependence--, thus rendering the pre-developed mathematical tools inapplicable to our setting. \n\n4. Layer normalization amounts to considering the dynamics on the unit sphere on $\\mathbb{R}^d$. In ongoing work, we have observed, both theoretically and numerically, that clustering persists in a sense much alike the results presented in the present paper. The proof techniques are however different due to the intrinsic symmetries entailed by the sphere, and thus involve significant additional developments. \n\n---\n\n**Questions.**\n\n1.  It is not clear how the limiting convex polytope depends on the initial distribution of tokens even without the consideration of different embeddings. \nWe believe that this could be an exciting avenue for future research. \n\n2. We do not believe that the vertices of the limiting polytope have a clear meaning as they are not straightforwardly related to the initial distribution, as alluded to in previous answers. Slight changes in the initial distribution can change the limiting polytope. The mean is not preserved.  We thank the Reviewer for the last observation which is very compelling--indeed, should one manage to gather a clearer understanding of how the limiting polytope depends on the initial distribution and on $Q^\\top K$, the repeated iteration of self-attention could perhaps be replaced in training and this could be seen as a possible convex relaxation. \n\n3. This is a possibility, we thank the Reviewer for this observation. What we had in mind was to potentially use transformers as a plain clustering algorithm in the spirit of the mean-shift algorithm, or perhaps for binary classification tasks, by considering the inputs of the entire dataset as an input sequence.\n \n4. Yes, there is a natural potential, or more precisely a natural Lyapunov function $\\mathscr{L}$, which appears in the proof of Lemma C.8. The function $\\mathscr{L}$ is non-increasing when $V=-I_d$, which corroborates the global decay of the norms of the points toward $0$ (Theorem C.5). When $V=I_d$, the function $\\mathscr{L}$ is non-decreasing along the flow of the original dynamics (15), and indeed all points tend to diverge toward infinity when evolved through this dynamics, as a corollary of Theorem 3.1 after removing the rescaling.\n\n5. With an infinite number of tokens, there is also shrinking of the convex hull as in Proposition C.2. The limiting shape is convex, but it is not necessarily a convex polytope, it may be for instance a sphere.\n\n6. There is empirically one leader in each cluster, but it is difficult to know which one it is without running the full dynamics. A limiting low-rank self-attention matrix allows to identify not only clusters, but also leaders within the clusters by reading the rows and columns. When we establish clustering, except for the case $d=1$ we are not able to prove convergence toward a low-rank self-attention matrix, although this is what we observe in practice.\n\n7. The outcome of [DCL21] is that without skip connections, the output of the transformers dynamics converges doubly exponentially to a rank-1 matrix. Therefore, [DCL21] highlights the necessity of skip-connections for transformers to be efficient. \n[DCL21] proves clustering toward a unique cluster for transformers without skip-connections; our work proves that clustering occurs for transformers with skip-connection, which are the ones used in practice."}]

Author rebuttal:
We thank the reviewers for their thoughtful feedback. We echo their concerns with regard to several assumptions we had made on the weight matrices for our analysis.

All in all, our goal was to consider the simplest setting of transformers amenable to rigorous mathematical analysis. To enhance the validity of the clustering theorems, we have conducted a few numerical experiments which violate the assumptions made for the development of the theory and observed that the clustering pattern persists. Namely, 
-  In Figure 1, we replicate the experiment illustrating Theorem 3.1 in the setting where $Q^\top K$ is a random matrix with entries sampled from $\mathrm{Unif}([-1,1])$. The clustering pattern persists even for this choice of weights.
- In Figure 2, we replicate the experiment illustrating Theorem 5.1 in the setting where $Q^\top K$ is a random matrix with entries sampled from $\mathrm{Unif}([-1,1])$. The same clustering pattern appears.
- In Figure 3, we add an additional MLP layer in the setup of Theorem 4.1. More precisely, we use a 2-layer neural network: we apply a component-wise activation function (either ReLU or Tanh), and then multiply by a time-independent weight matrix $W$. We again see clustering, the pattern depending on the weight matrix $W$ and on the activation function. 
    In particular, in the first row (corresponding to $W=I_d$ and ReLU), we see that the particles first evolve as to reach the upper right quadrant $\mathbb{R}^d_{>0}$ (due to the ReLU). Once they reach it, 
    every particle eventually follows one of three hyperplanes determined by the spectrum of $V$ and the projection onto $\mathbb{R}^d_{>0}$.
    In the other two cases, all the particles appear to collapse to $0$.

All experiments were conducted following the setup presented in Appendix F of the Supplementary Material. They will be added and documented in the Supplementary Material of the camera-ready version.
We comment on the difficulty of generalizing the theory to more general scenarios in the point by point replies to individual reviewers. 

Meta-review: This paper studies the asymptotic behavior of a sequence of tokens processed by a self-attention network and viewed as interacting particle systems. More precisely, Eq. (1) gives a differential equation for the evolution of the tokens $x(t)$ as a function of a self-attention matrix $P(t)$ (this result was derived in prior work). The main contribution of the paper is to show that the matrix $P(t)$ converges to a low-rank matrix with $0,1$ entries (this is for instance illustrated in Figure 2). Effectively, $P(\infty)$ shows the existence of clusters between tokens. This is an interesting perspective to study self-attention, which is still not well understood from an analytical perspective.

The reviewers are in general positive about the paper, they for instance find it clearly written, with an interesting contribution and the potential to lead to novel research directions. I therefore recommend acceptance. Please include the changes discussed in the rebuttal to the camera-ready version.
Decision: Accept (poster)