Title: A Phase Transition between Positional and Semantic Learning in a Solvable Model of Dot-Product Attention
Venue: NeurIPS 2024
Accepted: yes
Keywords: replica method

statistical physics

phase transition

high-dimensional limit

attention layer
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-Product Attention

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.

1 Introduction

Recent years have seen an upheaval in our ability to learn and implement complex tasks from textual data. Instrumental in these advances is the use of self-attention layers [1], which provide an efficient method of extracting information from sentences - both the information encoded in the ordering (i.e. positions ) of the words, and that encoded in the meaning (i.e. semantics ) of the words. In theory, attention layers can learn to leverage both types of information, by having tokens attend to each other based on their respective positions (a mechanism called positional attention in [2]) and/or respective meanings (henceforth referred to as semantic attention ). In this paper, we aim at a theoretical understanding of the emergence of these different mechanisms in attention layers, and the transitions therebetween.

Many empirical studies have provided evidence for the emergence of specific algorithmic mechanisms (abilities) in the learning of language models that lead to qualitative improvements of the model capabilities [3, 4, 5]. By reverse-engineering trained models into human-interpretable components [6, 7, 8] a growing body of work on mechanistic interpretability aims to empirically understand which precise algorithmic mechanisms a neural network learns. Such investigations have demonstrated that attention layers are able to implement a wide range of different algorithms, even for the same task, using both positional and semantic attributes of the inputs. We offer a particularly simple

y

∗

= (1

-

ω

)sftm

(

1

d

xQ

∗

Q

∗

T

x

T

)

x

+

ω

Ax

ϵ

t

(

ˆ Q

)

positional

semantic

A

B

C

x

∈

R

L

×

d

,

x

ℓ

∼ N

(0; Σ)

ˆ y

= sftm

(

1

d

(

x

+

p

)

ˆ

Q ˆ Q

T

(

x

+

p

)

T

)

x

student with positional encodings

sample complexity

few

many

positional

semantic

teacher mix

training loss diαerence

positional

semantic

is lower

is lower

Two local minima exist in the training loss▷

y

∗

= (1

-

ω

)sftm

(

1

d

xQ

∗

Q

∗

T

x

T

)

x

+

ω

Ax

ˆ y

= sftm

(

1

d

(

x

+

p

)

ˆ

Q ˆ Q

T

(

x

+

p

)

T

)

x

teacher◁target

tokenized data

length

L

dimension

d

0.5

0.0015

0.4

0.0010

0.0005

0.3

一

3

0.0000

△Et

0.2

-0.0005

0.1

-0.0010

-0.0015

0.0

0.5

1.0

1.5

2.0

α

illustration of this idea in Appendix A, where we show that for a sequence modelling task involving counting different algorithmic mechanisms co-exist, each corresponding to a distinct local minimum of the empirical loss. Out of these possibilities, the precise implementation that a model learns during training is jointly affected by its architecture [9], the training procedure itself [7, 4, 10] and the available data [11, 12, 13]. It remains an open question how to theoretically characterise the conditions under which a given behaviour emerges in the model leading to said qualitative improvements.

From a theoretical perspective, even the nature of this type of algorithmic emergence is unclear. It is not known whether it is simply a fast but smooth change in performance or whether the emergence is due to a sharp boundary between fundamentally different regimes of learning [14]. In our work, we take inspiration from physics, where a similar theoretical question about the nature of phase transitions was posed a century ago for models of interacting particles, such as the famous Ising model describing ferromagnetism [15, 16]. In the limit of infinitely many particles, it was shown that it is possible to theoretically deduce sharp discontinuities in some properties of interest (e.g. the magnetization of a magnet), delineating qualitatively very different regimes (e.g. magnetized or not). While mathematically, a large size limit needs to be considered to confirm the existence of sharp phase transitions, this asymptotic theory usually closely matches simulations, even for relatively moderate finite sizes.

In the theory of feed-forward fully connected neural networks, phase transitions in the network's generalisation ability as more training samples are available were studied as early as in [17, 18], and their existence was proven mathematically rigorously in [19]. In those works, the limit of many particles corresponds to taking the number of training samples and the dimensionality of the data to infinity at a fixed ratio. These theories rely on the property that macroscopic quantities of interest, such as the test error, become concentrated and deterministic in the high-dimensional limit. A dimension-free set of equations is then derived which predicts these deterministic quantities. Since then, a plethora of works along these lines both in statistical physics of phase transitions and in the theory of feed-forward neural networks have continued to study these phenomena, see e.g. [20, 21, 22, 23].

In this work, for the first time, we bring this type of study to the analysis of neural networks with attention layers. While several previous theoretical studies of the attention mechanism considered some type of high-dimensional limit [24, 25, 26, 2], none of them identified a phase transition between different types of mechanisms that are implemented by the attention layer. Simultaneously, the finite-dimensional real-world models that are the focus of works in mechanistic interpretability do not lend themselves to a tractable definition of a high-dimensional limit, which is necessary to identify a phase transition theoretically, as explained above. We aim to fill this gap by introducing and analysing a tractable model that permits such a sharp high-dimensional characterisation for attention layers (see Fig. 1).

In particular our contributions are:

We describe a model with a single self-attention layer with tied, low-rank query and key matrices. On Gaussian input data and realizable labels, we show that this model exhibits a phase transition in terms of sample complexity between a semantic and a positional mechanism.

As the main technical result, we analyse this model in the asymptotic limit where the embedding dimension d of the tokens and the number n of training samples grow proportionally. We provide a tight closed-form characterization of the test error and training loss achieved at the minima of the non-convex empirical loss. Using this high-dimensional characterization, we locate the positional-semantic phase transition, thus providing the first theoretical result about the emergence of sharp phase transitions in a model of dot-product attention.

We contrast the performance of the dot-product attention layer with that of a linear model, which can only implement positional mechanisms, and show how the former outperforms the latter once it learns the semantic mechanism, highlighting the advantage of the attention architecture for this task, when there is a sufficient amount of training data.

In Section 2, we discuss further related work. Section 3 defines the general version of a solvable model of tied low-rank attention, and Section 4 provides a tight characterization of the global minimum of its empirical loss. Section 5 analyses a concrete instance of dot-product attention and demonstrates that in this case the global minimum corresponds to either a semantic or positional mechanism, depending on the training data and task, with a phase transition between them. We conclude with a discussion of the limitations of our analysis in Section 6.

2 Related work

Theory of attention Attention models have been the object of sizeable theoretical scrutiny in recent years, with a growing body of work investigating various aspects such as their expressivity [25, 27, 28], inductive bias [29, 30, 31], training dynamics [2, 32, 33, 34], and in-context learning [26, 35, 36, 37, 38]. [39] and [25] analyze models with frozen non-trainable queries and keys, under the lens of signal propagation in such frozen models, or of their expressivity. [2] similarly studies the learning of the value matrix and positional encodings only, fixing keys and queries to identity, and shows how a transformer with a single attention head can learn spatial structure with a purely positional attention mechanism. The works of [29, 37] analyze the learning of a single layer of attention, with trainable queries and keys, assuming linear or ReLu activations - instead of the standard softmax activation. [33] provide convergence bounds for non-linear transformer models with a single attention layer, with trainable queries and keys. Because these studies are not tight, they do not allow to capture sharp changes in the behaviour of attention mechanisms such as phase transitions, and cannot for the same reason provide a theoretical description of the sudden emergence of new algorithmic mechanisms. A first tight description was provided in [24], in the context of learning a high-dimensional graphical model with a single layer of factored attention, leveraging its formal equivalence to a linear and convex learning problem. On the other hand, this model does not exhibit any emergent phenomenon. [38] analyze how induction head mechanisms can be learnt from gradient descent on the population loss, i.e. when an infinite amount of data is available. The present manuscript conducts a tight analysis of the non-convex learning of a non-linear attention model with trainable tied queries and keys from a finite train set, thereby allowing the description of sharp phase transitions in sample complexity in the behaviour and performance of the model.

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, e.g. recurrent architectures [40], where the positional information is provided implicitly by the order in which the input is processed. While some transformers can leverage implicit positional information through causal masks in training [41, 42, 43], in principle a dot product attention layer requires an explicit encoding of positional information as it views the input sequence in parallel, as a bag of words [1]. Several works experimentally explore different types of positional encodings with the goal of improving the downstream task performance [44, 45]. In this work, we provide a tractable model to quantify the generalization error of a single layer of attention in the presence of positional encodings.

Theory of phase transitions in neural networks In supervised learning with feed-forward fully connected neural networks, phase transitions in sample complexity were identified in settings where the data consists of random Gaussian samples, and the labels are generated by a target neural networks with random weights. For a single-layer perceptron and a variety of teacher weights distributions and activation functions, a discontinuity of the optimal test error as the number of samples increases was established in [18, 17, 19]. For a two-layer neural network, [21, 20] evidenced a specialization threshold in the sample complexity below which linear regression matches the optimal test accuracy, and above which a strictly better accuracy can be reached. To our awareness, phase transitions in neural networks with attention layers have not been studied theoretically yet.

3 Tied low-rank attention model

Input data model We consider a model of embedded sentences with uncorrelated (1-gram) words. More precisely, a sentence x ∈ R L × d , where L is the sentence length and d represents the embedding dimension, consists of L tokens { x ℓ } 1 ≤ ℓ ≤ L independently drawn from a Gaussian distribution x ℓ ∼ N (0 , Σ ℓ ) with covariance Σ ℓ ∈ R d × d . In the following, we denote the probability distribution of x as p x . Note that while this sentence model does not involve in itself statistical correlations between tokens, the task (target function) will entail interactions between different tokens. While more general data models involving inter-token correlations can also be readily analyzed such analyses come at the price of much more intricate analytical formulae. We thus choose for clarity to restrict the discussion to this simple instance, which already displays rich phenomenology, as will be explored in Section 5. We defer a discussion and an analytical treatment of the general case to Appendix C.

Target function The target function (teacher) is assumed to be of the form

for a function T : R L × t → R L × L . The term T [ 1 / √ d xQ ⋆ ] ∈ R L × L in (1) should be interpreted as the target attention matrix, which mixes the tokens of the input x . This attention matrix is parametrized by the target weights Q ⋆ ∈ R d × r t .

Tied attention We consider the learning of the target (1) by a single attention layer

In (2), p ∈ R L × d is a fixed matrix, corresponding to positional encodings, and Q ∈ R d × r s is a trainable weight matrix. We denote subsequently p ℓ ∈ R d the ℓ -th row of p . Like the target (1), the parametric function (2) takes the form of a data-dependent attention matrix S [ 1 / √ d ( x + p ) Q ] ∈ R L × L mixing the tokens of the input x . Note that, compared to the usual attention mechanism [1] employed in practice, (2) corresponds to setting the value weights to identity, and - since (2) is parametrized by a single matrix 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 C a 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 µ , y ( x µ ) } n µ =1 with n independently sampled sentences { x µ } n µ =1 , and the associated labels { y ( x µ ) } n µ =1 , is available. The target (1) can be learnt by carrying out an empirical risk minimization:

The performance of the resulting trained model f ˆ Q is measured at test time by the mean squared error (MSE)

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 α = n / d (henceforth referred to as the sample complexity) stays of order Θ d (1) . We further assume the sentence length L , the ranks r s , r t of the weights Q , Q ⋆ , and the norm of the positional embeddings ∥ p ∥ , to be Θ d (1) . This limit has been considered in a stream of previous works (e.g. [46, 47, 48]) and allows to derive closed-form characterization of the ERM problem (3), which we present in the next section. It also exhibits a particularly rich learning phenomenology which we further explore in Section (5). Finally, let us comment briefly on the assumption that r s = Θ d (1) , which in words implies that the weight matrix Q is low-rank . While primarily motivated by technical limitations here, it is worth noting that low-rank weights are also considered in machine learning practice, in the context of model compression [49] or fine-tuning [50].

The main technical result of the present work is a closed-formed characterization of the test MSE (4) and training loss (3) achieved in the high-dimensional limit when training the model (2) via the empirical risk minimization of (3).

Assumption 4.1. The covariances { Σ ℓ } L ℓ =1 admit a common set of eigenvectors { e i } d i =1 . We further note { λ ℓ i } d i =1 the eigenvalues of Σ ℓ . The eigenvalues { λ ℓ i } d i =1 and the projection of the positional embedding { p ℓ } L ℓ =1 and the teacher columns { Q ⋆ j } r t j =1 on the eigenvectors { e ⊤ i p ℓ } i,ℓ , { e ⊤ i Q ⋆ j } i,j are assumed to admit a well-defined joint distribution ν as d → ∞ - namely, for γ = ( γ 1 , ..., γ L ) ∈ R L , π = ( π 1 , ..., π r t ) ∈ R r t and τ = ( τ 1 , ..., τ L ) ∈ R L :

In words, Assumption 4.1 guarantees that all parameters of the problem admit well-defined limits in the considered asymptotic limit, with the further assumption that the covariances { Σ ℓ } L ℓ =1 of the different tokens can be jointly diagonalized. We are now in a position to state the main technical result of the present work.

Result 4.2. Under Assumption 4.1, in the limit n, d →∞ , ∥ p ∥ , n / d , L, r s , r t = Θ d (1) , the summary statistics

concentrate in probability, and are solutions of the set of finite-dimensional self-consistent equations

In (7) , U = { u ℓ } L ℓ =1 and Ξ = { ξ ℓ } L ℓ =1 , with u ℓ ∼ N ( ξ ⊤ ℓ q -1 / 2 ℓ θ ℓ , ρ ℓ -θ ⊤ ℓ q -1 ℓ θ ℓ ) and ξ ℓ ∼ N (0 , I r s ) , and · ⊗ 2 denotes the outer product of a vector with itself. Finally, the resolvents { prox(Ξ , U ) ℓ } L ℓ =1 are defined as the minimizers of the Moreau envelope

We noted Z ∈ R L × r s (resp. U ∈ R L × r t ) the matrix whose rows are z ℓ (resp. u ℓ ) and:

In the same limit, the test error (4) converges in probability to

where the average bears on h ∈ R L × r s , h ⋆ ∈ R L × r t with independent rows with statistics

Finally, the training loss ϵ t converges in probability to

Result 4.2 provides a tight asymptotic characterization of the test error ϵ g (10) and the training loss ϵ t (12), as a function of a finite set finite-dimensional summary statistics { q ℓ , V ℓ , m ℓ , θ ℓ } L ℓ =1 , thereby providing a finite-dimensional description of the high-dimensional learning problem (3). These summary statistics are further characterized in closed form by the set of equations (7), in terms of the solution of a low-dimensional minimization problem (91). Intuitively, this low-dimensional problem may be viewed as a form of an effective loss averaged over the finite training set. In practice, the solution of the self-consistent equations (7) can be found by numerically iterating the equations until convergence. The resulting summary statistics can then be used to evaluate the expressions (10) and (12) to reach the asymptotic limits of the test error and train loss. In Section 5 we evaluate exactly these functions to understand the different minima in the models empirical loss landscape. Similar sharp asymptotic characterizations have been derived in the literature for other neural network architectures trained with ERM, in particular generalized linear models (see e.g. [51, 52, 48]) and auto-encoders [47].

The derivation of Result 4.2 is provided in Appendix B, and is exploiting a mapping of the model (2) to a (variant of) a Generalized Linear Model (GLM) [53, 54]. The summary statistics characterized by the equations (7) (often called state evolution [55] in this context) asymptotically describe the fixed points of a Generalized Approximate Message Passing (GAMP) algorithm [56], which we state in Appendix B. The fixed points of GAMP in turn correspond to critical (zero-gradient) points of the non-convex empirical loss landscape (3). Therefore, while Result (4.2) is stated as a characterization of the global minimum of (3), which is the main concern of the present work, solutions of (7) also describe local minima and saddles.

This strategy has been used in many recent work to study asymptotics of a large number of highdimensional problems, see e.g. [57, 58, 59, 60, 48]. We note, however, that we importantly assume the point-wise convergence of GAMP. While we believe that this point can be rigorously justified, it would require a considerable amount of work -in particular, the usual rigorous tools used in recent works fall short because of the non-convexity of the loss- and we leave this point for future studies (see the discussion in Appendix C, where we also provide an alternative derivation using the replica method from statistical physics [61]). Finally, we mention that while Result 4.2 is presented for an ℓ 2 regularization of the empirical loss (3) for clarity, similar results can be reached for generic convex regularizers, and are presented in Appendix C. In the following section, we explore the phenomenology uncovered from the study of the equations (7) of Result 4.2, for the special case of dot-product attention.

0.0

0.5

1.0

1.5

2.0

c

0.0010

0.0005

0.0000

0.0005

t

0.0

0.5

1.0

1.5

2.0

c

0.0

0.2

0.4

0.6

0.8

1.0

/

2

m

/

2

0.0

0.5

1.0

2.0

c

l

0.002

0.003

0.004

g

min

g

5 Positional-to-semantic phase transition

Rank one dot-product attention In the following, we turn to a special case of tied low-rank attention (2) - namely a dot-product attention layer, which is the example from Fig. 1:

As in (2), we allow for positional encodings p in the dot-product attention parametrization (13). We further consider a specific case of target attention matrix (1) of the form

with A ∈ R L × L a fixed matrix. In (14), the parameter ω ∈ [0 , 1] tunes the relative strength of the dot-product term and the fixed matrix term, and interpolates between a fully positional and a fully semantic task:

For ω = 0 , the target reduces to the first dot-product term, and is purely semantic, in that the i, j -th element of the score matrix softmax( 1 / d xQ ⋆ Q ⊤ ⋆ x ⊤ ) only depends on the tokens x i , x j and not explicitely on their respective placements i, j inside the sentence. To learn satisfyingly the target, the learning model thus has to learn a semantic attention matrix.

For ω = 1 , the target reduces to the second fixed term A in (14). The attention matrix A associated thereto is purely positional, in the sense that A ij is a function of i, j but not of x i , x j . To complete the learning task, a positional mechanism then needs to be learnt.

The parameter ω thus allows to tune the amount of semantic/positional content in the target (14), and thus the extent to which the task requires the model to implement semantic attention (small ω s) or rather positional attention (large ω s). In the following, for definiteness, we further assume r s = r t = 1 and set Q ⋆ to be a fixed random Gaussian vector drawn from N (0 , I d ) , and choose the positional encodings p 1 = -p 2 = 1 d . Finally, for simplicity, we consider sentences with two tokens L = 2 and isotropic token covariances Σ 1 = Σ 2 = σ 2 d .

✶ Semantic and positional mechanisms The summary statistics θ ℓ , m ℓ describing the global minimizer of the empirical loss minimization (3) of the dot-product attention (13) on the target

0.0

0.2

0.4

0.6

0.8

1.0

1.2

m

/

2

0.0

0.2

0.4

0.6

0.8

/

2

semantic

positional

10 1

10 2

d

10 1

10 2

0.0

0.5

1.0

1.5

2.0

0.1

0.2

0.3

0.4

0.5

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

t

0.0

0.5

1.0

1.5

2.0

0.1

0.2

0.3

0.4

0.5

0.0020

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

0.0020

g

g

(14) are captured alongside the corresponding test error (4) and training loss (3), by Result 4.2. The solution of the system of equations (7) is not unique, and different stable fixed points describe different corresponding critical points of the non-convex empirical loss landscape (3). In practice, we notably find two solutions of (7), corresponding to two minima associated with different mechanisms implemented by the dot-product attention (13) when approximating the target (14):

-Positional solution One solution of (7) correspond to vanishing overlap θ = 0 between the trained weights ˆ Q and the semantic target weights Q ⋆ , and non-zero m > 0 between the trained weights ˆ Q and the positional embedding p 1 = -p 2 . Consequently, the argument of the dot-product attention ˆ Q ( x + p ) has a sizeable token-independent -thus positional- contribution ˆ Qp , alongside a token-dependent semantic part ˆ Qx . Because of the positional terms, the resulting learnt attention attention matrix softmax( 1 / d ( x + p ) ˆ Q ˆ Q ⊤ ( x + p ) ⊤ ) implements a partly positional mechanism. -Semantic solution Another solution of the system of equations (7) is associated with a vanishing overlap m = 0 between the learnt weights ˆ Q and the positional embeddings, and a finite overlap θ > 0 with the target weights Q ⋆ . Therefore the resulting learnt attention matrix softmax( 1 / d ( x + p ) ˆ Q ˆ Q ⊤ ( x + p ) ⊤ ) ≈ softmax( 1 / d x ˆ Q ˆ Q ⊤ x ⊤ ) is largely semantic.

While the system of self-consistent equations (7) may admit other solutions, we did not find solutions with lower training loss than the two aforedescribed fixed points. Which of these solution corresponds to the global minimum - and thus the solution of the optimization (3)- depends on the sample complexity α and the positional/semantic parameter ω (14), as we describe in the following subsection.

Positional-to-semantic phase transition For a fixed parameter ω in (14), an analysis of equations (7), further detailed in Appendix C, reveals that for a sizeable range of ω , in the probed setups, there exists a threshold α c for the sample complexity so that

For α < α c , the global minimum of (3) corresponds to a positional mechanism, and is described by the positional solution of (7) of Result 4.2 with θ = 0 , m > 0 .

For α > α c , the global minimum of (3) corresponds to a semantic mechanism, and is described by the semantic solution of (7) of Result 4.2 with θ > 0 , m = 0 .

The dot-product attention thus displays a phase transition in sample complexity from a positional mechanism to a semantic mechanism , implementing the simpler positional mechanism when having access to small amounts of data, and only learning the semantic content of the target (14) when presented sufficient data. The critical sample complexity α c generically grows with the positionality ω of the target function (14), as the semantic content - i.e. the first term of (14)- is less apparent for larger ω , and thus requires larger amounts of data to be identified and approximated by the dot-product attention (13). An example for ω = 0 . 3 is given in Fig. 2. In Fig. 3 (center) the difference

in training loss ∆ ϵ t between the positional and semantic solutions of (7) is represented, alongside the difference in training loss at convergence experimentally reached by gradient descent. For small (resp. large) sample complexity α < α c (resp. α > α c ), the training loss of the positional (resp. semantic) minimum is lower, and thus corresponds to the global minimum.

Experimentally, the positional minimum can be reached for α < α c via gradient descent by initializing the weights Q of the attention (13) close to the positional embedding p 1 . By the same means, the semantic minimum can be reached from an initialization at the teacher weights Q ⋆ (14). Henceforth, we refer with a slight abuse to the minimum experimentally reached from a positional (resp. semantic) initialization as the positional (resp. semantic) minimum, even when it is not global. Note that importantly the semantic initialization is informed in nature, in that it necessitates the knowledge of the target parameters Q ⋆ . Note that even though the minima we characterize analytically are fixed points of gradient descent, a precise analysis of the dynamics of gradient descent from an agnostic (random) initialization, and ascertaining whether the optimizer reaches the global minimum, is an interesting question but falls out of the scope of the present manuscript - which is an analysis of the loss landscape. We however conduct numerical experiments from a random initialization of Q in Appendix E.4, and show that the dynamics may reach either of the local minima, or get stuck in a different one.

In Fig. 2, we compare our analytical characterizations for different metrics at the global 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 a positional (resp. semantic) initialization for α < α c (resp. α > α c ), displaying overall good agreement. In Fig. 3 (left) and Appendix E.1 we further verify that in the scaling limit of our analysis, namely n, d →∞ for α = O (1) , the agreement improves with growing n, d .

The dot-product attention (13) thus implements a semantic mechanism when learning from sufficient amounts of data. The learning of the semantic mechanism by the dot-product attention at sample complexities α > α c corresponds to a noticeable drop in the generalization MSE as can be observed in Fig. 2, right. But just how essential is the learning of semantic mechanism in the ability of the dot-product attention to generalize well? We explore this question in the following subsection, by comparing the dot-product attention (13) to a purely positional attention model.

Purely positional baseline In this subsection, for the same target (14), we contrast the dot-product attention model (13), analyzed in the previous subsections, to the baseline given by a linear layer

with a trainable weight matrix W ∈ R L × L . As for the dot-product attention (13), we consider the case where the weights ˆ W are learnt by minimizing the empirical risk

The model (15) is a natural counterpart to the dot-product architecture (13). In (15), the attention matrix is parametrized by a single fully-trainable matrix W , instead of being parametrized as a dot-product attention as in (13). A seminal difference in the two parametrizations is that while the elements of softmax( 1 / d xQQ ⊤ x ⊤ ) can depend on the input tokens x , and therefore express semantic information, the elements W ij of W can only depend on the positions i, j . The model (15) can thus only implement positional mechanisms , while the dot-product attention (13) can implement both linear and semantic mechanisms, as discussed above. Finally, observe that the model (15) is closely related to the one analyzed by [24] in another asymptotic limit. The following result characterizes the test error achieved by the purely positional model (15):

Result 5.1. In the same asymptotic limit as Result (4.2) , the learnt weights ˆ W trained by minimizing the empirical risk (16) coincide with the minimizer of the population risk, and thus admit the compact expression

where the average bears over a finite-dimensional matrix h ∈ R L × t with independent rows h ℓ with statistics h ℓ ∼ N (0 , ρ ℓ ) , where ρ ℓ was defined in (6) in Result (4.2) . We remind that T [ 1 / √ d xQ ⋆ ]

corresponds to the target score matrix (1) . Finally, the test MSE 1 / dL E x ∥ y ( x ) -f ˆ W ( x ) ∥ 2 achieved by the trained dense linear model f ˆ W (15) admits the asymptotic characterization

The MSE achieved by the baseline (15) when learning the target (14) is plotted in Fig. 2 (right) as the orange solid line, alongside the MSE achieved by the dot-product attention (13) discussed in previous subsections. Remarkably, in the setup of Fig. 2, in the positional regime α < α c when the dot-product attention relies on a positional mechanism θ = 0 , m > 0 to approximate the target, the dot-product attention (13) is outperformed by the purely positional attention (15) ϵ g > ϵ lin g . In contrast, in the semantic regime α > α c where the dot-product attention learns the semantic mechanism, there exists a sample complexity α l ≥ α c above which ϵ g < ϵ lin g , i.e. the dot-product attention (13) outperforms the dense linear baseline (15). This threshold value α l is plotted for various positionality strengths ω in Fig. 3, alongside the positional-to-semantic threshold α c . Interestingly, we observe α l ≥ α c in all probed settings, temptingly suggesting the natural interpretation that the dot-product attention needs to learn the semantic mechanism first (at α = α c ) in order to then be able to outperform the best positional approximation f ˆ W (at α = α l ). This highlights the importance of the semantic mechanism, enabled by the dot-product parametrization (13), in learning targets with semantic content such as (14).

6 Limitations

Compared to the original transformer [1] we consider a simplified model : the query and key matrices in our model are sharing weights and are of a low rank, as a value matrix we use the identity, and we employ only one head and a single layer. Further, our data model is limited to Gaussian data with sentences of 1-grams. Concerning the analysis, our characterization holds only in the high-dimensional limit , but we show that even in the finite case experimental values lie close to the theoretical prediction. Since the analysis only concerns the minima of the models loss landscape the implications for the dynamics of learning algorithms, e.g. gradient decent, are limited. This shows in our numerical experiments where we need to initialize GD close to the minima in order to arrive at them. It is yet unclear if there are scaling limits of learning algorithms which would reliably find the lower or a specific one of the minima from a random initialization.

Conclusion

We explored the interplay between positional and semantic attention, through the prism of tied low-rank self-attention in high dimensions. In a theoretically controlled setting, we characterized the global optimum of the empirical loss, when learning a target attention layer. This global optimum was found to correspond to either a positional or a semantic mechanism, with a phase transition between the two mechanisms occurring as the sample complexity increases. We believe the present asymptotic analysis of the inner workings of attention mechanisms opens up exciting research directions. Considering alternative attention architectures (including a readout network after the attention layer, or considering cross-attention) or training procedures (such as masked language modelling, or training with causal masks), are some possible extensions which will hopefully pave the way towards a satisfactory theoretical comprehension of attention mechanisms. Finally, elucidating under which conditions either minimum can be reached by a given optimizer from a random initialization constitutes an important future research avenue.

References

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , volume 30, 2017.

Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems , 35:37822-37836, 2022.

Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682 , 2022.

Aaditya K 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.

Taylor Webb, Keith J Holyoak, and Hongjing Lu. Emergent analogical reasoning in large language models. Nature Human Behaviour , 7(9):1526-1541, 2023.

Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning , pages 11080-11090. PMLR, 2021.

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 Hatfield-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.

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

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.

Angelica Chen, Ravid Schwartz-Ziv, Kyunghyun Cho, Matthew L Leavitt, and Naomi Saphra. Sudden drops in the loss: Syntax acquisition, phase transitions, and simplicity bias in mlms. arXiv preprint arXiv:2309.07311 , 2023.

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.

Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In The Twelfth International Conference on Learning Representations , 2024.

Sanjeev Arora and Anirudh Goyal. A theory for emergence of complex skills in language models. arXiv preprint arXiv:2307.15936 , 2023.

Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models a mirage? Advances in Neural Information Processing Systems , 36, 2024.

Ernst Ising. Contribution to the theory of ferromagnetism. Z. Phys , 31(1):253-258, 1925.

Lars Onsager. Crystal statistics. i. a two-dimensional model with an order-disorder transition. Physical Review , 65(3-4):117, 1944.

Haim Sompolinsky, Naftali Tishby, and H Sebastian Seung. Learning from examples in large neural networks. Physical Review Letters , 65(13):1683, 1990.

Géza Györgyi. First-order transition to perfect generalization in a neural network with binary synapses. Physical Review A , 41(12):7097, 1990.

Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences , 116(12):5451-5460, 2019.

Benjamin Aubin, Antoine Maillard, Florent Krzakala, Nicolas Macris, Lenka Zdeborová, et al. The committee machine: Computational to statistical gaps in learning a two-layers neural network. Advances in Neural Information Processing Systems , 31, 2018.

Henry Schwarze. Learning a rule in a multilayer neural network. Journal of Physics A: Mathematical and General , 26(21):5781, 1993.

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 , 33:11071-11082, 2020.

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.

Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. Optimal inference of a generalised Potts model by single-layer transformers with factored attention. arXiv preprint arXiv:2304.07235 , 2023.

Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. Advances in Neural Information Processing Systems , 36, 2024.

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.

Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR, 2022.

Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics , 8:156-171, 2020.

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.

Davoud Ataee 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.

Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308.16898 , 2023.

Emmanuel Abbe, Samy Bengio, Enric Boix-Adsera, Etai Littwin, and Joshua Susskind. Transformers learn through gradual rank increase. Advances in Neural Information Processing Systems , 36, 2024.

Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015 , 2023.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in Neural Information Processing Systems , 36:71911-71947, 2023.

Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representations. arXiv preprint arXiv:2310.10616 , 2023.

Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning , pages 19565-19594. PMLR, 2023.

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 2023.

Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. Advances in Neural Information Processing Systems , 36, 2024.

Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on transformers. arXiv preprint arXiv:2312.10794 , 2023.

Jeffrey L. Elman. Finding structure in time. Cognitive Science , 14(2):179-211, 1990.

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.

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.

Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-Attention with Relative Position Representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers) , pages 464-468. Association for Computational Linguistics, 2018.

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 generalization of transformers. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers) , pages 1889-1903, 2023.

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.

Hugo Cui and Lenka Zdeborová. High-dimensional asymptotics of denoising autoencoders. Advances in Neural Information Processing Systems , 36, 2024.

Bruno Loureiro, Gabriele Sicuro, Cédric Gerbelot, Alessandro Pacco, Florent Krzakala, and Lenka Zdeborová. Learning gaussian mixtures with generalized linear models: Precise asymptotics in high-dimensions. Advances in Neural Information Processing Systems , 34:1014410157, 2021.

Yen-Chang Hsu, Ting Hua, Sungen Chang, Qian Lou, Yilin Shen, and Hongxia Jin. Language model compression with weighted low-rank factorization. arXiv preprint arXiv:2207.00112 , 2022.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

Benjamin Aubin, Florent Krzakala, Yue Lu, and Lenka Zdeborová. Generalization error in high-dimensional perceptrons: Approaching bayes error with convex optimization. Advances in Neural Information Processing Systems , 33:12199-12210, 2020.

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):015019, 2023.

John Ashworth Nelder and Robert WM Wedderburn. Generalized linear models. Journal of the Royal Statistical Society Series A: Statistics in Society , 135(3):370-384, 1972.

Peter McCullagh. Generalized linear models . Routledge, 2019.

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):115-144, 2013.

Sundeep Rangan, Philip Schniter, Erwin Riegler, Alyson K Fletcher, and Volkan Cevher. Fixed points of generalized approximate message passing with arbitrary matrices. IEEE Transactions on Information Theory , 62(12):7464-7474, 2016.

Mohsen Bayati and Andrea Montanari. The lasso risk for gaussian matrices. IEEE Transactions on Information Theory , 58(4):1997-2017, 2011.

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.

Melikasadat Emami, Mojtaba Sahraee-Ardakan, Parthe Pandit, Sundeep Rangan, and Alyson Fletcher. Generalization error of generalized linear models in high dimensions. In International Conference on Machine Learning , pages 2892-2901. PMLR, 2020.

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 Transactions on Information Theory , 69(3):1824-1852, 2022.

Giorgio Parisi. Toward a mean field theory for spin glasses. Physics Letters A , 73(3):203-205, 1979.

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.

Mohsen Bayati and Andrea Montanari. The dynamics of message passing on dense graphs, with applications to compressed sensing. IEEE Transactions on Information Theory , 57(2):764-785, 2011.

Lenka Zdeborová and Florent Krzakala. Statistical physics of inference: Thresholds and algorithms. Advances in Physics , 65(5):453-552, 2016.

Lucas Clarté, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Theoretical characterization of uncertainty in high-dimensional linear classification. Machine Learning: Science and Technology , 4(2):025029, 2023.

Giorgio Parisi. Order parameter for spin-glasses. Physical Review Letters , 50(24):1946, 1983.

Erwin Bolthausen. An iterative construction of solutions of the tap equations for the sherringtonkirkpatrick model. Communications in Mathematical Physics , 325(1):333-366, 2014.

Marc 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.

Yehoram Gordon. On milman's inequality and random subspaces which escape through a mesh in rn. In Geometric Aspects of Functional Analysis: Israel Seminar (GAFA) 1986-87 , pages 84-106. Springer, 1988.

Christos Thrampoulidis, Samet Oymak, and Babak Hassibi. The gaussian min-max theorem in the presence of convexity. arXiv preprint arXiv:1408.4837 , 2014.

Mihailo Stojnic. Meshes that trap random subspaces. arXiv preprint arXiv:1304.0003 , 2013.

Mihailo Stojnic. Upper-bounding l1-optimization weak thresholds. arXiv preprint arXiv:1303.7289 , 2013.

Hugo Cui. High-dimensional learning of narrow neural networks. arXiv preprint arXiv:2409.13904 , 2024.

Appendix: Table of Contents

B

B

C

C

D D

B

B

B

B

B

B

C

C

D

D

B

B

B

B

Positional

42

6

6

6

7

6

6

6

6

7

6

41

6

7

6

6

6

6

7

7

6

6

42

6

7

7

6

6

6

6

6

6

6

43

6

6

6

6

6

7

7

6

6

7

42

6

6

7

6

6

7

6

6

7

7

42

7

6

6

6

6

7

7

7

6

7

42

7

6

7

7

6

6

7

7

7

6

41

6

6

6

6

6

6

7

6

6

6

42

7

7

6

6

7

7

6

6

6

6

42

Example Sequence #1

L

M D M M M A

H

H

B

L

M

D

M

M

M

A

H

H

B

42

6

6

6

7

6

6

6

6

7

6

41

6

7

6

6

6

6

7

7

6

6

42

6

7

7

6

6

6

6

6

6

6

43

6

6

6

6

6

7

7

6

6

7

42

6

6

7

6

6

7

6

6

7

7

42

7

6

6

6

6

7

7

7

6

7

42

7

6

7

7

6

6

7

7

7

6

41

6

6

6

6

6

6

7

6

6

6

42

7

7

6

6

7

7

6

6

6

6

42

Example Sequence #2

O N M

L

K

J

I

H G

F

O

N

M

L

K

J

I

H

G

F

42

6

6

6

7

6

6

6

6

7

6

41

6

7

6

6

6

6

7

7

6

6

42

6

7

7

6

6

6

6

6

6

6

43

6

6

6

6

6

7

7

6

6

7

42

6

6

7

6

6

7

6

6

7

7

42

7

6

6

6

6

7

7

7

6

7

42

7

6

7

7

6

6

7

7

7

6

41

6

6

6

6

6

6

7

6

6

6

42

7

7

6

6

7

7

6

6

6

6

42

Example Sequence #3

B

B

C

C

D D

B

B

B

B

B

B

C

C

D

D

B

B

B

B

Semantic

15

15

3

3

3

3

15

15

15

15

15

15

3

3

3

3

15

15

15

15

5

5

28

28

5

5

5

5

5

5

5

5

28

28

5

5

5

5

5

5

5

5

6

6

28

28

5

5

5

5

5

5

6

6

28

28

5

5

5

5

15

15

3

3

3

3

15

15

15

15

15

15

3

3

3

3

15

15

15

15

15

15

3

3

3

3

15

15

15

15

15

15

3

3

3

3

15

15

15

15

L

M D M M M A

H

H

B

L

M

D

M

M

M

A

H

H

B

33

8

6

8

8

8

7

8

8

7

4

19

4

19

19

19

4

5

5

4

7

8

35

8

8

8

7

7

7

7

4

19

4

19

19

19

4

5

5

4

4

19

4

19

19

19

4

5

5

4

4

19

4

19

19

19

4

5

5

4

7

8

9

8

8

8

30

8

8

9

6

6

6

6

6

6

5

27

27

6

6

6

6

6

6

6

5

27

27

6

7

7

7

7

7

7

8

8

8

34

O N M

L

K

J

I

H G

F

O

N

M

L

K

J

I

H

G

F

34

8

7

7

7

6

6

7

9

8

7

37

7

7

7

7

7

7

7

7

7

7

34

8

7

5

8

8

7

8

6

8

8

33

7

7

7

8

7

8

8

7

7

7

35

7

7

7

7

8

7

7

6

7

7

35

8

7

8

7

5

8

7

8

8

7

34

8

6

8

6

9

8

7

8

6

7

36

6

7

7

7

7

7

7

7

7

7

36

7

8

7

7

7

8

7

7

8

7

34

0.2

0.4

0.6

0.8

1.0

attention value

0.0

A The Histogram Task: An algorithmic toy example with a positional and semantic solution

A.1 Phenomenology

In this appendix, we demonstrate that for a simple counting task two qualitatively different solutions exist in the loss landscape of a simple transformer using a dot-product attention layer with positional encodings - a practical example that further motivates our theoretical investigations in the main text. One solution corresponds to a dot-product attention matrix which is largely independent of the tokens making up the input sequence, and another strongly varies based on the tokens (and thus the semantic content of the input). Both solutions achieve a test accuracy close to 100%.

The training task is a sequence-to-sequence counting task, referred to as the histogram task in [6]. Given an input sequence x = [ x 1 , x 2 , · · · , x L ] of length L of tokens from a fixed alphabet, the goal is to return a sequence y = [ y 1 , y 2 , · · · , y L ] , where each token y i is the number of occurrences of the token x i in x . In Fig. 1, 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 i thus take values up to the maximum count L .

We encode the input using token embeddings and absolute positional encodings which are trained jointly with the model weights. As an architecture we consider a small transformer made up of a single layer of dot-product attention, followed by a fully connected hidden layer and with learned embeddings for both tokens and positions. For each output position, it generates logits for the L possible classes of the output alphabet; training is done using the cross-entropy loss. (Further details are provided in the remainder of this appendix).

We conduct experiments where we set two different sections of the model's weights to zero at the initialization of training -removing either the model's access to positional or semantic informationand keeping the weights frozen throughout training with the Adam optimizer. After convergence, we check that the resulting configurations of weights are stable in the unconstrained loss landscape, i.e.

without frozen weights. More precisely, we ascertain that these weights only change marginally when further trained with SGD on the unconstrained loss, and that the qualitative behaviour of the attention layer is retained. Our experiments demonstrate that the loss landscape of the transformer has at least two qualitatively different local minimizers (or close to minimizers), subsequently referred to as the semantic and positional solution.

We inspect the learnt attention matrix for different input sequences in Fig. 4. The positional solution corresponds to a learnt attention matrix whose i, j -th component only depend on the positions i, j , and little on the tokens occupying these positions. The attention matrix is thus almost independent of the input sequence. In fact, the attention matrix is similar to the identity. In this case, the attention layer simply serves to aggregate the other tokens uniformly, and the fully connected layer learns the counting.

In contrast, the attention matrix learnt at the semantic solution displays larger i, j -th component if the tokens at position i and j are identical. In other words, identical tokens attend more to each other. This mechanism hence does not rely on the positions, but rather on the semantic content of the tokens. Both solutions and associated attention matrices thus correspond to feasible algorithms which ultimately allow the transformer to solve the downstream task.

Our experimental exploration gives compelling evidence that different stable solutions exist in the empirical loss landscape of simple transformers, which correspond to different algorithmic solutions to a given task. However, it remains an interpretation of an experiment and does not allow for a precise characterization of their behaviour or of the conditions under which they are established. This stands in contrast to the example in the main paper: the model treated there is simpler, but still presents similar phenomenology and can be analyzed theoretically.

In the remainder of this Appendix we discuss the implementation details such as architecture and training procedure which enabled us to exhibit the different minima for the histogram task.

The code for reproducing the results is available at github.com/SPOC-group/positional-and-semanticattention.

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 = [ x 1 , x 2 , · · · , x L ] we sample x i i.i.d. and uniformly from a set of tokens T of size T , which we set to 15 in our experiments. For visualization purposes we use capitalized letters as tokens. To obtain the target y = [ y 1 , y 2 , · · · , y L ] , we set y i = ∑ L j =1 1 ( x i = x j ) , where 1 ( b ) is 1 if the boolean statement b is true and zero otherwise. Therefore, y i ∈ { 1 , 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 E token = [ t A , t B , · · · ] ∈ R T × d and E pos = [ p 1 , p 2 , · · · ] ∈ R L × d . We convert the input sequence x = [ x 1 , x 2 , · · · , x L ] ∈ T L into an embedded sequence ˜ x ∈ R L × d of the same length where ˜ x i = t x i + p i . This input sequence is then fed into the first layer of the transformer.

Formally, we have

with W 1 ∈ R d × h , b 1 ∈ R h , W 2 ∈ R h × c , b 2 ∈ R c . The final prediction is obtained using the argmax on the logits. We have that the score matrix A and the dot-product attention mechanism is

where Q,K,V ∈ R d × d and the softmax is applied row-wise. Also, for an x ∈ R d we define

where γ, β ∈ R d . We define the empirical loss for a dataset D = { x µ , y µ } n µ =1 as the average cross entropy loss for all output tokens C , i.e.

Because of the way in which the histogram dataset is created for a 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 E token , E pos , Q, K at initialization, and also freeze these zero weights during training. Note that this procedure was used in the literature to study architectural biases of varying model architectures, e.g. by [62] for convolutional neural networks.

With · we denote the initialization that is taken as the default Pytorch initialization for a linear layer. For both semantic and positional initialization, we overwrite this initialization with zeros as follows with i = 1 , · · · , T and j = 1 , · · · , L

where 0 d/ 2 is the all-zero vector of size d/ 2 . For the positional initialization we additionally set

and for the semantic initialization we additionally set

where 0 d/ 2 × d/ 2 is the all-zero matrix of size d/ 2 by d/ 2 . In either case, the attention mechanism only has access to the semantic or positional part of the model.

Training all weights except the ones frozen to zero using Adam on a dataset of size n = 35 , 000 , we obtain a > 99 . 8% test accuracy on a test set of size n = 15 , 000 for both datasets. We call these parameter configurations θ pos and θ sem respectively (the local minima as referred to in the main text).

In Fig. 4, we show that the attention layer of the two models behaves in qualitatively different ways, by showing the activations of the matrix A for different input sequences.

A.5 Unfreezing model parameters and SGD convergence

Finally, we want to verify that the parameter configurations θ pos and θ sem we obtained using the special initialization and frozen training are stable configurations in the unconstrained parameter space of the model. To show this, we perturb the given parameter configuration slightly using additive Gaussian i.i.d. noise with a scale of 0 . 001 . Subsequently, we run SGD on this perturbed configuration in the unconstrained parameter space for another 100 epochs. We call the parameter configurations obtained after this step ˜ θ pos and ˜ θ sem .

While this further training leads the training loss to further decrease and the test accuracy to slightly increase, each models qualitative behaviour before and after extra training remains unchanged, as shown in Fig. 5 and 6, for the examples also used in Fig. 1. The absolute difference in norm between ˜ θ pos and θ pos (respectively ˜ θ sem and θ sem ) of parameters is also small (see Table 1).

This evidences, that both resulting parameter configurations are in flat regions of the full parameterization of the transformer model, and that these two regions are further qualitatively different in the same sense that the original ones were. We use this stability in the unconstrained loss landscape as a justification to refer to them as 'local minima'.

B

B

C

C

D D

B

B

B

B

B

B

C

C

D

D

B

B

B

B

pos

40

7

6

7

7

7

7

7

7

7

6

41

7

6

6

6

6

6

7

7

6

6

43

6

6

6

6

6

6

7

7

6

7

41

7

7

7

6

6

7

7

7

6

7

40

7

7

7

7

6

7

6

7

6

7

40

7

7

7

7

6

7

6

7

7

7

41

7

7

6

7

6

6

7

7

7

7

41

7

6

6

7

7

6

7

7

7

7

41

6

6

7

6

6

6

6

6

6

7

42

Example Sequence #1

L

M D M M M A

H

H

B

L

M

D

M

M

M

A

H

H

B

40

7

6

7

7

7

7

7

7

7

6

41

7

6

6

6

6

6

7

7

6

6

43

6

6

6

6

6

6

7

7

6

7

41

7

7

7

6

6

7

7

7

6

7

40

7

7

7

7

6

7

6

7

6

7

40

7

7

7

7

6

7

6

7

7

7

41

7

7

6

7

6

6

7

7

7

7

41

7

6

6

7

7

6

7

7

7

7

41

6

6

7

6

6

6

6

6

6

7

42

Example Sequence #2

O N M

L

K

J

I

H G

F

O

N

M

L

K

J

I

H

G

F

40

7

6

7

7

7

7

7

7

7

6

41

7

6

6

6

6

6

7

7

6

6

43

6

6

6

6

6

6

7

7

6

7

41

7

7

7

6

6

7

7

7

6

7

40

7

7

7

7

6

7

6

7

6

7

40

7

7

7

7

6

7

6

7

7

7

41

7

7

6

7

6

6

7

7

7

7

41

7

6

6

7

7

6

7

7

7

7

41

6

6

7

6

6

6

6

6

6

7

42

Example Sequence #3

B

B

C

C

D D

B

B

B

B

B

B

C

C

D

D

B

B

B

B

pos

42

6

6

6

7

7

6

7

6

6

6

43

6

6

6

6

6

6

7

6

6

6

43

6

6

6

6

6

6

7

7

6

6

42

7

7

6

6

6

7

7

6

7

7

40

7

7

6

6

7

7

6

7

6

6

42

7

6

6

7

6

7

6

7

7

6

41

7

7

6

7

6

6

6

7

7

7

42

7

6

7

7

6

6

6

7

7

6

42

6

6

7

6

6

6

6

6

6

7

42

L

M D M M M A

H

H

B

L

M

D

M

M

M

A

H

H

B

42

6

6

7

7

6

6

7

7

6

6

43

7

6

6

6

6

6

7

6

6

6

43

6

6

7

6

6

6

7

7

6

7

42

7

7

7

6

6

7

7

7

6

7

41

7

7

6

7

6

7

6

6

6

7

41

7

6

7

7

6

7

6

7

6

6

41

7

7

6

7

6

6

6

7

7

7

42

7

6

6

7

6

6

6

7

7

6

42

6

6

7

6

6

6

6

6

6

7

42

O N M

L

K

J

I

H G

F

O

N

M

L

K

J

I

H

G

F

41

7

6

7

7

7

6

7

6

7

6

43

6

6

6

6

6

6

7

7

6

6

43

6

6

6

6

6

6

7

7

6

6

42

7

7

7

6

6

7

7

7

6

7

41

7

7

6

7

6

7

6

6

6

7

41

7

6

7

6

6

6

7

6

7

6

41

7

7

6

7

6

6

6

7

7

7

42

7

6

6

6

7

6

6

6

7

6

42

6

6

6

6

6

7

6

6

6

7

42

0.2

0.4

0.6

0.8

1.0

attention value

0.0

B Derivation of Result 4.2

In this Appendix, we provide a 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.4 shows that fixed points of GAMP correspond to critical (zero-gradient) points of the empirical loss landscape (3), thus establishing that equations 7 of Result 4.2 describe fixed points of GD.

B.1 Notations

For simplicity, we place ourselves in the setting r s = 1 explored in Section 5 of the main text, but allow the length L of the sentences to be arbitrary, and allow a generic learning model S (2), i.e. not necessarily the dot-product attention model analyzed in Section 5. The case r s ≥ 2 follows identical derivation steps, modulo the replacement of all variables by tensor objects. We provide another alternative derivation of Result 4.2 in full generality in Appendix C, using the replica method from statistical physics. Let us note { X ℓ } 1 ≤ ℓ ≤ L a series of L n × d matrices, with X ℓ corresponding to the ℓ -th rows (tokens) of each input sentence x µ stacked vertically, and normalized by √ d . We denote ˜ X ℓ ≡ X ℓ + P ℓ , where P ∈ R n × d is the matrix with all rows equal to the ℓ -th positional encoding p ℓ . Let us further define ρ ∈ R n × L × L the tensor corresponding to the sequence of n matrices 1 d x µ ( x µ ) ⊤ ∈ R L × L . Finally, let us denote T ∈ R n × L × L the tensor so that the µ -th row of T satisfies y ( x µ ) = T µ x µ , see equation (1). In other words, T corresponds to the concatenation of the target attention matrices.

B

B

C

C

D D

B

B

B

B

B

B

C

C

D

D

B

B

B

B

sem

15

15

3

3

3

3

15

15

15

15

15

15

3

3

3

3

15

15

15

15

5

5

29

29

5

5

5

5

5

5

5

5

29

29

5

5

5

5

5

5

5

5

6

6

30

30

5

5

5

5

5

5

6

6

30

30

5

5

5

5

15

15

3

3

3

3

15

15

15

15

15

15

3

3

3

3

15

15

15

15

15

15

3

3

3

3

15

15

15

15

15

15

3

3

3

3

15

15

15

15

Example Sequence #1

L

M D M M M A

H

H

B

L

M

D

M

M

M

A

H

H

B

31

8

7

8

8

8

7

7

7

7

3

20

4

20

20

20

3

3

3

4

6

7

39

7

7

7

7

7

7

6

3

20

4

20

20

20

3

3

3

4

3

20

4

20

20

20

3

3

3

4

3

20

4

20

20

20

3

3

3

4

8

7

8

7

7

7

36

8

8

7

6

5

7

5

5

5

5

28

28

6

6

5

7

5

5

5

5

28

28

6

6

7

7

7

7

7

7

6

6

39

Example Sequence #2

O N M

L

K

J

I

H G

F

O

N

M

L

K

J

I

H

G

F

37

7

7

6

7

8

7

8

7

8

7

38

7

7

7

7

7

8

6

7

7

8

38

7

6

7

8

6

7

6

5

8

9

34

6

7

9

8

8

7

8

8

8

7

33

8

8

8

6

7

7

7

7

7

7

34

8

8

9

7

6

8

8

8

8

8

35

7

7

7

8

9

7

7

7

8

7

35

7

6

7

8

8

7

5

8

8

6

34

7

9

7

7

8

6

8

8

6

7

34

Example Sequence #3

B

B

C

C

D D

B

B

B

B

B

B

C

C

D

D

B

B

B

B

sem

15

15

3

3

3

3

15

15

15

14

15

15

3

3

3

3

15

15

15

14

5

5

29

29

5

5

5

5

5

5

5

5

29

29

5

5

5

5

5

5

5

5

6

6

30

30

5

5

5

5

5

5

6

6

30

30

5

5

5

5

15

15

3

3

3

3

15

15

15

14

15

15

3

3

3

3

15

15

15

14

15

15

3

3

3

3

15

15

15

14

15

15

3

3

3

3

15

15

15

14

L

M D M M M A

H

H

B

L

M

D

M

M

M

A

H

H

B

32

8

7

8

8

8

7

7

7

7

4

19

5

19

19

19

4

3

3

4

6

7

38

7

7

7

7

7

7

6

4

19

5

19

19

19

4

3

3

4

4

19

5

19

19

19

4

3

3

4

4

19

5

19

19

19

4

3

3

4

8

7

8

7

7

7

36

8

8

7

6

6

7

6

6

6

5

28

28

6

6

6

7

6

6

6

5

28

28

6

6

7

7

7

7

7

7

7

7

36

O N M

L

K

J

I

H G

F

O

N

M

L

K

J

I

H

G

F

36

7

7

6

7

8

7

8

7

8

7

37

7

7

7

7

7

8

6

7

7

8

37

7

6

7

8

7

7

6

6

8

9

34

6

7

9

8

8

7

8

7

8

7

34

7

8

8

6

7

7

7

7

7

7

35

8

7

8

7

6

7

8

7

7

8

37

7

7

7

8

9

7

7

7

8

7

35

7

6

7

8

8

7

5

8

8

6

34

7

9

7

7

8

6

8

8

6

7

34

0.2

0.4

0.6

0.8

1.0

attention value

0.0

Before detailing the derivation, we first highlight a simplifying observation. Note that a loss item can be expanded as

where we used that with high probability in the considered asymptotic limit, for all 1 ≤ µ ≤ n ,

Since the first term of (27) does not depend on the weights Q , it can be without loss of generality substracted from the loss. Without loss of generality, one can thus consider the equivalent empirical risk minimization problem

The risks (29) and (3) are equivalent, and we shall use the former in the following.

Finally, for arguments T ∈ R L × L , ρ ∈ R L × L , ω ∈ R L , V ∈ R L × L we introduce the resolvent prox( T, ρ, ω, V ) ≡

Note that the latter part of the bracketed term corresponds to the simplified loss (27) derived in the beginning of Appendix C, which is the one we shall without loss of generality consider in the present appendix. For ease of presentation, we place ourselves under Assumption 4.1, where all the input covariances { Σ ℓ } ℓ are codiagonalizable. In the following, without loss of generality, we thus assume them diagonal, by placing ourselves in the common basis { e i } 1 ≤ i ≤ d of Assumption 4.1.

B.2 AMPalgorithm

We are now in a position to state the AMP algorithm:

Algorithm 1 GAMP

The GAMP algorithm can be derived in standard fashion from the Belief Propagation (BP) algorithm, see e.g. [56, 63] or [64] for an overview. Compared to the standard GAMP iterations for Generalized linear models, one needs to account for the fact that there exist different sources of data X ℓ (corresponding to the ℓ -th tokens of each input sentence), and for the fact that the output of the equivalent GLM are R L × L - valued attention matrices. In the following subsection, we show that the fixed points of GAMP 1 correspond to critical points of the empirical loss (3), i.e. fixed points of Gradient Descent (GD), allowing to connect Result 4.2 to our numerical experiments using GD.

B.3 State evolution

In this section we show that the dynamics of the GAMP Algorithm 1 are tracked by the summary statistics of Result 4.2. In particular, the equations (7) describe the statistics of the GAMP fixed points. To see this, it is convenient to take as a starting point the relaxed Belief Propagation (rBP) equations, which are a step upstream in the derivation of the GAMP iterations, and which are asymptotically equivalent- see e.g. [64] for a review or e.g. [65], Appendix A, for a detailed walkthrough. The rBP equations read

As conventional, we note · µ the version of a variable · µ → i where the summation also encompasses the index i , and · i the version of a variable · i → µ where the summation also encompasses the index µ . Note that in all cases above the two variables differ by at most Θ d ( 1 / √ d ) .

Concentration of ( V t µ → i ) ℓκ We first study the statistics of V t µ → i , A t i → µ , remembering that the data ˜ x µ ℓi ≡ ( x µ ℓ ) i / √ d + ( p ℓ ) i / √ d in the notation of the main text, with ( x µ ℓ ) i = Θ d (1) , ( p ℓ ) i = Θ d ( 1 / √ d ) . Replacing in the rBP updates:

̸

̸

̸

̸

Algorithm 2 rBP

end for

return Estimator ˆ Q

Distribution of ω t ℓ,µ → i Let us first introduce the teacher local field

Like e.g. [65], Appendix A, we first ascertain the joint distribution of h µ,ℓ , ω t ℓ,µ → i with respect to the data. These variables have mean

and respective variance

Distribution of b t i → µ Let us ascertain the distribution of b t i → µ .

̸

̸

̸

̸

̸

̸

leading asymptotically to

where the expectations bear over ξ ℓ ∼ N (0 , 1) and h ℓ ∼ N ( ξ ℓ θ t ℓ / √ q t ℓ , ρ ℓ -( θ t ℓ ) 2 / q t ℓ ) . The variance is given by

Concetration of A t i → µ Similarly to the derivation for V t µ → i , A t i → µ concentrates to

Recovering Result 4.2 Wrapping up, we now massage these equations to recover equations (7) from Result 4.2 of the main text. Starting from (31):

Next, for q t ℓ (34):

For θ t ℓ (34):

Finally for m t ℓ (33):

For ˆ m t ℓ (38):

while for ˆ θ t ℓ (38):

Now turning to ˆ q t ℓ :

Finally, for ˆ V t ℓ (40):

Summary : State evolution equations The state evolution equations asymptotically describing the dynamics of the GAMP algorithm 1 thus read

which exactly recovers equations (7) of Result 4.2 of the main text, for the case r s = 1 considered in the present Appendix. Again, we mention that the case r s ≥ 2 should follow straightforwardly with the exact same derivation steps, using tensor variables (see e.g. [52]). This subsection has thus established that the equations (7) (with time indices) describe the summary statistics capturing the dynamics of GAMP iterations 1. In particular, (7) describe the fixed points of GAMP. The next subsection further shows that the (stable) fixed points of GAMP correspond to critical (zero-gradient) points of the empirical landscape (3), i.e. fixed points of gradient descent. Finally, we provide in Appendix C an alternative derivation of the state evolution equations (7)(49), using the replica method from statistical physics [61, 66].

B.4 Fixed points of GAMP are fixed points of GD

In this subsection, we show that fixed points of GAMP 1, as asymptotically described by (7) in Result 4.2, correspond to critical (zero gradient) points of the empirical landscape (3). Again, we present the

result for r s = 1 for clarity, the generalization to r s ≥ 2 being straightforward (see e.g. [52]). In the previous notations, let us denote the (simplified, see (27)) empirical loss as

where we introduced the shorthands

i.e. respectively the simplified empirical loss (3) and the regularization, as functions with matrix arguments. The empirical minimization problem (3) can thus be written compactly as

with the critical (zero-gradient) condition being given by

Let us choose a diagonal definite A ∈ R d × d , and a sequence { V µ } 1 ≤ µ ≤ n of symmetric definite L × L matrices. Group them into a block diagonal matrix ˇ V ∈ R Ln × Ln , so that the µ -th block of ˇ V corresponds to V µ . It shall prove useful to further introduce the matrices ˇ ˜ X ∈ R d × nL (resp. ˇ ∂L ( ˇ ˜ X ) ∈ R nL ), defined as the concatenation of the matrices ˜ X 1 , ..., ˜ X L (resp. ∂ 1 L, ..., ∂ L L ), viewed as n blocks of length L . Then without loss of generality the zero-gradient condition can be rewritten as

Similarly to [52], let us introduce where

So that

This can be written in terms of a resolvent as

which corresponds to (30). Similarly, we denote

In the particular case of an ℓ 2 regularization g ( · ) = λ / 2 ∥·∥ 2 , note that

The zero-gradient condition can now be rewritten as

One is now in a position to expand the concatenated variables ˇ · into a sequence of L n -dimensional parameters. For u = prox(ˇ ω ) , ˇ ω let us denote u µℓ ( 1 ≤ µ ≤ n, 1 ≤ ℓ ≤ L ) the ℓ -th component of the µ -th block. Introduce

Denote f ℓ ≡ ( f µℓ ) 1 ≤ µ ≤ n ∈ R n , ω ℓ ≡ ( ω µℓ ) 1 ≤ µ ≤ n ∈ R n . The system of equations (63) can then be rewritten as (further redefining b ← -Ab ):

We used the assumption that g ( · ) is an ℓ 2 regularization. Finally, introducing ˆ Q = prox g ( A -1 b ) = ( λ I d + A ) -1 b , on reaches

which corresponds to the fixed-point equations of GAMP (Algorithm 1). This finishes to show the correspondence between the fixed points of GAMP and the critical points of the empirical landscape (3). To summarize, we have shown that equations (7) describe the zero-gradient points of the empirical loss landscape (3), i.e. fixed points of GD.

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, 58, 59, 60, 48], there exist technical difficulties in adapting these rigorous arguments to the present setting, and a fully rigorous proof would warrant sizable work. While we leave this challenging task for future work, we wish to discuss how it can be potentially achieved. The first task would require the proof of point-wise convergence of GAMP, as indeed, the identification of the GAMP estimates with the one of the extrema of the loss function requires to be at the fixed point of the iteration. This difficulty, discussed in detail in, e.g. [67, 60, 48], can be in principle addressed by computing the convergence criterion from the state evolution equations (see [67] the discussion in Lemma 7 in [60]), a criterion sometimes called the "replicon" in the context of replica theory [68].

Provided the replicon criterion is satisfied, all converging fixed point described by our theory thus correspond rigorously to fixed point of the loss. The last task would be to prove that the minimum of the loss is indeed the fixed point we found with minimum energy. A potential strategy to prove this would be to use the Gordon-Minimax approach of [69]. While it is used in many situations for convex problems (e.g. [70, 71, 72]), only one side would be required for our (non-convex) problem thanks to the GAMP matching bound. We hope that our results would provide inspiration for further research in this direction.

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 [61, 66] in its replica-symmetric formulation. First observe that for any test function ϕ ( ˆ Q ) of the minimizer ˆ Q of (3),

where we denoted R [ Q ] the empirical loss (3), and

the normalization factor, also known as the partition function in statistical physics. We remind that D refers to the training set. In order to access key summary statistics and learning metrics associated to ˆ Q , it is therefore reasonable to seek to compute the generating function associated to the measure (67), namely E ln Z . Such computations can be addressed using the replica method from statistical physics [61, 66], building on the identity

The backbone of the derivation thus lies in the computation of E Z s .Below, we detail the derivation for a generic convex regularizer g : R d → R + and later specialize to the case of ℓ 2 regularization. The replicated partition function thus reads

Introduce the local fields

and the overlaps

with rows m ℓ a . These fields have statistics

Thus

where we decomposed the replicated free entropy into the trace, entropic and energetic potentials Ψ t , Ψ Q , Ψ y . Note that all exponents are scaling with d →∞ . Therefore the integral in (76) can be computed using a Laplace saddle-point approximation.

C.1 Replica-Symmetric ansatz

We have thus rephrased the analysis of the measure (67) as a optimization problem over the order parameters { q ℓ ab , θ ℓ a , m a } , and the associated conjugate variables. However, these still represent

2 L ( s 2 +1) + 2 s variables, and s → 0 . In order to make progress, we assume that the maximizer is of replica-symmetric (RS) form [66, 61]

The RS ansatz holds in a number of machine learning settings, notably for convex problems and Bayes-optimal settings, see e.g. [64] for a review. In the present setting, since the empirical loss (3) is non-convex, we emphasize that the RS ansatz constitutes a heuristic technical assumption of our analysis.

C.2 Entropic potential

We now turn to the entropic potential Ψ w . It is convenient to introduce the variance order parameter

The entropic potential can then be expressed as

Therefore

For a matrix Ξ ∈ R r × d and tensors A , B ∈ R r × d ⊗ R r × d , we denoted ( Ξ ⊙ A ) kl = ∑ ij Ξ ij A ij,kl and ( A ⊙ B ) ij,kl = ∑ rs A ij,rs B rs,kl .

C.3 Energetic potential

The computation of the energetic potential Ψ y is rather standard and follows the same lines as in e.g. [20], yielding

C.4 Zero-temperature limit

We now take the limit β →∞ . Rescaling

The entropic potential then reduces to

where we defined the entropic Moreau enveloppe

The energetic potential can be similarly recast into a more compact form

where the Moreau envelope is defined as

C.5 Replica free entropy

One finally reaches an expression for the replica free entropy as

C.6 Saddle-point equations : general regularizer

The extremization of the free entropy (92) yields, similarly to [47], the following system of selfconsistent equations on the summary statistics:

where the proximals prox g and prox ℓ respectively refer to the arginf in Q (resp. x ℓ ) of the envelopes M g (89) (resp. M 91).

C.7 Saddle-point equations : ℓ 2

We now specialize the saddle-point equations (93) to the case of an ℓ 2 regularizer g ( · ) = 1 / 2 ∥·∥ the entropic potential admits the simple form

The replica free energy thus reads

leading to the saddle point equations

which finishes to recover (7). Let us finally mention that the update equations (7) for the summary statistics (6) do not describe the dynamics of gradient descent, but rather that of an Approximate Message Passing algorithm [63], which we elicit in Appendix B for completeness.

C.8 test MSE

The generalization performance is measured by the test error

Expliciting this expression in terms of the correlated Gaussian variables xQ ⋆ , xQ allows to straightforwardly show that ϵ g admits the sharp asymptotic characterization in terms of the summary statistics characterized by (97):

where the average bears on X ∈ R L × r , Y ∈ R L × t with independent rows with statistics

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 ℓ 2 regularization g = 1 / 2 ∥·∥ 2 . First note that the training loss ϵ t can be expressed as

Where Φ( β ) is the free entropy at finite temperature. The trace potential Ψ t bears no explicit dependence on β . On the other hand,

Thus

Finally, going through the same rescaling steps to take the β →∞ limit,

By the same token, it is straightforward to see that

We used the self-consistent equations (7) to identify the term in underbrace. Putting everything together,

This constitutes a sharp asymptotic characterization of the training loss ϵ t as a function of the summary statistics characterized in Result 4.2.

For completeness, we finally explicit the connection between ϵ t and the negative free entropy (i.e. the free energy in statistical physics). We go back to massage the expression for the free entropy

In other words, the training loss is equal to the zero-temperature free energy.

C.10 Extensions

For the sake of clarity and definiteness, we stated the technical results for in the simplest instance, under notably the symplifying assumptions of independent tokens, tied key and query weight matrices, as exposed in Section 3 of the main text. These two assumptions can in fact be relaxed. In this final subsection of Appendix C, we discuss how the characterization of Result 4.2 can be generalized to more complex data distributions and attention architectures.

Data model Let us consider a data distribution allowing for statistical correlations between different tokens. A simple model meeting this criterion, proposed in [73], consists in assuming every token (row) x ℓ is drawn from a Gaussian mixture with K ℓ with respective probabilities ρ ℓ,k ,

In (108), the random variable c ℓ ∈ [ K ℓ ] indicates the index of the cluster the ℓ -th token is drawn from. Note that we generically allow for different rows to be sampled from distinct mixtures. In order to introduce correlations between different tokens, we further assume that the variables c = ( c 1 , ..., c L ) follow a generic joint law ρ ( c ) . Intuitively, the different models can model different types of words (e.g. verbs, pronouns, adjectives), and the nature of the words, as reflected by the variable c ℓ , is correlated across tokens. We assume that ∥ µ ℓ,k ∥ = Θ d (1) , i.e. are of the same order as the positional encodings.

Assumption C.1. Finally, similarly to Assumption 4.1, let us assume that the set of matrices {{ Σ ℓ,k } K ℓ k =1 } L ℓ =1 admits a common set of eingevectors { e i } d i =1 with eigenvalues { λ ℓ,k i } d i =1 . The eigenvalues { λ ℓ,k i } ℓ,k,i and the projection of the cluster means (to which the positional encodings are added) { µ ℓ,k + p ℓ } ℓ,k and the teacher columns { Q ⋆ i } r t i =1 on these eigenvectors are assumed to admit a well-defined joint distribution ν as d → ∞ - namely, for γ = ( γ ℓ,k ) ℓ,k , π = ( π 1 , ..., π t ) ∈ R r t , τ = ( τ ℓ,k ) ℓ,k :

Architecture Again in the interest of further generality, we alleviate the architectural restrictions on the attention (2), and allow for untied key and query weight matrices. More precisely, let us consider function of the form

parameterized by two set of weights Q,K ∈ R d × r s . In order to work with more compact expression, let us introduce the total weight matrix W ∈ R d × ( r s + r s ) , obtained as the horizontal concatenation of Q,K . Then (110) can be written compactly as

where for any argument h ∈ R L × ( r s + r s ) , wiewed as a 1 × 2 block matrix with two blocks h 1 , h 2 ∈ R L × r , G ( h ) = S ( h 1 , h 2 ) .

Asymptotic characterization A sharp asymptotic characterization of the test error and train loss for this more general model can be derived along the same lines, as presented in Appendix C (see also [73]). For the sake of conciseness, we only report the final result.

Result C.2. Under Assumption C.1, in the limit n, d →∞ , ∥ p ∥ , n / d , L, r s , r t = Θ d (1) , the summary statistics

concentrate in probability, and are solutions of the set of finite-dimensional self-consistent equations

The averages bears over Ξ ∈ R L × 2 r s , Y ∈ R L × r t , with rows ξ ℓ ∼ N (0 , I 2 r s ) independently and y ℓ ∼ N ( ξ ⊤ ℓ q -1 / 2 ℓ,c ℓ θ ℓ,c ℓ ρ ℓ,c ℓ -θ ⊤ ℓ,c ℓ q -1 ℓ,c ℓ θ ℓ,c ℓ ) , conditionally on c . prox c ℓ corresponds to the arginf of the minimization defining the Moreau envelope M ( c, Y, Ξ) , as

We introduced m c ⋆ ∈ R L × r t , whose rows are

where ρ c Σ ∈ R L × L is diagonal with elements

and the shorthand

Test error The test error admits the sharp asymptotic characterization

where, conditioned on the class assignments c , the average bears on h ∈ R L × 2 r s , Y ∈ R L × r t with independent rows with statistics

Finally, the train loss ϵ t converges in probability to

Trainable value matrix Finally, let us remark that it should be possible to leverage a similar derivation to further accommodate trainable value matrices V , provided they are low rank, i.e. in R d × r V with r V = Θ d (1) as d →∞ . Similarly, introducing the stacked weights W = ( Q | K | V ) ∈ R d × (2 r s + r V ) allows to perform the asymptotic analysis. Since this architectural modification then implies that f now takes values in R L × r V , one cannot immediately analyze this architecture for the task considered in the present manuscript, as the labels are valued in R L × d , and the setting and questions explored would need to be entirely modified. We thus postpone a thorough analysis of the effect of a trained value matrix to future investigations.

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)

where we use the shorthand notation A ⋆ ( x ) ≡ T [ 1 / √ d xQ ⋆ ] for the target attention score matrix (1). The expression for the risk can be asymptotically simplified as

Therefore the learnt weight ˆ W is simply equal to

Naming h = 1 / √ d xQ ⋆ the argument of the last term, the matrix h possesses independent rows with statistics

where we remind that ρ ℓ ≡ 1 / d Q ⊤ ⋆ Σ ℓ 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 → ∞ and L = Θ d (1) , training W , i.e. a number L 2 = Θ d (1) of parameters on the empirical loss for n ≪ L 2 data points is equivalent

asymptotically to directly training on the population loss. Expliciting the corresponding test MSE, one reaches the sharp asymptotic characterization

We close this Appendix by giving a few examples for definiteness.

D.1 Examples

Purely positional target Let us consider the case where the target is purely positional, i.e. T [ x ] = A for all x . This corresponds to the ω = 1 limit of the target (14). It then follows from Result 5.1 that the dense linear layer recover perfectly the target weights ˆ W = A .

Target (14) For the target discussed in the main text (14),

with h and having independent rows h ℓ ∼ N (0 , ρ ℓ ) .

E Supplementary Experiments

E.1 Empirical scaling of α = d/n

In the following we verify that our experiments are consistent with the scaling behaviour predicted from the theory. We jointly increase d and n for a fixed value of α . In Fig. 7 we indeed observe the expected behaviour for an exemplary value of α = 2 . The same holds for the summary statistics θ and m , which concentrate as d and n jointly grow, shown in Fig. 3 (left) in the main text.

10 1

10 2

d

0.002

0.003

0.004

0.005

0.006

0.007

g

10 1

10 2

d

positional

semantic

empirircs

theory

E.2 Alternative hyperparameters

We provide supplementary results for different parameter settings. Fig. 8 on the left shows more slices from the phase diagram that appears in the main Fig. 3. For the experimental section of the main text, we chose a specific A for definiteness. In the following, we present the same results for a different A with a stronger off-diagonal and a higher rank,

In Fig. 9 we present the analogous simulations to Fig. 3 (center, right). While the global phenomena match the previous example, the details of the transitions location differ.

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

0.0015

0.0020

0.0025

0.0030

0.0035

0.0040

0.0045

g

0.2

0.3

0.7

semantic

positional

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

0.001

0.002

0.003

0.004

0.005

g

0.1

0.01

0.001

semantic

positional

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

0.0175

0.0200

g

0.3

0.5

0.7

semantic

positional

0.0

0.5

1.0

1.5

2.0

0.1

0.2

0.3

0.4

0.5

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

t

0.0

0.5

1.0

1.5

2.0

0.1

0.2

0.3

0.4

0.5

0.0020

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

0.0020

g

g

E.3 More Complex Architectural Choices

As discussed previously in Section C.10, the theory developed in this work can be relaxed to incorporate arbitrary statistical correlations between tokens, changes in the (low-)rank student, generally untied key and query matrices and the presence of a low-rank trainable value matrix. This surplus in complexity comes at the cost of more intricate and heavier equations.

In the present section, we report experiments in the four settings above, but not the theoretical predictions. Fig. 10-13 compare the change in architecture with the setting considered in the main text, for a given ω = 0 . 5 and varying sample complexity α . The exhaustive experiments for varying ω are shown in Fig. 14 - which immediately shows that the general idea of the phase transition from positional to semantic solution is consistent, but that the shape of the phase transition curve is influenced by the architecture.

For correlated inputs, the transitions seems to moves to larger values of α (Fig. 10). For the case of a student with a rank-2 query matrix, and a teacher with a rank-1 query matrix, we consider three possible initilisations of the columns of Q when finding minima of GD empirically.

Initializing both semantically, both positional, or one of them in each way, as shown in Fig. 11. Interestingly, the dual initialization seems to do approximately as well as the best of the two other initializations. The optimization generally seems more complex, as the phase diagram in the upper right corner of Fig. 14 are not entirely clean, and more noisy than empirical results from other architectures.

When we compare tied and untied Q and K weights in Fig. 12 we observe that the additional parameters we need to learn in the case of an extra K come at a small cost, as the transition to the semantic solution is moved to the right.

Finally, for a value matrix (Fig. 13) the results seem to largely resemble the transition line we observed without the value matrix.

0.0

0.5

1.0

1.5

2.0

0

1

2

3

4

t

unmodified

0.0

0.5

1.0

1.5

2.0

correlated inputs (

=0.1)

positional init

semantic init

0.0

0.5

1.0

1.5

2.0

0

1

2

3

4

t

unmodified (

r

student

=1)

0.0

0.5

1.0

1.5

2.0

higher rank student (

r

student

=2)

positional init

semantic init

both init

0.0

0.5

1.0

1.5

2.0

0

1

2

3

4

t

unmodified

0.0

0.5

1.0

1.5

2.0

untied

Q

,

K

positional init

semantic init

0.0

0.5

1.0

1.5

2.0

0

1

2

3

4

5

6

t

unmodified

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

with value matrix

V

R d

×

d

positional init

semantic init

0.0

0.5

1.0

1.5

2.0

0.1

0.2

0.3

0.4

0.5

correlated (

=0.1)

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

t

0.0

0.5

1.0

1.5

2.0

0.1

0.2

0.3

0.4

0.5

higher rank student (

r

student

=2)

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

t

0.0

0.5

1.0

1.5

2.0

0.1

0.2

0.3

0.4

0.5

untied

Q

,

K

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

t

0.0

0.5

1.0

1.5

2.0

0.1

0.2

0.3

0.4

0.5

with value matrix

V

R d

×

d

0.0015

0.0010

0.0005

0.0000

0.0005

0.0010

0.0015

t

E.4 Uninformed initialization and training via Adam

In or experiments, to obtain the empirical results, we initialize the GD optimizer in an informed fashion, i.e. initializing Q ⋆ of the student with r = 1 as either p 1 (positional) or Q ⋆ (semantics). GD then converges in the two local optima described by our theory.

Since our theory only ascertains that these solutions predicted are indeed fixed points of GD for large sizes, this does not have direct implications for other types of optimization algorithms. In Fig. 15 we show that indeed running the Adam optimizer from an uninformed initialization may lead one to either of the local minima for d = 100 . For larger d we observe the semantic minimum is reached less often than the positional minimum, and a considerable number of times the algorithm simply does not find either of them.

0.0

0.2

0.4

0.6

0.8

1.0

1.2

2

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

/

2

informed initialization, GD, d=1000

semantic

positional

0.0

0.2

0.4

0.6

0.8

1.0

1.2

2

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

/

2

random initialization, Adam, d=100

m

/

m

/

F NeurIPS Paper Checklist

1. 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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: There is a dedicated section on limitations.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [No]

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 in the main text and in 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The full code is 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 a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 to 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 (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code will be made public on github as it is available with the supplementary material.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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. If only a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The concrete details of how many samples per batch etc. are contained in the supplementary material.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of 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, or 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [No]

Justification: The compute needed is negligable, so it is not specified.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

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.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: It is foundational research.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: We do not use external original material.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

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 is not available online, the authors are encouraged to reach out to the asset'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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is 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. Crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nUnderstanding 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.\n\nWEAKNESSES:\nBecause 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\n\nQUESTIONS:\nDo 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\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: 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).

STRENGTHS:
* Provides theoretical analyses of a model of self-attention, finding a closed-form solution.
* Asymptotic analysis demonstrates a phase-shift between two minima, one which relies on position and one which does not.
* 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.
* While toy models are indeed useful objects to study in general, there are unclear connections between some of the assumptions (e.g. independent samples of individual tokens, and the low-rank attention) and actual language modeling practice.

QUESTIONS:
* 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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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$.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nAs 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.\n\nQUESTIONS:\nCan 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?\n\nLIMITATIONS:\nLimitations are identified and discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}