The emergence of clusters in self-attention dynamics

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.

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. 2. 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. 3. 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)." 4. Please strengthen the contribution with layer normalization.

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). 2. 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?) 3. 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. 4. 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? 5. Please comment on the shrinking convex hull/ polytope (since it is finite) (Appendix C.1.1). What about infinite number of tokens? 6. 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? 7. Please comment on the effects of skip and non-skip connections [DCL21].

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Authors addressed the limitation of fixed $Q,K,V$, simple structures for $V$. Authors mentioned multi-head attention as future work.

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". They 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. Across 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.

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. The 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. Furthermore, 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. The paper is very math-heavy for a general machine learning audience, and as such it is often hard to follow.

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? 2. 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? 3. 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

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.

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: i) the convergence guarantees of self-attention matrix to a low rank matrix. ii) the characterization of convergence of initial inputs to final outputs as a function of query, key and value matrix, In 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.

Strengths

This is a very well written paper. It studies a very important problem of understanding the core of transformers, i.e., self attention. Originality: The paper has quite an original idea and approach to solve the problem. Quality: This is a high quality paper that can inspire important developments on the path of demistifying transformers. Clarity: The paper is a joy to read. Significance: 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. 1. 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. 2. 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? 3. 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. 4. I would appreciate if authors could add some more intuition on what is geometric meaning of Definition 1. 5. 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.

Questions

Please see the weakness section.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

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.

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). The authors first study the one dimension case and show that the self-attention matrix converges to a low-rank boolean matrix. They 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. The 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. All the theoretical results are numerically illustrated.

Strengths

1. The paper is very clear and well written. 2. The overall contribution is strong, as it is the first paper provably showing the emergence of clusters in self-attention only Transformers. 3. More specifically, each theorem, (well summarized in Table. 1) is an interesting result of clustering in interacting particles systems. 4. Most assumptions on (Q, K, V) are clearly discussed. 5. 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. 2. 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$ ? Typo: l. 242, I believe a $Q$ is missing in the quadratic form.

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 ? 2. 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 ? 3. 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 ?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Yes.

Official Comment by Reviewer 8epq

I would like to thank the authors for their rebuttal; they have addressed my questions. I am still positive towards this paper and I think its contribution is significant. I am keeping my score to 6. I hope the authors incorporate some of the discussion in the camera-ready version of the paper.

Thanks

I thank the authors for these clarifications. I would appreciate if the authors can revise the paper in the light of the above discussion. I am happy to maintain my score for the paper.