Emergence of meta-stable clustering in mean-field transformer models

We thank the reviewer for their helpful comments and suggestions. We provide answers to the points raised below. - **Q1**: The paper "Dynamic Metastability in the Self-Attention Model" is indeed connected to our work in its aim to rigorously explore metastability, as initially suggested in "A Mathematical Perspective on Transformers". However, these two papers approach the problem from complementary perspectives. Geshkovski et al. focus on well-separated configurations as initial conditions, with an interest in the behavior as $\beta\to \infty$ and relatively small token counts $N$ compared to the feature dimension $d$. In contrast, our work initiates from a uniformly distributed configuration, with a particular focus on the limit $N \to \infty$ while keeping $\beta$ finite. This leads us to investigate fundamentally different dynamical regimes and employ distinct mathematical frameworks. Bridging these perspectives would undoubtedly advance the mathematical understanding of transformer architectures, and we are interested in exploring these challenges in future work. - **Q2**: The convergence statements in our paper (see Theorem 4.2, 4.3 and 4.5) provide qualitative results highlighting the clustering phenomenon in the $N \to \infty$ limit. It is possible to extend our estimates in order to provide quantitative results showing clustering (with high probability) in the finite $N$ case. However, these estimates are typically technically quite challenging and we opted not to include them in the interest of conciseness and clarity (see also the response to Q3 by reviewer MWwW). Finally, we have understood this reviewer’s question as referring to discreteness in terms of $N$ (number of tokens). If discrete is interpreted as “having a discrete number of layers” (i.e. discrete time in our setup), the validity of our method is justified by Euler’s method for a sufficiently large number of layers and appropriate scaling of the parameters. We hope that we interpreted the reviewer’s comment correctly, but we would be happy to discuss this (and all other points) further during the discussion period. - **Q3**: We expect transformers to sequentially operate in all three phases presented in the paper, given enough “time” (depth of the architecture) and large enough $N$. One can, however, classify these phases based on different measures of operational importance. If we examine token movement in terms of spatial displacement, the majority of this occurs during a general version of phase 3, where tokens converge closely to unstable configurations. If we instead consider the duration (in terms of time/layer count) that tokens remain in the neighborhood of certain configurations, such as saddles in the dynamical landscape, then phases 1 and 2 dominate, as they involve longer periods. - **Q4**: We agree on the importance of justifying the omission of the remainder term $R$. To address this, starting from Proposition C.17 in the appendix, we provide a rigorous analysis of the growth of $R$ using a Lagrangian approach. This approach enables us to bound the evolution of $R$ in negative Sobolev norms, demonstrating that by the end of the quasi-linear phase, the distance between the original solution $\mu_N$ and the dominant mode $f$ vanishes as $N \to \infty$. In particular, $R$ becomes negligible at every finite time horizon following the quasi-linear phase. - **Q5**: As the reviewer suggests, the limit $\alpha \to 0$ is specifically considered because when $N \to +\infty$, the definition provided on lines 321-322 tends toward zero in probability. This is due to the fact that $(\hat{\rho\_0})\_{k\_{\text{max}}} = O(\|\rho_0\|_{H^{-1}})$, which justifies our asymptotic treatment of this parameter in the mean-field context. Attention should be paid, however, to the fact that $f^\alpha$ is only a technical tool in our analysis, representing (a sequence of) leading order approximations of the measure at the exit of the linear phase for increasing values of $N$.

We appreciate the reviewer’s feedback and we respond to their questions below. - **Q1**: The existence of a meta-stable phase implies that the network retains a richer and more complex representation of the data because it avoids the immediate collapse into a single cluster, a phenomenon often associated with a loss of expressivity. This collapse has been identified in trained Transformer models and is viewed as a negative property by practitioners due to its limitation on the network's ability to distinguish between different tokens or features in the data [1,2,3]. In contrast, the metastability phenomenon ensures the persistence of multiple clusters for extended periods, allowing the model to maintain diverse intermediate representations. This diversity supports the encoding of more complex relationships and structures in the data, contributing to the model's overall ability to learn more sophisticated patterns. Our work aims to rigorously describe and prove the appearance of this desirable metastability behavior. - **Q2**: The main takeaway for practitioners is that even in our focused analysis of transformer dynamics, the core hyperparameters (such as the number of tokens $N$, the number of layers, the temperature, and the feature dimension $d$) are deeply interdependent. Our findings highlight that adjusting these parameters independently may disrupt the delicate balance required for the model to fully express its dynamics and achieve optimal performance, emphasizing the importance of tuning these parameters in a coordinated way. In particular, the fact that the duration of the metastable phase scales as $O(\ln(N))$ may suggest, in a regime of large context window size in LLMs, a choice of depth of the transformer architecture allowing to maintain ideal expressivity and performance. **References:** 1. Dong, Yihe, Jean-Baptiste Cordonnier, and Andreas Loukas. "Attention is not all you need: Pure attention loses rank doubly exponentially with depth." International Conference on Machine Learning. PMLR, 2021. 2. Zhai, Shuangfei, et al. "Stabilizing transformer training by preventing attention entropy collapse." International Conference on Machine Learning. PMLR, 2023. 3. Bao, Han, Ryuichiro Hataya, and Ryo Karakida. "Self-attention Networks Localize When QK-eigenspectrum Concentrates." arXiv preprint arXiv:2402.02098 (2024).

We thank the reviewer for their detailed review and the points raised, which have improved the paper. We address each of them in the following responses. - **W1+W2+W3**: Thank you for your careful reading and for pointing this out, we apologize for these oversights. To clarify, the dominant mode developed in the first phase continues to grow exponentially in time and its periodicity is maintained over time intervals of length $O(\ln N)$. We updated the wording at lines 69-70 to reflect this more precise characterization, and made the suggested corrections at lines 504 and 233. - **Q1**: The delicate aspect of this assumption lies in the fact that a $t \to \infty$ limit is taken: without this assumption, our results show that even in the presence of small perturbations of the uniform measure (activating potentially all Fourier modes), the linear phase naturally isolates a dominant mode, allowing us to recover equation (6) and, under Assumption 3, convergence to $k_{max}$ clusters (in the joint limit $N \to \infty$ and $t_N \to \infty$). In the $t \to \infty$ limit, we proceed to discuss possible extensions of our Assumption 3 to the multimodal setting: - In general, for an initial condition with a distribution that combines multiple modes, the dynamics restrict to distributions with periodicity given by the greatest common divisor $k_{gcd}$ of the modes. Proving that all the initial conditions perturbing the uniform measure with $k_{gcd}$ periodicity converge to $k_{gcd}$ equally spaced clusters is, however, harder than proving that Assumption 3 holds, since the set of initial conditions on which the convergence property should be proven is strictly larger than the ones in Assumption 3. Note that it is easy to construct fixed points of the dynamics that are composed of multiple, non-identical clusters. It is unclear to us what the basins of attractions of these stationary distributions are and, as a consequence, it is difficult to develop an intuition on the veracity of the more general version of Assumption 3 considered in this paragraph (more difficult than in the case of Assumption 3 as written in the paper, which was corroborated by numerical simulations). - On the other hand, one could ask themselves the converse question: are there initial conditions perturbing the uniform measure with modes that do *not* have $k^*$ periodicity but still leading to $k^*$ clusters as $t \to \infty$, i.e., “creating” periodicity? While we strongly suspect that this phenomenon cannot occur in our system (if at all, it can only occur in a set of initial conditions of measure 0), a proof of this fact would require a more in-depth understanding of the PDE dynamics. The *direct* implications of the above on our analysis would be relatively limited, given that in the linear phase we know that only one mode is generically selected by the dynamics. However, indirectly, a more in-depth understanding of the structure of the basins of attraction of this PDE would have far-reaching implications in the study of its dynamics, for instance in allowing to prove that Assumption 3 actually holds.

- **Q1 + Q2 + W**: We sincerely appreciate the reviewer's thoughtful comments regarding the theoretical assumptions in our work. While our model makes several simplifying assumptions, these choices enable rigorous analysis of fundamental transformer dynamics that would be difficult to isolate, at present, in more complex settings. We note that the case $Q = K = V = \lambda Id$ is not the only situation in which our results apply. Indeed, our results extend immediately to the case $Q^T K = \lambda Id$. Furthermore, the gradient flow structure is preserved under the assumption $Q^T K = V$ for a general choice of $V$, suggesting our analytical tools could be extended to these settings. Our interest in developing a theoretical analysis of clustering phenomena is also motivated by empirical observations in the reference work [1], which performed experiments on more realistic transformer architectures [2], and is further supported by a long line of research on token uniformity, oversmoothing, and rank collapse (see, for example, the references in the answer to reviewer WDSo or in [1]). This suggests that, while the mechanisms we analyze theoretically may not capture all aspects of transformer behavior, they isolate key components that are fundamental to the dynamics of transformers. We have briefly summarized this discussion in the conclusions, specifically on lines 522–532, of the new version of the paper for clarity. While we are actively investigating extensions including more general weight matrices and MLPs, we believe that, by clearly isolating specific dynamical phenomena in a tractable setting, our current theoretical contribution could provide a foundation for understanding how additional complexities interact with these base dynamics, thereby representing an important stepping stone toward the open problem of analyzing more general architectures. **References:** 1. Geshkovski, Borjan, et al. "A mathematical perspective on transformers." arXiv preprint arXiv:2312.10794 (2023). 2. Lan, Z. "Albert: A lite bert for self-supervised learning of language representations." arXiv preprint arXiv:1909.11942 (2019).

Thank you for your thoughtful and detailed review. We appreciate your insights and feedback. Below, we address each of your questions and comments. - **Q1**: You are correct, and we appreciate your careful reading. The coefficient $\frac{1}{N}$ in equation (USA-MF) is indeed redundant, as it is implicitly included in the empirical distribution. This typo has been corrected at lines 148-149. The (SA-MF) representation remains appropriate due to the definition of $Z_{\beta, \mu}$ which ensures normalization. - **Q2 + W3**: We understand the concern about the well-definedness of $T_1$​ when the parameter becomes negative. This point is addressed by noting that $\rho_0 = \mu_N - \mu_{\infty}$ implicitly depends on $N$ and that, as a result of the Central Limit Theorem, when $N \to \infty$, the $H^{-1}$ norm of $\rho\_0$ scales as $N^{-1/2}$ (Lemma E.3). Consequently, even with large values of $N$, this scaling ensures that the linear phase remains well-defined. We have clarified this point in the revised version at lines 308-309 and explicitly specified that the dependence of $\rho_0$ on $N$ was omitted for simplicity at lines 302-303. - **Q3**: Our approach does not provide a specific convergence rate for the $W_1$​ distance to keep the proof focused and readable. Instead, we rely on a qualitative argument: weak convergence on compact metric spaces implies $W_1$ convergence. However, a quantitative bound for the convergence rate could indeed be provided: the steps in our proof could yield an estimate in a negative Sobolev norm setting, showing a polynomial rate in $\frac{1}{N}$. We expect that by adapting techniques from [1], one could derive from such results a quantitative upper bound in $W_1$. Nonetheless, we believe that the technical details required by such an adaptation would obscure the main insights of the paper, and we opted to omit such a bound in the interest of readability and conciseness. - **W1**: While our model simplifies certain aspects, these assumptions allow us to isolate key phenomena in transformer dynamics, such as metastability and clustering, which would be difficult to study rigorously in more complex settings. We believe our analysis provides a foundation for future work, both theoretical and empirical, on more general transformer architectures. We elaborate further on these considerations in our response to Reviewer bdQd and have added a comment in the conclusions of the revised version of the paper (lines 522–532). - **W2 + W4**: We appreciate your feedback on the notation, particularly around Theorems 4.2 and 4.3. We have now incorporated the definition of $(\hat{\rho_0})\_{k\_{max}}$ as the Fourier coefficient of $\rho\_0$ with index $⁡k_{max}$​ (lines 314-315). Since we present every statement in that section for $d=2$, we have also clarified the identification of $S^1$ with the interval $[0, 2\pi]$ in defining $\mu_{\text{cluster}}$​ (lines 296-297). Concerning Theorem 3.1, we have emphasized on lines 62-63 that it justifies and supports the more refined metastability analysis carried out in the rest of the paper. However, we stress that this bound is later combined, during the nonlinear dynamical phase 3, with the more refined estimates we provide to complete the dynamical picture presented in our paper. **References**: 1. Peyre, Rémi. "Comparison between W2 distance and Ḣ− 1 norm, and localization of Wasserstein distance." ESAIM: Control, Optimisation and Calculus of Variations 24.4 (2018): 1489-1501.

Thank you for the clarification and for revising the manuscript. My concerns have been addressed, and each statement is now clear and understandable. While I still find the model somewhat restrictive, I acknowledge that this work provides an interesting interpretation of the transformer model. Therefore, I have increased the score to 8.