A Phase Transition between Positional and Semantic Learning in a Solvable Model of Dot-Product Attention

Authors: hugo  cui
freya  behrens
lenka  zdeborová
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-08 07:52:33

Abstract: Many empirical studies have provided evidence for the emergence of algorithmic mechanisms (abilities) in the learning of language models, that lead to qualitative improvements of the model capabilities. Yet, a theoretical characterization of how such mechanisms emerge remains elusive. In this paper, we take a step in this direction by providing a tight theoretical analysis of the emergence of semantic attention in a solvable model of dot-product attention. More precisely, we consider a non-linear self-attention layer with trainable tied and low-rank query and key matrices. In the asymptotic limit of high-dimensional data and a comparably large number of training samples we provide a tight closed-form characterization of the global minimum of the non-convex empirical loss landscape. We show that this minimum corresponds to either a positional attention mechanism (with tokens attending to each other based on their respective positions) or a semantic attention mechanism (with tokens attending to each other based on their meaning), and evidence an emergent phase transition from the former to the latter with increasing sample complexity. Finally, we compare the dot-product attention layer to a linear positional baseline, and show that it outperforms the latter using the semantic mechanism provided it has access to sufficient data.

Main paper: 
 introduction: 
 Recent years have seen an upheaval in our ability to learn and implement complex tasks from textual data. Instrumental in these advances is the use of self-attention layers [1] , which provide an efficient method of extracting information from sentences -both the information encoded in the ordering (i.e. positions) of the words, and that encoded in the meaning (i.e. semantics) of the words. In theory, attention layers can learn to leverage both types of information, by having tokens attend to each other based on their respective positions (a mechanism called positional attention in [2] ) and/or respective meanings (henceforth referred to as semantic attention). In this paper, we aim at a theoretical understanding of the emergence of these different mechanisms in attention layers, and the transitions therebetween.Many empirical studies have provided evidence for the emergence of specific algorithmic mechanisms (abilities) in the learning of language models that lead to qualitative improvements of the model capabilities [3, 4, 5] . By reverse-engineering trained models into human-interpretable components [6, 7, 8] a growing body of work on mechanistic interpretability aims to empirically understand which precise algorithmic mechanisms a neural network learns. Such investigations have demonstrated that attention layers are able to implement a wide range of different algorithms, even for the same task, using both positional and semantic attributes of the inputs. We offer a particularly simple illustration of this idea in Appendix A, where we show that for a sequence modelling task involving counting different algorithmic mechanisms co-exist, each corresponding to a distinct local minimum of the empirical loss. Out of these possibilities, the precise implementation that a model learns during training is jointly affected by its architecture [9] , the training procedure itself [7, 4, 10] and the available data [11, 12, 13] .It remains an open question how to theoretically characterise the conditions under which a given behaviour emerges in the model leading to said qualitative improvements.From a theoretical perspective, even the nature of this type of algorithmic emergence is unclear. It is not known whether it is simply a fast but smooth change in performance or whether the emergence is due to a sharp boundary between fundamentally different regimes of learning [14] . In our work, we take inspiration from physics, where a similar theoretical question about the nature of phase transitions was posed a century ago for models of interacting particles, such as the famous Ising model describing ferromagnetism [15, 16] . In the limit of infinitely many particles, it was shown that it is possible to theoretically deduce sharp discontinuities in some properties of interest (e.g. the magnetization of a magnet), delineating qualitatively very different regimes (e.g. magnetized or not). While mathematically, a large size limit needs to be considered to confirm the existence of sharp phase transitions, this asymptotic theory usually closely matches simulations, even for relatively moderate finite sizes.In the theory of feed-forward fully connected neural networks, phase transitions in the network's generalisation ability as more training samples are available were studied as early as in [17, 18] , and their existence was proven mathematically rigorously in [19] . In those works, the limit of many particles corresponds to taking the number of training samples and the dimensionality of the data to infinity at a fixed ratio. These theories rely on the property that macroscopic quantities of interest, such as the test error, become concentrated and deterministic in the high-dimensional limit. A dimension-free set of equations is then derived which predicts these deterministic quantities. Since then, a plethora of works along these lines both in statistical physics of phase transitions and in the theory of feed-forward neural networks have continued to study these phenomena, see e.g. [20, 21, 22, 23] .In this work, for the first time, we bring this type of study to the analysis of neural networks with attention layers. While several previous theoretical studies of the attention mechanism considered some type of high-dimensional limit [24, 25, 26, 2] , none of them identified a phase transition between different types of mechanisms that are implemented by the attention layer. Simultaneously, the finite-dimensional real-world models that are the focus of works in mechanistic interpretability do not lend themselves to a tractable definition of a high-dimensional limit, which is necessary to identify a phase transition theoretically, as explained above. We aim to fill this gap by introducing and analysing a tractable model that permits such a sharp high-dimensional characterisation for attention layers (see Fig. 1 ).In particular our contributions are:• We describe a model with a single self-attention layer with tied, low-rank query and key matrices. On Gaussian input data and realizable labels, we show that this model exhibits a phase transition in terms of sample complexity between a semantic and a positional mechanism.• As the main technical result, we analyse this model in the asymptotic limit where the embedding dimension d of the tokens and the number n of training samples grow proportionally. We provide a tight closed-form characterization of the test error and training loss achieved at the minima of the non-convex empirical loss. Using this high-dimensional characterization, we locate the positional-semantic phase transition, thus providing the first theoretical result about the emergence of sharp phase transitions in a model of dot-product attention.• We contrast the performance of the dot-product attention layer with that of a linear model, which can only implement positional mechanisms, and show how the former outperforms the latter once it learns the semantic mechanism, highlighting the advantage of the attention architecture for this task, when there is a sufficient amount of training data.In Section 2, we discuss further related work. Section 3 defines the general version of a solvable model of tied low-rank attention, and Section 4 provides a tight characterization of the global minimum of its empirical loss. Section 5 analyses a concrete instance of dot-product attention and demonstrates that in this case the global minimum corresponds to either a semantic or positional mechanism, depending on the training data and task, with a phase transition between them. We conclude with a discussion of the limitations of our analysis in Section 6. 
related work: 
 Theory of attention Attention models have been the object of sizeable theoretical scrutiny in recent years, with a growing body of work investigating various aspects such as their expressivity [25, 27, 28] , inductive bias [29, 30, 31] , training dynamics [2, 32, 33, 34] , and in-context learning [26, 35, 36, 37, 38] . [39] and [25] analyze models with frozen non-trainable queries and keys, under the lens of signal propagation in such frozen models, or of their expressivity. [2] similarly studies the learning of the value matrix and positional encodings only, fixing keys and queries to identity, and shows how a transformer with a single attention head can learn spatial structure with a purely positional attention mechanism. The works of [29, 37] analyze the learning of a single layer of attention, with trainable queries and keys, assuming linear or ReLu activations -instead of the standard softmax activation. [33] provide convergence bounds for non-linear transformer models with a single attention layer, with trainable queries and keys. Because these studies are not tight, they do not allow to capture sharp changes in the behaviour of attention mechanisms such as phase transitions, and cannot for the same reason provide a theoretical description of the sudden emergence of new algorithmic mechanisms. A first tight description was provided in [24] , in the context of learning a high-dimensional graphical model with a single layer of factored attention, leveraging its formal equivalence to a linear and convex learning problem. On the other hand, this model does not exhibit any emergent phenomenon. [38] analyze how induction head mechanisms can be learnt from gradient descent on the population loss, i.e. when an infinite amount of data is available. The present manuscript conducts a tight analysis of the non-convex learning of a non-linear attention model with trainable tied queries and keys from a finite train set, thereby allowing the description of sharp phase transitions in sample complexity in the behaviour and performance of the model.To combine the positional and semantic information in textual or general sequential data, a variety of models and input encodings have been explored. Many approaches are based on autoregressive models, e.g. recurrent architectures [40] , where the positional information is provided implicitly by the order in which the input is processed. While some transformers can leverage implicit positional information through causal masks in training [41, 42, 43] , in principle a dot product attention layer requires an explicit encoding of positional information as it views the input sequence in parallel, as a bag of words [1] . Several works experimentally explore different types of positional encodings with the goal of improving the downstream task performance [44, 45] .In this work, we provide a tractable model to quantify the generalization error of a single layer of attention in the presence of positional encodings.In supervised learning with feed-forward fully connected neural networks, phase transitions in sample complexity were identified in settings where the data consists of random Gaussian samples, and the labels are generated by a target neural networks with random weights. For a single-layer perceptron and a variety of teacher weights distributions and activation functions, a discontinuity of the optimal test error as the number of samples increases was established in [18, 17, 19] . For a two-layer neural network, [21, 20] evidenced a specialization threshold in the sample complexity below which linear regression matches the optimal test accuracy, and above which a strictly better accuracy can be reached. To our awareness, phase transitions in neural networks with attention layers have not been studied theoretically yet. x ℓ ∼ N (0, Σ ℓ ) with covariance Σ ℓ ∈ R d×d .In the following, we denote the probability distribution of x as p x . Note that while this sentence model does not involve in itself statistical correlations between tokens, the task (target function) will entail interactions between different tokens. While more general data models involving inter-token correlations can also be readily analyzed such analyses come at the price of much more intricate analytical formulae. We thus choose for clarity to restrict the discussion to this simple instance, which already displays rich phenomenology, as will be explored in Section 5. We defer a discussion and an analytical treatment of the general case to Appendix C.The target function (teacher) is assumed to be of the formEQUATIONfor a function 1) should be interpreted as the target attention matrix, which mixes the tokens of the input x. This attention matrix is parametrized by the target weights Q ⋆ ∈ R d×rt .T : R L×t → R L×L . The term T [ 1 / √ dxQ ⋆ ] ∈ R L×L in (We consider the learning of the target (1) by a single attention layerEQUATIONIn (2), p ∈ R L×d is a fixed matrix, corresponding to positional encodings, and Q ∈ R d×rs is a trainable weight matrix. We denote subsequently p ℓ ∈ R d the ℓ-th row of p. Like the target (1), the parametric function (2) takes the form of a data-dependent attention matrixS [ 1 / √ d(x + p)Q] ∈ R L×Lmixing the tokens of the input x. Note that, compared to the usual attention mechanism [1] employed in practice, (2) corresponds to setting the value weights to identity, and -since ( 2) is parametrized by a single matrix Qtying the key and query weights. While the assumption of tied weights is not strictly necessary, it makes for simpler and more interpretable analytical characterizations, and is thus considered in this work for clarity. We provide in Appendix C a full analysis of the untied architecture for completeness.We study the learning of the attention layer (2), when a training set D = {x µ , y(x µ )} n µ=1 with n independently sampled sentences {x µ } n µ=1 , and the associated labels {y(x µ )} n µ=1 , is available. The target (1) can be learnt by carrying out an empirical risk minimization:EQUATION)The performance of the resulting trained model f Q is measured at test time by the mean squared error (MSE)EQUATION)4 Closed-form characterization of the training lossHigh-dimensional limit We analyze the learning problem (3) in the limit where the embedding dimension d and the number of training samples n jointly tend to infinity, while their ratio α = n /d (henceforth referred to as the sample complexity) stays of order Θ d (1) . We further assume the sentence length L, the ranks r s , r t of the weights Q, Q ⋆ , and the norm of the positional embeddings ∥p∥, to be Θ d (1). This limit has been considered in a stream of previous works (e.g. [46, 47, 48] ) and allows to derive closed-form characterization of the ERM problem (3), which we present in the next section. It also exhibits a particularly rich learning phenomenology which we further explore in Section (5) . Finally, let us comment briefly on the assumption that r s = Θ d (1) , which in words implies that the weight matrix Q is low-rank. While primarily motivated by technical limitations here, it is worth noting that low-rank weights are also considered in machine learning practice, in the context of model compression [49] or fine-tuning [50] . EQUATIONIn words, Assumption 4.1 guarantees that all parameters of the problem admit well-defined limits in the considered asymptotic limit, with the further assumption that the covariances {Σ ℓ } L ℓ=1 of the different tokens can be jointly diagonalized. We are now in a position to state the main technical result of the present work. Result 4.2. Under Assumption 4.1, in the limit n, d → ∞, ∥p∥, n /d, L, r s , r t = Θ d (1), the summary statisticsEQUATIONconcentrate in probability, and are solutions of the set of finite-dimensional self-consistent equationsEQUATIONIn (7), U = {u ℓ } L ℓ=1 and Ξ = {ξ ℓ } L ℓ=1 , with u ℓ ∼ N (ξ ⊤ ℓ q -1 /2 ℓ θ ℓ , ρ ℓ -θ ⊤ ℓ q -1 ℓ θ ℓ )and ξ ℓ ∼ N (0, I rs ), and • ⊗2 denotes the outer product of a vector with itself. Finally, the resolvents {prox(Ξ, U ) ℓ } L ℓ=1 are defined as the minimizers of the Moreau envelopeM(Ξ, U ) = inf z1,...,z L L ℓ=1 Tr V -1 ℓ x ℓ -q 1 /2 ℓ ξ ℓ -m ℓ ⊗2 +Tr S(Z)ρ Σ S(Z) ⊤ -2Tr T(U )ρ Σ S(Z) ⊤ .We noted Z ∈ R L×rs (resp. U ∈ R L×rt ) the matrix whose rows are z ℓ (resp. u ℓ ) and:EQUATIONIn the same limit, the test error (4) converges in probability toEQUATION)where the average bears on h ∈ R L×rs , h ⋆ ∈ R L×rt with independent rows with statisticsEQUATIONFinally, the training loss ϵ t converges in probability toEQUATION)Result 4.2 provides a tight asymptotic characterization of the test error ϵ g (10) and the training loss ϵ t (12) , as a function of a finite set finite-dimensional summary statistics {q ℓ , V ℓ , m ℓ , θ ℓ } L ℓ=1 , thereby providing a finite-dimensional description of the high-dimensional learning problem (3). These summary statistics are further characterized in closed form by the set of equations (7) , in terms of the solution of a low-dimensional minimization problem (91). Intuitively, this low-dimensional problem may be viewed as a form of an effective loss averaged over the finite training set. In practice, the solution of the self-consistent equations (7) can be found by numerically iterating the equations until convergence. The resulting summary statistics can then be used to evaluate the expressions (10) and (12) to reach the asymptotic limits of the test error and train loss. In Section 5 we evaluate exactly these functions to understand the different minima in the models empirical loss landscape. Similar sharp asymptotic characterizations have been derived in the literature for other neural network architectures trained with ERM, in particular generalized linear models (see e.g. [51, 52, 48] ) and auto-encoders [47] .The derivation of Result 4.2 is provided in Appendix B, and is exploiting a mapping of the model (2) to a (variant of) a Generalized Linear Model (GLM) [53, 54] . The summary statistics characterized by the equations (7) (often called state evolution [55] in this context) asymptotically describe the fixed points of a Generalized Approximate Message Passing (GAMP) algorithm [56] , which we state in Appendix B. The fixed points of GAMP in turn correspond to critical (zero-gradient) points of the non-convex empirical loss landscape (3) . Therefore, while Result (4.2) is stated as a characterization of the global minimum of (3), which is the main concern of the present work, solutions of (7) also describe local minima and saddles. This strategy has been used in many recent work to study asymptotics of a large number of highdimensional problems, see e.g. [57, 58, 59, 60, 48] . We note, however, that we importantly assume the point-wise convergence of GAMP. While we believe that this point can be rigorously justified, it would require a considerable amount of work -in particular, the usual rigorous tools used in recent works fall short because of the non-convexity of the loss-and we leave this point for future studies (see the discussion in Appendix C, where we also provide an alternative derivation using the replica method from statistical physics [61] ). Finally, we mention that while Result 4.2 is presented for an ℓ 2 regularization of the empirical loss (3) for clarity, similar results can be reached for generic convex regularizers, and are presented in Appendix C. In the following section, we explore the phenomenology uncovered from the study of the equations (7) of Result 4.2, for the special case of dot-product attention. 2) trained using gradient descent, initialized previously for the overlaps. All experiments were performed in d = 1, 000 with the Pytorch implementation of full-batch gradient descent, for T = 5, 000 epochs and learning rate η = 0.15. All points are averaged over 24 instances of the problem each.Question: Does the paper discuss the limitations of the work performed by the authors?Justification: There is a dedicated section on limitations.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [No] • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
positional-to-semantic phase transition: 
 Rank one dot-product attention In the following, we turn to a special case of tied low-rank attention (2) -namely a dot-product attention layer, which is the example from Fig. 1 :EQUATION)As in (2), we allow for positional encodings p in the dot-product attention parametrization (13) . We further consider a specific case of target attention matrix (1) of the formEQUATIONwith A ∈ R L×L a fixed matrix. In (14) , the parameter ω ∈ [0, 1] tunes the relative strength of the dot-product term and the fixed matrix term, and interpolates between a fully positional and a fully semantic task:• For ω = 0, the target reduces to the first dot-product term, and is purely semantic, in that the i, j-th element of the score matrix softmax( 1 /dxQ ⋆ Q ⊤ ⋆ x ⊤ ) only depends on the tokens x i , x j and not explicitely on their respective placements i, j inside the sentence. To learn satisfyingly the target, the learning model thus has to learn a semantic attention matrix.• For ω = 1, the target reduces to the second fixed term A in (14) . The attention matrix A associated thereto is purely positional, in the sense that A ij is a function of i, j but not of x i , x j . To complete the learning task, a positional mechanism then needs to be learnt.The parameter ω thus allows to tune the amount of semantic/positional content in the target (14) , and thus the extent to which the task requires the model to implement semantic attention (small ωs) or rather positional attention (large ωs). In the following, for definiteness, we further assume r s = r t = 1 and set Q ⋆ to be a fixed random Gaussian vector drawn from N (0, I d ), and choose the positional encodings p 1 = -p 2 = 1 d . Finally, for simplicity, we consider sentences with two tokens L = 2 and isotropic token covariancesΣ 1 = Σ 2 = σ 2 1 d .Semantic and positional mechanisms The summary statistics θ ℓ , m ℓ describing the global minimizer of the empirical loss minimization (3) of the dot-product attention (13) on the target The solution of the system of equations ( 7) is not unique, and different stable fixed points describe different corresponding critical points of the non-convex empirical loss landscape (3) . In practice, we notably find two solutions of (7) , corresponding to two minima associated with different mechanisms implemented by the dot-product attention (13) when approximating the target ( 14):-Positional solution One solution of (7) correspond to vanishing overlap θ = 0 between the trained weights Q and the semantic target weights Q ⋆ , and non-zero m > 0 between the trained weights Q and the positional embedding p 1 = -p 2 . Consequently, the argument of the dot-product attention Q(x + p) has a sizeable token-independent -thus positional-contribution Qp, alongside a token-dependent semantic part Qx. Because of the positional terms, the resulting learnt attention attention matrix softmax( 1 /d(x + p) Q Q⊤ (x + p) ⊤ ) implements a partly positional mechanism.-Semantic solution Another solution of the system of equations ( 7) is associated with a vanishing overlap m = 0 between the learnt weights Q and the positional embeddings, and a finite overlap θ > 0 with the target weights Q ⋆ . Therefore the resulting learnt attention matrix softmax(1 /d(x + p) Q Q⊤ (x + p) ⊤ ) ≈ softmax( 1 /dx Q Q⊤ x ⊤) is largely semantic.While the system of self-consistent equations ( 7) may admit other solutions, we did not find solutions with lower training loss than the two aforedescribed fixed points. Which of these solution corresponds to the global minimum -and thus the solution of the optimization (3)-depends on the sample complexity α and the positional/semantic parameter ω (14), as we describe in the following subsection.Positional-to-semantic phase transition For a fixed parameter ω in ( 14), an analysis of equations ( 7), further detailed in Appendix C, reveals that for a sizeable range of ω, in the probed setups, there exists a threshold α c for the sample complexity so that • For α < α c , the global minimum of (3) corresponds to a positional mechanism, and is described by the positional solution of (7) of Result 4.2 with θ = 0, m > 0. • For α > α c , the global minimum of (3) corresponds to a semantic mechanism, and is described by the semantic solution of (7) of Result 4.2 with θ > 0, m = 0.The dot-product attention thus displays a phase transition in sample complexity from a positional mechanism to a semantic mechanism, implementing the simpler positional mechanism when having access to small amounts of data, and only learning the semantic content of the target (14) when presented sufficient data. The critical sample complexity α c generically grows with the positionality ω of the target function (14) , as the semantic content -i.e. the first term of ( 14)-is less apparent for larger ω, and thus requires larger amounts of data to be identified and approximated by the dot-product attention (13) . An example for ω = 0.3 is given in Fig. 2 . In Fig. 3 (center) the difference in training loss ∆ϵ t between the positional and semantic solutions of ( 7) is represented, alongside the difference in training loss at convergence experimentally reached by gradient descent. For small (resp. large) sample complexity α < α c (resp. α > α c ), the training loss of the positional (resp. semantic) minimum is lower, and thus corresponds to the global minimum.Experimentally, the positional minimum can be reached for α < α c via gradient descent by initializing the weights Q of the attention (13) close to the positional embedding p 1 . By the same means, the semantic minimum can be reached from an initialization at the teacher weights Q ⋆ (14) . Henceforth, we refer with a slight abuse to the minimum experimentally reached from a positional (resp. semantic) initialization as the positional (resp. semantic) minimum, even when it is not global. Note that importantly the semantic initialization is informed in nature, in that it necessitates the knowledge of the target parameters Q ⋆ . Note that even though the minima we characterize analytically are fixed points of gradient descent, a precise analysis of the dynamics of gradient descent from an agnostic (random) initialization, and ascertaining whether the optimizer reaches the global minimum, is an interesting question but falls out of the scope of the present manuscript -which is an analysis of the loss landscape. We however conduct numerical experiments from a random initialization of Q in Appendix E. 4 , and show that the dynamics may reach either of the local minima, or get stuck in a different one.In Fig. 2 , we compare our analytical characterizations for different metrics at the global mimimumthe summary statistics θ, m (middle), and the test MSE (right)-, with the corresponding experimental estimates, obtained by optimizing (3) with the Pytorch implementation of gradient descent, from a positional (resp. semantic) initialization for α < α c (resp. α > α c ), displaying overall good agreement. In Fig. 3 (left) and Appendix E.1 we further verify that in the scaling limit of our analysis, namely n, d → ∞ for α = O(1), the agreement improves with growing n, d.The dot-product attention ( 13) thus implements a semantic mechanism when learning from sufficient amounts of data. The learning of the semantic mechanism by the dot-product attention at sample complexities α > α c corresponds to a noticeable drop in the generalization MSE as can be observed in Fig. 2 , right. But just how essential is the learning of semantic mechanism in the ability of the dot-product attention to generalize well? We explore this question in the following subsection, by comparing the dot-product attention (13) to a purely positional attention model.In this subsection, for the same target (14) , we contrast the dot-product attention model (13) , analyzed in the previous subsections, to the baseline given by a linear layerEQUATIONwith a trainable weight matrix W ∈ R L×L . As for the dot-product attention (13), we consider the case where the weights Ŵ are learnt by minimizing the empirical riskEQUATION)The model ( 15) is a natural counterpart to the dot-product architecture (13) . In (15) , the attention matrix is parametrized by a single fully-trainable matrix W , instead of being parametrized as a dot-product attention as in (13) . A seminal difference in the two parametrizations is that while the elements of softmax( 1 /dxQQ ⊤ x ⊤ ) can depend on the input tokens x, and therefore express semantic information, the elements W ij of W can only depend on the positions i, j. The model (15) can thus only implement positional mechanisms, while the dot-product attention (13) can implement both linear and semantic mechanisms, as discussed above. Finally, observe that the model ( 15) is closely related to the one analyzed by [24] in another asymptotic limit. The following result characterizes the test error achieved by the purely positional model (15):Result 5.1. In the same asymptotic limit as Result (4.2), the learnt weights Ŵ trained by minimizing the empirical risk (16) coincide with the minimizer of the population risk, and thus admit the compact expressionEQUATIONwhere the average bears over a finite-dimensional matrix h ∈ R L×t with independent rows h ℓ with statistics h ℓ ∼ N (0, ρ ℓ ), where ρ ℓ was defined in (6) in Result (4.2). We remind that T[ 1 / √ dxQ ⋆ ]corresponds to the target score matrix (1). Finally, the test MSE 1 /dLE x ∥y(x) -f Ŵ (x)∥ 2 achieved by the trained dense linear model f Ŵ (15) admits the asymptotic characterizationEQUATIONThe MSE achieved by the baseline ( 15) when learning the target ( 14) is plotted in Fig. 2 (right) as the orange solid line, alongside the MSE achieved by the dot-product attention (13) discussed in previous subsections. Remarkably, in the setup of Fig. 2 , in the positional regime α < α c when the dot-product attention relies on a positional mechanism θ = 0, m > 0 to approximate the target, the dot-product attention ( 13) is outperformed by the purely positional attention (15) ϵ g > ϵ lin g . In contrast, in the semantic regime α > α c where the dot-product attention learns the semantic mechanism, there exists a sample complexity α l ≥ α c above which ϵ g < ϵ lin g , i.e. the dot-product attention ( 13) outperforms the dense linear baseline (15) . This threshold value α l is plotted for various positionality strengths ω in Fig. 3 , alongside the positional-to-semantic threshold α c . Interestingly, we observe α l ≥ α c in all probed settings, temptingly suggesting the natural interpretation that the dot-product attention needs to learn the semantic mechanism first (at α = α c ) in order to then be able to outperform the best positional approximation f Ŵ (at α = α l ). This highlights the importance of the semantic mechanism, enabled by the dot-product parametrization (13) , in learning targets with semantic content such as (14) . 
limitations: 
 Compared to the original transformer [1] we consider a simplified model: the query and key matrices in our model are sharing weights and are of a low rank, as a value matrix we use the identity, and we employ only one head and a single layer. Further, our data model is limited to Gaussian data with sentences of 1-grams. Concerning the analysis, our characterization holds only in the high-dimensional limit, but we show that even in the finite case experimental values lie close to the theoretical prediction. Since the analysis only concerns the minima of the models loss landscape the implications for the dynamics of learning algorithms, e.g. gradient decent, are limited. This shows in our numerical experiments where we need to initialize GD close to the minima in order to arrive at them. It is yet unclear if there are scaling limits of learning algorithms which would reliably find the lower or a specific one of the minima from a random initialization.We explored the interplay between positional and semantic attention, through the prism of tied low-rank self-attention in high dimensions. In a theoretically controlled setting, we characterized the global optimum of the empirical loss, when learning a target attention layer. This global optimum was found to correspond to either a positional or a semantic mechanism, with a phase transition between the two mechanisms occurring as the sample complexity increases. We believe the present asymptotic analysis of the inner workings of attention mechanisms opens up exciting research directions. Considering alternative attention architectures (including a readout network after the attention layer, or considering cross-attention) or training procedures (such as masked language modelling, or training with causal masks), are some possible extensions which will hopefully pave the way towards a satisfactory theoretical comprehension of attention mechanisms. Finally, elucidating under which conditions either minimum can be reached by a given optimizer from a random initialization constitutes an important future research avenue. A The Histogram Task: An algorithmic toy example with a positional and semantic solutionL M D M M M A H H B L M D M M M A H H BIn this appendix, we demonstrate that for a simple counting task two qualitatively different solutions exist in the loss landscape of a simple transformer using a dot-product attention layer with positional encodings -a practical example that further motivates our theoretical investigations in the main text.One solution corresponds to a dot-product attention matrix which is largely independent of the tokens making up the input sequence, and another strongly varies based on the tokens (and thus the semantic content of the input). Both solutions achieve a test accuracy close to 100%.The training task is a sequence-to-sequence counting task, referred to as the histogram task in [6] .Given an input sequence x = [x 1 , x 2 , • • • , x L ] of length L of tokens from a fixed alphabet, the goal is to return a sequence y = [y 1 , y 2 , • • • , y L ],where each token y i is the number of occurrences of the token x i in x. In We encode the input using token embeddings and absolute positional encodings which are trained jointly with the model weights. As an architecture we consider a small transformer made up of a single layer of dot-product attention, followed by a fully connected hidden layer and with learned embeddings for both tokens and positions. For each output position, it generates logits for the L possible classes of the output alphabet; training is done using the cross-entropy loss. (Further details are provided in the remainder of this appendix).We conduct experiments where we set two different sections of the model's weights to zero at the initialization of training -removing either the model's access to positional or semantic informationand keeping the weights frozen throughout training with the Adam optimizer. After convergence, we check that the resulting configurations of weights are stable in the unconstrained loss landscape, i.e.without frozen weights. More precisely, we ascertain that these weights only change marginally when further trained with SGD on the unconstrained loss, and that the qualitative behaviour of the attention layer is retained. Our experiments demonstrate that the loss landscape of the transformer has at least two qualitatively different local minimizers (or close to minimizers), subsequently referred to as the semantic and positional solution.We inspect the learnt attention matrix for different input sequences in Fig. 4 . The positional solution corresponds to a learnt attention matrix whose i, j-th component only depend on the positions i, j, and little on the tokens occupying these positions. The attention matrix is thus almost independent of the input sequence. In fact, the attention matrix is similar to the identity. In this case, the attention layer simply serves to aggregate the other tokens uniformly, and the fully connected layer learns the counting.In contrast, the attention matrix learnt at the semantic solution displays larger i, j-th component if the tokens at position i and j are identical. In other words, identical tokens attend more to each other. This mechanism hence does not rely on the positions, but rather on the semantic content of the tokens. Both solutions and associated attention matrices thus correspond to feasible algorithms which ultimately allow the transformer to solve the downstream task.Our experimental exploration gives compelling evidence that different stable solutions exist in the empirical loss landscape of simple transformers, which correspond to different algorithmic solutions to a given task. However, it remains an interpretation of an experiment and does not allow for a precise characterization of their behaviour or of the conditions under which they are established. This stands in contrast to the example in the main paper: the model treated there is simpler, but still presents similar phenomenology and can be analyzed theoretically.In the remainder of this Appendix we discuss the implementation details such as architecture and training procedure which enabled us to exhibit the different minima for the histogram task.The code for reproducing the results is available at github.com/SPOC-group/positional-and-semanticattention.We use the histogram task as proposed in [6] . We consider sequences of fixed length L = 15. For every input sequence x = [x 1 , x 2 , • • • , x L ]In order to read the input sequence using a transformer we learn an embedding of dimension d for each of the T tokens and L positions, stored in the matricesE token = [t A , t B , • • • ] ∈ R T ×d and E pos = [p 1 , p 2 , • • • ] ∈ R L×d . We convert the input sequence x = [x 1 , x 2 , • • • , x L ] ∈ T L into an embedded sequence x ∈ R L×d ofthe same length where xi = t xi + p i . This input sequence is then fed into the first layer of the transformer. Formally, we haveEQUATIONwithW 1 ∈ R d×h , b 1 ∈ R h , W 2 ∈ R h×c , b 2 ∈ R c. The final prediction is obtained using the argmax on the logits. We have that the score matrix A and the dot-product attention mechanism isEQUATIONwhere Q, K, V ∈ R d×d and the softmax is applied row-wise. Also, for an x ∈ R d we defineEQUATIONwhere γ, β ∈ R d . We define the empirical loss for a dataset D = {x µ , y µ } n µ=1 as the average cross entropy loss for all output tokens C, i.e.EQUATIONBecause of the way in which the histogram dataset is created for a fixed input length L, it follows C = L.To obtain the positional and semantic minima, we set some weights to zero in E token , E pos , Q, K at initialization, and also freeze these zero weights during training. Note that this procedure was used in the literature to study architectural biases of varying model architectures, e.g. by [62] for convolutional neural networks.With • we denote the initialization that is taken as the default Pytorch initialization for a linear layer.For both semantic and positional initialization, we overwrite this initialization with zeros as followsEQUATIONwhere 0 d/2 is the all-zero vector of size d/2. For the positional initialization we additionally setEQUATIONand for the semantic initialization we additionally setEQUATION)where 0 d/2×d/2 is the all-zero matrix of size d/2 by d/2. In either case, the attention mechanism only has access to the semantic or positional part of the model.Training all weights except the ones frozen to zero using Adam on a dataset of size n = 35, 000, we obtain a > 99.8% test accuracy on a test set of size n = 15, 000 for both datasets. We call these parameter configurations θ pos and θ sem respectively (the local minima as referred to in the main text).In Fig. 4 , we show that the attention layer of the two models behaves in qualitatively different ways, by showing the activations of the matrix A for different input sequences.Finally, we want to verify that the parameter configurations θ pos and θ sem we obtained using the special initialization and frozen training are stable configurations in the unconstrained parameter space of the model. To show this, we perturb the given parameter configuration slightly using additive Gaussian i.i.d. noise with a scale of 0.001. Subsequently, we run SGD on this perturbed configuration in the unconstrained parameter space for another 100 epochs. We call the parameter configurations obtained after this step θpos and θsem .While this further training leads the training loss to further decrease and the test accuracy to slightly increase, each models qualitative behaviour before and after extra training remains unchanged, as shown in Fig. 5 and 6 , for the examples also used in Fig. 1 . The absolute difference in norm between θpos and θ pos (respectively θsem and θ sem ) of parameters is also small (see Table 1 ).This evidences, that both resulting parameter configurations are in flat regions of the full parameterization of the transformer model, and that these two regions are further qualitatively different in the same sense that the original ones were. We use this stability in the unconstrained loss landscape as a justification to refer to them as "local minima". θ init Adam ---→ θ θ SGD --→ θ accuracy θ ∥θ init -For simplicity, we place ourselves in the setting r s = 1 explored in Section 5 of the main text, but allow the length L of the sentences to be arbitrary, and allow a generic learning model S (2), i.e. not necessarily the dot-product attention model analyzed in Section 5. The case r s ≥ 2 follows identical derivation steps, modulo the replacement of all variables by tensor objects. We provide another alternative derivation of Result 4.2 in full generality in Appendix C, using the replica method from statistical physics. Let us note {X ℓ } 1≤ℓ≤L a series of L n × d matrices, with X ℓ corresponding to the ℓ-th rows (tokens) of each input sentence x µ stacked vertically, and normalized by √ d. We denote Xℓ ≡ X ℓ + P ℓ , where P ∈ R n×d is the matrix with all rows equal to the ℓ-th positional encoding p ℓ . Let us further define ρ ∈ R n×L×L the tensor corresponding to the sequence of n matrices 1 d x µ (x µ ) ⊤ ∈ R L×L . Finally, let us denote T ∈ R n×L×L the tensor so that the µ-th row of T satisfies y(x µ ) = T µ x µ , see equation (1) . In other words, T corresponds to the concatenation of the target attention matrices. Before detailing the derivation, we first highlight a simplifying observation. Note that a loss item can be expanded asEQUATIONwhere we used that with high probability in the considered asymptotic limit, for all 1 ≤ µ ≤ n,EQUATIONSince the first term of (27) does not depend on the weights Q, it can be without loss of generality substracted from the loss. Without loss of generality, one can thus consider the equivalent empirical risk minimization problemEQUATIONThe risks ( 29) and (3) are equivalent, and we shall use the former in the following.Finally, for argumentsEQUATIONNote that the latter part of the bracketed term corresponds to the simplified loss (27) derived in the beginning of Appendix C, which is the one we shall without loss of generality consider in the present appendix. For ease of presentation, we place ourselves under Assumption 4.1, where all the input covariances {Σ ℓ } ℓ are codiagonalizable. In the following, without loss of generality, we thus assume them diagonal, by placing ourselves in the common basis {e i } 1≤i≤d of Assumption 4.1.We are now in a position to state the AMP algorithm:Algorithm 1 GAMP Inputs : { Xℓ ∈ R n×d } L ℓ=1 , T ∈ R n×L×L , ρ ∈ R n×L×L Initialize Q0 =∼ N (0, I d ), ĉ0 = I d , {f 0 ℓ = 0 n } L ℓ=1 for t ≤ t max do ∀1 ≤ ℓ, κ ≤ L, V t ℓκ = ( Xℓ ⊙ Xκ )ĉ t ∀1 ≤ ℓ ≤ L, ω t ℓ = Xℓ Qt - L κ=1 V t ℓκ f t-1 κ ∀1 ≤ ℓ ≤ L, f t ℓ = κ (V -1 ) ℓκ (prox(T, ρ, ω t , V t ) κ -ω t κ ) ∀1 ≤ ℓ, κ ≤ L, g t ℓκ = ∂ ω ℓ f t κ A t = - L ℓ,κ=1 ( Xℓ ⊙ Xκ ) ⊤ g t ℓκ b t = L ℓ=1 X⊤ ℓ f t ℓ + A t ⊙ Qt Qt+1 = (λI d + A t ) -1 b t ĉt+1 = (λI d + A t ) -1 end for return Estimator QThe GAMP algorithm can be derived in standard fashion from the Belief Propagation (BP) algorithm, see e.g. [56, 63] or [64] for an overview. Compared to the standard GAMP iterations for Generalized linear models, one needs to account for the fact that there exist different sources of data X ℓ (corresponding to the ℓth tokens of each input sentence), and for the fact that the output of the equivalent GLM are R L×L -valued attention matrices. In the following subsection, we show that the fixed points of GAMP 1 correspond to critical points of the empirical loss (3), i.e. fixed points of Gradient Descent (GD), allowing to connect Result 4.2 to our numerical experiments using GD.In this section we show that the dynamics of the GAMP Algorithm 1 are tracked by the summary statistics of Result 4.2. In particular, the equations ( 7) describe the statistics of the GAMP fixed points.To see this, it is convenient to take as a starting point the relaxed Belief Propagation (rBP) equations, which are a step upstream in the derivation of the GAMP iterations, and which are asymptotically equivalent-see e.g. [64] for a review or e.g. [65] , Appendix A, for a detailed walkthrough. The rBP equations readAs conventional, we note • µ the version of a variable • µ→i where the summation also encompasses the index i, and • i the version of a variable • i→µ where the summation also encompasses the index µ.Note that in all cases above the two variables differ by at mostΘ d ( 1 / √ d).Concentration of (V t µ→i ) ℓκ We first study the statistics of V t µ→i , A t i→µ , remembering that the dataxµ ℓi ≡ (x µ ℓ )i / √ d + (p ℓ )i / √ d in the notation of the main text, with (x µ ℓ ) i = Θ d (1), (p ℓ ) i = Θ d ( 1 / √ d).Replacing in the rBP updates:EQUATIONAlgorithm 2 rBPInputs : { Xℓ ∈ R n×d } L ℓ=1 , T ∈ R n×L×L , ρ ∈ R n×L×L Initialize ∀1 ≤ µ ≤ n, 1 ≤ i ≤ d, Q0 i→µ = 0, ĉ0 i→µ = 1, {f 0 ℓµ→i = 0} L ℓ=1 for t ≤ t max do ∀1 ≤ ℓ, κ ≤ L, 1 ≤ µ ≤ n, 1 ≤ i ≤ d, (V t µ→i ) ℓκ = j̸ =i (x µ ℓj )(x µ κj )ĉ t j→µ ∀1 ≤ ℓ, 1 ≤ µ ≤ n, 1 ≤ i ≤ d, ω t ℓ,µ→i = j̸ =i xµ ℓ,i Qj→µ ∀1 ≤ ℓ, 1 ≤ µ ≤ n, 1 ≤ i ≤ d, f t ℓ,µ→i = κ (V -1 µ→i ) ℓκ (prox(T µ , ρ µ , ω t µ→i , V t µ→i ) κ -ω t κ,µ→i ) ∀1 ≤ ℓ, κ ≤ L, 1 ≤ µ ≤ n, 1 ≤ i ≤ d, g t ℓκ,µ→i = ∂ ω ℓ f t κµ→i ∀1 ≤ µ ≤ n, 1 ≤ i ≤ d, A t i→µ = - L ℓ,κ=1 ν̸ =µ (x ν ℓi )(x ν κi )g t ℓκ,ν→i ∀1 ≤ µ ≤ n, 1 ≤ i ≤ d, b t i→µ = L ℓ=1 ν̸ =µ x ν ℓi f t ℓ,ν→i ∀1 ≤ µ ≤ n, 1 ≤ i ≤ d, Qt+1 i→µ = (λI d + A t i→µ ) -1 b t i→µ ∀1 ≤ µ ≤ n, 1 ≤ i ≤ d, ĉt+1 i→µ = (λI d + A t i→µ ) -1 end for return Estimator Q Distribution of ω t ℓ,µ→iLet us first introduce the teacher local fieldEQUATIONLike e.g. [65] , Appendix A, we first ascertain the joint distribution of h µ,ℓ , ω t ℓ,µ→i with respect to the data. These variables have meanEQUATION)and respective varianceEQUATION)Let us ascertain the distribution of b t i→µ .EQUATIONleading asymptotically toEQUATIONwhere the expectations bear over ξ ℓ ∼ N (0, 1) and h ℓ ∼ N (ξ ℓ θ t ℓ/ q t ℓ , ρ ℓ -(θ t ℓ ) 2 /q t ℓ ). The variance is given byEQUATIONSimilarly to the derivation for V t µ→i , A t i→µ concentrates toEQUATIONRecovering Result 4.2 Wrapping up, we now massage these equations to recover equations ( 7) from Result 4.2 of the main text. Starting from (31):EQUATIONNext, for q t ℓ (34):EQUATIONFor θ t ℓ (34):EQUATION)Finally for m t ℓ (33):EQUATIONFor mt ℓ (38) :EQUATIONwhile for θt ℓ (38) :EQUATIONNow turning to qt ℓ :EQUATIONFinally, for V t ℓ (40):EQUATIONSummary : State evolution equations The state evolution equations asymptotically describing the dynamics of the GAMP algorithm 1 thus readEQUATIONEQUATION)which exactly recovers equations ( 7) of Result 4.2 of the main text, for the case r s = 1 considered in the present Appendix. Again, we mention that the case r s ≥ 2 should follow straightforwardly with the exact same derivation steps, using tensor variables (see e.g. [52] ). This subsection has thus established that the equations ( 7) (with time indices) describe the summary statistics capturing the dynamics of GAMP iterations 1. In particular, (7) describe the fixed points of GAMP. The next subsection further shows that the (stable) fixed points of GAMP correspond to critical (zero-gradient) points of the empirical landscape (3), i.e. fixed points of gradient descent. Finally, we provide in Appendix C an alternative derivation of the state evolution equations (7)(49), using the replica method from statistical physics [61, 66] .In this subsection, we show that fixed points of GAMP 1, as asymptotically described by (7) in Result 4.2, correspond to critical (zero gradient) points of the empirical landscape (3) . Again, we present the result for r s = 1 for clarity, the generalization to r s ≥ 2 being straightforward (see e.g. [52] ). In the previous notations, let us denote the (simplified, see ( 27)) empirical loss asL({ Xℓ Q} ℓ ) + g(Q) (51)where we introduced the shorthandsEQUATIONi.e. respectively the simplified empirical loss (3) and the regularization, as functions with matrix arguments. The empirical minimization problem (3) can thus be written compactly asEQUATION)with the critical (zero-gradient) condition being given byEQUATION)Let us choose a diagonal definite A ∈ R d×d , and a sequence {V µ } 1≤µ≤n of symmetric definite L × L matrices. Group them into a block diagonal matrix V ∈ R Ln×Ln , so that the µth block of V corresponds to V µ . It shall prove useful to further introduce the matrices X ∈ R d×nL (resp. ∂L( X) ∈ R nL ), defined as the concatenation of the matrices X1 , ..., XL (resp. ∂ 1 L, ..., ∂ L L), viewed as n blocks of length L. Then without loss of generality the zero-gradient condition can be rewritten asEQUATION)Similarly to [52] , let us introduceEQUATION)This can be written in terms of a resolvent asEQUATIONwhereEQUATIONwhich corresponds to (30) . Similarly, we denoteEQUATIONSo thatEQUATIONIn the particular case of an ℓ 2 regularization g(EQUATION)The zero-gradient condition can now be rewritten asEQUATIONOne is now in a position to expand the concatenated variables • into a sequence of L n-dimensional parameters. For u = prox(ω), ω let us denote u µℓ (1 ≤ µ ≤ n, 1 ≤ ℓ ≤ L) the ℓ-th component of the µ-th block. IntroduceEQUATIONDenotef ℓ ≡ (f µℓ ) 1≤µ≤n ∈ R n , ω ℓ ≡ (ω µℓ ) 1≤µ≤n ∈ R n .The system of equations ( 63) can then be rewritten as (further redefining b ← -Ab):EQUATION)We used the assumption that g(•) is an ℓ 2 regularization. Finally, introducingEQUATION)which corresponds to the fixed-point equations of GAMP (Algorithm 1). This finishes to show the correspondence between the fixed points of GAMP and the critical points of the empirical landscape (3). To summarize, we have shown that equations ( 7) describe the zero-gradient points of the empirical loss landscape (3), i.e. fixed points of GD.While the connection between the GAMP fixed point and the extrema of the loss is sound, and has been at the roots of many rigorous results for convex losses, see e.g. [57, 58, 59, 60, 48] , there exist technical difficulties in adapting these rigorous arguments to the present setting, and a fully rigorous proof would warrant sizable work. While we leave this challenging task for future work, we wish to discuss how it can be potentially achieved. The first task would require the proof of point-wise convergence of GAMP, as indeed, the identification of the GAMP estimates with the one of the extrema of the loss function requires to be at the fixed point of the iteration. This difficulty, discussed in detail in, e.g. [67, 60, 48] , can be in principle addressed by computing the convergence criterion from the state evolution equations (see [67] the discussion in Lemma 7 in [60] ), a criterion sometimes called the "replicon" in the context of replica theory [68] .Provided the replicon criterion is satisfied, all converging fixed point described by our theory thus correspond rigorously to fixed point of the loss. The last task would be to prove that the minimum of the loss is indeed the fixed point we found with minimum energy. A potential strategy to prove this would be to use the Gordon-Minimax approach of [69] . While it is used in many situations for convex problems (e.g. [70, 71, 72] ), only one side would be required for our (non-convex) problem thanks to the GAMP matching bound. We hope that our results would provide inspiration for further research in this direction.In the Appendix we provide an alternative derivation of Result 4.2, which sharply characterizes the global minimum of the empirical loss (3), using the heuristic replica method from statistical physics [61, 66] in its replica-symmetric formulation. First observe that for any test function ϕ( Q) of the minimizer Q of (3),EQUATIONwhere we denoted R[Q] the empirical loss (3), andEQUATION)the normalization factor, also known as the partition function in statistical physics. We remind that D refers to the training set. In order to access key summary statistics and learning metrics associated to Q, it is therefore reasonable to seek to compute the generating function associated to the measure (67) , namely E ln Z. Such computations can be addressed using the replica method from statistical physics [61, 66] , building on the identityEQUATION)The backbone of the derivation thus lies in the computation of EZ s .Below, we detail the derivation for a generic convex regularizer g : R d → R + and later specialize to the case of ℓ 2 regularization. The replicated partition function thus readsEQUATION)Introduce the local fieldsh a ≡ xQ a √ d ∈ R L×r , h ⋆ ≡ xQ ⋆ √ d ∈ R L×t (71)and the overlapsEQUATIONwith rows m ℓ a . These fields have statisticsEQUATION)ThusEZ s = dmd mdθd θdqdq e -d a ℓ [ mℓ⊤ a m ℓ a +Tr(θ ℓ a θℓ⊤ a )]-d ℓ 1≤a≤b≤sTr(q ℓ ab qℓ⊤ ab )EQUATION)where we decomposed the replicated free entropy into the trace, entropic and energetic potentials Ψ t , Ψ Q , Ψ y . Note that all exponents are scaling with d → ∞. Therefore the integral in (76) can be computed using a Laplace saddle-point approximation.We have thus rephrased the analysis of the measure (67) as a optimization problem over the order parameters {q ℓ ab , θ ℓ a , m a }, and the associated conjugate variables. However, these still represent 2L(s 2 + 1) + 2s variables, and s → 0. In order to make progress, we assume that the maximizer is of replica-symmetric (RS) form [66, 61] q ℓ ab = (r ℓ -q ℓ )δ ab + q ℓ (77)EQUATION)EQUATION)EQUATIONThe RS ansatz holds in a number of machine learning settings, notably for convex problems and Bayes-optimal settings, see e.g. [64] for a review. In the present setting, since the empirical loss ( 3) is non-convex, we emphasize that the RS ansatz constitutes a heuristic technical assumption of our analysis.We now turn to the entropic potential Ψ w . It is convenient to introduce the variance order parameterEQUATIONThe entropic potential can then be expressed asEQUATIONThereforeEQUATIONFor a matrix Ξ ∈ R r×d and tensors A, B ∈ R r×d ⊗ R r×d , we denoted(Ξ ⊙ A) kl = ij Ξ ij A ij,kl and (A ⊙ B) ij,kl = rs A ij,rs B rs,kl .The computation of the energetic potential Ψ y is rather standard and follows the same lines as in e.g. [20] , yieldingEQUATIONC.4 Zero-temperature limitWe now take the limit β → ∞. RescalingEQUATIONThe entropic potential then reduces toEQUATIONwhere we defined the entropic Moreau enveloppeM g (Ξ) ≡ inf Q 1 2 L ℓ=1 Vℓ ⊗Σ ℓ 1/2 Q- L ℓ=1 Vℓ ⊗Σ ℓ -1 L ℓ=1 ( √ d mℓ p ⊤ ℓ + θℓ Q ⊤ ⋆ Σ ℓ )+ L ℓ=1 Ξ ℓ ⊙(q ℓ ⊗Σ ℓ ) 1 2 2 + g(Q) .(89) The energetic potential can be similarly recast into a more compact formEQUATIONwhere the Moreau envelope is defined asEQUATIONOne finally reaches an expression for the replica free entropy asΦ = 1 2 L ℓ=1 Tr Vℓ q ℓ -Tr qℓ V ℓ - L ℓ=1 m⊤ ℓ m ℓ - L ℓ=1Tr θ⊤ ℓ θ ℓ -EQUATIONThe extremization of the free entropy (92) yields, similarly to [47] , the following system of selfconsistent equations on the summary statistics:EQUATIONwhere the proximals prox g and prox ℓ respectively refer to the arginf in Q (resp. x ℓ ) of the envelopes M g (89) (resp. M 91).We now specialize the saddle-point equations (93) to the case of an ℓ 2 regularizer g(•) = 1 /2∥•∥ the entropic potential admits the simple formEQUATIONThe replica free energy thus readsΦ = 1 2 L ℓ=1 Tr Vℓ q ℓ -Tr qℓ V ℓ - L ℓ=1 m⊤ ℓ m ℓ - L ℓ=1Tr θ⊤ ℓ θ ℓ -αE y,ξ M(y, ξ)EQUATIONleading to the saddle point equations                 qℓ = αE Ξ,Y V -1 ℓ prox ℓ -q 1 2 ℓ ξ ℓ -m ℓ ⊗2 V -1 ℓ Vℓ = θℓ θ ⊤ ℓ q -1 ℓ -αE Ξ,Y V -1 ℓ prox ℓ -q 1 2 ℓ ξ ℓ -m ℓ ξ ⊤ ℓ q -1 2 ℓ mℓ = αE ξ,η V -1 ℓ prox ℓ -q 1 2 ℓ ξ ℓ -m ℓ θℓ = αE ξ,η V -1 ℓ prox ℓ -q 1 2 ℓ ξ ℓ -m ℓ y ℓ -ξ ⊤ ℓ q -1 /2 ℓ θ ℓ ⊤ ρ ℓ -θ ⊤ ℓ q -1 ℓ θ ℓ -1                            q ℓ = dν(γ, τ, π)γ ℓ λI r + L κ=1 γ κ Vκ -1 L κ=1 γ κ qκ + L κ=1 mκ τ κ + γ κ θκ • π ⊗2 λI r + L κ=1 γ κ Vκ -1 V ℓ = dν(γ, τ, π)γ ℓ λI r + L κ=1 γ κ Vκ -1 m ℓ = dν(γ, τ, π)τ ℓ λI r + L κ=1 γ κ Vκ -1 L κ=1 mκ τ κ + γ κ θκ • π θ ℓ = dν(γ, τ, π)γ ℓ λI r + L κ=1 γ κ Vκ -1 L κ=1 mκ τ κ + γ κ θκ • π π ⊤ .(97) which finishes to recover (7) . Let us finally mention that the update equations ( 7) for the summary statistics (6) do not describe the dynamics of gradient descent, but rather that of an Approximate Message Passing algorithm [63] , which we elicit in Appendix B for completeness.The generalization performance is measured by the test errorEQUATIONExpliciting this expression in terms of the correlated Gaussian variables xQ ⋆ , xQ allows to straightforwardly show that ϵ g admits the sharp asymptotic characterization in terms of the summary statistics characterized by (97):EQUATIONwhere the average bears on X ∈ R L×r , Y ∈ R L×t with independent rows with statisticsEQUATIONWe finally turn to the training loss. It is reasonable to expect, from statistical physics, that the training loss should be equal to the free energy -Φ at zero temperature. We provide below an alternative derivation, for simplicity in the case of ℓ 2 regularization g = 1 /2∥•∥ 2 . First note that the training loss ϵ t can be expressed asEQUATIONWhere Φ(β) is the free entropy at finite temperature. The trace potential Ψ t bears no explicit dependence on β. On the other hand,EQUATIONThusEQUATION) Finally, going through the same rescaling steps to take the β → ∞ limit,EQUATIONBy the same token, it is straightforward to see thatlim β→∞ ∂ β (βΨ y ) = -E Y,Ξ     M(Y, Ξ) - 1 2 L ℓ=1 Tr     V -1 ℓ x ℓ -q 1 /2 ℓ ξ ℓ -m ℓ ⊗2 qℓ V ℓ         (105)We used the self-consistent equations (7) to identify the term in underbrace. Putting everything together,EQUATIONEQUATIONThis constitutes a sharp asymptotic characterization of the training loss ϵ t as a function of the summary statistics characterized in Result 4.2.For completeness, we finally explicit the connection between ϵ t and the negative free entropy (i.e. the free energy in statistical physics). We go back to massage the expression for the free entropyEQUATIONIn other words, the training loss is equal to the zero-temperature free energy.For the sake of clarity and definiteness, we stated the technical results for in the simplest instance, under notably the symplifying assumptions of independent tokens, tied key and query weight matrices, as exposed in Section 3 of the main text. These two assumptions can in fact be relaxed. In this final subsection of Appendix C, we discuss how the characterization of Result 4.2 can be generalized to more complex data distributions and attention architectures.Data model Let us consider a data distribution allowing for statistical correlations between different tokens. A simple model meeting this criterion, proposed in [73] , consists in assuming every token (row) x ℓ is drawn from a Gaussian mixture with K ℓ with respective probabilities ρ ℓ,k ,EQUATIONIn ( 108), the random variable c ℓ ∈ [K ℓ ] indicates the index of the cluster the ℓ-th token is drawn from. Note that we generically allow for different rows to be sampled from distinct mixtures.In order to introduce correlations between different tokens, we further assume that the variables c = (c 1 , ..., c L ) follow a generic joint law ρ(c). Intuitively, the different models can model different types of words (e.g. verbs, pronouns, adjectives), and the nature of the words, as reflected by the variable c ℓ , is correlated across tokens. We assume that ∥µ ℓ,k ∥= Θ d (1), i.e. are of the same order as the positional encodings.Assumption C.1. Finally, similarly to Assumption 4.1, let us assume that the set of matrices {{Σ ℓ,k } K ℓ k=1 } L ℓ=1 admits a common set of eingevectors {e i } d i=1 with eigenvalues {λ ℓ,k i } d i=1 . The eigenvalues {λ ℓ,k i } ℓ,k,i and the projection of the cluster means (to which the positional encodings are added) {µ ℓ,k + p ℓ } ℓ,k and the teacher columns {Q ⋆ i } rt i=1 on these eigenvectors are assumed to admit a well-defined joint distribution ν as d → ∞ -namely, for γ = (γ ℓ,k ) ℓ,k , π = (π 1 , ..., π t ) ∈ R rt , τ = (τ ℓ,k ) ℓ,k :EQUATIONArchitecture Again in the interest of further generality, we alleviate the architectural restrictions on the attention (2), and allow for untied key and query weight matrices. More precisely, let us consider function of the formEQUATIONparameterized by two set of weights Q, K ∈ R d×rs . In order to work with more compact expression, let us introduce the total weight matrix W ∈ R d×(rs+rs) , obtained as the horizontal concatenation of Q, K. Then (110) can be written compactly asEQUATIONwhere for any argument h ∈ R L×(rs+rs) , wiewed as a 1 × 2 block matrix with two blocksh 1 , h 2 ∈ R L×r , G(h) = S(h 1 , h 2 ).Asymptotic characterization A sharp asymptotic characterization of the test error and train loss for this more general model can be derived along the same lines, as presented in Appendix C (see also [73] ). For the sake of conciseness, we only report the final result.Result C.2. Under Assumption C.1, in the limit n, d → ∞, ∥p∥, n /d, L, r s , r t = Θ d (1), the summary statisticsEQUATIONconcentrate in probability, and are solutions of the set of finite-dimensional self-consistent equations EQUATION                                             q ℓ,k = dν(γ, τ, π)γ ℓ,k λI 2rs + κ j γ κ,j Vκ,j -1   κ j mκ,j τ κ,j + γ κ,j θκ,j π ⊗2 + κ j γ κ,j qκ,j   λI 2rs + κ j γ κ,j Vκ,k -1 V ℓ,k = dν(γ, τ, π)γ ℓ,k λI 2rs + κ j γ κ,j Vκ,jThe averages bears over Ξ ∈ R L×2rs , Y ∈ R L×rt , with rows ξ ℓ ∼ N (0, I 2rs ) independently andy ℓ ∼ N (ξ ⊤ ℓ q -1 /2 ℓ,c ℓ θ ℓ,c ℓ ρ ℓ,c ℓ -θ ⊤ ℓ,c ℓ q -1 ℓ,c ℓ θ ℓ,c ℓ ), conditionally on c. prox c ℓ corresponds to the arginf of the minimization defining the Moreau envelope M(c, Y, Ξ), asEQUATION)We introduced m c ⋆ ∈ R L×rt , whose rows areEQUATIONand the shorthandEQUATIONwhere ρ c Σ ∈ R L×L is diagonal with elementsEQUATION)Test error The test error admits the sharp asymptotic characterizationEQUATIONwhere, conditioned on the class assignments c, the average bears on h ∈ R L×2rs , Y ∈ R L×rt with independent rows with statisticsEQUATION)Finally, the train loss ϵ t converges in probability toEQUATIONTrainable value matrix Finally, let us remark that it should be possible to leverage a similar derivation to further accommodate trainable value matrices V , provided they are low rank, i.e. in R d×r V with r V = Θ d (1) as d → ∞. Similarly, introducing the stacked weights W = (Q|K|V ) ∈ R d×(2rs+r V ) allows to perform the asymptotic analysis. Since this architectural modification then implies that f now takes values in R L×r V , one cannot immediately analyze this architecture for the task considered in the present manuscript, as the labels are valued in R L×d , and the setting and questions explored would need to be entirely modified. We thus postpone a thorough analysis of the effect of a trained value matrix to future investigations.In this Appendix we derive the asymptotic characterization of the learning performance of the dense linear baseline (15) , as stated in Result 5.1. Consider the empirical risk minimization ( 16)EQUATION)where we use the shorthand notation A ⋆ (x) ≡ T[ 1 / √ dxQ ⋆ ] for the target attention score matrix (1). The expression for the risk can be asymptotically simplified asEQUATIONTherefore the learnt weight Ŵ is simply equal toEQUATIONNaming h = 1 / √ dxQ ⋆ the argument of the last term, the matrix h possesses independent rows with statisticsEQUATIONwhere we remind that ρ ℓ ≡ 1 /dQ ⊤ ⋆ Σ ℓ Q ⋆ . Therefore, the trained weights Ŵ obtained by minimizing the empirical loss also coincide with the minimizer of the population loss. Intuitively, this follows from the fact that in the asymptotic limit considered n, d → ∞ and L = Θ d (1) , training W , i.e. a number L 2 = Θ d (1) of parameters on the empirical loss for n ≪ L 2 data points is equivalent asymptotically to directly training on the population loss. Expliciting the corresponding test MSE, one reaches the sharp asymptotic characterizationEQUATION)We close this Appendix by giving a few examples for definiteness.Purely positional target Let us consider the case where the target is purely positional, i.e. T[x] = A for all x. This corresponds to the ω = 1 limit of the target (14) . It then follows from Result 5.1 that the dense linear layer recover perfectly the target weights Ŵ = A.Target (14) For the target discussed in the main text (14) ,EQUATIONwith h and having independent rows h ℓ ∼ N (0, ρ ℓ ).E.1 Empirical scaling of α = d/nIn the following we verify that our experiments are consistent with the scaling behaviour predicted from the theory. We jointly increase d and n for a fixed value of α. In Fig. 7 we indeed observe the expected behaviour for an exemplary value of α = 2. The same holds for the summary statistics θ and m, which concentrate as d and n jointly grow, shown in Fig. 3 (left) in the main text. We provide supplementary results for different parameter settings. Fig. 8 on the left shows more slices from the phase diagram that appears in the main Fig. 3 . For the experimental section of the main text, we chose a specific A for definiteness. In the following, we present the same results for a different A with a stronger off-diagonal and a higher rank,EQUATION)In Fig. 9 we present the analogous simulations to Fig. 3 (center, right). While the global phenomena match the previous example, the details of the transitions location differ. As discussed previously in Section C.10, the theory developed in this work can be relaxed to incorporate arbitrary statistical correlations between tokens, changes in the (low-)rank student, generally untied key and query matrices and the presence of a low-rank trainable value matrix. This surplus in complexity comes at the cost of more intricate and heavier equations.In the present section, we report experiments in the four settings above, but not the theoretical predictions. Fig. 10 -13 compare the change in architecture with the setting considered in the main text, for a given ω = 0.5 and varying sample complexity α. The exhaustive experiments for varying ω are shown in Fig. 14 -which immediately shows that the general idea of the phase transition from positional to semantic solution is consistent, but that the shape of the phase transition curve is influenced by the architecture.For correlated inputs, the transitions seems to moves to larger values of α (Fig. 10 ). For the case of a student with a rank-2 query matrix, and a teacher with a rank-1 query matrix, we consider three possible initilisations of the columns of Q when finding minima of GD empirically.Initializing both semantically, both positional, or one of them in each way, as shown in Fig. 11 . Interestingly, the dual initialization seems to do approximately as well as the best of the two other initializations. The optimization generally seems more complex, as the phase diagram in the upper right corner of Fig. 14 are not entirely clean, and more noisy than empirical results from other architectures.When we compare tied and untied Q and K weights in Fig. 12 we observe that the additional parameters we need to learn in the case of an extra K come at a small cost, as the transition to the semantic solution is moved to the right.Finally, for a value matrix (Fig. 13 ) the results seem to largely resemble the transition line we observed without the value matrix. We compare different initializations of the higher-rank student. Positional is when both columns of the student matrix Q are initialized using the positional strategy. We do the same for the semantic strategy. The 'both' initialization initializes one column using the positional strategy and one using the semantic strategy. Experiments are repeated 5 times per data point with d = 1000. In or experiments, to obtain the empirical results, we initialize the GD optimizer in an informed fashion, i.e. initializing Q ⋆ of the student with r = 1 as either p 1 (positional) or Q ⋆ (semantics). GD then converges in the two local optima described by our theory.Since our theory only ascertains that these solutions predicted are indeed fixed points of GD for large sizes, this does not have direct implications for other types of optimization algorithms. In Fig. 15 we show that indeed running the Adam optimizer from an uninformed initialization may lead one to either of the local minima for d = 100. For larger d we observe the semantic minimum is reached less often than the positional minimum, and a considerable number of times the algorithm simply does not find either of them. 2 for the sample complexity α = 2. The student parameter Q is obtained via either (left) positional and semantic informed initialization and (right) GD training from a random initialization are compared. Each point represents a single run. For the informed GD, we used the same optimization parameters as in Fig. 2 (24 runs per initialization). For Adam we trained on the same data, but for 2, 500 epochs with learning rate η = 0.01 (showing 140 runs). Justification: Yes, it just matches.• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [No] Justification: The compute needed is negligable, so it is not specified.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: It is a theoretical work that aims to advance understanding only.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [No] Justification: It is a foundational work that aims to advance our understanding of the model internals.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: It is foundational research.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: We do not use external original material.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: The paper presents an asymptotic result on a simplified model of attention, and use it to show that, given certain assumptions, a sharp transition between a more order-dependent and a more "semantic" type of attention mechanism will emerge.

As with all theoretical papers based on simplified models, it is not entirely clear the extent to which the presented insights generalize to real-life models, but all reviewers agree that this is an exciting paper worth being presented at the conference.

Reviews: review: paper_summary: paper_summary: The authors state an asymptotic result characterizing the test MSE and training loss in a simplified single-layer model of dot product attention. They apply this result to study a special case in which the target attention function contains a tradeoff parameterized by $\omega$ between positional (i.e. dependent only on index location) and semantic (i.e. input-dependent) terms. Analysis of the solution characterized in the theoretical result shows that for a fixed $\omega$, there is a sharp boundary in terms of the sample complexity $\alpha$ (ratio of sample size to embedding dimension) between a semantic vs. positional parameter as the global minimum. These results are corroborated by an empirical analysis illustrating the distinct minima and comparing the empirical difference in train loss at semantic vs positional minima to the theoretical prediction as a function of $\alpha$ and $\omega$.
strength_weakness: strength_weakness: As noted by the authors in the limitations, the stated result applies to a simplified model both in terms of the structure of the attention function and in terms of the input data. 

It is unclear how the solutions of (7) were found in practice, as discussed in Section 5, or how it was determined that the global minimum was among this pair of fixed points for the values of $(\alpha, \omega)$ studied in the experiments.
questions: questions: Can the authors comment on why the simplifications in the attention model (value weights set to identity, key and query weights tied) are required for their result?

The positional-to-semantic phase transition detailed in Section 5 is discovered as a consequence of Theorem 4.2 applied to a specific model (Eq (14)). Have the authors considered other models or aspects of attention-based learning that could be studied through the same lens?
limitations: limitations: Limitations are identified and discussed in the paper.
paper_summary: paper_summary: This paper introduces a simplifed model of attention and analyzes it theoretically, showing that there exists a phase transition between a paradigm where attention is based mostly on position to one where it is not (which they call "semantic").  I will confess to not being an expert on the methods used and so did not follow the main results (which take about a page just to state) and proofs in detail.  They strike me, however, as genuinely useful and insightful, albeit with a caveat or two about some of the assumptions needed to get them to work (e.g. sequential independence).
strength_weakness: strength_weakness: * Very dense mathematically, so hard to follow for a reader not intimately familiar with the literature to which it contributes.
* While toy models are indeed useful objects to study in general, there are unclear connections between some of the assumptions (e.g. independent samples of individual tokens, and the low-rank attention) and actual language modeling practice.
questions: questions: * How much do you think the results depend on the nature of the data? I'm thinking in particular of the fact that words are drawn _independently_: could this be part of why positional information becomes irrelevant, since the distribution at each position is the same?  Do you have any expectations for slightly more realistic settings (even, e.g. independent $n$-grams instead of unigrams)?
limitations: limitations: Yes
paper_summary: paper_summary: The authors consider a simplified attention network with shared Query and Key matrices trained with an MSE loss and show a sharp phase transition exists when training this network. In the high-dimensional limit, this paper provides a closed form solution to the training and test loss and shows that a phase transition exists in terms of sample complexity where the model goes from solving the solution with positional information (locations of the tokens in the sentence) to semantic information (content of the tokens in the sentence). The authors show that this theoretical result shows the advantage of the attention mechanism over a fully-connected network for this task with sufficient data.
strength_weakness: strength_weakness: Because the results rely on the model reaching the minimum, it is unclear how well these results extend to randomly initialized networks. This is more a problem for studying training dynamics, though, which this paper does not aim to address.

some parts of the paper are not well motivated or presented. In particular 4.1 and 4.2 may only reach a small audience without more description.

The authors make many simplifications on the attention mechanism
questions: questions: Do the authors have speculations about how changes to the current architecture would affect the phase transition? For example, if the value matrix was not the identity, would the model be likely to transition to the semantic solution faster?
limitations: limitations: Yes

Comments: title: Thank you for the reply
comment: Thanks to the authors for their comments. These additional details address the main points of my review and the corresponding (minor) updates to the manuscript will further help contextualize the main result. I am happy to support this paper for acceptance.
Thanks!  I really appreciate these clarifications and am looking forward to reading the generalization Appendix in a camera-ready version.
Thank you for the detailed reply. I appreciate the pointer to appendix D.3 which is helpful. Besides that, I definitely agree this would be out of the scope of the current work. 

P2: Thank you the extra explanation will be helpful. Again, I am sympathetic that it simply isn't possible to catch everyone up in such a short space, I have found that adding the extra explanations has been worth it, though.

I think the authors properly address the concerns, and I'm a bit more confident after the followup discussion and reading the other reviews