A Phase Transition between Positional and Semantic Learning in a Solvable Model of Dot-Product Attention
Authors: Hugo Cui
Freya Behrens
Florent Krzakala
Lenka Zdeborova
Venue: NeurIPS 2024 spotlight (2024)
Keywords: replica method
statistical physics
phase transition
high-dimensional limit
attention layer
Primary area: learning_theory
TL;DR: We provide a tight asymptotic analysis of the learning of an attention layer, and evidence a phase transition from a positional to a semantic attention mechanism with sample complexity.

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.

A Phase Transition between Positional and Semantic Learning in a Solvable Model of Dot-Produce Attention
Hugo Cui
Freya Behrens
Statistical Physics of Computation laboratory EPFL, Lausanne, Switzerland
Statistical Physics of Computation laboratory EPFL, Lausanne, Switzerland
Florent Krzakala
Lenka Zdeborová
Information, Learning and Physics laboratory EPFL, Lausanne, Switzerland
Statistical Physics of Computation laboratory EPFL, Lausanne, Switzerland
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 step in this direction by providing tight theoretical analysis of the emergence of semantic attention in a solvable model of dot-product attention. More precisely, we consider 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 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 each other based on their meaning), and evidence 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 has access to sufficient data.
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 (. 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 leverage both types of information, by having tokens attend to each other based on their respective positions (a mechanism called positional atlention in (2] and/or respective meanings (henceforth referred to as semantic attention). this paper, we aim at 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 qualitative improvements of the model capabilities 3 5) By reverse-enge enneinernngn trained models into human- interrrrareree components [6] 8 growing body of work on mechanistic interpretabilityy aims to empirically understand which precise algorithmic mechanisms neural network learns. Such investigations have demonstrated that attention layers are able implement a wide range different algorithms, even for the same task, using both positional and semantic attributes of the inputs. We offer particularly simple
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

A tolemboed data
B
x Riie x.- (0.22)
posit ional
0:0510
lower
teocler mix
seacher /tacher/target stargt
(1 waJafam 22Q.0 x0,''11)1 WAX
4.1
semantic
student with positional encodings
8 aftm (Hix= (x +
Two local mininus the training lons.
0.0
0.5
3.0
1.5
203
sample compleeiiy
Figure A phase transition in a toy model of attention. (A) We investigate a tied low-rank attention model in a teacher-student setting. The teacher mixes the L individual tokens of dimension according to a semantic (as function the token's content x) and a positional (as function the token's S position) attention matrix. The student can only use positional encodings p to fit the positional properties of the teacher. (B) Schematic view of the loss landscape of the teacher, which contains both positional and semantic minimum. (C) We find that in the asymptotic high-dimensional limit and as a function of the sample complexity and the composition of the teacher, the global minimum switches, constituting a phase transition between positional and semantic learning.
illustration of this idea in Appendix A where we show that for sequence modelling task involving counting different algorithmic mechanisms co-exist, each corresponding 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 IO] and the available data [1] 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 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 century ago for models of interacting particles, such as the famous Ising model describing ferromagnetism [15] [6]. 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 magnet), delineating qualitatively very different regimes (e.g. magnetized or not). While mathematically, a large size limit needs be considered 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 more training samples available were studied early in [I7] [8]. and their existence was proven mathematically rigorously in [9]. 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 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 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 244 25 26 2, none of them identified a phase transition between different types of mechanisms that are implemented by the attention layer. Simultaneously, real-world models that are the focus of works in mechanistic interpretabilityy do not lend themselves to 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 sharp high-dimensional characterisation for attention layers (see Fig.
the
2

---

In particular our contributions are:
We describe a model with a single self-attention layer with tied, low-rank query and key ma- trices. On Gaussian input data and realizable labels, we show that this model exhibits 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 embed- ding dimension d of the tokens and the number n of training samples grow proportionally. We provide a right closed-forn characterization of the test eтror and training loss achieved at the minima of the non-convex empirical loss. Using this high-dimensional characterization, we locate the positional-semantiee phase transition, thus providing the first theoretical result about the emergence of sharp phase transitions in model of dot-product attention.
We contrast the performance of the dot- 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 sufficient amount of training data.
In Section we discuss further related work. Section 3 defines the general version of a solvable model of tied low-rank attention, and Dection[proviin 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 miniumm corresponds to either a semantic or positional mechanism, depending on the training data and task, with phase transition between them. We conclude with discussion of the limitations of our analysis in Section 6
2 Related work
Theory of attention Attention models have been the object of sizeable theoretical scrutiny in recent years, with growing body of work investigating various aspects such as their expressivity
25
27
28,, inductive bias 299 30 31, training dynamics 2 32 33 34] and in-context learning
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 single layer of attention, with trainable queries and keys, assuming linear or ReL.u 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 theoretical description of the sudden emergence of new algorithmic mechanisms. learning high- nigh-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.
26
35
36
37
first tight description was provided in (24, in the context of
(38] analyze how induction head mechanisms can be learnt from gradient descent on the population loss, i.e. when an infinite amount of data available. The present manuscript conducts tight analysis of the non-convex learning of non-linear attention model with trainable tied queries and keys from finite train set, thereby allowing the description of sharp phase transitions in sample complexity in the behaviour and performance of the model.
Positional encodings 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, c.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 training [1
43 J in principle
a dot product attention layer requires an explicit encoding of positional information as it views the input sequence in parallel, as bag of words [. 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.
3

---

Theory of phase transitions in neural networks In supervised learning with feed-forward fully connected neural networks, phase transitions sample complexity were identified in settings where the data consists of random Gaussian samples, and the labels are generated by target neural networks with random weights. For a single-layer perceptron and a variety of teacher weights distributions and activation functions, discontinuity of the optimal test error as the number of samples increases was established [88 19]. For a two-layer neural network, 20) evidenced 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.
3
Tied low-rank attention model
Input data model We consider a model of embedded sentences with uncorrelated (I-gram) words. More precisely, sentence IR 1xddd, where is the sentence length and d represents the embedding dimension, consists of L tokens f independently drawn from a Gaussian distribution It ~ (0, Σe) with covariance Σg E Rdxd In the following, we denote the probability distribution of x Pa- 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. Section [5]
will be explored in We defer a discussion and an analytical treatment of the general case to Appendix C
Target function
The target function (teacher) is assumed to be of the form
. = T vd QQ x,
for function _x LLLL The term T [3/Vxx..] RLxL in (I) should be interpreted the target attention matrix, which mixes the tokens of the input æ. This attention matrix is parametrized by the target weights Q. ddd..
Tied attention
We consider the learning of the target (T) by single attention layer
1 va (2 + Q)Q p).
fa(x) S
In Q E RLxd is fixed matrix, corresponding to positional encodings, and Q E RRdx., trainable weight matrix. We denote subsequently Ps E Re the f-th row of p. Like the target (.. the parametric function (2) takes the form of data-dependent attention matrix + + E L mixing the tokens of the input æ. Note that, compared to the usual attention mechanism II employed in practice, 2 corresponds setting the value weights to identity, - (2) is parametrized by a single matrix Q- tying 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 Cl full analysis of the untied architecture for completeness.
Empirical risk minimization We study the learning of the attention layer 2 when a training set D = {x* {[==)) with n independently sampled sentences {x* 1p=1 and the associated labels (y(*"))p= is available. The target can be learnt by carrying out empirical risk minimization:
X 1011
Q
argmin
Σ
2d [lv((**) [1
QeRaxr
The performance of the resulting trained model få is measured at test time by the mean squared error (MSE)
y(z) -f f0((0)
Eg = dL
E Ex-pp
(4)
4

---

4 Closed-form characterization of the training loss
High-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 a = "/4
(henceforth referred to as the sample complexity) stays of order Ox(1). We further assume the sentence length L, the ranks rs, of the weights Q, Q.. and the norm of the positional embeddings ||p], to be (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. (a(1), which in words implics 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]
The main technical result of the present work is (4) and training loss (3) achieved in the high- nigh-dimensional limit when training the model (2) via the empirical risk minimization of 3
characterization of the test MSE
Assumption 4.1. The covariances {EE)E- admit common set of eigenvectors {ei We
further note {2] the eigenvalues of E.. The eigenvalues {2() and the projection of the positional embedding {peYé=1 and the teacher columns {Q on the eigenvectors {e, Petic, namely, for {e] Qihis are assumed admit well-defined joint distribution as d 7 (?,,.., ...)).))))))))) ,R,,, (T.,.... (T.)) E R and T (Ti..., TL.) R.:
d L
EI 01 (4f--0) (Videe Pt - Tt) PE t=
II j=1
(e, Qj - RT) d+00 v (T)
(5)
d
In words, Assumption guarantees that all parameters of the problem admit well-defined limits in the considered usymptotic limit, with the further assumption that the covariances {E of the different tokens can be jointly diagonalized. We are now in position to state the main technical result of the present work.
Result 4.2. Under Assumption statistics
4.1 in the limit n, d oc, 11pl|, 1/d, I Ts, Oa(1), the summary
P: = Q.E.Q.
E R
E ER
qe d
d
Q Pt ER' mg d
T E,Q, R
0. d concentrate in probabilit, and solutions of the of fiite-dimenssonnl self-consistent equations
(6)
qe [d T, T)e
AI,
Σ
Yx
V..)
Σ 1ng
MaTa +r 7) -7)
4 (AUI Σnm V.))
x=1
E
Ve f du(y, T, T)76
Adr,
Σ V/
k=1
L
-1
L
me = f dul(, TT
All, + Σ
Vn
Σ
mr + On T
x=1
x=1
1
0:- f du(), T. T)
Alr, +
Σ
The V/
Σ Tx Tx + IrBm x TT.
602
d: a.,,uVE (prox(E,U); /r Er mg
V
Ve 0.07 qe" -aBE, V; (prox(E. (E, U)e -02 Ee mi)
E7 4e
(8)
me VV (prox(E, U), u² m²)
0. VB (prox(E, U)2 / Ee me)
(we- - ET qe ==1)))) (pe
02 0f 100)
5

---

In -we and = (E) with we NEt and _2 denotes the outer product of vector with itself. Finally, the resolvents {prox(E, U)eké= are defined as the minimizers of the Moreau envelope
0e, PL ve-00 02 96 0z) and & ~ (O
48
82
M) = inf
Tr Ve
(Me-Ge Ec-me)
31
We noted Z E BLXx, (resp. U BLxxs the matrix whose rows are ze (resp. ue) and:
PI = diag
diag
E RLLL
(9)
t=1
In the same limit, the test error converges in probability to
Ey - - Tishiisshh + En Tr [psTh]]]]
- TrSh""
(10)
where the average bears on h E IRLx,, h² E BLxx with independent rows with statistics
(he, kb) N
ml
qe 0
O2
]
(11)
Finally, the training loss E converges in probability to
L
(x+ (M Tri Vel 2 7/I m0)) 7 reat ()) (12)
)00)
& = aBay
Result 122 provides tight asymptotic characterization of the test error Eg TT) and the training loss er (12., as function of a finite set finte-dimennional summary statistics (ge, Ve, m. 0r) f= thereby providing a finite-dimensional description of the high-dimensional learning problem 33. These summary statistics are further characterized closed form by the set of equations 7 in terms of the solution of a low-dimensional minimization problem (I). Intuitively, this low-dimensional problem may be viewed solution of the self-consistent equations 7 can be found by numerically iterating the equations until convergence. The resulting summary statisties can then be used to evaluate the expressions we evaluate
form of effective loss averaged over the finite training set. In practice, the
(10) and (12) to reach the asymptotic limits of the test error and train loss. In Section [ 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, particular generalized linear models (see e.g. [1 52 48 and
auto-encoders 47).
The derivation of Result 4.2 is provided in Appendix B and is exploiting mapping of the model (2)
to (variant of) a Generalized Linear Model (GLM) [33 54]. The summary statistics characterized by the equations 7 (often called state evolution 551 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-graa sero--gaddeet) 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 also describe local minima and saddles.
This strategy has been used in many recent work to study asymptotics of large number of high- dimensional problems, see e.g. [87]58 60) 60 48). We note, however, that we importantly assume the point-wise convergence of GAMP. While 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 [1]). Finally, we mention that while Result 44. is presented for an &2 regular- ization of the empirical loss 33 for clarity, similar results can be reached for generic convex regulariz- ers, and are presented in Appendix C In the following section, we explore the phenomenology uncov- ered from the study of the equations of Result 4.2 for the special case of dot- dot-product attention.

---

0,0005
10-
8/
0.004
45
0.0000
0.6
0.003
-0.0005
0.2
0.002
0.0010
0.0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0 1.5
2.0
0.0
1.0
2.0
Figure 2: Mixed positional/semanteee teacher for 0.3. Setting is
1, 2,
(00.6, 0.4), (0.4, 0.6)). Σ1 = Σg 0.251a. P1 = 1g = p2 and Q (0, (left) Solid lines:
difference in training loss Ae, between the semantic and positional solutions of (7) in Result 4.2 Markers:
difference training loss at convergence achieved by training the model (2) using gradient
descent initialized resp. Q. and at p1. Marker color in Fig. (center) overlap 0 between the learnt weights Q and the target weights Q. overlap between the learnt weights Q and the positional emhedding P1- Solid lines represent the theoretical characterization of these two summary statistics provided by Result (.2 Only the solution of ( corresponding the lowest found training loss is represented (i.c. the positional solution for < and the semantic otherwise). Markers represent experimental measures of these quantities, for gradient descent at convergence. Gradient descent was initialized at 1ı for a < Or and at Q. for a > aee. (right) We show the MSE achieved by the dense linear as cg (Result 4.2), and MSE achieved by the dense linear baseline ein (15) (Result (5 Markers indicate the MSE experimentally reached by the model 2 trained using gradient descent, initialized previously for the overlaps. All experiments were performed in 1, 000 with the Pytorch implementation of full-batch gradient descent, for T = 5, 000 epochs and learning rate r) = 0.15. All points are averaged over 24 instances of the problem each.
Positional-to-semaniic phase transition
Rank one dot-product attention In the following, we turn to special case of tied low-rank attention namely dot- product attention layer, which the example from Fig.
1 Vd
p = softmax (1
x (x
(x +p)
(13)
As in 2 we allow for positional encodings in the dot-product attention parametrization further consider specific case of target attention matrix | of the form
13 We
= (1 - wssfftmax ( =
T
=Q.
(14)
with RLXL a fixed matrix. In (4). the parameter w E [0, 1] tunes the relative strength of the dot-product term and the fixed matrix term, and interpolates between fully positional and fully semantic task:
For w 0, the target reduces to the first dot-product term, and is purely semantic, in that the i j. element of the score matrix softmax (aaaaaxx///n....... QTx::" only depends on the tokens Ti, B. and not explicitely on their respective placements i, 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;j is a function of i, but not of æ; To complete the learning task, positional mechanism then needs to be learnt.
The parameter i 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 ws) or rather positional attention (large ws). In the following. for definiteness, we further assume T, Tt 1 and set Q, to be fixed random Gaussian vector drawn from (000, La), and choose the positional encodings P1 = P2 = 1a. Finally, for simplicity, we consider sentences with two tokens L 2 and isotropic token covariances Σ = Eg 1d-
Semantic and positional mechanisms The summary statistics Or, me describing the global
minimizer of the empirical loss minimization B of the dot-product attention (13 on the target

---

0.5
5.0022
0.5020
5.4
0,0020
0.4
0.0015
0.0035
0.3
0.0030
00.000
0.2
-0.0205
0.0030
6.3
-30010
0.0325
-3.0015
0.0320
00 04
2.0
3.0
Figure 3: Phase transition between semantic and positional training loss. Setting and experiments were performed identical to Fig. 2 (left) Scaling d and n jointly for a = 1.5 concentrates for 0 and m, in different locations for the positional and semantic local minima each. We show 30 runs for each E 110, 15, 23, 36, 56, 87, 135, 209, 323, 500). (center) The color map represents the difference in training loss at convergence when training the model (2) using the Pytorch implementation of full-batch gradient descent, respectively from an initialization at P1 or at Q.. The green dashed line represents the theoretical prediction for the threshold (w( above which the semantic solution of 77) in Result (.2 has lower loss than the positional solution. (right) The color map represents the difference in test MSE at convergence when training the attention model (13) using the Pytorch implementation of full-batch gradient descent initialized at Q.. and the dense linear baseline (15). The red dashed lines indicate the theoretical prediction ffollowing from Result 4.2 and Result 15
for the threshold sample complexity (w( above which the dot- not-product attention outperforms the
baseline (15
(14) are captured alongside the corresponding test error (4) and training loss (3). by Result 4.2
The solution of the system of equations 77 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 (14b:
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 p1 = -P2. Consequently, the argument of the dot-product attention (x + p) has a sizeable token- ooken-independent - -thus positional- contribution Qp, alongside a deeendde semantic part Qx. Because of the positional terms, the resulting learnt attention attention matrix softmax(+/a( + QQQ (æ + +)) implements a partly positional mechanism.
-Positional solution One solution of 7 correspond to vanishing overlap
Semantic solution Another solution of the system of equations (7) is associated with a vanishing overlap m = between the learnt weights Q and the positional embeddings, and finite overlap 0 > 0 with the target weights Q. Therefore the resulting learnt attention matrix softmax (oltmax(+/a( + QQ (æ T) R softmax(H/arQQ x T) 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 com- plexity a and the positional/semaniic parameter w 14 as we describe in the following subsection.
-oriioomnn----ssmntiieee phase transition For fixed parameter in (4), an analysis of equations 77. further detailed in Appendix C reveals that for a sizeable range of w, in the probed setups, there exists threshold ae for the sample complexity so that
For a < Og. the global minimum of 33) corresponds to a positional mechanism, and is described by the positional solution of 7 of Result 4.2 with 0 0, m O.
For > Oce the global minimum of 3 corresponds to a semantic mechanism, and is described by the semantic solution of of Result 4.2 with 0 > 0, m 0.
The dot-product attention thus displays phase transition in sample complexity from 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 a, generically grows with the positionality
of the target function (14), as the semantic content - i.e. the first term of (14p- is less apparent for larger w, and thus requires larger amounts of data to be identified and approximated by the dot-product attention
13 An example for w 0.3 is given in Fig, In Fig. 3 (center) the difference
8

---

in training loss de 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 < Oe (resp. > ot.). the training loss of the positional (resp. semantic) mínimum is lower, and thus corresponds to the global minimum.
Experimentally, the positional minimum can be reached for a < Or via gradient descent by initializing the weights Q of the attention 113 close to the positional embedding p1. By the same means, the semantic minimum can be reached from an initialization at the teacher weights Q (T4. Henceforth, we refer with 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 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 an analysis of the loss landscape. We however conduct numerical experiments from 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 mimimum the 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 positional (resp. semantic) initialization for < a (resp. a > a₂), displaying overall good agreement. In Fig. (left) and Appendix F. we further verify that in the scaling limit of our analysis, namely n, d - oc for a = O()) 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 a > Oke 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- not-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.
Purely positional baseline In this subsection, for the same target (1).. we contrast the dot-product attention model 133. analyzed in the previous subsections, to the baseline given by a linear layer
Sw (z)) W x
(15)
with a trainable weight matrix W E RXI. As for the dot- not-product attention 13 we consider the case where the weights W are learnt by minimizing the empirical risk
W argmin Σ Ell((=") - (w()) WERLKL
(16)
The model [15) is a natural counterpart to the dro-pur architecture (3) In (5). the attention matrix is parametrized by single fully-trainable matrix W, instead of being parametrized dot-product attention as in (I3). A seminal difference in the two parametrizations is that while the elements of (aatmaxx(///aaQ0 TTQTa Ta can depend the input tokens æ., and therefore express semantic information, the elements Wij of W can only depend on the positions j.j. The model (15) can thus only implement positional mechanisms, while the dot-produse attention (13) can implement both linear and semantic mechanisms, as discussed above. Finally, observe that the model (15) is closely related the one analyzed by (24) in another asymptotic limit. The following result characterizes the test error achieved by the purely positional model (1)::
Result 5.1. In the same asymptotic limit as Result (4.2). the learnt weights W trained by minimizing the empirical risk (16) coincide with the minimizer of the population risk, and thus admit the compact expression
E=T Lvd
=
xQ.
TAT]
(17)
where the average bears over a fiite--diennional matrix h E Lxt with independent rows by with statistics he ~ (O, pt). where pe was defined in 66) in Result 4. We remind that T [./ax..]]

---

corresponds to the target score matrix by the trained dense linear model fw
T
Finally, the test MSE //LLE (lv([l))())))))((()))))) fw (x))= achieved
admits the asymptotic characterization
1 Tr W ps: WT + T. 9 T.
Er TT[[[M)TTT(]]]]] - En Tr W PM
(18)
The MSE achieved by the baseline (5) when learning the target (14) plotted in Fig. (right) as the orange solid line, alongside the MSE achieved by the dot- dot-product attention (13) discussed in previous
subsections. Remarkably, in the setup of Fig. in the positional regime a < ar when the dot-product
attention relies on positional mechanism O 0. >> O to approximate the target, the dot-product attention (13p is outperformed by the purely positional attention (15) 6g > <g contrast, in the semantic regime or > OC where the dot-product attention learns the semantic mechanism, there exists a sample complexity aj 2 Ore above which Eg < clin i.c. the dot-product attention (13) outperforms the dense linear baseline (15). This threshold value oq is plotted for various positionality strengths in Fig. 3 alongside the positional- to-smalll-- threshold aee. Interestingly, we observe aj 2 a, in all probed settings, temptingly suggesting the natural interpretation that the dot-product attention needs to learn the semantic mechanism first (at a = az)) in order to then be able to outperform the best positional approximation f (at az). This highlights the importance of the semantic mechanism, enabled by the dot-product parametrization (13 3 in learning targets with semantic content such as (14).
6 Limitations
Compared to the original transformer 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 limited to Gaussian data with sentences of 1 1-grams. Concerning the analysis, our characterization holds only in the nigh-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 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 them. It is yet unclear if there are scaling limits of learning algorithms which would reliably find the lower or specific one of the minima from random initialization.
Conclusion
We explored the interplay between positional and semantic attention, through the prism of tied low-rank self-attention in high dimensions. In theoretically controlled setting, characterized the global optimum of the empirical loss, when learning a target attention layer. This global optimum was found to correspond either a positional or 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 random initialization constitutes an important future research avenue.
Acknowledgements
We thank Luca Biggio, Federica Gerace, and Matteo Vilucchio for insightful discussions. We acknowledge funding from the Swiss National Science Foundation grant SNFS OperaGOST (grant number 200390), and SMArtNet (grant number 212049).
10

---

References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
[2] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37836, 2022.
[3] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07060 2022.
[4] Aaditya Singh, Stephanie C.Y. Chan, Ted Moskovitz, Erin Grant, Andrew M Saxe, and Felix Hill. The transient nature of emergent in-context learning in transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[5] Taylor Webb, Keith Holyoak, and Hongjing Lu. Emergent analogical reasoning in large language models. Nature Human Behaviour, 79):1522-1541, 2023.
[6] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference Machine Learning, pages 11080-11090 PMLR, 2021.
[7] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hattield-Dodds Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022.
[8] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
[9] Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks. Advances in Neural Information Processing Systems, 36, 2024.
[10] Angelica Chen, Ravid Schwartz-Ziv, Kyunghyun Cho, Matthew L Leavitt, and Naomi Saphra. Sudden drops the loss: Syntax acquisition, phase transitions, and simplicity bias in mlms. arXiv preprint arXi::2309.07311, 2023.
[11] Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in Neural Information Processing Systems, 36, 2024.
[12] Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In The Twelfth International Conference Learning Representations, 2024.
[13] Sanjeev Arora and Anirudh Goyal. A theory for emergence of complex skills in language models. arXiv preprint arXin:23071159:: 15936,, 2023.
[14] Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models mirage? Advances in Neural Information Processing Systems, 36, 2024.
[15] Ernst Ising. Contribution to the theory of ferromagnetism. Z. Phys, 31(1):253-258, 1925.
[16] Lars Onsager. Crystal statistics. a two-dimensional model with an order-disorder transition. Physical Review, 65(3-4):117, 1944.
[17] Haim Sompolinsky, Naftali Tishby, and Sebastian Seung. Learning from examples in large neural networks. Physical Review Letters, 65(13):1683, 1990.
11

---

[18] Géza Györgyi. First-order transition to perfect generalization in a neural network with binary synapses. Physical Review A, 41((2):7097 1990.
[19] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 1 116(12):5451-466) 2019.
[20] Benjamin Aubin, Antoine Maillard, Florent Krzakala, Nicolas Macris, Lenka Zdeborová, et al. The committee machine: Computational to statistical gaps in learning two-layers neural network. Advances Neural Information Processing Systems, 31, 2018.
[21] Henry Schwarze. Learning rule in multilayer neural network. Journal of Physics A: Mathematical and General, 26(21):5781, 1993.
[22] Antoine Maillard, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase retrieval in high dimensions: Statistical and computational phase transitions. Advances in Neural Information Processing Systems,
2020.
[23] Adriano Barra, Giuseppe Genovese, Peter Sollich, and Daniele Tantari. Phase transitions in
restricted boltzmann machines with generic priors. Physical Review E, 96(4):042156, 2017.
[24] Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. Optimal inference of generalised Potts model by single-layer transformers with factored attention. arXiv preprint 23Xi:23304.02235. 2023.
[25] Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What single attention layer learn? study through the random features lens. Advances Neural Information Processing Systems, 36, 2024.
[26] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
[27] Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and vari- able creation in self-attention mechanisms. In International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
[28] Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 2020.
[29] Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In International Conference on Machine Learning, pages 19050-19088 PMLR, 2022.
[30] Davoud Atace Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. Advances in Neural Information Processing Systems, 36:48314 48362, 2023.
[31] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transform- ers as support vector machines. arXiv preprint arXiv:2308. 16898, 2023.
[32] Emmanuel Abbe, Samy Bengio, Enric Boix-Adsera, Etai Littwin, and Joshua Susskind. Trans- formers learn through gradual rank increase. Advances Neural Information Processing Systems, 36, 2024.
[33] Hongkang Li, Meng Wang. Sijia Liu, and Pin-Ya Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.00 06015, 2023.
[34] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamies and token composition 1-layer transformer. Advances in Neural Information Processing Systems, 36:71911-719999 2023.
12

---

[35] Tianyu Guo, Wei Hu, Song Mei, Huan Wang. Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-contett beyond simple functions? case study on learning with representations. arXiv preprini arXiv:2310. 10616, 2023.
[36] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pages 19565-19594, PMLR, 2023.
[37] Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.099) 2023.
[38] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
[39] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on transformers. arXiv preprint arXiv:2312.0079, 2023.
[40] Jeffrey I. Elman. Finding structure in time. Cognitive Science, 14(2):179-211, 1990.
[41] Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer language models without positional encodings still learn positional information. In Findings of the Association for Computational Linguistics: EMNLP 2022, pages 1382-1390 December 2022.
[42] Koustuv Sinha, Amirhossein Kazemnejad, Siva Reddy, Joelle Pineau, Dieuwke Hupkes, and Adina Williams. The curious case of absolute position embeddings. In Findings of the Association for Computational Linguistics: EMNLP 2022, pages 4449-4472 2022.
[43] Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[44] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-Attention with Relative Position Representations. In Proceedingy of the 2018 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), pages 464-468 Association for Computational Linguistics, 2018
[45] Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length gen- eralization of transformers. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Pupers), pages 1889-1903, 2023.
[46] Francesca Mignacco, Florent Krzakala, Yue Lu, Pierfrancesco Urbani, and Lenka Zdeborova. The role of regularization in classification of high-dimensional noisy gaussian mixture. In International conference on machine learning, pages 6874-6883. PMLR, 2020.
[47] Hugo Cui and Lenka Zdeborová. High-dimensional asymptotics of denoising autoencoders. Advances in Neural Information Processing Systems, 36, 2024.
[48] Bruno Loureiro, Gabriele Sicuro, Cédric Gerbelot, Alessandro Pacco, Florent Krzakala, and Lenka Zdeborová. Learning gaussian mixtures with generalized linear models: Precise asymp- totics in high-dimension.. Advances in Neural Information Processing Systems, 34:10144 10157, 2021.
[49] Yen-Chang Hsu, Ting Hua, Sungen Chang, Qian Lou, Yilin Shen, and Hongxia Jin. Language model compression with weighted low-rank factorization. arXiv preprint 22X0:2207.000112 2022.
[50] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li. Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.00688 2021.
[51] Benjamin Aubin, Florent Krzakala, Yue Lu, and Lenka Zdeborová. Generalization enror in high-dimensional perceptrons: Approaching bayes error with convex optimization. Advances in Neural Information Processing Systems, 33:12199-122100 2020.
13

---

[52] Elisabetta Cornacchia, Francesca Mignacco. Rodrigo Veiga, Cédric Gerbelot, Bruno Loureiro, and Lenka Zdeborová. Learning curves for the multi-class teacher-student perceptron. Machine Learning: Science and Technology, 4(1 (0119, 2023.
[53] John Ashworth Nelder and Robert WM Wedderbum. Generalized linear models. Journal of the Royal Statistical Society Series A: Statistics in Society, 135(3):370-335) 1972.
[54] Peter McCullagh. Generalized linear models. Routledge, 2019.
[55] Adel Javanmard and Andrea Montanari. State evolution for general approximate message passing algorithms, with applications to spatial coupling, Information and Inference: A Journal of the IMA, 2(2): 12)):1551144 15-144, 2013.
[56] Sundeep Rangan, Philip Schniter, Erwin Riegler, Alyson Fletcher, and Volkan Cevher. Fixed points of generalized upproximate message passing with arbitrary matrices. IEEE Transactions on Information Theory, 62(12))746,-747, 2016.
[57] Mohsen Bayati and Andrea Montanari. The lasso risk for gaussian matrices. IEEE Transactions on Information Theory, 58(4)):9977-2017 2011.
(58) David Donoho and Andrea Montanari. High dimensional robust m-estimation: Asymptotic variance via approximate message passing. Probability Theory and Related Fields, 166:935-969 2016.
[59] Melikasadat Emami, Mojtaba Sahraee- Ardakan, Parthe Pandit, Sundeep Rangan, and Alyson Fletcher. Generalization error of generalized linear models in high dimensions. International Conference on Machine Learning, pages 2892-2901. PMLR, 2020.
[60] Cedric Gerbelot, Alia Abbara, and Florent Krzakala. Asymptotic errors for teacher-student convex generalized linear models (or: How to prove kabashima's replica formula). IEEE Trunsuctions on Information Theory, 69(3):1824-1885 2022.
[61] Giorgio Parisi. Toward mean field theory for spin glasses. Physics Letters A, (3(3):203-208 1979.
[62] Stéphane d'Ascoli, Levent Sagun, Giulio Biroli, and Joan Bruna. Finding the needle in the haystack with convolutions: on the benefits of architectural bias. In Advances in Neural Information Processing Systems, volume 32, 2019.
[63] Mohsen Bayati and Andrea Montanari. The dynamics of message passing on dense graphs, with applications to compressed sensing. IEEF Transactions on Information Theory, 77(2):764-775 2011.
[64] Lenka Zdeborová and Florent Krzakala. Statistical physics of inference: Thresholds and algorithms. Advances in Physics, 6565):453-552, 2016.
[65] Lucas Clarté, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Theoretical characteri- zation uncertainty in high-dimensional linear classification. Machine Learning: Science and Technology, 4(2):025029, 2023.
[66] Giorgio Parisi. Order parameter for spin-glasses. Physical Review Letters, (0024):1946, 1983.
[67] Erwin Bolthausen. An iterative construction of solutions of the tap equations for the sherrington- kirkpatrick model. Communicutions in Mathematical Physics, 325(1):333-366 2014.
[68] Mare Mézard, Giorgio Parisi, and Miguel Angel Virasoro. Spin glass theory and beyond: An Introduction to the Replica Method and Its Applications, volume 9, World Scientific Publishing Company, 1987.
[69] Yehoram Gordon. On milman's inequality and random subspaces which escape through a mesh in In Geometric Aspects of Functional Analysis: Israel Seminar (GAFA) 1986-87 87, pages 84-106. Springer, 1988.
[70] Christos Thrampoulidis, Samet Oymak, and Babak Hassibi. The gaussian min-max theorem in the presence of convexity, arXiv preprint arXiv:1408. 4837, 2014.
14

---

[71] Mihailo Stojnic. Meshes that trap random subspaces. arXiv preprint arXiv:1304.0003, 2013.
[72] Mihailo Stojnic. Upper-bounding 11-optimization weak thresholds.
arXiv preprint
arXiv:1303.7283: 7289, 2013.
[73] Hugo Cui. digh--immnssioal learning of narrow neural networks. a4Xiv:2409. 13904, 2024.
arXiv preprint
15

---

Appendix: Table of Contents
The Histogram Task: An algorithmic toy example with a positional and semantic solution
17
A.1
Phenomenology
17
A.2
Dataset
18
Model
18
A.4
Training procedure and freezing model parameters
19
A.5
Unfreczing model parameters and SGD convergence
19
Derivation of 4.2
20
Notations
20
B.2
AMP algorithm
21
State evolution
22
B.4
Fixed points of GAMP are fixed points of GD
25
B.5
Towards a rigorous proof of result
27
Derivation of Result[4.2 with the replica method
27
C.1
ansatz
28
Entropic potential
................................
29
Energetic potential
................................
30
Zero-temperature limit
30
Replica free entropy
30
C.6
Saddle-point equations: general regularizer
31
C.7
saddle-point equations: b
8
test MSE
32
C.9
Training loss
32
10 Extensions
D Derivation of 5.1
36
D.1 Examples
37
Supplementary Experiments
37
E.1
Empirical scaling of d/n
37
Alternative hyperparameters
More Complex Architectural Choices
38
B,4
4
Uninformed initialization and training via Adam
40
NeurlPS Paper Checklist
42
16

---

Foanele Sesutece
Exampl Sequence #? N M H
Examalle Segureen
D B
M K J
D
Postlona
M
M
M
A
H
B
M
N MANMLKI K HIGEF
D
Sema
M
27
Figure 4: Several solutions exist for the histogram task. Elements of attention matrices for the histogram task for local minima in the empirical loss landscape. We generated a dataset of sequences by sampling each token of the sequence i.idd. from the uniform distribution over all tokens. The target of a a given input sequence x = [A, D. D, C| is the number of occurence of each token in the complete sequence, i.e. y [1, 2, 2, 1]. Models were trained with their respective frozen initialization using n = 35, 000 samples and the Adam optimizer. Top Row: The attention matrix of the positional solution is largely independent of the specific input sequence. Boltom Row: The attention matrices from the semantic solution vary based on the input token. Red squares highlight the elements of Ag where T; = Tz.
A
The Histogram Task: An algorithmic toy example with a positional and semantic solution
A.1 Phenomenology
In this appendix, demonstrate that for simple counting task two qualitatively different solutions exist in the loss landscape of simple transformer using a dot-product attention layer with positional encodings- 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 ssquencc-to-sequeee Given an input sequence æ = [x1, 22, LL of length L of tokens from a fixed alphabet, the goal is return sequence y [y1, V2. V2), where each token y; is the number of occurrences of the token x; in x. In Fig. we show an example where we consider sequences where the tokens are from the fixed alphabet X (A, B, C....) } of size |X| 15. In this setting, for instance, the sequence x = [A, B, B, C, A, B) should be mapped to its histogram sequence y = [2, 3, 3, 1, 2,3). When the input data is limited to length L, the output elements y/ thus take values up the maximum count L.
counting task, referred to the histogram task in [6].
We encode the input using token embeddings and absolute positional encodings which are trained jointly with the model weights. As an architecture we consider small transformer made up of single layer of dot-produe 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 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 's access to positional or semantic information- and 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.c.
17

---

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. The positional solution corresponds to a learnt attention matrix whose i,j- 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 the semantic solution displays larger i,j-th component the tokens at position i and j are identical. 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 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, remains an interpretation of an experiment and does not allow for a precise characterization of their behaviour of the conditions under which they 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 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 attention
-ed-saanntic
A.2 Dataset
We use the histogram task as proposed in (6. We consider sequences of fixed length L 15. For every input sequence x = [z1, 22 ... FL we sample z; i.i.d. and uniformly from a set of tokens T of size T, which we set 15 in experiments. For visualization purposes
capitalized
letters as tokens. To obtain the target y [y1. 1/2 ... VL], we set y, E/= 1(x; = xz), where (b) is the boolean statement b true and zero otherwise. Therefore, y: E f1,2. L).
A.3
Model
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 matrices Etokenn (B ...] RTxd and g P2,,,22... embedded sequence & E IR 1-xd of the same length where d tz. + Pc. This input sequence is then fed into the first layer of the transformer.
E gLxd We convert the input sequence [x1, X2, xr E T into an
Formally, we have
logits,(x) = Re
(X)) + b1) + b2 E R
(19)
with W1 E dd by E ARA, W2 E Rixe by E R. The final prediction is obtained using the argmax on the logits. We have that the score matrix A and the dot- dot-product attention mechanism is
Allention(x)
A(x)Vx
(20)
XQKTXT
A(x) Softmax
(21)
Vd
where Q,K.V V dddd and the softmax is applied wiwe.ise. Also, for x E Rd we define
E[xr]
Loyer\Norm(z)=
VVar(x+ +E *7+
(22)
18

---

where 3 E I We define the empirical loss for dataset D = {xx" y 1p-= as the average cross entropy loss for all output tokens C, i.e.
C(D)
-1
Wf (oggs
(23)
Because of the way in which the histogram dataset is created for fixed input length L, it follows C= L.
A.4 Training procedure and freezing model parameters
To obtain the positional and semantic minima, we set some weights to zero in Etoken Evos Q, K 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, c.g. by [62 for convolutional neural networks.
With we denote the initialization that is taken as the default Pytorch initialization for linear layer. For both semantic and positional initialization, overwrite this initialization with zeros follows with = 1,... and j = =,,-.. ,,,
d/2 0d/2 Pj
(0a/2
'd/2
(24)
where Oa/2 is the all-zero vector of size d/2. 2. For the positional initialization we additionally set
Q
1d/2xa/2
0 01/2xd/2
K =
'd/2xd/2
02/2xd/2
(25)
0a/2xd/2
0d/2xd/2
0a/2xd/2
and for the semantic initialization additionally set
00////xa// 04/2xa/2
0d/2xd/2
04/2xd/2
Od/2xd/2
02/2xd/2
Q =
K- =
(26)
"d/2xd/2
0a/2xa/2
-4/2xd/2
where 0d/2xd/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 zero using Adam obtain a > 99.8% test accuracy on a test set of size n = 15, 000 for both datasets. We call these parameter configurations 0pee and Osem respectively (the local minima as referred to in the main text).
dataset of size 35, 000,
In Fig. 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.
A.5 Unfreezing model parameters and SGD convergence
Finally, we want to verify that the parameter configurations 6pos and Osem we obtained using the special initialization and frozen training stable configurations in the unconstrained parameter space 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 Opee and 0 0sem-
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 θ 8pos and 0 pos (respectively 8 and @sem) of parameters is also small (see Table 1}.
This evidences, that both resulting parameter configurations are flat regions of the full parameteri- zation 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 justification to refer to them "local minima".
19

---

finit Adam,
0 SGD, 0
accuracy 0
||
0|12
||0 0|1
positional
0.99885 0.0004
158.85 + 4.35
0.61 + 0.13
semantic
0.99895 + 0.0012
163.69 + 8.88
0.46 + 0.24
Table 1: Parameter configurations for different types of initializations for the histogram task. Opti- mizer parameters as in the accompanying code. Average over 10 runs using the same dataset.
20
M
10
H
-16
-42
Figure 5: Comparison of the attention layer activations for different sequences for ppos and 6 poss.
B Derivation of Result 4.2
In this Appendix, we provide detailed derivation of Result 4.2 of the main text. In subsection
B.2
we introduce a Generalized Approximate Message Passing algorithm (GAMP) [56]. Subsection
B.3
then establishes that equations (7) of Result 4.2 track the dynamics of summary statistics describing the GAMP algorithm. In particular, the equations 7 describe the fixed points of GAMP. Finally, subsection B. shows that fixed points of GAMP correspond critical (zero-gradee (rooogrrrient)) points of the empirical loss landscape (3). thus establishing that equations of Result 4.2 describe fixed points of GD.
B.1 Notations
For simplicity, we place ourselves in the setting r' explored in Section of the main text, but allow the length L of the sentences to be arbitrary, and allow a generic learning model (2). i.e. not necessarily the dot- not-product attention model analyzed in Section 5 The case r, 2 follows identical derivation steps, modulo the replacement all variables by tensor objects. We provide another alternative derivation of Result in full generality in Appendix C using the replica method from statistical physics. Let us note ( (Xcheees a series of n x matrices, with X, corresponding to the E-th rows (tokens) of each input sentence r stacked vertically, and normalized by Vd. We denote Xe = X2 + Pt, where P E Raxd is the matrix with all rows equal to the f-th positional encoding Pe. Let us further define p E RExLxL the tensor corresponding to sequence of matrices f,c (z E R Finally, let us denote T E x LXL the tensor so that the fI-th th row of satisfies y(**)) TH see equation (T) In other words, T corresponds to the concatenation of the target attention matrices.
20

---

Eoamplle Segueves
Same Seguecee N
Baamele Segree K
D B
M
M
D
M
H
20 299
O
HIG
M
Figure 6: Comparison of the attention layer activations for different sequences for 0sem and 6sem-
Before detailing the derivation, we first highlight a simplifying observation. Note that a loss item can be expanded as
1 d s(x*) - Vi (xt + +D)
xt
2 -11x(+*))1
Trs
Va (at
+
pr:
(o* + +11
Va
TTT
2.Q. Prs
(x + +n)
(27)
where we used that with high probability in the considered asymptotic limit, for all S μ Sn n
T- (ar + p) p) = x(x + +p) pr.
(28)
Since the first term of (27 does not depend on the weights Q, can be without loss of generality substracted from the loss. Without loss of generality, one can thus consider the equivalent empirical risk minimization problem
Q QeRaXT
argmin
Trs Ja (w" +2)4] pns (o" +ri@] -2nIT Q. 2.JoR.Jess (27 +p)))
X
1Q11..
2d
2
(29)
The risks
(29 and 3 are equivalent, and we shall use the former in the following.
Finally, for arguments T e R.xt, , E LL R.,, R glxx we introduce the resolvent prox(T, p, V)
arginf [R(
(xt wz) (en( wx) 2 [<[xpp]]]TT]] Tr [[t]] []
(30)
Note that the latter part of the bracketed term corresponds to the simplified loss (7 derived in the beginning of Appendix C which is the shall without loss of generality consider in the present appendix. For case presentation, we place ourselves under Assumption 4.1 where all the input covariances {Ee) codiagonalizable. In the following, without loss of generality, thus assume them diagonal, by placing ourselves in the common basis (e, (ehisia of Assumption 4.1
B.2
AMP algorithm
We are now in position to state the AMP algorithm:
21

---

Algorithm 1 GAMP
Inputs: {X, E B |En, E Raxtil, E RuxLxl
Initialize Q0= =~ (0,,,d)), = La
for t I faax do
V1 L L, Vix =(x Xx)e
V1 S S L. w/ = X1Q2 - - Vintth
VI 5 L. E(V-"))ee
p. V()x wh)
V x 4, Dex =wwefi
A' Σ (Xe (xxxx))Ta Tee
(,==1
6- Σ X f/ A O
Q+1 = (ATla + 1
2+ = (Ala + A')
end for
return
Estimator
The GAMP algorithm can be derived in standard fashion from the Belief Propagation (BP) algorithm, or
see e.g.
56 63
64 for an overview. Compared the standard GAMP iterations for General-
ized 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 f. valued attention matrices. In the following subsection, we show that the fixed points of GAMP [ 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.
B.3 State evolution
In this section we show that the dynamics of the GAMP Algorithm I are tracked by the summary statistics of Result 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 step upstream in the derivation of the GAMP iterations, and which are asymptotically cquivalent- see e.g. [64] for a review or e.g. 65 Appendix A, for a detailed walkthrough. The rBP equations read
As conventional, we note "n the version of a variable "u++ where the summation also encompasses the index i, and + the version of a variable "H-H where the summation also encompasses the index Jr. Note that in all cases above the two variables differ by at most (2(v/z))
u-vi A5 remembering that the data Concentration of ee We first study the statistics of VI
3/ = (=t)i//a + /oel/V/in in the notation of the main text, with (x²²); = Oa(1), (pt): (./v/7)). Replacing the rBP updates:
(V (Ve =
Σ (3H))))
Σ
(1)))
+ +
(at)) + (e k) +
E(De),
d
den (1.e (oa
(a(1)))
@x(1/4)
=de (EE)) = V
(31)
22

---

Algorithm 2 rBP
Inputs: {Xe E [ Raxixi, PEE Ruxiix,
Initialize v1 S S n, isisd, 0i-p== = =0, =1, -1, 03
for 1 S mmx do
VI f
(Vi =
371
2 M.L.155S -/ = Σ H O
MIS4SSsSmiiis fi (VV ViVyiii (prox([TTnn Pas wa Va-sde wt a
WS..sS5.1155 d, Pim.p-1 = Owe Oueshup-s
L
NISHSm11is444 = Σ (3
MISSmIIis4 6/ = Σ Σ
(114
M 0 = (AIa + Al
10
V1 S / Sn G4p (AIa + A )-1
end for
return Estimator Q
Distribution of wi Let us first introduce the teacher local field
huee E
(32)
Like c.g. (65), Appendix A, we first ascertain the joint distribution of hp w/ w with respect the data. These variables have mean
Pr Q' Vd
E B|w2
=
m
(33)
and respective variance
Viw p- w kp Spmden
Σ
(((E)) Spndengh
(34)
1.j
[[llueeeene 8pmden
Σ
(E = Spu
(35)
7
Elhpeewh = Spnden
1
(Eee)
O3 Spndentt
(36)
is
Distribution of 6/-
Let us ascertain the distribution of 69 bpp
69 6/-a = Σ
(i)) (V))-1 (prox(Te, Pv. twh dno Veniee - wh
8 Vg
Sprox(Tu.pa Jn.
ΣΣ 1/va((z)) + (pe).)
prox(T((hu-ma)))) Pur twhwessi Ve
8 V$#
+ +1/38 > prrp::/00..oox(//////aabbbbeebe Pur twho-sstrs Vonsle,
(37)
23

---

leading asymptotically to
p = <(vdpr) pCn PE, {m + Ve tWitwee
=m)
Qi Σ De 0Ennss, PE, {m Vihtnatss (V (38)
=0
where the expectations bear over & NO, 1) and he ~ N Vil (0t) The variance is given by
by = ōig ((E)) [.,, pr, {m + Voheahe (w))))
(39)
E
Concetration of A/-
Similarly the derivation for V V Ales concentrates to
Ainp = Σ
[HH::O.,prox(([[]] Pz. {m + (V - -E)
(40)
V
Wrapping up. we now massage these equations to recover equations from Result 4.2 of the main text. Starting from (3I)
Recovering Result 4.2
1 Σ (E) Vt à
A Σ V
1---- (Ex) (Ex)
-1
/ dull)
(+ (^+Σ
Σ
K
7/c
(41)
Next, for of 34):
((I (I(vala))o +0;((w)in
(A+I) ++
of
1 y d
E(Eei
+ (Ex)
(*+
(Ex)
((I mibilo Tn. + "+
) tmid (*+ Σ 1
2
/ du(t)
mm
(42)
For 0/(34):
(x+ (A+>) Σ 1z (2irr)
-1
02 02= - 7 E(Eiliis
+ Σ + @2(())al)))
+ +0a(1)
Σ mi To 0.TiTTT) (++ -++ 2
= =fdu(7)
(43)
Finally for ml (33
(x+ Σ 1
-1
m² [(Vape):
Σ (Vd(pa)em +1 ([W))))
+oa(1)
Σ E'm Tr + (ariatt)) ( (A+ E.E'"* ve-lgn
= =fduly
(44)
For m 38)
m = aEHHE Vt prox. V -mi
ml
(45)
24

---

while for di 38
0 aBH, The prox: v Vatte m3
The
0 = aBu Vi PC 101)²//i
proxe V mà
(46)
Now turning to 0Z:
H.E ( (v.DrOss - V mi) )
(47)
Finally, for V (40
V = H.H
12, prox; -
Vi
a. V7
de (prox, - Voi5o me)
02
he Vai5e
-s)
=,,
(proxe Volle -ms)
V
V
( 101)
0,61
(prox, Vatte-me)es me)
(48)
V
Summary State evolution equations dynamics of the GAMP algorithm thus read
The state evolution equations asymptotically describing the
V f du(r)))
x Σ 1
( M+Σ V4- 7m
-2
of du((,
Σ m/ Tx + OnTTTT
+ 7u
(49)
-1
0f = dulr, T, meme
Σ M Tr. + 0 7kTh
Σ Ve-i Vn
k
x
m f dd T)T
Σ m - TT + 1x
A + Σ Σinーlag
V 0 of
aEHHI Vo (prox, VoLle -
o aEH (
( prox, - V - mt) 27
(50)
0f EH H1
hi he-01 0//oTte
proxe v T.ce - m²)
102)+/03
m HHH
prox.
Votte mi
which exactly recovers equations 7 of Result 4.2 of the main text, for the case r. considered in the present Appendix. Again, we mention that the case r, 2 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
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.c. fixed points of gradient descent. Finally, we provide in Appendix C an alternative derivation of the state evolution equations (74)) using the replica method from statistical physics 61 66).
B.4 Fixed points of GAMP are fixed points of GD
In this subsection, we show that fixed points of GAMP as asymptotically described by 7 in Result correspond to critical (zero gradient) points of the empirical landscape (3). Again, we present the
25

---

result for r, = 1 for clarity, the generalization to 2 2 being straightforward (see e.g. (52)) In the previous notations, let us denote the (simplified, see (27) empirical loss as
+(1x:0))+((()
(51)
where we introduced the shorthands
L(the) E R")
2 ( "((((1))"
(52)
sQ)) = 2 1111?
(53)
i.c. respectively the simplified empirical loss 3 and the regularization, as functions with matrix arguments. The empirical minimization problem 3 can thus be written compactly as
argmin QERa
(((())))++()))
(54)
with the critical veeo-ggradien) condition being given by
Σ xx) +g(()
(55)
Let us choose a diagonal definite A E Rdxd, and a sequence (Vitisusa of symmetric definite L x L matrices. Group them into block diagonal matrix V E R LinxLn, so that the μ- th block of V corresponds Vμ- shall prove useful to further introduce the matrices X RddxL (resp. iT((X) IR=L), defined the concatenation of the matrices Xi,.., X (resp. O T, Or. L.), viewed as n blocks of length L. Then without loss of generality the zero-gradient condition can be rewritten as
X (VŌL(X) - + -
(56)
Similarly to 52 let us introduce
w VŌL(X) -xQ.
(57)
This can be written in terms of resolvent as
XQ = prox(
(58)
where
(
prox(w) E IR" nL = argmin
|| -1((+((+(
(59)
aER=L 2
which corresponds
30
Similarly, we denote
b= a=(()) +Q
(60)
So that
prox, argmin xER 1 b11 bl12-1+9(z))
(61)
In the particular case of an & regularization g(.) = A/211-112 note that
prox, (b) (AIa A) ' Ab.
(62)
The zero-gradient condition can now be rewritten as
TTV- (prox(() - w) A(b prox, (b) Xprox, (b) = prox(()
(63)
26

---

One is now in position to expand the concatenated variables into a sequence of L n- -n-dimensionn parameters. For u = prox (w), let us denote une (1 S p S n, S l S L) the e-th component of the fu-th block. Introduce
fue = E(V (prox(((()))))) - (pp)).
(64)
K
Denote ft = (fue)igasn E IR" we = (Wae) E R The system of equations 63 can then be rewritten as (further redefining Ab):
Σ E fe b- A(AIa + A)-
(65)
+ AA) Venfu WE
We used the assumption that g()) is an la regularization. Finally, introducing Q = proxx(A-1b) =
(ATa + 0, on reaches
- b- AQ
(66)
Viaaa
which corresponds to the fixed-point equations of GAMP (Algorithm D correspondence between the fixed points of GAMP and the critical points of the empirical landscape (3). To summarize, we have shown that equations describe the zero-gradient points of the empirical loss landscape i.c. fixed points of GD.
This finishes to show the
B.5 Towards a rigorous proof of result 4.2
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 588 55 60 48), there exist technical difficulties in adapting these rigorous arguments the present setting, and fully rigorous proof would warrant sizable work. While we leave this challenging task for future work, we wish to discuss how 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 Lemma in 60 a criterion sometimes
called the "replicon" in the context of replica theory [68].
Provided the replicon criterion 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-Minng Minimax approach of [69]. While it is used in many situations for convex problems (c.g. (1 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.
C Derivation of Result 4.2 with the replica method
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 formulation. First observe that for any test function ((Q) of the
61 66) in its replica-symmette
minimizer Q of (3).
x()) lim ED (Qo(Q)e-8RR)))) [R[Q]]
1 0x z
/
(67)
where we denoted R[Q] the empirical loss Bp. and
Z dQe-
(68)
27

---

the normalization factor, also known as the partition function in statistical physics. We remind that D refers the training set. In order access key summary statisties and learning metrics associated to Q. is therefore reasonable to seek to compute the generating function associated the measure (67). namely E In Z. Such computations can be addressed using the replica method from statistical physics 61 66]. building on the identity
In 7 lim -
(69)
The backbone of the derivation thus lies in the computation of EE Below, we detail the derivation for a generic convex regularizer g : R² + R4 and later specialize to the case of l2 regularization. The replicated partition function thus reads
E a(Qa)
EZ
dQae
a=1
Tri ta (ctt))))]
(x+))) (= -2THT =2 p= +/2
(=+t)_)))
e
#=1
(70)
Introduce the local fields
RQ Vd
hi* xQ. Vd
hta
E Lx
E R
(71)
and the overlaps
pQa vd
ma
BLxr
(72)
with rows m These fields have statistics
2eQ0 =
En dEN
Lab
(73)
d
E.Q.
E=(4((h2)]] = der
= Pt
(74)
d
Q L.Q. 0r
E, [1( = SEx
(75)
d
Thus
-dE m (T )
dI Σ
Tr
EZ =
dmdfinddddddde e
150500s.
/
-B
E ETTRONTT isasbs.
dQue
4m +m T 1 1ossh ad
- E +m
(76)
eoddy
where decomposed the replicated free entropy into the trace, entropic and energetic potentials v, V', y Note that all exponents are scaling with - oc. Therefore the integral in (76) can be computed using Laplace saddle-point approximation.
C.1 Replica-Symmettre ansatz
We have thus rephrased the analysis of the measure (67 as a optimization problem over the order parameters {g of ). and the associated conjugate variables. However, these still represent
28

---

2L(s² + + 2s variables, and + 0. In order to make progress, we assume that the maximizer is of aaplica-symmetie (RS) form 66 61
Lab (re qe)dah + 98
(77)
m me
(78)
Of O2
(79)
Tab (fe/2 + qe)) +
(80)
ri rîte.
(81)
0 0g
(82)
The RS ansatz holds in number of machine learning settings, notably for convex problems and Bayes-optimal settings, see e.g. [64] for review. 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.
C.2 Entropic potential
We now turn to the entropic potential V.. is convenient to introduce the variance order parameter
V = + de-
(83)
The entropic potential can then be expressed
BsdV g
dQae -0 $ E (1
/ DE
/ aQe
Pa(() 1 PO ( Vamert +0.4) + {
-p((2) t0ol & FesE.Joo soo 2 (Vāmeer + +.2..:)++ & =00())))))
(
= Ez / dee
(84)
Therefore
NV 1 /
BB In / dQe
V. oQ+ ( (Vamepi 2.)+ 200(SE)) -
(85)
For a matrix 3 E Rd and tensors A, B. E Rr @ Brxd we denoted (E A) E, 3d AAA and (A @ B 1g Ers Brakk.
29

---

C.3 Energetic potential
The computation of the energetic potential yy is rather standard and follows the same lines as in e.g. (20), yielding
BV
J 20x8
dYDZ
DE
II
8 ye (pe 0g qr (g) (1)))zz - 03 Se)
Rlxr
8=1
( $c)
Ve (=c-0ž xe
In
dX
- - -2 Ty
JRixe
del (2xVo)
(=1
ve o7 o2 5r)
(ps- (0c-0 20,) (v- (
{
dY
DE
JLLx
JRLxr
det (2T 0 0e)
t=1
=Bv;=
( to)
Vi (ar-az Ex)
x In
Jaxr
dx
e
--DTss/otm/ppolpotm -2 -1
c
det (2
8=1
(86)
C.4 Zero-temperature limit
We now take the limit B + cc. Rescaling
Bi, - Ve,
- Vi + Ve,
Bin + me:
30,+ + Or,
B t 40
(87)
3
The entropic potential then reduces to
V w
B= Tr
Σ
Σ
(Vamapl + +0Q..2)) + Σ Se o (ge x Ex)+
2d
(2)
(88)
where we defined the entropic Moreau enveloppe
M2 inf Q 2 heE.) 1/2 " ((a-((
VSE.) (( (Vamup?
+
(89)
The energetic potential can be similarly recast into a more compact form
Vy = M(Y
(90)
where the Moreau envelope is defined
M(Y -inl
Tr V
(Ix
1/2ge
827 mp)
T+(s())0=88))
(Tx(((())=(()))
X
2
(91)
C.5 Replica free entropy
One finally reaches expression for the replica free entropy as
I.
Tr Vege - -TraVV)) Σ
O 0, B=, Σ
mh
me
1=1
E=1
E=1
L
((Vatann++ -Vaiupp++0....0 +0:Q.E.) + Σ 3
627
+
Es Tr
Σ
V & 3
o (ee 8 E,)
- aF -aEy.MMHH.
2d
F=11
F=1
(92)
30

---

C.6 Saddle-point equations: general regularizer
The extremization of the free entropy 92 yields, similarly to 47 consistent equations on the summary statistics:
the following system of self-
Vi Es prox, @ (ee & Ex) O Ex) E
qe à E= prox, Exprox
(93)
my = Vd E= [prox gPt]
0e va Es proxx EcQ.]
1 |prox;- q2 E2 m²)
② aBz,yVi V
Vi-
Ve = 0.0% 98
a.EExVV; Vi-
(prox,
36 Ee- ml
1
Et a8
(94)
me = alBe,n
(prox; - q/ &7 - mm)
0e alLE Vi=1
(prox,- q8 Ee me)
(ye E, gr
101) (pe-07 qe 02)
where the proximals prox, and prox, respectively refer to the arginf in Q (resp. xze) of the envelopes Mg 89 (resp. M 91.
C.7 Saddle-point equations &
We now specialize the saddle-point equations (33 entropic potential admits the simple form
the case of an regularizer g(-) 1/211-11 the
VQ
Tr
AI, o
Σ
V/
Σe
Σ
qe 30 E₁
Σ (Vanintt U.Q.
2d
Tr
(~+
x+ - x(10)
m Xfau+ +
(Vitment at cu0 sees ( (Vanuri as ao Exes
2d
d+ 1
2 /
( x (^+
Σ reVe
Σ
4
reqe +
Σ Temr + ye -T
82
(95)
The replica free energy thus reads
①
1
Tr Vege -T
Σ
m m
T
0r d: a.E.sM(y.E))
2
1=1
1m1
)
+ //2
1
do(()
x+
Σ
ne
Σ
Σ
Tainee + 700e
E=1
(=1
(96)
31

---

leading to the saddle point equations
82
de a.',, Vi¹²
pprxx
9i ce m
V
Ve 0.07 qe
(prox,
oi Ee me Ee 41
me = a.Bg.VV V
(proxo
me
0.- - aB
prox;
1 q6 E3 - me)
(ye - 5 q1 1/00)) T
(pe -00-0/10 00 (1)))
q f du(,,
1. 1I, + Σ
The Vx
-1
Σ
L
In +
Σ
1
MrT + Tn :x)
02
(AL,
( x=1
Σ vnli
1
-1
x=1
1=1
-==
Vi [ du(y, TT
AI Σ
Nex Vi
x=1
my dd((),,
Xl, + Σ
7 Vn
Σ Tx + 7r00x TT
x=1
L
-1
0. = d,
All + Σ
Vhe
Vn
Σ MATN Tx + Tw0m T..
A=1
(97)
which finishes to recover 77) 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 633. which we clicit in Appendix for completeness.
C.8 test MSE
The generalization performance is measured by the test error
c9 = -
EpE T [
Vd 2 x
( (x +)
(x (x+
(98)
Expliciting this expression in terms of the correlated Gaussian variables xQ., 2Q allows to straight- forwardly show that eg admits the sharp asymptotic characterization terms the summary statistics characterized by 97
1 Eg T, Ex Tr S[X + 7,EY |
1
T(s((,)sTT(Y'', (99)
where the average bears on E gLxr, Y glxt with independent rows with statistics
(x),8e) ~ N
m
q8
0r.
]
(100)
0
C.9 Training loss
We 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 &g regularization g 1/211-12.. First note that the training loss e, can expressed as
1
6, =
lim 03 9
In ((3)
(101)
((1))
Where ((3) is the free entropy at finite temperature. The trace potential V bears no explicit dependence on 3. On the other hand,
BVQ
In det BXI, & Tla
Σ
V
& Σ
E=1
Tr
BAL, о ll
Σ
11 Ee
Ee & E3
Σ (Vdiuud U.Q. Ee
2d
(102)
32

---

Thus
X Tr
1
-1
0)(3Vi) = 2
BAI & Ia +
Σ
Vi Σt
[=1
(BAL. o
L Σ
-2
1
1 (Vdniust + +,00% Ee
Tr
BAL,
V1 8
Σ
Σ
de & E, +
2d
t=1
(103)
Finally, going through the same rescaling steps to take the / - 00 limit,
(MI Σ V
-2
L
62
li (3V)) B-000
2d Tr
AT, A
Ta
Σ de 00 Eg +
Σ
(Vdiert + D.
(104)
By the same token, it straightforward to see that
lim y)( Ey M(Y 1 Σ Tr
me Ve 7 to identify the term in underbrace.
Ve
21 qi
/
Ee-
002
(105)
8-+00
We used the self-consistent equations together,
Putting everything
(
L
82
- = lim 0z =
{ du(( d Tt
A (A+
Σ
VeVe
Σ rele +
Σ 1=1
Teme + 0e.
+000
8=1
L
aBy. (M(, + Σ
1
Tr Ve]..
(106)
2
8=1
This constitutes sharp asymptotic characterization of the training loss e, as function of the summary statistics characterized in Result [4.
For completeness, we finally explicit the connection between e, and the negative free entropy (i.c. the free energy in statistical physics). We go back to massage the expression for the free entropy
Tr Vege Tr
Σ
m m
Tr 0g - a(y,
2
1=1
L
-1
62
/ du((;:)
+ 2
A (x+E
VeVe
Σ
TEge +
Σ
Teme + 7600 a
8=1
1=1
1=1
(
1
002
Tr Le Ve + Veld - /du() Σ 1=1
Σ
reVe
Σ
Tede +
Σ
Teine + 7 -
t=1
8=11
1=1
aBy,
Tr|ae Vil + / de( x Σ
yeVe
$ Vere->
E 10/))
( 1 roett ( 4
remeet )or))
t=1
-aBy=M(Y,3)
4
1
82
A
(2(7,)) Tr
A+
Σ
TeVe
Σ
Geqe +
Σ
Tems 0:
2
(=1
1=1
1=1
1
-BEr (M +
TrlgeVi Vi]
(107)
33

---

In other words, the training loss is equal to the zero- pero-temperature free energy.
C.10 Extensions
For the sake of clarity and definiteness, stated the technical results for in the simplest instance, under notably the symplifying assumptions of independent tokens, tied key and query weight matrices, exposed in Section 3 of the main text. These assumptions 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 consider data distribution allowing for statistical correlations between different tokens. A simple model meeting this criterion, proposed in [73], consists in assuming every token (row) zy is drawn from Gaussian mixture with K₃ with respective probabilities pe,k.
K2 Σ pe WN De,a)-
Be, CE
(108)
In (108. the random variable ce [Kg] 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 introduce correlations between different tokens, we further assume that the variables c e= (CI,.... .... Cr.) follow a generic joint law p(c). Intuitively, the different models can model different types of words (e.g. verbs, pronouns, adjectives), and the nature of the words, reflected by the variable ez, is correlated across tokens. We assume that || Ox(1), i.e. are of the same order as the positional encodings.
Assumption C.1. Finally, similarly Assumption 4.1 let us assume that the set of matrices HE Ke 1 admits common set of eingevectors {e:já= with eigenvalues { Ji=1 The eigenvalues { Jek, and the projection of the cluster means (to which the positional encodings are added) {pe,K + pe be, and the teacher columns {Q1) on these eigenvectors are assumed to admit well well-defined joint distribution v as + DC namely, for 7 (Ne,,)e, & (KI,..,;T, Te) E R"
=T,
Ke IIII
8 (1. 7exx) (vde, (Ne, + +P) TEK
II 8 (eT - -T,) d (T,T,)) j=1
7 ==
(109)
Architecture Again the interest of further generality, we alleviate the architectural restrictions on the attention ( and allow for untied key and query weight matrices. More precisely, let consider function of the form
fx. = [ (a p)Q.
val (a (/(0=
p)k
(x + p).
(110)
parameterized by two of weights Q, K e gdxr, In order to work with more compact expression, let us introduce the total weight matrix W E (dx((r.+r.)) obtained as the horizontal concatenation of Q, K. Then (110) can be written compactly as
fk.2(z) - (w))))
[
(a + p)w (x (c+p) p),
(111)
where for any argument h (lx((ttr)), wiewed 2 block matrix with two blocks h1 h2 Lixr, G(h) = S(h1,h2)
Asymptotic characterization A sharp asymptotic characterization of the test error and train loss for this more general model can derived along the same lines, as presented in Appendix (see also [73). For the sake of conciseness, only report the final result.
34

---

Result C.2. Under Assumption C.I in the limit d + oc, llpll, 5/d, L. Ob(1). the summary statistics
PE, Q2 Er
Dr,W
E "
qe,k d
E 122xx22,,
d
W WHERK + pe)
W MTA
me, =
E E
02 =
E xr
(112)
d
d
concentrate in probabilit,, and are solutions of the set of fuitteddimens self-consistent equations
de. = alE &e (proxs Vk,kEe mp,k) e,k me.k) gek
Ve, 0, 1 allieden Vi (prox;-
Re
gek
mer = aB,dee RE,r Vi.k
(proxe - gi,55 me,k)
0c,4 N.. Bg Vel (prxxf g. mex)
(ve- 02. ke) (pe. 0f -1
(113)
-1
9r.k f (,, Te, Allar. + +ΣΣ Noa Disd
82
Σ Tr.b + OnoT
+ Σ Σ AnJ
Alara EΣ Tni Tmo Vark
Ve, J dv(7), T,
Al2r, + +ΣΣ Mhar Vnnd
me, = d,
Al2r,
+ 6
+EEns. Pn. Vr.j
Σ Σ m JRi + Tn. InsORITT K 3
0.. du((,, T)A
AT2r. + Emnn
The Ennnia Vans
ΣΣ MagTas Jnj + O 7
(114)
The averages bears over E E Y E BLix,, with rows & (N.0.22.. independently and 0' yt NET qe,c Pe,c,PE,c TT,,, 0e,e)) conditionally on prox corresponds to the arginf of the minimization defining the Moreau envelope M(c, Y,E), as
1/2 Ex mg,02 Xx - -2,
82
+ eYY + +m X,ex
prox arginl
Tr
(115)
2
1=1
We introduced m R.xx,, whose rows
HL Q. Vd
(m
(116)
and the shorthand
Cly,x,0 -[ - TT [[y]ep
(117)
where P& TRLXX is diagonal with elements
(p$))e /
dull,, N NVt.er
(118)
35

---

Test error The test ervor admits the sharp asymptotic characterization
1 Eg - - [Tr[TTh"]SSTIN + [T -2Tr []]
(119)
where, conditioned on the class assignments c, the average bears on 2x2,, Y E R Rlx, with independent rows with statistics
(he. h7) ~ N (mea))
(
9t.cr 0gces
Or,ee PE.a
]
(120)
m
Finally, the train loss e converges in probability to
(>ler.
-2
a = / du (7. 5.
Allr. +
ΣΣ Teek Vaikeaik
[ me, taanra
+ 1m..
L
Kr
+ ((BoysMMcxX,3)
Tr( Ve..o.
(121)
2
Trainable value matrix Finally, let remark that it should be possible leverage similar derivation to further accommodate trainable value matrices V, provided they are low rank, i.e. in Rdxrv with rv Oa(1) as d cc. Similarly, introducing the stacked weights W = (QKKIV) (dx((r,trr)) allows perform the asymptotic analysis. Since this architectural modification then implies that f now takes values in g_xrV, one cannot immediately analyze this architecture for the task considered in the present manuscript, as the labels are valued in 12dxd and the setting and questions explored would need to be entirely modified. We thus postpone a thorough analysis of the effect of trained value matrix to future investigations.
D Derivation of Result 5.1
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)
R(W)= Σ [A" )x* - Wx
(122)
where we use the shorthand notation A T[3/Vx[Q2] for the target attention score matrix (T) The expression for the risk can be asymptotically simplified
R(W)
1 n p=1
| (x*
(al (= -AA (A"("") - + +h.c) + WpeW + +o(//a)
E. (A'
E. +h.c) W
E. (A (z) - E. (x)) (A E. [A)( + | (w E. -W-E, [A))))))
(123)
Therefore the learnt weight W is simply equal to
W E., [A*())) = E. [T(/VarQ.]]-
(124)
Naming = 1/vanQ, the argument of the last term, the matrix h possesses independent rows with statistics
h4 ~ NN(0, pt)
(125)
where we remind that P: pe I/4QI E.Q.. Therefore, the trained weights W. 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 - oo and = @x(1), training W, i.e. number 2 Oa(1) of parameters on the empirical loss for n < 2 data points is equivalent
36

---

asymptotically to directly training on the population loss. Expliciting the corresponding test MSE, one reaches the sharp asymptotic characterization
Jin y
Tr W p E. TT(TT]]sTT[]]
1 En Tr piTT"TT
(126)
We close this Appendix by giving few examples for definiteness.
D.1 Examples
Purely positional target Let consider the case where the target is purely positional, i.e. T/ = A for all z. This corresponds to the w = 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
W (1 ) softmmaaa ) WA
(127)
with h and having independent rows pr).
E
Supplementary Experiments
E.1
Empirical scaling of a = d/n
In the following we verify that our experiments consistent with the scaling behaviour predicted from the theory. We jointly increase d and n for a fixed value of a. In Fig. [w indeed observe the expected behaviour for exemplary value of 2. The same holds for the summary statisties and m, which concentrate as d and n jointly grow, shown in Fig. (left) in the main text.
positional
semantic empliness
0.006
theory
0.005
0.084
0.003
d
Figure 7: Scaling d and jointly for 1.5 approaches the theoretical prediction of the generaliza- tion error of the positional and semantic local minima. Experimental settings as in Fig 2 with 70 runs per datapoint.
E.2 Alternative hyperparameters
provide supplementary results for different parameter settings. Fig. 8 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 specific A for definiteness. the following, we present the same results for different A with stronger off-diagonal and higher rank,
0.3 0.7
(128)
0.8 0.2
In Fig. we present the analogous simulations to Fig. (center, right). While the global phenomena match the previous example, the details of the transitions location differ.
37

---

0.0
083
6.04
5810
6100
1002
080
Figure 8: Alternative Parameters. Mixed positional/semantic teacher for w = 0.3. Settings is To L, A ((0.6, 0.4), (0.4, Σ1 ΣS 0.25 p1 1d P2 and Q. ~ (2). While keeping all other settings the same, we vary from left to right: The target positionality w, the student regularizer À and the standard deviation (which is 0.5 v Eii Experiment settings in Fig.
0.5
0.0015
0.5
0.0020
0.0010
0.0015
0.4
0.0010
0.0005
0.3
0,,0000
3
0.0000
-0.0005
0.2
-0.0005
0.0010
0.0010
0.1
-0.0015
0.0015
-0.0020
1.0
1.5
2.0
1.0
1.5
2.0
a
Matrix. r 2 P1. Q. (O0 Ia) independently. Here, use definite matrix A from (128). which differs form the one used in the main text. Experiments were conducted as in Fig
Figure 9: Alternative
Eg 0.25 = = and
E.3 More Complex Architectural Choices
As discussed previously Section C.10.
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. |0 compare the change in architecture with the setting considered in the main text, for a given w 0.5 and varying sample complexity a. The exhaustive experiments for varying are shown in Fig. 14 which immediately shows that the general idea of the phase transition from positional 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 larger values of (Fig. 10). For the case of a student with 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. Interestingly, the dual initialization seems to do approximately 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.
well the best of the two other
When we compare tied and untied Q and K weights in Fig. 12 we observe that the additional parameters we need learn in the case of an extra K come at a small cost, as the transition the semantic solution is moved to the right.
38

---

Finally, for value matrix (Fig. 13 the results seem to largely resemble the transition line we observed without the value matrix.
0.10 Seroadg-14
Figure 10: Uncorrelated vs. correlated inputs (w= = 0.3, 0.5, A= 0.001). In the main text we sample all datapoints R=xx such that the columns are independent. We compare this setting with a correlated data structure with hidden latent: We sample three vectors u, ~ No 1). and set the first column x to be (ou y0)/ vo + and the second one (ow + 70)/ /0 o Experiments are repeated 5 times per data point with d = 1000.
unmod sed Draaee
tigher rank studen: Sraser 21
20
2.6
3.
23
Figure 11: Rank 1 vs. rank 2 student (w 0.3, 0.5, = 0.001). 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 using the semantic strategy. Experiments are repeated times per data point with d = 1000.
urmodifed
urteed O.K
1.5 25
Tied vs. independent weight Q, K (w = 0.3, a= 0.5, - 0.001). We compare the student setting from the paper, where the query and key matrices are bound to each other with the setting where we set them independently. We initialize with them both being either close to the positional close to the semantic initialization. The phase transition for the semantic minimum dominating moves to the left, i.e. more samples are now needed. Experiments are repeated 5 times per data point with 1000.
Figure 12:
39

---

unmozilied
macile CA
podtioaal K0:
No value matrix vs. value matrix. (w 0.3, 0.5, 0.001). We compare the setting from the main text with adding value matrix, i.e. trainable parameter V Raxd This applied to every embedding before they are averaged over using the attention matrix. We ran the experiment times for each with d 500, and rescaled the training compare the experiments in the main, where d = 1000.
Figure 13:
correlated (y= 0.1)
higher rank student (Fuudeer
0.5
0.4
0.4
0.3
0.3
AF
0.0005
0.0005
-000100
0.1
-0,0010
-0.0015
-0.0015
0.0
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
untied Q.
with value matrix VERaxa
0.0015
0.0010
0.0005
0,0005
0.3
0.3
0.0000
0,0000
0.2
-0.0005
-0.0005
-00000
-000010
-0.0015
-0.0015
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
Figure 14: Phase diagrams of the difference in training loss between the semantic and positional solutions from Fig. 10 13 in terms of and the sample complexity The green dashed line represents the theoretical phase transition which we obtained from the architecture as considered in the main text.
E.4 Uninformed initialization and training via Adam
In experiments, to obtain the empirical results, we initialize the GD optimizer in an informed fashion, i.e. initializing Q- of the student with r either p1 (positional) or Q+ (semantics). GD then converges in the two local optima described by our theory.
40

---

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. [3] we show that indeed running the Adam optimizer from an wninformed initialization may lead one to either of the local minima for 100. For larger d we observe the semantic minimum reached less often than the positional minimum, and considerable number of times the algorithm simply does not find either them.
informed intidlizaton, GD, d=1000
random iriialization Adam, d=100
18
BNTA//:
1.8
57
377
122
0.2
0.1
0.0
2.5
0.0
3.5
m/3
Figure 15: Comparing GD and Adam. Settings as in Fig. 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 single run. For the informed GD, we used the same optimization parameters as in Fig. 2 (24 runs per initialization). For Adam we trained the same data, but for 2, 500 epochs with learning rate rj 0.01 (showing 140 runs).
41

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, it just matches.
Guidelines:
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: There is 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 not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [No]
42

---

Justification: At least not fully- - our theoretical results are on the level of rigor of theoretical physics and hence we provide assumptions and formal derivations (which are all provided in the Appendix of this work), but not proofs. We verify the match between the theory and empirical simulations both the main text and the Appendix. We verify the deterministic concentration of properties in the high-dimensional limit through empirical simulations in the main text and the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
The full code available in the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
43

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: The code will be made public on github as it available with the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
wublic//giide//Cdeaabbimsss www.aaammmiioooooooooooooooeeiie for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//nips //iiss.c//pule
guides ProddeSubmissiooool....o
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: The concrete details how many samples per batch etc. are contained the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: Error bars are provided.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
44

---

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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification:
The compute needed is negligable, so 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 compute required for each of the individual experimental runs as well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: :ttpp:::///eurreeeee hwwrips........ hctpppblic//thicccciiiieeel.llll
Answer: [Yes]
Justification:
It is a theoretical work that aims to advance understanding only.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. 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.
45

---

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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
is foundational research.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
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 dataset.
The authors should state which version of the asset used and, possible, include 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.
46

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The code is well documented and accessible.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This work is foundational theoretical research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This work foundational theoretical research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
47

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
48


Reviews:
[{"review_id": "DsP2WIdfkt", "reviewer": "Reviewer_uCdk", "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$.", "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?\n\nThe 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 are identified and discussed in the paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper contributes an original result on the learning theory of attention models, along with a creative and insightful application of this result to a setting that contrasts positional versus semantic solutions. The result appears significant as a theoretical lens through which to characterize the loss landscape of an attention model, and it may have applications beyond the specific positional-vs-semantic target model studied in this paper. The exposition and empirical illustration of the result are clear.", "weaknesses": "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. \n\nIt 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their appreciation of our work and their constructive questions, which we address below.\n\n> the stated result applies to a simplified model [...] in terms of [...] the input data.\n\nWe have indeed chosen for the sake of clarity to present the results for the simplest instance of input data distribution --namely uncorrelated Gaussian tokens--, for which the trained attention can learn both positional and semantic solutions. The analysis can however be extended to include statistical correlations between the tokens, and cover more structured Gaussian mixture token distributions, see also the answer to reviewer QM4T, at the price of much heavier equations. For this reason, we chose not to discuss this general case for the sake of clarity. We generically expect that introducing correlations leaves the phenomenology qualitatively unchanged, but shifts the phase transition towards higher sample complexities (i.e. increases $\\alpha_c$), as we illustrate in Fig. 1 of the attached pdf. We will include these additional theoretical discussions, alongside the additional figure, in the final revision of the manuscript.\n\n> 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?\n\nThe analysis can actually be extended to include trainable value weights, as well as untied key and query weights, provided they remain low-rank. On the other hand, these extensions come at the price of more cumbersome equations -- for instance, the size of the summary statistics matrices of Result 4.2 triple, and the expression of the Moreau envelope (l.168) becomes sizeably more intricate. For these reasons, we chose for the sake of clarity and conciseness to present our result under these various simplifications. We shall include a detailed discussion on how the analysis can be generalized to these cases in the camera-ready version of the manuscript.\n\nWe illustrate in attached pdf how these architectural changes affect the phase transition : Fig. 3 shows how untying the key and query weights keeps the phase transition but shifts it to towards higher sample complexities (i.e. larger $\\alpha_c$). Fig. 4 shows how appending a trainable value matrix leads to no sensible change to the phase transition.\n\n> 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)$\n studied in the experiments.\n\nWe will include further discussion on these two points in the final revision of the manuscript. The solutions of (7) were found by numerically iterating the set of self-consistent equations (7) until convergence, for various initializations. The training loss of the different fixed points thus reached was then evaluated using equation (12) of Result 4.2. We found in all examined settings, both in the theory and in experiments, that the lowest training loss was achieved by one of the two fixed points (positional or semantic) discussed in Section 5, and thus concluded that the global minimizer belongs to this pair of fixed points.\n\n> 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?\n\nAs the reviewer correctly surmises, Result 4.2 provides a very versatile playground to explore other models of attention-based learning, beyond the model considered in the present manuscript. Other aspects of attention mechanism which can be directly explored through the lens of Result 4.2. include the use of causal masks, and cross-attention mechanisms. A thorough study theoreof however warrants separate works, and fall out of the scope of the current manuscript, whose focus is on the interplay between positional and semantic learning. We will however include a highlight of these other aspects in the conclusion section of the revised manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jPQW0e3v5F", "reviewer": "Reviewer_QM4T", "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).", "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": "Yes", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* Provides theoretical analyses of a model of self-attention, finding a closed-form solution.\n* Asymptotic analysis demonstrates a phase-shift between two minima, one which relies on position and one which does not.\n* The first analysis of this type to an attention layer, instead of just a feed-forward layer.", "weaknesses": "* Very dense mathematically, so hard to follow for a reader not intimately familiar with the literature to which it contributes.\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their reading of our work, and the many interesting questions, which we answer below.\n\n> Very dense mathematically, so hard to follow for a reader not intimately familiar with the literature to which it contributes.\n\nWe will take advantage of the additional page allowed in the camera-ready version to provide further intuition-building discussion below the statement of Result 4.2, alongside further context and discussion on how our contribution fits in the broader literature in the related works section.\n\n> 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.\n\nWe thank the reviewer for raising this important point. The assumption of independent tokens was actually made for the sake of clarity and conciseness of presentation, and can be relaxed in the analysis to include generic statistical correlations between the tokens, as we further detail in our answer to the following question.\n\nLow-rank attention weights have been considered in practice in language modeling in the context of model compression, see for example [a]. In this approach, the weights of large language models are approximated by low-rank matrices to reach a smaller model, easier to fine-tune and deploy. The idea to train low-rank weights, at least at the finetuning stage, also underlies the celebrated LoRA scheme [b], which allows for the resource-efficient yet effective fine-tuning of large language models.\n\n[a] Hsu et al, language model compression with weighted low-rank factorization, ICLR 2022.\n\n[b] Hu et al, LoRA: Low-Rank Adaptation of Large Language Models, ICLR 2022.\n\n> 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 \nn-grams instead of unigrams)?\n\nThe reviewer's intuition is right that the amount of correlation between tokens affects the phase transition. In Fig.1 of the attached pdf (see global rebuttal), we provide additional numerical experiments showing how introducing more correlation between tokens shifts the phase transitions to higher sample complexities, i.e. increases $\\alpha_c$. On an intuitive level, this is because each data point carries less semantic information (as correlations make tokens more redundant), and therefore more data points are needed to identify the semantic content, and learn a semantic mechanism. By the same token, we generically expect that the phase transition happens at higher sample complexities for n-grams than for unigrams.\n\nWe would also like stress that the analysis can, in fact, be extended to cover arbitrary statistical correlations between different tokens, albeit at the price or heavier equations (see also the answer to Reviewer uCdk). For these reasons, we have chosen for the sake of clarity not to discuss the effect of these correlations, and focus on the uncorrelated case, which already yields a very rich phenomenology in terms of semantic versus positional learning. We will however include a discussion of this generalization in Appendix A of the camera-ready version, alongside the aforedescribed figure.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F1ET7nTpQc", "reviewer": "Reviewer_wGcv", "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.", "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": "Yes", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Understanding the properties that lead to phase transitions in neural networks, and more broadly understanding training dynamics in transformer models is an important area of research. This work provides the first theoretical result showing phase transitions existing in attention mechanisms from learning, and thus opens the door for more work on the learning dynamics of transformers. I think the findings of this paper are important for interpretability research, but since it is not my area, I'm not able to strongly recommend it one way or the other.\n\nThe design for the task and properties resulting in the phase transition are clear and simple. The authors are able to empirically test their results and find that models training on data nearly match the theory. The purely positional baseline provides a nice comparison for the dot product attention mechanism which *can* modulate between positional and semantic information as a result of the phase transition. Other empirical results are well justified and presented clearly.", "weaknesses": "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.\n\nsome 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.\n\nThe authors make many simplifications on the attention mechanism", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful reading of our work and insightful comments, which we address below.\n\n> 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.\n\nWe completely agree that further understanding which of the characterized minima is reached by a given optimizer under given conditions (in other words, the dynamics) is of great interest. As evidenced in [9] in a related setting, the answer to this question can depend on many optimization hyperparameters. While we give some elements of answer in Appendix D.3, we believe a thorough answer to this question warrants a careful empirical investigation, and theoretical analysis, of the training dynamics. This significant research endeavour however falls out of the scope of this first work-- whose primary focus is indeed to provide the first theoretical analysis of a phase transition in the learning of attention models, as appreciated by all three reviewers. We will however further emphasize the importance of this future direction in the conclusion of the revised manuscript.\n\n> 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.\n\nWe will take advantage of the allowed extra page in the camera-ready version to include further intuition-building discussion below 4.1 and 4.2, to describe in more detail the content and consequences of these technical statements.\n\n> The authors make many simplifications on the attention mechanism\n\nThe results are indeed stated under four main simplifying assumptions, namely (a) uncorrelated tokens, (b) tied key and query weights (c)  value weights set to identity and no readout (d) all weights are low-rank . Assumptions (a,b,c) can in fact be relaxed, and the analysis can be extended to incorporate arbitrary statistical correlations between tokens, generically untied key and query weights, and accomodate a low-rank trainable value matrix -- at the price of much more intricate and heavier equations (see also our answers to Reviewers uCdk and QM4T). For this reason, we have chosen for the sake of clarity to restrict the discussion to a simpler case, already exhibiting the rich phenomenology of semantic versus positional learning. We will highlight these generalizations in Appendix A of the revised manuscript, and include a pointer thereto in the main text.\n\nAssumption (d) is on the other hand required for the analysis. Note however that weights are enforced to be low-rank in a number of practical settings in language modeling, notably in the context of model compression [a] or finetuning [b], to ease the costs induced by finetuning or deploying large language models, see also our answer to reviewer QM4T.\n\n[a] Hsu et al, language model compression with weighted low-rank factorization, ICLR 2022.\n\n[b] Hu et al, LoRA: Low-Rank Adaptation of Large Language Models, ICLR 2022.\n\n> 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?\n\nAfter the reviewer's question, we ran some preliminary experiments that suggest that including a learnable value matrix leads to no sensible change to the phase transition. These experiments are illustrated in Fig. 4 of the attached pdf, where we also include experiments illustrating the effect of other architectural changes (see also the answers to reviewers uCdk and QM4T). We will include all these figures, and a discussion thereof, in the final version of the manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Thank you all for taking the time to read and review our work. 

In this global response we post the pdf file containing plots of some preliminary experiments that help to clarify several questions raised in the reviews. 
We refer to this global file in the separate responses to each reviewer.

Precisely, we show how the phase transition between positional and semantic minima moves empirically for:
1. Uncorrelated vs. correlated inputs,
2. Rank 1 vs. rank 2 student,
3. Tied vs. independent weight Q, K,
4. No value matrix vs. value matrix.

We will provide more extensive versions of the same experiments in the appendix of a camera-ready version.

Meta-review: 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.
Decision: Accept (spotlight)