Title: The emergence of clusters in self-attention dynamics
Venue: NeurIPS 2023
Accepted: yes
Keywords: Transformers

Self-Attention

Clustering

Interacting Particle Systems

Continuous 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

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. Cluster locations are determined by the initial tokens, confirming context-awareness of representations learned by Transformers. Using techniques from dynamical systems and partial differential equations, we show that the type of limiting object that emerges depends on the spectrum of the value matrix. Additionally, in the one-dimensional case we prove that the selfattention 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.

1 Introduction

The introduction of Transformers in 2017 [VSP ` 17] marked a turning point in the AI revolution, powering breakthroughs in natural language modeling and computer vision. With remarkable empirical 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 an astonishing rate, the need to understand their inner workings is becoming a pressing scientific challenge. In this work, we make a first step in this direction by describing the geometry of learned representations.

To provide a transparent presentation of our findings, we take a leaf out of the literature on continuoustime dynamics such as neural ordinary differential equations (ODEs) [CRBD18, Wei17, HR17]. By viewing layers as a time variable, this formalism has emerged as a flexible mathematical framework to implement and study ResNets [HZRS16a] as particular discrete-time versions of a parametrized dynamics of the form

Here θ is the trained parameter of a neural network and f θ is characterized by the precise architecture of the ResNet 1 . In turn, an input (e.g., an image) x p 0 q P R d is mapped to its representation x p T q .

Unlike neural ODEs and ResNets, the representation map of Transformers is not solely a function of an individual input x p 0 q P R d but rather of a set/sequence p x 1 p 0 q , . . . , x n p 0 qq of n ě 1 d -dimensional tokens. These tokens then evolve in time by interacting with each other per the selfattention mechanism. Namely, following [SABP22], we view tokens as particles, and the transformer dynamics as an interacting particle system of the form

for any i P r n s , where P ij p t q are the entries of a n ˆ n stochastic matrix P p t q , given by

Here the matrices Q (Query), K (Key), and V (Value) are learned from data. Note that Q,K need not be square. The n ˆ n matrix P p t q is called self-attention matrix . The wording attention stems precisely from the fact that P ij p t q captures the attention given by token i to token j relatively to all tokens /lscript P r n s . 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 a 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 to self-attention dynamics of (1). 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 representations of a trained Transformer by studying the asymptotic behavior of a sequence of tokens p x 1 p t q , . . . , x n p t qq as they evolve through the layers of a transformer architecture using the dynamics (1). In this setup, a Transformer is completely described by the weight matrices p Q,K,V q obtained during training. Note that we assume that these three matrices are timeindependent . While this assumption is motivated by mathematical convenience, it is worth noting that such weight-sharing

scenarios are in fact used in practice-see, e.g., ALBERT [LCG ` 20]-as they drastically reduce the number of parameters of a network.

With parameters p Q,K,V q 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 t x 1 p t q , . . . , x n p t qu , appropriately rescaled, tends to a clustered configuration as t Ñ8 . Our theoretical findings justify the empirical observation made in [VSP ` 17] that leaders appear in a 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 ' 1 and V ą 0 , the self-attention matrix P p t q converges to a low-rank matrix with entries 0 and 1 as t Ñ`8 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 as past empirical work) indicate that the result may extend to higher dimensions for almost all initial sequences of tokens.

(ii) In Section 3 we first focus on the case V ' I d as a natural canonical choice that enables us 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 n . This elucidates the clustering phenomenon. (See Fig. 1.) When V ' ´ I d , all tokens following the dynamics (1) collapse to 0 .

(iii) We build on these results and in Section 4 consider the case wherein V is only assumed to have a 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 at most three hyperplanes which are determined by the corresponding eigenvector.

(iv) In Section 5 we complete the results of Sections 3 and 4 by addressing the case where the leading eigenvalue has multiplicity. This results in clustering toward the vertices of a convex polytope in some directions, and a 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).

Remark 1 (Discrete time) . While we focus on the idealized setting of self-attention dynamics in continuous-time, this is solely done for convenience and all of our methods are straightforwardly applicable to the discrete-time setting. (See also Remark 4.) The discrete-time analog of (1) with time-step ∆ t ą 0 (equal to 1 in practice) is simply the forward Euler iteration

Notation. We denote by x¨ , ¨y and } ¨ } the Euclidean dot product and norm respectively, and we use the shorthand r n s : ' t 1 , . . . , n u . For any matrix M P R d ˆ d , we order its eigenvalues (repeated according to multiplicity) by decreasing order of modulus: | λ 1 p M q| ě . . . ě | λ d p M q| . We denote by } M } op the /lscript 2 -operator norm of the matrix M , equal to the largest singular value of M . Given a set S Ă R d , we define the distance of a point x P R d to S as dist p x, S q : ' inf s P S } x ´ s } , and by conv p S q the convex hull of S .

Related work

Our study and results build on several different lines of work, and we draw some parallels in what follows.

Analysis of attention-based models. Given the widespread use of Transformers in natural language processing, there has been a surge of interest in understanding the function and significance of attention layers within these models. In [YBR ` 20], the authors show that when treated as discrete-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-diffusion equations to slightly improve the performance of Transformers by employing a Strang-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 a 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 to process sequences, because each self-attention layer contains n 2 products of the form x Qx i , Kx j y (for i, j P r n s ). The empirical observation that the selfattention matrix P is close to a low rank matrix-see [LWLQ22, Section 4.4] for references-is 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 [WLK ` 20, 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 a rich cluster structure emerges.

Compared to the usual BERT , ALBERT [LCG ` 20] uses parameter-sharing 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 O p n 2 q complexity of the original Transformer, but, quoting [LCG ` 20], it "significantly reduce[s] the number of parameters for BERT without seriously hurting performance, thus improving parameter-efficiency. 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".

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 a 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/clustering of oscillators [Kur75, ABV ` 05], the model which is most similar to (1) is the Krause model [Kra00]

which is non-symmetric in general ( a ij ‰ a ji ), much like (1). When φ is compactly supported, it has been shown in [JM14] that the particles x i p t q assemble in several clusters as t Ñ `8 . 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-Smale model [CS07]. These models may also exhibit a 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 patterns. Much more remains to be discovered and this work is a 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 ` 20] 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 n which is inherent to Transformers, resulting in a non-negligible decrease in the cost of training.

As a warm-up to deriving complete geometric representations of the dynamics, our first result shows, in the simple 1 d case that P p t q indeed converges exponentially fast toward a matrix which is typically both Boolean and low-rank (see Fig. 3). Although there are clear obstructions to a 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).

0

5

10

15

20

25

30

35

P

of

n

ˆ

n

matrices having the form

P

σ

1

»

-

-

-

-

-

-

-

-

-

1

0

. . .

0

.

.

.

.

.

.

.

.

.

.

.

.

1

0

. . .

0

˚

˚

. . .

˚

0

. . .

0

1

.

.

.

.

.

.

.

.

.

.

.

.

0

. . .

0

1

fi

ffi

ffi

ffi

ffi

ffi

ffi

ffi

ffi

ffi

P

σ

2

-

fl

˚

To set this up, we introduce the set

To set this up, we introduce the set P of n ˆ n matrices having the form illustrated in Fig. 2, where the asterisks denote arbitrary 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.

the last one. Theorem 2.1 (Self-attention matrix converges to a low-rank Boolean matrix) . Let d ' 1 . Suppose that the scalars p Q,K,V q satisfy V ą 0 and QK ą 0 . For any initial sequence of pairwise distinct tokens p x 1 p 0 q , . . . , x n p 0 qq P R n , there exists some P ˚ P P such that the self-attention matrix P p t q defined in (2) converges to P ˚ as t Ñ`8 .

Theorem 2.1 that the scalars p Q,K,V q satisfy The proof may be found in Appendix B. The rate of convergence toward P ˚ is in fact doubly exponential in t for coefficients

The proof may be found in Appendix B. The rate of convergence toward P is in fact do exponential in t for coefficients outside the row of asterisks in 4. The proof the theorem also re that for almost all initial sequences of pairwise distinct tokens, P ˚ is actually of rank 1 or 2 , i.e. row of asterisks is equal to either e 1 ' p 1 , 0 , . . . , 0 q or e n ' p 0 , . . . , 0 , 1 q . The interpretation of Theorem 2.1 is that in the 1 d case, at most three tokens capture the attention of all tokens except at most one. Typically, these leading tokens are those carrying the largest amount of information. This is also illustrated in Fig. 4. Since the tokens x i here evolve on R , the right-most and left-most ones (which typically tend toward ˘8 ) capture the attention of all the others.

tokens p x 1 p 0 q , . . . , x n p 0 qq P R n , there exists some P ˚ P P such that the self-attention matrix defined in (2) converges to P ˚ as t Ñ`8 . outside the row of asterisks in 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 e 1 ' p 1 , 0 , . . . , 0 q or e n ' p 0 , . . . , 0 , 1 q .

The interpretation of Theorem 2.1 is that in the all tokens except at most one. Typically, these

1

d

case, at most three tokens leading

capture the attenti tokens are those carrying the largest amou

here evolve on

R

, the right-

) capture the attention of all the others.

information. This is also illustrated in Fig. 3. Since the tokens

x

i

and left-most ones (which typically tend toward

˘8

0

5

10

15

20

25

30

35

t

= 0

.

0, rank= 11

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

t

= 3

.

0, rank= 23

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

t

= 5

.

0, rank= 14

0

5

10

15

20

25

30

35

0

5

10

15

20

25

30

35

t

= 10

.

0, rank= 2

Figure 2: An illustration of the asymptotics of

Q

K

1

and

V

1

P

p

t

q

entailed by Theorem 2.1 for

n

. (See Appendix F for details on computing.) Increasing

'

40

tokens,

n

has no effe

this behavior of

P

p

t

q

-see Fig. 10.

t

= 0

.

0

t

= 2

.

0

t

= 9

.

0

'

'

'

the self-attention matrix

[VSP

`

'

P

p

t

2

.

q

2

where pair

0

and

Q

ą

ď

2

structu

This interaction echoes findings illustrated in the original p

17]-for instance, Figures 3-5 therein.

3 Clustering toward vertices of convex polytopes

5 In the rest of the paper, we seek to taxonomize various clustering results for the solutions to (4) when t Ñ`8 , depending the sign and the multiplicity of the eigenvalues of V . We begin by focusing

on what may appear to be the most natural 2 case V ' I d , as is also done in [SABP22]. In fact, we demonstrate (theoretically and numerically) later on, clustering is a generic phenomenon which holds under much less restrictive assumptions.

The transformer dynamics considered in (1) does not contain a layer normalization mechanism typically encountered in practice [VSP ` 17]. In absence of such a device, tokens may diverge to infinity as in Theorem 2.1. In fact, the norm of the tokens x i p t q typically diverges exponentially toward `8 for any d : this is expected, by analogy with the non-trivial solutions to 9 y p t q ' y p t q .

To remedy this situation, we take inspiration from the solution y p t q ' e tV y p 0 q to 9 y p t q ' V y p t q . Namely, for any i P r n s we consider the rescaled tokens z i p t q : ' e ´ tV x i p t q , which solve

The initial condition remains the same: x i p 0 q ' z i p 0 q for any i P r n s . More importantly, the coefficients of the self-attention matrix for the rescaled tokens z i p t q are the same as those for the original tokens x i p t q . Whence, the conclusion of Theorem 2.1 also applies to the dynamics (4). We see this rescaling of tokens as a 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 z i p t q then gives information on the dynamics of the original tokens x i p t q by virtue of the relation x i p t q ' e tV z i p t q .

We are now able to state the main result of this section on the case V ' I d . The following theorem shows that the tokens z i p t q evolving per dynamics (4) converge to the boundary of a convex polytope as t Ñ`8 . We present here a simplified but weaker version of our result for convenience, and refer the reader to Theorem C.1 in the appendix for a complete statement.

Theorem 3.1 (Convergence to points on the boundary of a convex polytope) . Suppose V ' I d and Q J K ą 0 . Then, for any initial sequence of tokens t z i p 0 qu i Pr n s Ă R d , there exists a convex polytope K Ă R d such that for any i P r n s , z i p t q converges either to 0 or to some point on B K as t Ñ`8 .

The convex polytope K is completely determined by the initial sequence of tokens, and Q J 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 z i p t q ( i P r n s ) 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 B 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.

²

-

5

0

5

-

5

0

5

-

5

0

5

t

= 0

.

0

-

5

0

5

-

5

0

5

-

5

0

5

t

= 1

.

0

-

5

0

5

-

5

0

5

-

5

0

5

t

= 2

.

0

-

5

0

5

-

5

0

5

-

5

0

5

t

= 5

.

0

Recall that the points x i p t q ' e t z i p t q when V ' I d follow the original dynamics (1). Akin to Theorem 2.1, this result also shows the emergence of a set of leaders (given by the vertices of K )

attracting all tokens as t grows. It has been experimentally observed (first in [VSP ` 17]) that in trained Transformers, tokens focus their attention on local leaders in a way that seems to reproduce the syntactic and semantic structure of sentences.

The proof of Theorem 3.1 is postponed to Appendix C, and 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) 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 to increase in magnitude until they themselves reach this boundary (Step 2 in the proof). This occurs due to the time-rescaling.

Remark 2. Assuming Q J K ą 0 does not seem to be essential for our conclusions; instead, it 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 3 (Rate of convergence) . Although Theorem 3.1 (as well as Theorems 4.1 and 5.1 stated below) does not specify a rate of convergence toward B K , we expect (and observe through numerics) that convergence happens very quickly-after few layers, most tokens are already clustered. What "few layers" means here necessarily depends on 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 in 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 R ' I d ` V ∆ t , and assume that R is invertible (which is the case for sufficiently small ∆ t ), then z i p k ∆ t q ' R ´ k x i p k ∆ t q : ' z r k s i satisfies

The proofs of Theorems 2.1, C.5, 3.1, 4.1, and 5.1 carry through with straightforward modifications.

Let us provide some comments on the proof of Theorem 3.1 in the discrete-time setting, for the sake of completeness. First of all, Proposition C.2 holds intuitively because for all integers i P r n s and k ě 1 ,

We then define the candidate set of limit points as in (37) , and Claim 1 holds without any change in the statement or in the proof. Then, just as in Steps 2 and 3 in the proof of C.1, we can first show that if z r k s i is not already near some point in the candidate limit set, it will keep moving toward the boundary of the convex polytope. Finally, 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) .

4 Clustering toward hyperplanes

While being a natural example to consider, value matrices found empirically are much more general than V ' I d , which we considered in the previous section. We now turn our attention to a significantly more general setting of value matrices, which we formalize as follows.

Definition 1. We call p Q,K,V q a good triple if the two following conditions are satisfied:

the eigenvalue of V with largest modulus is real, positive, and simple; namely,

x Q ϕ 1 , K ϕ 1 y ą 0 for any ϕ 1 P R d lying on the line ker p V ´ λ 1 p V q I d q .

The second condition simply states that the quadratic form x Q ¨ , K ¨y is positive definite along the eigenspace associated to the leading eigenvalue of V . Note also that if all entries of V are positive,

the first condition is automatically satisfied by virtue of the Perron-Frobenius theorem. In fact, this assumption is generic. On the one hand, it is satisfied by some 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 in dimension d ' 128 -this proportion is known to vanish as d Ñ8 , albeit very slowly [RS14].

Our clustering result in the setting of good triples can be summarized as follows: the coordinate x z i p t q , ϕ 1 } ϕ 1 } y of any token z i p t q along the eigenspace spanned by ϕ 1 converges, as t Ñ`8 , toward one among possibly 3 real scalars. Consequently, all the tokens z i p t q converge toward one among at most three parallel hyperplanes; see Fig. 6 for an illustration.

Theorem 4.1 (Convergence toward ď 3 hyperplanes) . Assume that p Q,K,V q is a good triple in the sense of Definition 1. Then, for any initial sequence of tokens t z i p 0 qu i Pr n s Ă R d , there exist at most three parallel hyperplanes in R d such that for any i P r n s , the distance of the solution z i p t q to (4) to one of these hyperplanes converges to 0 as t Ñ`8 .

0

10

-

10

0

10

t

= 0

.

0

0

10

-

10

0

10

t

= 1

.

0

0

10

-

10

0

10

t

= 5

.

0

0

10

-

10

0

10

t

= 15

.

0

The proof may be found in Appendix D. The important role played by λ 1 p V q in the dynamics may be seen in (4): the component of z i p t q along ϕ 1 determines the size of e tV z i p t q in the exponent appearing in (4). The tokens z j p t q attracting other tokens z i p t q are those for which this component along ϕ 1 is 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 ` 17]).

Furthermore, Theorem 4.1 can also be interpreted in more classical machine learning terms. On the one hand, it can be seen as an instance of K -flats clustering [BM00, Vid11]-points in the input sequence are clustered, based on their intrinsic similarity, to at most 3 "flats" of dimension d ´ 1 . On the other hand, it ensures that for a good triple p Q,K,V q , (4) generates a linearly separable representation of tokens.

Beyond a single direction?

Numerical experiments (e.g., Fig. 7) indicate that a similar phenomenon emerges for more complex V . We formulate following conjecture which is a natural generalization of Theorem 4.1.

Conjecture 4.2 (Codimension conjecture) . Let k ě 1 be the number of eigenvalues of V with positive real part. Then there exist at most three parallel Euclidean subspaces of R d of codimension k such that for any i P r n s , the distance of z i p t q to one of these subspaces converges to 0 as t Ñ`8 .

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 as a combination of

-

5

0

5

-

5

0

5

-

5

0

5

t

= 0

.

0

-

10

0

10

-

10

0

10

-

5

0

5

t

= 5

.

0

-

40

-

20

0

20

-

25

0

25

-

10

0

10

t

= 10

.

0

-

100

-

50

0

50

-

100

0

100

-

40

-

20

0

20

t

= 15

.

0

(a) Conjecture 4.2: low-dimensional case.

0

.

0

2

.

5

5

.

0

7

.

5

10

.

0

12

.

5

15

.

0

t

-

0

.

2

0

.

0

0

.

2

Positive limits for clustered coordinates

0

.

0

2

.

5

5

.

0

7

.

5

10

.

0

12

.

5

15

.

0

t

-

0

.

75

-

0

.

50

-

0

.

25

0

.

00

0

.

25

0

.

50

Negative limits for clustered coordinates

(b) Conjecture 4.2: high-dimensional case.

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 is normal-so that its eigenvectors form an orthonormal basis-so we call such a V paranormal.

Definition 2. We call p Q,K,V q a good triple with multiplicity if the following conditions hold:

(i) Q J K is positive definite: Q J K ą 0 ;

(ii) V is paranormal: there exist two linear subspaces F , G Ă R d which are invariant under V , and such that F ' G ' R d , V | F ' λ Id for λ ą 0 , and ρ p V | G q ă λ , where ρ p¨q denotes the spectral radius (the maximal modulus of eigenvalues).

An example of such a V 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 λ 1 with multiplicity) . Suppose that p Q,K,V q is a good triple with multiplicity in the sense of Definition 2. Then, for any initial sequence t z i p 0 qu i Pr n s Ă R d , there exists a bounded convex polytope K Ă F such that setting H : ' pB K Yt 0 uq ˆ G , for any i P r n s , we have dist p z i p t q , H q Ñ 0 as t Ñ`8 .

6 Outlook

Several important directions regarding the mathematical theory of Transformers remain unexplored. An important extension of our work would amount to studying multi-headed Transformersborrowing the notation from Remark 4, they amount to:

-

5

0

5

5

0

5

-

5

0

5

t

= 0

.

0

-

5

0

5

5

0

5

-

50

-

25

0

t

= 5

.

0

-

5

0

5

5

0

5

-

500

0

t

= 10

.

0

-

5

0

5

5

0

5

-

10000

-

5000

0

t

= 15

.

0

-

-

-

-

For each h P r H s (corresponding to a different head ), the weight matrices Q h , K h , V h are constant. Proofs regarding clustering or convergence of the self-attention matrix for such dynamics is an 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] 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, DMS-2022448, and CCF-2106377. Y.P. is supported in part by the MIT-IBM Watson AI Lab.

References

[ABV ` 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-plane clustering. Journal of Global optimization , 16:23-32, 2000.

[Che95] Yizong Cheng. Mean shift, mode seeking, and clustering. IEEE Transactions on Pattern Analysis and Machine Intelligence , 17(8):790-799, 1995.

[CHH ` 16] Junghee Cho, Seung-Yeal Ha, Feimin Huang, Chunyin Jin, and Dongnam Ko. Emergence of bi-cluster flocking for the Cucker-Smale model. Mathematical Models and Methods in Applied Sciences , 26(06):1191-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 , 52(5):852-862, 2007.

[DCL21] Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention is not 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. Funktsional'nyi Analiz i ego Prilozheniya , 13(2):48-58, 1979.

[Gol13] François Golse. Mean field kinetic equations. Course of Polytechnique , 2013.

[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] 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 on Learning Representations , 2022.

[HZRS16a] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on 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] Ulrich 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.

[Kur75] Yoshiki Kuramoto. Self-entrainment of a population of coupled non-linear oscillators. In International Symposium on Mathematical Problems in Theoretical Physics: January 23-29, 1975, Kyoto University, Kyoto/Japan , pages 420-422. Springer, 1975.

[LCG ` 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. In 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 a multi-particle 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. AI 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 , 117(40):24652-24663, 2020.

[PR13] Benedetto Piccoli and Francesco Rossi. Transport equation with nonlocal velocity in Wasserstein spaces: convergence of numerical schemes. Acta Applicandae Mathematicae , 124:73-105, 2013.

[PR16] Benedetto Piccoli and Francesco Rossi. On properties of the generalized Wasserstein distance. Archive for Rational Mechanics and Analysis , 222:1339-1365, 2016.

[PRT15] Benedetto Piccoli, Francesco Rossi, and Emmanuel Trélat. Control to flocking of the kinetic Cucker-Smale model. SIAM Journal on Mathematical Analysis , 47(6):46854719, 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 E 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] Tamás Vicsek, András Czirók, Eshel Ben-Jacob, Inon Cohen, and Ofer Shochet. Novel type of phase transition in a system of self-driven particles. Physical Review Letters , 75(6):1226, 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:2166521674, 2020.

[VSP ` 17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz 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] E 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 arXiv:2006.04768 , 2020.

[YBR ` 20] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations , 2020.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFor $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.\n\nSTRENGTHS:\nFollowing [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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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].\n\nLIMITATIONS:\nAuthors addressed the limitation of fixed $Q,K,V$, simple structures for $V$. Authors mentioned multi-head attention as future work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nOverall, 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nPlease see the weakness section.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}