Attention layers provably solve single-location regression
Authors: Pierre Marion
Raphaël Berthier
Gérard Biau
Claire Boyer
Venue: ICLR 2025 Poster (2025)
Keywords: theory of neural networks
attention
Transformer
statistical learning theory
optimization
first-oder optimization
gradient flow
Primary area: learning theory
TL;DR: We introduce single-location regression, a novel regression task, and show that attention solves this task.

Abstract: Attention-based models, such as Transformer, excel across various tasks but lack a comprehensive theoretical understanding, especially regarding token-wise sparsity and internal linear representations. To address this gap, we introduce the single-location regression task, where only one token in a sequence determines the output, and its position is a latent random variable, retrievable via a linear projection of the input. To solve this task, we propose a dedicated predictor, which turns out to be a simplified version of a non-linear self-attention layer. We study its theoretical properties, by showing its asymptotic Bayes optimality and analyzing its training dynamics. In particular, despite the non-convex nature of the problem, the predictor effectively learns the underlying structure. This work highlights the capacity of attention mechanisms to handle sparse token information and internal linear structures.

Published as a conference paper at ICLR 2025
ATTENTION LAYERS PROVABLY SOLVE SINGLE-
LINGLE-LOCATION LOCATION REGRESSION
Pierre Marion
Raphaël Berthier
Institute of Mathematics
Sorbonne Université, Inria
HPFL
Centre Inria de Sorbonne Université
Lausanne, Switzerland
Paris, France
pierre ma www.mmanerllll.omm ch
raphael www.beerrtteerrrbee.eee
Gérard Biau
Claire Boyer
Sorbonne Université
Université Sacia-Saclay
Institut Universitaire de France
Institut Universitaire de France
Paris, France
Orsay, France
gerard biau@upmc. fr
ABSTRACT
Attention-based models, such as Transformer, excel across various tasks but lack a comprehensive theoretical understanding, especially regarding token-wise sparsity and internal linear representations. To address this gap, we introduce the single- location regression task, where only one token sequence determines the output, and its position is a latent random variable, retrievable via a linear projection of the input. To solve this task, we propose dedicated predictor, which turns out to be simplified version of a non-linear self-attention layer. We study its theoretical properties, by showing its asymptotic Bayes optimality and analyzing its training dynamics. In particular, despite the non-convex nature of the problem, the predictor effectively Iearns the underlying structure. This work highlights the capacity of attention mechanisms to handle sparse token information and internal linear structures.
INTRODUCTION
Attention-basee models (Bahdanau al., 2015), such as Transformer (Vaswani et al., 2017), have achieved unprecedented performance in various learning tasks, including natural language processing (NLP), e.g., text generation (Bubeck al., 2023), translation (Luong et al., 2015), sentiment analysis (Song et al., 2019; Sun et al., 2019; Xu et al., 2019), and audio/speech analysis (Bahdanau et al., 2016). These developments have led to many architectural and algorithmic variants of attention- models (see the review by Lin et al., 2022). At high level, the success of attention has been linked to its ability to manage long-range dependencies in input sequences (Bahdanau et al., 2015; Vaswani al., 2017), since it consists in computing pairwise dependence between input tokens according to their projection in learned directions, independently of their location in the sequence.
On the theoretical front, however, deeper understanding of attention-based neural networks is still in its infancy. This limited progress is due both to the complexity of the architectures and to the disturbing diversity of relevant tasks. common approach introduce a simplified task that models certain features of real-world tasks, followed by demonstrating simplified version of the attention mechanism capable of solving the task. Prominent examples of this pattern include studying in-context learning with linearized attention (Ahn et al., 2023; von Oswald et al., 2023; Zhang et al., 2024), topic understanding with single-layer attention and alternate minimization scheme (Li et al., 2023b), learning spatial structure with positional attention (Jelassi et al., 2022), next-token prediction with latent bigram (Bietti et al., 2023; Tian et al., 2023) or causal graph (Nichani et al., 2024) structures, and sparse token selection (Wang et al., 2024). We refer to Appendix for additional discussion on some of these related works. While these works shed light on some abilities of Transformer, they do not encompass all the characteristics of tasks where
tackle these challenges is to

---

Published as a conference paper at ICLR 2025
Transformer performs well, in particular in NLP. Two features of particular interest, which to knowledge have not been addressed in previous theoretical studies on Transformer, are token-wiss wise sparsity, where relevant information is contained in limited number of tokens, and internal linear representations, which are interpretable representations of the input constructed by the model.
Contributions. To understand why attention is a suitable architecture for addressing these features, we introduce single-location regression, a novel statistical task where attention-based predictors excel (Section 2). In nutshell, this task is a regression problem with a sequence of tokens as input. The key novelty is that only one token determines the prediction, and the location of this token is a latent random variable that changes based on the input sequence. Consequently, solving the task requires first identifying the location of the relevant token, which can be done by learning latent linear projection, followed by performing regression on that token.
To tackle this problem, we propose a dedicated predictor, which turns out to be simplified version of a non-linear self-attention layer. We show that this attention-based predictor is asymptotically Bayes optimal, whereas more standard linear regressors fail to perform better than the null predictor. We then analyze the training dynamics of the proposed predictor, when trained to minimize the theoretical risk by projected gradient descent. Despite the non-convexity of the problem and the non-linearity of this transformer-basee method, we show that the learned predictor successfully retrieves the underlying structure of the task and thus solves single-location regression.
Organization. Section 2 presents the mathematical framework of single-location regression, fol- lowed by motivations from language processing. Section 3 is dedicated defining our predictor and explaining its connection with attention. We then move on to the mathematical study, from both statistical (Section 4) and optimization (Section 5) points of view. Section 6 concludes the paper.
SINGLE-LOCATION REGRESSION TASK
In this section, we describe statistical task, and connect it to language processing motivations.
2.1 STATISTICAL SETTING
We consider regression scenario where the inputs are sequences of I. random tokens' (X1,.... Xx.) taking values in Rd. The output Y IR is assumed to be given by
== == v- +,,
(P (Plearn)
.. L} and, conditionally Jo. where Jo is latent discrete random variable f1, ....
X)o N( ///kk,,,1, 22 Xe ~ Ia for e +
In the above formulation, (N, Σ) denotes the normal distribution with expectation μ and covariance matrix Σ, and la is the identity matrix of size d x d. All vectors are considered as column matrices, and the noise term & assumed to be centered random variable independent of X and Jn, with finite second-order moment ² Conditionally on Jo. the tokens (Xy)ıęJεL are assumed to be independent.
The parameters of the regression problem (Pearn) are the unknown vectors k* and v*, both assumed to be on the unit sphere gd-1 in dimension d, i.e., ||A* = ||e 1|2 = 1. The output is determined by a specific token in the sentence, indexed by the discrete random variable Ju on (1..... L). This token can be detected via its mean, which is proportional to k²*, contrarily to the others which have zero mean. Once X,, is identified, the prediction formed as linear projection in the direction v*. Therefore, the originality and difficulty of this task lies in the fact that the response Y is linearly related to a single informative token XJ.- whose location varies from sequence sequence-in in this sense, the problem is sparse, but with a random support.
knee-jerk reaction would be to fit a linear model to the pair (XI X Y). One might also consider tackling the problem with classical statistical approaches dedicated to sparsity, such
For the sake of simplicity, we interchangeahly use the terms "token" and "embedding", although they have slightly different meanings in the NL. community.
2

---

Published as a conference paper at ICLR 2025
a Lasso estimator group-Lasso technique (Hastie et al., 2009). However, we will see (in Section 4), all linear predictors fail due to the unknown and changing location of Jo- We note in addition that [[]] d when f * Jo. while X [2 d/2 + pdd. Therefore, Jo
choosing -0 = 1/2 implies that tokens have the same squared norm in expectation, whether they are discriminatory of not. This shows that any approach based comparing the magnitude of the tokens does not yield meaningful results. Ultimately, it is necessary to implement a more sophisticated approach, capable of taking into account the characteristics of the problem.
2.2 LANGUAGE PROCESSING MOTIVATION
The structure of the task (Peearn) motivated by natural language processing (NLP), and more specifically by two features, token- wiie sparsity and internal linear representations, as we detail next.
1.00
0.95
What have become sweetest friend?
Everyone know goes 111 the end.
0.75
0.09
What have become sweatest friend?
Accnncy
0.6
Everyone I know goes well in the end,
0,50
Birds flying high, you know feel?
0.4
And, I'm feeling (godd]
Birds flying high, you know how feal? And, I'm feeling (awful)
Initinl embeddiings
Train
Test
OOD OOD
tolens
(a) Examples of input-output pairs. The input is text
(b) Accuracy of logistic regression on embeddings of [CLS] tokens in the hidden layers of a pretrained Trans- former model. Initial embeddings of [CLS] (at layer 0) are context-aware, they have pure-chance accuracy of 50%. In hidden layers, the [CLS] token contains a representation of the sentence that achieves high scores and is robust to out-of diitrribbtttttiinn changes in token distribution and sentence structure.
containing two sentences (c.g., a question and an an-
swer). and the task is perform sentiment analysis
only for the second sentence. The Y output is sym- bolized here by a color code, where green (resp. red) corresponds to positive (resp. negative) feelings. The relevant information is sparse, typically concentrated in single token: changing the grey token flips the output.
A simple sentiment analysis task with synthetic data, which exemplifies (a) token-wise sparsity and (b) internal linear representations. We refer to Appendix for details on the experiment.
Figure 1:
Token-wise sparsity. In language tasks, the relevant information is often contained in few tokens, where recall that tokens correspond small text units (typically, words subwords), which are embedded in Rd using a learned dictionary. This sparsity is revealed by the success of sparse attention (Martins & Astudillo, 2016; Niculae & Blondel, 2017: Correia et al., 2019; Child et al., 2019; Jaszezur et al., 2021; Kim et al., 2022; Farina et al.. 2024), which is competitive with full attention while attending to fewer tokens. As an illustration, we consider simple sentiment analysis task in Figure la, and observe that changing one token flips the output. This is modeled in (Picarn) by having the output Y depend on a single token Jo, whose location furthermore varies with the input.
Internal linear representations. Linear projections of internal representations of Transformer (a.k.a. linear probing) contain interpretable information (Bolukbasi et al., 2021; Burns et al., 2023; Li al., 2023a). Such a linear structure is also present in the learned token embeddings that are fed as input to language models (Mikolov et al., 201 2013a;b; Bolukbasi et al., 2016; Nanda et al., 2023; Wen-Yi & Mimno, 2023). our task (Plear). ), the two directions k** and o* have to be learned by the model in order to solve the task. Figure gives an example of possible such directions for the toy task described above. While this illustration relies on initial embeddings, similar structures also appear in the intermediate representations of Transformer. This is shown in Figure 1b, where observe that pretrained Transformer architectures indeed build internal representations that are sufficient to solve the task with linear classifier.
3

---

Published as a conference paper at ICLR 2025
Ja =
How
aze
doing sweetheart?
To
the
least,
devastated.
embeddingg
Postionnl encodings
embeddings
X2
X2
Alignment E XJK. k'
Alignmant XJ
O
Y
Figure 2: Modeling of an NLP task within our statistical setting (Plear)). The token embeddings X1,...,Xr Xr, constructed by adding the embeddings of each word and positional encoding. For illustration purposes, we assume that each token corresponds to a word, and that the positional encoding solely depends on the part of the sentence (before or after the question mark), which differs from usual practice. Then, let the direction k* encode both the notion of sentiment and the position in the second part of the sentence. Thus only the last token of the sentence is aligned (positively) with k*, and we have Jo L. As for v*, it encodes whether the word is associated with positive or negative sentiment. Note that several tokens are positively or negatively aligned with v², but the output Y only depends on the token Jo. This illustrates the interest of having two latent directions k* and v* one that filters the informative token one that aligns with the output Y.
acknowledge that our statistical task presents limitations such as fixed sequence length, indepen- dent tokens, and output depending only on single token. More complex models could be considered, but at significant technical cost. Moreover, as argued above, our problem (Plearn) preserves interesting aspects of NLP tasks, which makes it relevant for theoretical study of Transformer. Furthermore, an original statistical task requiring the implementation of a customized estimation strategy. is precisely in this context that attention models prove their effectiveness,
show next.
AN ATTENTION-BASEE PREDICTOR TO SOLVE THE REGRESSION TASK
In this section, we propose predictor adapted to the problem (Plearn) and discuss its connection with attention. In order make our point as clear as possible, the construction divided into three steps. We represent the input sequence in matrix format X E Lxd where X = (X, ... (Xz)
Step 1: An oracle non-differentiable predictor. the vectors (k*,"*) E (Sd-1)2 were known, then natural procedure to solve the task (Pearn) would be to predict Y from X via
T(X) (Xu") (x) X 0x) where jx arg max (Xk*) 18EL
(I)
The arg max part detects the location Ju by exploiting the fact that all X₂ have zero mean except X while the Xu* part exploits the linear relationship Y XJ + &. In a more compact format, this ideal predictor can be rewritten as T(X) = Σke 1arg (xxxxxx)==(xx"))) which is a linear regression in the direction u² with non-differentiable weights depending k*
Step 2: trainable predictor. In practice, the vectors k* and V* are unknown and must be estimated from the data. In addition, the non-differentiabiiiy of the arg max function poses significant
4

---

Published as a conference paper at ICLR 2025
optimization challenges. To solve this problem, the most common approach in machine learning to replace arg max with a softmax function with inverse temperature > 0, i.c., for 2 (Z1,..., 2L) E RR, [softmax(A=)], = e² /tin />3 ente, This leads us to the model
T siooft R,0) (X)
Σ (sottmaxxxX))) softmax (1
Kk
Lx1
Lx1
where k v E gd-1, and the superseript 'soft' used to indicate the presence of the softmax function.
Step 3: The final predictor. The softmax nonlinearity. by inducing a coupling between all tokens, significantly complicates the mathematical analysis. To alleviate this difficulty, we replace it by the component-wise nonlinear function erf(z) = /x Jo e dt, which is differentiable, increasing on R, and such that erf(-oo) -1 and erf( 1. We are therefore led our operational model
_2
T ex) (AX T X Σ [er((AXTh)XTT.
where the erf function is applied component-wise. The choice of this activation function enables closed-form expectations for functions of Gaussian random variables (see, e.g., Lemma 18). Note that the role of softmax in attention is an open question the community. Several empirical papers investigate simplifying softmax into component- wise nonlinearity (Qin et al., 2022; Shen et al., 2023; Wortsman et al., 2023; Ramapuram et al., 2024), and have observed a similar performance. These works emphasize the importance of the normalization À when replacing softmax, which also find out to play an important role (see Corollary 2 and Section 5).
Connection attention. turns out that our estimation method finds a natural interpretation in terms of attention models. To see this, consider a model consisting of single attention layer with single head (Vaswani et al., 2017)
T(Q.K.V.)) (X) softmax
XQ K X
XV
OT LxP pxo
Exp PxL
where the dimensions p, E N* are hyperparameters of the model, the softmax function is applied row by row, Q, K. V E R and O e Roxp are the regular query, key, value, and output matrices, and À is usually taken to be 1/vp. In practice, the attention head is added to X via skip connection, which enforces = d. In a nutshell, K detects which tokens are relevant in the sentence, V encodes the regression coefficient, and Q encodes where to store the information.
In a supervised context, it is classical in practice to concatenate in first position an additional token [CLS] to the tokenized sentence X (see, e.g., Devlin et al., 2019). In this context, only the first coordinate of the output is used for the prediction task. Thus, we focus on the first row of (4), corresponding to the embedding of [CLS), namely
T.Q.
(X)1 = softmax X
(5)
with X Q E RIXP, where X(CTs] Id denotes the embedding of the [CLS] token.
important to note that only considering the first output coordinate is a mathematically valid simplification for a single attention layer, but not when multiple layers are stacked, all coordinates of the attention output contribute. Nevertheless, even in this latter more realistic case, the [CLS] token-op or the similar concepts of attention sinks and registers has been empirically shown to play a crucial role (Clark et al., 2019; Darcet et al., 2024; Xiao et al., 2024). This is also confirmed by our experiment Figure 1b, where we show that the [CLS] token in pretrained Transformer architectures stores internal representation of the sentence that is sufficient solve simple NLP tasks with linear classifier. This further motivates the need to understand how information is stored in this token.
It turns out that there is direct connection between the model T (s) (x) defined in (2) and the attention model T T..KKVV.O (VO) (X), described (5). To see this, take o = 1, to adapt the model (5) for univariate regression, and set p reasonable assumption given both empirical and theoretical
5

---

Published as a conference paper at ICLR 2025
evidence suggesting that Transformer parameter matrices low-rank (Aghajanyan et al., 2021; Kajitsuka & Sato, 2024). Then, let Q E Rd be any vector with positive correlation with XJCLSI (for instance it suffices to take Q X(Css)), and O 1. We then deduce that
T( (X)1 = T 2
(vo.. (X)
In other words, the attention layer (5) matches the considered predictor in (2) with a softmax inverse temperature proportional to the scalar product between ((CLS) and Q. Thus, our results, in particular the study the training dynamics in Section 5, can be seen as a model of how Transformer builds internal representations of the input during training. This is also supported by numerical experiments showing that Transformer layers behave similarly to our predictor (see Appendix E).
RISK OF THE ORACLE AND OF THE LINEAR PREDICTORS
Now that we have constructed our predictor (k,) (see Eq. (3)), a first key question is to assess its statistical performance. Recall that k, v E S are the two parameters of the model, and their purpose is approximate their theoretical counterparts k' and e* defined in (1). This begs in particular the question of the performance of the oracle predictor Tr(K","", ,²), To answer these questions, we introduce the risk of the predictor, which is measured by the mean squared error
Rx (ki 1) EE [(x T (x))
(6)
To proceed with the analysis, we make the following assumption. Assumption 1. The vectors k-, V* E gd-1 are orthogonal. i.e., TT
This assumption is made everywhere in the remainder of the paper, even though it is not reminded explicitly at each result. It is a relatively mild assumption in high- digeeniimnsional setting where any two independent vectors uniformly distributed the sphere are close being orthogonal.
Oracle predictor. Our first result characterizes the risk of the proposed transformer model (3) with oracle parameters (k*, u²). All the proofs of the paper are deferred to the Appendix.
Theorem 1. There exists function R, RR R such that, for any (k, (84-1)3
R2 (k, v) RR (K, (,,,V0, 0,,,0,mmP
where x = T kt, v -= uT= 0², 0: = v := k u*, and = kT v. A closed-form expression of RS, is given in Appendix In particular,
RA (kt,, = R (1,1,0,0,0)
+ + (d 2(1 2X222)
d
/d X2y2) + (I),))((0 =
= = erf (1
where, t, R.
(tt,)):= [erf G)] G~N((0,73).
(7)
This result fundamental for the analysis of gradient descent studied in the next section since reduces the dimension of the dynamical system defined by the optimization dynamics. Before delving into the optimization analysis, we study below the statistical optimality of the estimator R (k* v*) and its comparison with linear regression.
Asymptotic Bayes optimality. Let us start by observing that the Bayes risk associated with problem (Plearn) is larger than s², which follows from elementary properties of the conditional expectation (Le Gall, 2022, Chapter 11). Indeed, using the Pythagorean theorem, one easily shows that
E E|| > E[(Y E|Y X, Ju =[{"|=
8
Then, the following corollary to Theorem shows that the oracle predictor achieves the Bayes-optimal risk in the asymptotic scaling T. < 1/X2 < d.

---

e², and the linear predictor Published as a conference paper at ICLR 2025
Corollary 2. Assume joint asymptotic scaling where d and L o(d). Taking X such that Xvd - oo and XVT. O, we have
Rx(k", v²) + E²
Thus, in this asymplotic regime, the oracle predictor T x²) T(l"","))
is asymptotically Bayes optimal.
Note that Corollary holds any finite L E Nyo. but L may also tend to infinity, as long as L o(d). Let give intuition why this result holds and where the scalings of L and intervene. The oracle predictor can be decomposed as
Tuk (X) = XLO
erf erf(AX
X erf ( k')
(9)
j7 =0(1) =0(A)
EEYY Jo
=(Ava
With the scaling Xvd + oc, the argument of the first erf nonlinearity diverges to infinity with d. Thus it reaches the saturating part of erf, so the first term (9) converges E Ja). On the other hand, the argument of the erf nonlinearities inside the sum are of order X = o(1). Thus they are the linear part of erf. Therefore, the sum consists of L independent terms, each of magnitude As a consequence, by the central limit theorem, the whole sum is of order ((vV)). and we get
T .²") ) (X) I[ + @(vVL).
Due to the scaling AVT, - O, the second term decays to zero, and the oracle predictor implements the conditional expectation of given X and Jp. This the best that we can hope for: the predictor succeeds in inferring the latent variable Jo. then gives the best possible prediction of Y given X and Jo. We also see the crucial role played by the nonlinearity of erf, whose linear part acts for j A Jo and saturating part for Jo. In particular, the reasoning would not hold for linear activation.
Linear model. The asymptotic optimality of our oracle predictor is particularly striking in compari- son to the risk of the optimal linear predictor. More precisely, let
3 E arg min E [Y - (XT, ..... X[)
BERAL
be the optimal linear predictor for the regression task (Piearn). Its associated risk is R(3*) (X( (.T.,,")))))) Both the optimal predictor and its risk can be made explicit as follows.
Proposition 3. Let Pj = PJoo = j) for j E (1,..., L). Then the optimal linear predictor is parameterized by 8 = (b1e* ..... by with bj 1+ms TT. and its risk is
7Pp
23
R(B*) = 2 +
13(22
In particular,
R(8) 2 8² + 7 +1)
max Pj. j=1....!
This result calls for a few comments. If the number of tokens (meaning that one Pý is equal to while the others are equal to 0), then the learning problem (Pieamn) corresponds to a standard linear regression. In this case, R(3*)
= or if Jo is a constant location
(Xi...., Xr.) + (XT ...., x 3r achieves the Bayes risk. At the other end of the spectrum, case where Jn is uniform over {1.... L}. the formula for the risk of the linear predictor simplifies to R(8*)= =² +2 - T. When L + 20, this risk tends to e² + v² that is, the performance of the null predictor. In other words, the optimal linear predictor performs no better than always predicting zero. More generally, this conclusion is true in any limit where I. + and max - 0. This can be explained by the fact that the location of the relevant token for prediction is random, varying from sentence sentence. Unable leverage this latent information, the linear regressor balances all its coefficients, resulting in poor prediction performance. This stands in sharp contrast to Corollary 2, which shows that the oracle predictor T(Kk* .u"). is able to account for the complexity of
7

---

Published as a conference paper at ICLR 2025
10
20
Figure 3: Risk of the oracle predictor (Theorem 1, solid lines) and of the best linear predictor (Propo- sition 3, dashed lines), depending the dimen- sions d and L. The oracle predictor outperforms the linear predictor when scaling d. We take z² 0. 7- 1/v2, À 1/d and all P5 equal to 1/L.
50
10a
Risk
10-1
400
600
800
1000
d
the task, at least asymptotically. This is also illustrated by Figure 3, which compares the value of the risks given by Theorem and Proposition
Naturally, implementing the attention sreettoo-bbased oracle predictor T requires knowledge of the parameters k* and v². Our goal in the next section therefore to show that gradient descent is able to recover these parameters.
5 GRADIENT DESCENT PROVABLY RECOVERS THE ORACLE PREDICTOR
This section is devoted to the analysis of the optimization dynamics in (k: u) E (8d-1)² of the risk
RA(k v) - (x) 21 - crf(A T Xv)
2
We emphasize that R2 v) a theoretical risk, which depends on the distribution of the pair (X, Y) (defined in Section 2). In practice, an empirical version of this risk is minimized. As we show experimentally (see Figure 5), the stochastic dynamics induced by the empirical version of the risk are qualitatively similar to the deterministic dynamics of the theoretical risk. In the remainder of the article, we focus on the theoretical risk for simplicity, and leave the empirical risk for future research.
Our optimization method is the Projected (Riemannian) Gradient Descent (PGD), described below. Definition 1 (PGD), Given an initialization (kg. vo) E (8d- 12, a step size a > O, and an inverse temperature sequence (Ae) the sequence (kt, Ut 330 (S (gd-1)2 is recursively defined by
kx+1 Projga-1 a(Ia k V RA, (Kie, ve)),
(10)
1 Projga-1 a - da3a R, Ut),
where Projg--1 1 / |11|| denotes the Euclidean projection the unit sphere of Rd
The operators (L, - kg and (Ia - )z2t7) correspond to Riemannian gradient descent (Boumal, 2023, Section 4.3), meaning that compute the gradient of the risk the Riemannian manifold (5d-1)2. In other words, the gradient step is performed on the tangent space to the sphere at the current iterate. This is precaution we are taking because, in the analysis of the dynamics, we rely on an expression of the risk (6) that is valid only on this manifold. In addition, this ensures that the subsequent projection Sd-1 is always well-defined, despite the fact that the sphere is non-convex set, because iterates always avoid the pathological cases & 0 or V = O.
Experimentally, we observe in Figure 4a that PGD able to recover the oracle parameters (k* v*). Note that running the PGD iterates (10) involves computing the gradients RR (ke, v.) and V,RA, (ke, 14), which is non-trivial a priori. A direct approach using Monte Carlo simulations would require large number of sample points to reduce variance, which is computationally in- tractable in particular high-dimension, and in any case gives an approximate result. Instead, we leverage our closed form formula for RS from Theorem to get exact values for the gradients (up to numerical errors). Interestingly, we also observe in Figure 4a that v aligns with v² much faster than i aligns with k*. This is typical of two-timescale dynamics, which is common framework in analysis of convec learning dynamics (Heusel et al., 2017; Dagréou et al., 2022; Hong et al., 2023; Marion & Berthier, 2023; Berthier et al., 2024; Marion al., 2024).
8

---

Published as a conference paper at ICLR 2025
Delane ta dhe manifoll M
300
Exoose ttek
Altannent with
3,5
10-2
10-1
2
10-
M
0
30000
100000
10000
10000
50000
100000
Step
(a) From a random initialization on (8*-1)3.
Delumee the mattifll M
1.00
1.0
200
100000 risk
Alipamene with paranotess
3.5
10-
10-1
2
A]
-3.
10-
0
10-1
0
10000 Step
20000
0
L0000
20000
-1
10000
20000
Step
Step
(b) From random initialization on M (see Eq. (11)).
Figure 4: Convergence of PGD to the oracle parameters. Left; Excess risk as function of the number k*| and | = u*| with the oracle parameters. Middle right: Trajectories of к and y in two repetitions of the experiments. Each repetition corresponds color, the trajectory starts in the middle and ends the invariant manifold M. In all plots except the middle right ones, the experiment is repeated 30 times with independent random initializations, and 95% percentile intervals are plotted (but are not visible when the variance is too small). Parameters are 400, = 10, 7 = V1/2, and (a)
of Steps. Middle left: Alignment [s] = [
comner of the plot. Right: Distance to
X 1/(1 + 10 4t), (b) > 0.1. More details are given in Appendix E.
Moving on to the mathematical study, even with the formula for RS. a full analysis of the dynam- ics (10) is difficult. For instance, the dynamics (10) can be formulated in terms of the five variables of RR. but then one needs to study a 5-dimensional highly nonlinear dynamical system. In the following, we consider the case where the parameters are initialized on the submanifold of (Sd-1)2
= {(k: v) Sd- sd-1,k kT T = Tk* 0, v== =0- 0}.
(11)
We introduce this manifold on the one hand owing to the observation in Figure 4a (right) that the dynamics converge to this manifold even when initialized on the sphere, and on the other hand because this allows to reduce the problem to lower-dimensional subspace and to simplify the expression of the risk. Clearly, due to Assumption 1, the oracle parameters (k* u²) belong to M. A first key property of this manifold is invariance under the PGD dynamics.
Lemma 4. The manifold M is invariant under the PGD dynamics (10), in the sense that if (ki, ve) E M, then (k141, 1t+1) E M.
This lemma shows that, if the initialization is taken on the manifold, then it is enough to understand the dynamics the manifold conclude. Such analysis the manifold is tractable. This yields Theorem 5, our main result, which shows that the sequence (k,, Ut ):20 converges to the oracle values (k*, u*) (up sign) as + co, for any small enough step size, and a constant inverse temperature. Theorem Take a constant inverse temperature X₁ = X > 0. Then there exists TI >> 0 such that, for any step size & a. and for generic initialization (ko, vo) E M. (ke, vz) - =(k,, v*).
This result shows that, despite the non-convexity of the risk, the attention layer trained by PGD can recover the underlying structure the problem. Convergence to (k**, v*) or (-k",-u*) is not at problematic, since 7
(k"k ,=") - --=") by symmetry of the erf function. Furthermore, recovery is guaranteed for a generic initialization on M, in the sense that the pathological pairs (Ng, Uy) E M such that PGD fails to recover the oracle parameters Theorem are illustrated by Figure 4b. We observe that, due to roundoff errors, the dynamics are not exactly on the manifold but stay very close the manifold.
of Lebesgue measure zero. The results of

---

Published as a conference paper at ICLR 2025
We emphasize that the manifold M depends the unknown parameters k' and v", making it impractical to initialize directly on the manifold. If the initialization is not on M, more diverse phenomena possible. As already pointed out in Figure 4a, it is possible to obtain recovery of (k", v*) and convergence to the manifold M from a general initialization on the sphere. This suggests that our analysis the manifold is relevant, and completing the analysis for a general initialization is left for future work. However, we note that using decreasing inverse temperature sequence A is crucial for the recovery of (k+, u** when initialized out of M. Indeed, the best of our experiments, an deerationinndppeddee case (see Appendix E). This contrasts with the dynamics on the manifold proven in Theorem 5.
choice of À does not consistently lead to the recovery of k- and v* in this
To investigate these behaviors. a fruitful direction would be to investigate the (local) stability of the manifold M for the PGD dynamics. the manifold is indeed stable, one can hope to transfer the analysis on the manifold to dynamics initialized close to the manifold. Furthermore, recall that, in high dimension, random vectors on the sphere are close to being orthogonal. Thus, with high probability, uniform initialization in (8d-1)² falls in the neighborhood of the manifold M, that the local analysis should allow to conclude.
The proof of the theorem relies on detailed analysis of the dynamics of the PGD algorithm on the invariant manifold M. in particular the properties of its stationary points. These arguments, which lie at the intersection of dynamical systems and topology, are of independent interest. A key idea is to reduce the problem to a two-dimensional system depending only on k = k k-* and v = v 6*.
Finally, numerical experiments show that full Transformer layer is able solve the single-location regression task. Similarly to our simplified predictor, the weights align with the oracle parameters k* and v* This supports the connection drawn in Section 3 between our predictor and attention layers. We refer to Appendix E for details and plots, as well as experiments on multiple-location regression, a variant of single-location regression where the output depends on several tokens.
6
1,00
1.0
Distaane to the masiitdd
tisk
Aligznent with
cracer semunelees
0.75
0.5
10
Escoss
0,50
0.0
10-1
2
0.25
10-9
0
10-1
0
100000
200000
100000
200000
0
100000
200000
Step
Step
Step
Figure 5: Convergence of online stochastic PGD to the oracle parameters from a random initialization on
Left: Excess risk as function of the number of steps. Middle left: Alignment |s| = I and | = |u" Tu'l with the oracle parameters. Middle right: Trajectories к and color, the trajectory starts
in two repetitions of the experiment. Each repetition corresponds in the middle and ends at a corner of the plot. Right: Distance to the invariant manifold M. In all plots except the middle right one, the experiment is repeated 30 times with independent random initializations, and 95% percentile intervals are plotted. Parameters are d 80, L 10, 7 = V1/2. > 2/(1 + 10-4)), and batch size of 5. More details are given in Appendix E.
6 CONCLUSION
This paper introduced introduceddsinkke-oraalii regression, a novel statistical task where the relevant information in the input sequence is supported by single token. We analyzed the statistical properties and optimization dynamics of a natural estimator for this task, which can be seen as basic attention layer. We hope this work encourages further research into how Transformer architectures address sparsity and long-range dependencies, while simultaneously constructing internal linear representations of their input-aan an aspect with significant implications for interpretability. Beyond NLP, potential applications include problems connected to sparse sequential modeling such as anomaly detection in time series. A natural extension of our framework is when relevant information is spread across few input tokens rather than just one, which relates to multi-head attention. Future mathematical analyses should also consider extensions to general initialization schemes and stochastic dynamics. Our experiments (Figures 4a, 5, and Appendix E) yield encouraging results in all these directions.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
Authors thank Peter Bartlett Linus Bleistein, Alex Damian, Spencer Frei, and Clément Mantoux for fruitful discussions and feedback. P.M. is supported by Google PhD Fellowship.
REFERENCES
Armen Aghajanyan, Sonal Gupta, and Luke Zettlemoyer. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. In C. Zong, F. Xia, W. Li, and R. Navigli (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume I: Long Papers), pp. 7319-7328, Association for Computational Linguistics, 2021.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, Curran Associates, Inc., 2023.
volume 36, pp.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. Y. Bengio and Y. LeCun (eds.), 3rd International Conference on Learning Representations, 2015,
Dzmitry Bahdanau, Jan Chorowski, Dmitriy Serdyuk, Philémon Brakel, and Yoshua Bengio, End-to- end attention- based large vocabulary speech recognition. In 2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1945-4949, 2016.
Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. Foundations of Computational Mathematics, pp. 1-84, 2024.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Hervé Jégou, and Léon Bottou. Birth of transformer: A memory viewpoint. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 1560-1588. Curran Associates, Inc., 2023.
Tolga Bolukbasi, Kai-Wei Chang, James Zou, Venkatesh Saligrama, and Adam Kalai. Man is to computer programmer woman is to homemaker? Debiasing word embeddings. In D. Lee, M. Sugiyama, U. von Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29, pp. 4356-4364. Curran Associates, Inc., 2016.
Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Fernanda Viégas, and Martin Wattenberg. An interpretability illusion for BERT. arXiv:2104.07143, 2021.
Nicolas Boumal. An Introduction Optimization on Smooth Manifolds. Cambridge University
Press, Cambridge, 2023.
James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018. URL
http:// /ttp:///githbb com/ jax-ml jax.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with GPT-4. arXiv:2303.12712. 2023.
Collin Burns, Haotian Ye, Dan Klein, and Jacob Steinhardt. Discovering latent knowledge in language models without supervision. In The Eleventh International Conference on Learning Representations, 2023.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. 10Xi::1004.105000 2019.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bun Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
11

---

Published as a conference paper at ICLR 2025
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? An analysis of BERT's s attention. In T. Linzen, G. Chrupata, Y. Belinkov, and D. Hupkes (eds.). Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NL.,, pp. 276-286. Association for Computational Linguistics, 2019.
Gonçalo M. Correia, Vlad Niculae, and André F.T. Martins. Adaptively sparse transformers. In K. Inui, Jiang, V. Ng, and X. Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNI.--ICNIPP), pp. 2174-2184. Association for Computational Linguistics, 2019.
Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances Neural Information Processing Systems, volume 35, pp.
Curran Associates, Inc., 2022.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, 2024.
Richard D. De Veaux. Mixtures of linear regressions. Computational Statistics & Data Analysis, 8: 227-245, 1989.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-41866 Association for Computational Linguisties, 2019.
Mirko Farina, Usman Ahmad, Ahmad Taha, Hussein Younes, Yusuf Mesbah, Xiao Yu, and Witold Pedrycz. Sparsity transformers: A systematic literature review. Neurocomputing. 582:127468, 2024.
Trevor Hastie, Robert Tibshirani, and Jerome Friedman. The Elements of Statistical Learning. Data Mining, Inference, and Prediction. Springer, New York, 2 edition, 2009.
Bobby He and Thomas Hofmann. Simplifying transformer blocks. In The Twelfth International Conference on Learning Representations, 2024.
Bobby He, James Martens, Guodong Zhang, Aleksandar Botev, Andrew Brock, Samuel I. Smith, and Yee Whye Teh. Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation. In The Eleventh International Conference on Learning Representations, 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to local nash equilibrium. In Guyon, U. von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Mingyi Hong, Hoi-To To Wai, Zhaoran Wang. and Zhuoran Yang. two-timescale stochastic algorithm framework for bilevel optimization: Complexity analysis and application to actor-critic. SIAM Journal on Optimization, 33: 147-180, 2023.
Sebastian Jaszezur, Aakanksha Chowdhery, Afroz Mohiuddin, Lukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, and Jonni Kanerva. Sparse is enough in scaling transformers. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 9895-9907. Curran Associates, Inc., 2021.
Samy Jelassi. Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 37822-37836. Curran Associates, Inc., 2022.
Tokio Kajitsuka and Issei Sato. Are transformers with layer self-attention using low-rank weight matrices universal approximators? In The Twelfth International Conference on Learning Representations, 2024.
12

---

Published as a conference paper at ICLR 2025
Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 784 794. Association for Computing Machinery, 2022.
Kenneth Lange. Optimization. Springer, New York, 2 edition, 2013.
Jean- Saan-François Le Gall. Measure Theory, Probability, and Stochastic Processes. Springer Cham, 2022.
Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Inference-tinee intervention: Eliciting truthful answers from language model. A. Oh, T. Naumann, A. Glober- son, K. Saenko, M. Hardt, and S. Levine (eds.). Advances in Neural Information Processing Curran Associates, Inc., 2023a.
Systems, volume 36, pp. 41451-41530
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In A. Krause, E. Brunskill, K. Cho, B. Engelhardr, S. Sabato, and Scarlett (eds.), Proceedings the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 19689-19729, PMLR, 2023b.
Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transformers. Al Open, 3: 111-132, 2022.
Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention- -bsed neural machine translation. In L. Mårquez, C. Kallison-Burchhe and J. Su (eds.), Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pp. 1412-1421. Association for Computational Linguistics, 2015.
Pierre Marion and Raphaël Berthier. Leveraging the two timescale regime to demonstrate convergence of neural networks. In Oh, T. Naumann, A. Globerson, Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 11996-65029 Curran Associates, Inc., 2023.
Pierre Marion, Anna Korba, Peter Bartlett, Mathieu Blondel, Valentin De Bortoli, Arnaud Doucet, Felipe Llinares-López, Courtney Paquette, and Quentin Berthet. Implicit diffusion: Efficient optimization through stochastic sampling, XXiv:2402.00466 2024.
Andre Martins and Ramon Astudillo. From softmax to sparsemax: sparse model of attention and multi-label classification. In M.F. Balcan and K.Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings Machine Learning Research, pp. 1623. PMLR, 2016.
Peter McCullagh and John A. Nelder. Generalized Linear Models. Chapman & Hall, London, edition, 1983.
Tomas Mikolov, Quoc V. Le, and Ilya Sutskever. Exploiting similarities among languages for machine translation. arXiv: 1ttv:1339,. 1309.41066 418,, 2013a.
Tomas Mikolov, Wen-tau Yih, and Geoffrey Zweig. Linguistic regularities in continuous space word representations. In L.. Vanderwende, H. Daumé III, and K. Kirchhoff (eds.). Proceedings of the 2013 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 75,-751, Association for Computational Linguisties, 2013b.
Neel Nanda, Andrew Lee, and Martin Wattenberg. Emergent linear representations in world models of self-supervised sequence models. In Y. Belinkov, S. Hao, J. Jumelet, N. Kim, A. McCarthy, and H. Mohebbi (eds.). Proceedings of the 6th BlackboxNLP Workshop: Analyzing and Interpreting Neural Networks for NLP, pp. 16-30. Association for Computational Linguistics, 2023.
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In International Conference Machine Learning. PLMR, 2024.
13

---

Published as a conference paper at ICLR 2025
Vlad Niculae and Mathieu Blondel. A regularized framework for sparse and structured neural attention. In Guyon, U. von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwamathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel. M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12(85):2825-2830) 2011.
Mary Phuong and Marcus Hutter. Formal algorithms for transformers. arXiv:2207.0923 2022.
William Press, Saul Teukolsky, William Vetterling, and Brian Flannery. Numerical Recipes: The Art of Scientific Computing. Cambridge University Press, Cambridge, 3 edition, 2007.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arXiv:2202.08791 2022.
Jason Ramapuram, Federico Danieli, Eeshan Dhekane, Floris Weers, Dan Busbridge, Pierre Ablin, Tatiana Likhomanenko, Jagrit Digani, Zijin Gu, Amitis Shidani, et al. Theory, analysis, and best practices for sigmoid self-attention, arXiv:2409.04431 2024.
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study on ReLU and Softmax in Transformer. arXiv:2302.0644 2023.
Michael Shub. Global Stability of Dynamical Systems. Springer, New York, 1987.
Youwei Song, Jiahai Wang, Tao Jiang, Zhiyue Liu, and Yanghui Rao. Attentional encoder network for targeted sentiment classification. 1rXiv:1902.09314, 2019.
Charles M. Stein. Estimation of the mean of multivariate normal distribution. The Annals of Statistics, 9:1 9:1135-115,, 1981.
Chi Sun, Luyao Huang, and Xipeng Qiu. Utilizing BERT for aspect-based sentiment analysis via constructing auxiliary sentence. J. Burstein, C. Doran, and T. Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 380-385. Association for Computational Linguistics, 2019.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S. Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances Neural Information Processing Systems. volume 36, pp. 71911-71947. Curran Associates, Inc., 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, pp. 6000-6010. Curran Associates, Inc., 2017.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In A. Krause, E. Brunskill, K. Cho, B. Engelhardr, S. Sabato, and J. Scarlett (eds.). Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 35151-35174, PMLR, 2023.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D. Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. In International Conference on Machine Learning. PLMR, 2024.
Andrea W Wen-Ve Yi and David Mimno. Hyperpolyglot LLMs: Cross-lingual interpretability in token embeddings. In H. Bouamor, J. Pino, and K. Bali (eds.). Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pp. 1124-1131. Association for Computa- tional Linguistics, 2023.
14

---

Published as a conference paper at ICLR 2025
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac. Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger. Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-45. Association for Computational Linguistics, 2020,
Mitchell Wortsman, Jachoon Lee, Justin Gilmer, and Simon Kornblith. Replacing softmax with ReLU vision transformers. arXiv:2309.08886, 2023.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024.
Hu Xu, Bing Liu, Lei Shu, and Philip Yu. BERT post-training for review reading comprehension and aspect-based sentiment analysis. In Burstein, Doran, and T. Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume I (Long and Short Papers), pp. 2324-2335. Association for Computational Linguistics, 2019.
Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems, 34:17084-17097 2021.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett, Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-555 2024.
15

---

Published as a conference paper at ICLR 2025
Appendix
Organization of the Appendix. Section A presents the main steps of Theorem 5. The intermediate results of this proof, well as the other statements of the main text, are proven in Section B. Section C provides an expression the risk R beyond the manifold M that extends the one provided in Lemma 6 on M. Section D gives some useful technical lemmas. Experimental details and additional results are in Section E. Finally, Section F discusses additional related models.
Notation. In the whole Appendix, consider a constant inverse temperature schedule X, A as in Theorem 5. For this reason, it is not necessary to make explicit the dependence of Rx and R on A, and we use the lighter notations R and R instead.
A OUTLINE OF THE PROOF OF THEOREM 5
This section outlines the essential steps for the proof of Theorem 5, For clarity, the proofs are to be found in Appendix B, except the proof of Proposition I0.
Step 1: Invariant manifold & reparameterization. We first show that the risk R(k, v) has simpler expression when considered on the manifold M.
Lemma 6. The risk R(k, v) restricted to M has the form
R(k, U) = - erf \v 2(1 + 21222) 2x2 TTTT) + +>
d k kh x2p2)))7)
2
+ (L - 1)((0, +
where, for 1,7 E.7ER. RR,
<(1,72) = (1 (er
G G~N(0077))
This expression has two main consequences. First, we use it to prove that the manifold M invariant by PGD, according to Lemma 4. Second, we observe that the risk on the manifold depends on the variables (k, v) E Sd-1 x sd-1 only through the two scalar quantities
k = k k and = TT=-
This suggests studying the dynamics in terms of the reduced variables (k, v) E [-1, More precisely, in the following, we denote by R the risk function R reparameterized function of V), ie., we let
R (K,D) = -2 2222re (Ni
d 2(1 + 2 K + + T K. (2)) T
(. (L-1))
Note that, with slight abuse of notation, we use R denote both the function of five variables (R, v.O, p) (as in Theorem 1) and the function of only the first two variables (k., 2/). There should be no confusion, as both functions coincide on the manifold M where = P = denote the corresponding PGD iterates using this reparameterization by (Ke, 14)) ;= (k, k* 02 v*).
0. We also
With this notation, the following lemma reformulates the PGD iterations as an autonomous discrete dynamical system in terms of (Kt, Vi).
Lemma 7. When initialized on the manifold M, the PGD iterations (10) can be reformulated in terms of the autonomous discrete dynamical system
(6441,441) = g v4),
(12)
where the mapping 1- []1, is given by
s--a (x,
D - a (n, v))(1 v2)
g(6, v) =
(13)
11+ (R,
1 (n. v)) -22)
16

---

Published as a conference paper at ICLR 2025
1.00
1.
035
08
525
K
0.0
025
055
02
-100
-1.
-0.5
0.0
0.5
(a)
(b)
Figure 6: Dynamics in (k, v) on the manifold M. In (a), the fixed points of the dynamics are represented; the minimizers, saddle point, and maximizers respectively depicted in yellow, blue and red. In (b), the vector field (k, v) + - (r,, (1 x²), 8, (K, v1 - v²)) is displayed (the colormap corresponds the magnitude of the vector field).
Step 2: Analysis of the stationary points. Regarding the dynamies restricted the invariant manifold M, we can characterize the limit points of the PGD iterates as follows.
Proposition 8. For a sufficiently small step size a and for any (ko, vo) E M, the risk R is decreasing along the PGD iterates. Furthermore, the distance between successive PGD iterates tends zero, and, if (K, 2) is an accumulation point of the sequence of iterates (K4, then
- k )0 (K. v) = and (1 (1-v" (1-r²)a,R= (K, v) = =0
(14)
We stress that the system (14) of equations corresponds fixed points of the dynamics (12)-(13). We next solve this system of equations.
Proposition 9. The points (r, v) [-1, satisfying (14) are (n, v) (=1, and (n,, v) (0, O).
The identity (N,, p) = #1) corresponds the situation where the variables (k, are aligned (up to sign) with the targets (k* v²). As the next proposition shows, these are the only global minima of R
Proposition 10. The fixed points of the dynamics can be classified follows:
(i) The points (K, v) = (-1,1) and (1, -1) are global maxima of R- on 1-1, 1/2.
(ii) The points (k, 1/) = (1, 1) and (-1, are global minima of R on 1-1,
(iti) The point (R, v) = (0, 0) is saddle point R
s. on [11,
The fixed points of the dynamics well the vector field
(rs, v) - (O.R (k, x² 0, (k, u] v²))
are displayed Figure 6.
Step 3: Convergence to global minima. The convergence of the sequence of iterates (Ke, Mt) to global minimum shown in two stages. First, we show that the iterates converge to one of the five fixed points described in Proposition 10.
2This notation is used to designate any extreme point of the square [-1, 1)2, i.e., (N, (1, 1), (1, -1),
(-1,1) 1), and (-1,-1)
17

---

Published as a conference paper at ICLR 2025
Proposition 11. For sufficiently small step size a, the sequence of iterates (Ke, Pe)ezo converges to one of the five fixed points (0,,,,
Proof. According to Proposition 8, the distance between successive iterates (Ke, Mt) tends zero. Therefore, the set of accumulation points of the sequence (Kt, M Pe) is connected (Lange, 2013, Proposition 12.4.1). Since there a finite number of possible accumulation points (by Proposition 9), we deduce that the sequence has unique accumulation point. Furthermore, the sequence belongs to a compact. Thus, converges, and its limit is one of the five fixed points.
It remains to precisely characterize the limit of the sequence (Ke, To this aim, begin by showing key properties of the gradient mapping 9.
Proposition 12. For a sufficiently small step size a, the mapping g is a local diffeomorphism around (0, 0), whose Jacobian matrix has one eigenvalue in (0, 1) and one eigenvalue in (1)) Furthermore, is injective on 1-1,10 1|², differentiable, and its Jacobian non-degenerate.
These properties enable to apply the Center-Stable Manifold theorem (Shub, 1987, Theorem I11.7), a tool from dynamical systems theory, to deduce the next proposition.
Proposition 13. For a sufficiently small step size a, the set of initializations such that the sequence (Ka, vi)izo converges to (-1, 1). (1, -1), or (0,0) 0) has Lebesgue measure zero (with respect to the Lebesgue measure on the manifold M).
Combining Proposition 11 and Proposition 13, conclude that, provided the step size is chosen small enough, the sequence (Ke, Pr)eso almost surely converges to one of the minimizers, (1, 1) or (-1, -1). This convergence is almost sure with respect the Lebesgue measure on the manifold M. Indeed, Proposition 13 ensures that the pathological initializations converging towards a maximizer or a saddle point are of Lebesgue measure zero. This concludes the proof of Theorem
The use of the Center-Stable Manifold theorem is crucial to our proof. Unfortunately, this tool does not provide quantitative rates of convergence. Obtaining a rate a challenging task as would require quantifying the distance of the iterates to the saddle points of the risk (the dynamics is indeed slower near saddle points), which in turn requires other tools of analysis and potentially additional assumptions.
B PROOFS OF THE MAIN RESULTS
B.1 PROOF OF LEMMA 6 AND THEOREM I
We recall the formula for the risk
R.(k: v) BB Σ erx( x))xx))
t=1
and the data model
Y =
V=XJ, XJ v² +
where
X N ((V _k ,2aa,,
Jo P(fI.... L}) and
X N(O, La) for 4
In the above expression for the risk, we can condition on the value of Jo. Actually, the conditioned risk is independent of Jg. Thus in this section, we assume without loss of generality that Jo = 1 a.s.:
R(k: v) = B [x + erf(xx
Σ en((AXT()xXT 0)
(15)
where
{ Xi NVV
Xe la) for 8> 2.
18

---

Published as a conference paper at ICLR 2025
We rewrite this quantity in terms of multivariate standard Gaussian random variables. Using Assump- tion 1, we get
orf( x k xx]]++(()))))))))))) R xxx x87 Σenth exx [e u)
R(k, v) B ( uz +++++=---- -
1
t=2
where X1, X2...., XL NN This can be formulated in terms of the five scalar quantities i = ko ke, V = v 0².0= v T k", . = v² and p p= kJ.. u. Indeed, we have
R(k, u) = R (1c, v. v.0. p)
= (1G +
(v
d
/ 2 (o+)) 0 + GGi erf (x(V2
Σ G)
(16)
2
where
N (0,
1 V V
(o.
(17)
This last expression only involves the five parameters K; v, v): P. which play a role either explicitly in the function or as parameters of the covariance of the random variables. This proves the first statement of Theorem A computation of a closed-form formula for this expectation is given in Appendix C.
On the manifold M defined by f == 17 = p = O, we can simplify the expressions (16)-(17)
d R$ (M, 0,0, = E (707' + - -0-GGoorf( (( 2 not) /
G)af((AA))) "
GH G
C C nid )m,m
N (0, (1 1) ,h,...,.G Gt ~ NOO, 1), and ~ N' 82)
where
independent.
We first expand in & and obtain
d R (K, = 5² = - -Gferf -Gifrr (>(==++ Σ )jen((a())) )
1
We now expand the square, as follows:
R (k,,0,000000) 8² + r E [GT))]] - -27²B Gferl (2(( (N(V=*+*G*))
+ 2 2 (G)) ert2 (2V/2 (
E (ap Gierf( (2(11) dd
( + ((+) )
2
+
Σ GIG erf(A
f.m=2
We address each term in this separately.
Since G ~ ~N 1), 22E [[G]' )²] =
1 N (0,
0. ( 1)) is independent from G ~ NN 1), we have
Since
(x(1 / x
2 Gy orf
= ++++ne ++++)]]]]] = 22 . न E orf (V
19

---

Published as a conference paper at ICLR 2025
2--27 erf (ort(a(((// g +
Finally, using Lemma 18(i)), obtain
-2 Gi [ G'erf (VV Z" R r - - Xv d
2 I 21292
Since Gj, Gf widd NO0, 1), we have
2 (G)) (x()
d 2" + (c4))]]] - (GG)
erf2 ( (2(V2x+
=>-7(enn (e9 eer (2(1 _ 2 mC4
Using the definition of ( Eq. (7), we have
y²B Par) err (>(v d
2 x +++ = - (G B > (2
2nt + +
d = xm X222)
For f= I. (GT Gi, Gt), Gū, and G3 are independent. Thus d
E - Gienl (> (1
2 x rG4)) G)er((AGt)]
=E [o'
- Gjerf -Gjeff (2V x * { [or(A)))) (or =0,
where in the last step we use [[G 0.
Finally, to tackle the last term, we address the cases é + m and é = m separately. If l # m, Gg Gh G and Gh, are independent, we have
E|Gyer CO er((AAA )] [CC
0.
If f = m. G7, G mid. (O 1), we have B[(G)) (AG))] = " (AG+)]=<((0))
= ((0,)3)
Putting together these computations, we obtain
-2-² R U,K,,,,00000) = 8² + +2² - verl \v d
d + +02c
2 Zh, X322)
/1+ 2X2,2
+ (I,,1)((0,,)))
This proves Lemma 6. Taking x = proves Theorem
B.2 PROOF OF COROLLARY
Recall that, according to Theorem 1.
d 2(1 + 21222)
d _ x2() + 23) + (
Rk" (()),)))=""") = 22--27² erf (v/=
where, for t, y E IR,
(( E (trf2 (t +]]]
G~ G00,1.
We compute the limit of each term separately. First, we have
d
Avd d
(18)
2(1 2X272)
/2
20

---

Published as a conference paper at ICLR 2025
Therefore, the second term of (x(k", u*) tends -2y2. To handle the third term, note by Jensen's inequality that
orf 1 1> (2 (2 2 E
d
( ( 2
+ (t) > E erf (antAA/
f d + xnc)
Thus, by Lemma 18(i).
1 > (A 2. > exf (NV
d 2(1 + 22272)
- +1,
where we used (18). Thus the third term of R2(k", (²) converges 72. As for the fourth term, observe by Lemma 17 that
4 u² crf((u)
π
hence
0 < ((0,A)) < _
G
_
Since AVL O, get
(L- X²) = O(A2L) = o(1)
Putting everything together, we obtain
Rk w²) d=100 2 +
Since we already know by (8) that the Bayes risk is lower-bounded by e2, this proves that the Bayes risk is asymptotically equal to e², and that the oracle predictor is asymptotically Bayes optimal.
B.3 PROOF OF PROPOSITION 3
Let us first introduce a useful notation for the proof. If M is a block matrix, we denote by M[j its (i.j)-th block, and likewise, u is block vector, we denote by Ubj its j-th block. Next, note that
E[y = = + E ((o')xxx,)
= = + 7 11
8² + 7
since ||u* =1. Recall that
8+ E arg arg min E (Y -xT.xxX..,,,,).. x[)))]
BeRdn
is the optimal linear predictor. The classical formula for linear regression shows that
X1
8
IF
(X
E ((1*) X1o +
XL
XL
On the one hand, let
X1
M =
(X] ,.))).
Then [E[M] =
B [o]], and EEM [o] is a block-diagonal matrix, where, for j.j E {l,.... L}.
E|M[Jo o dz la + 8, 8j-5 d (k)))))).
Thus
E| = - (1- + * Ia 1)Ia + ps' ks (1:)).
21

---

Published as a conference paper at ICLR 2025
On the other hand, let
X1
((0*)
No
+ +5)
X1
Then
X1
X1
E ju]
E
:
X
U¹
E
E
:
X T
Ju
1²
X1
X1
PI
5k (K*))
Piu' ++
0²
?2
:
PL (1 La
(**()))TT)
PLU*
since, by Assumption A O.
Since [[ is block-diagonal matrix and Elu] is block vector, we get by standard computation rules for block matrices
1 (B[M) = B Blulns ps(7" 1)Ia+ + Pj 1h (k+))) ("*))))""pp''
ppp
Recall the Sherman- Morrison formula (Press al., 2007, Section 2.7.1), which states that for any vectors u, U E IR (Ia + uu)) - = (Ia - uu /(1 + u uTu) v. Applying this formula with orthogonal vectors, we obtain
p²p;
8 = (1+ 1)) "ppp"" 1 (T 1)
which shows the first formula of the proposition. Finally, the risk associated with the optimal linear predictor (XT,.. ..... x,.) + (XT ..... xT.) is given by
R(3*) = E| [X ... X[881
Bin
= 7 -2 (pi(v*) PL(O))T)
p3
3² + 72 7² 4 Σ
(19)
ps( 11 1)
This shows the formula for the risk given in the Proposition. To obtain the last bound, observe that, if 2² 2 1, we have 1 + (v - 1) > 1. If 72 & 1, since Py & 1, we have 1 - (> 1) >> = 22. Thus we obtain 1 - p, - 1) 3 min(1, Therefore,
R(8") > 22 max ( ,,72 Σ
pj
j=l
L
> 2 + 7- -max(74, 72)
Σ
Pj
max Pj j=
j=1
26 2 (7 + 72)
max Pj j= I.....
2 (22 j=.....
P
When all py are equal to 1/L, all terms in the sum are equal, and Eq. (19) simplifies to
+
R(B*) = 2² +2² Lna
2 +2²
1 1
212 2
L + 72
22

---

Published as a conference paper at ICLR 2025
B.4 PROOF OF LEMMA 4
As a first step in the proof, we prove the next lemma, which is the key towards the invariance property we are aiming at, in that it shows that, for point on the manifold M (defined by 0 = `/ - p 0), the gradient of the risk does not "push" the point outside of the manifold. Its proof leverages the expression of the risk as a function of five parameters derived in the previous section
Lemma 14. At any point (K. 0 T7, p) such that O = = = 0, we have 0₂ = 0₂ = 0,R = O.
Proof. We use Eq. (17) and change signs in the square function:
R (K. v, p.
=E [1GT + + d rG:)eer( erf (Vg + - -- (Gjer((a(t) '
2/200 + rG:
)
1=2
= = - 5- (1 (-0) + (-) + erl ( d
/ rGt)
Σ (r 9.
2
where
G
v
7))
N
0,
1
&~,,00,
G
G
i.d.
p
Thus
G
G G
-n -p 1 )
v
-1)
G G
N
0,
v
1
- =
Gk
iid.
As consequence,
R R,,,,,, n. p) = R² R^,,,((( U1 -0, -0,-11-11 -p).
Taking the partial derivative in 0, we are led
dgR^(R (R. v. - (R, v, -1,
At a point such that 0 p 0, this gives R= 0, 0) = - RO V,. 0) and thus O₂ (K, 0, = The proof for the other two derivatives R,, 0₂ is identical.
v.
We now complete the proof of Lemma 4. By the chain rule for total derivatives applied to R(k, v) R and then by Lemma 14, on the manifold M, we have
V&R = (0.R*) yk* + (8, )u" + (a,R5) Ju )**
(20)
and, similarly,
V.R (O, + (00 (0, (8,R
(21)
Recall the formulas for the PGD updates
h2 a k R R.ee, 1x)
h+1 Projgs 1 (k2 a(I - ke Rkke, 01))
|kk a( k V R (Kir, 11)
122 a(1 [ 1.4n V R R( Uz)
U241 = Projga- (vr a(I - 1, T VV N. Rikke, vc)) =
|0p a(I 1.4o V RR v2)
Let c = llbg - a(I kake V R(ke, 1z)1l1 and cu = -1102 - aa( m. V (ke,, 1z) Then,
(ke, v.) M,
(ot)Tees Tke
(0..)) k)) =0,
Ck
23

---

Published as a conference paper at ICLR 2025
(k*) 1t+1 (K*) v -
-1 -1 (0,R (No, vi)le" = -0,
Cr.
and
v2 ke a(0,RR 0. )00) T a((0 Kokk }1r))
T 1(41 1441
Cock
(0.R5 k )k* ( ULu U.U,
0.
CoCk
where we have omitted the dependence of (*x and R5) in (Ke, in the last expression for the ease of readability. Note that the last term is equal to zero since
((I (1--ki) ke }kr) T (I 1aia )o" = (k* k Tou()T(u=- v(v1) ) = 0.
This shows that (k441, E M.
B.5 PROOF OF LEMMA 7
By definition the PGD iterates and by (20)-211, one has
Kt - or o& uz (k+) - keki k
K2 - ((0,R5)( - k2)
K++1 = -ki41
VI+a ) k,,; 3
(1+ R2)
Vt nd, (K,, (1)()) 0, )ur
" a (1+ vz)
V/+1 = 0141
veo Jut
where we have used the Pythagorean theorem and the idempotent property projection matrices for the denominator.
B.6 PROOF OF PROPOSITION 8
In this proof, C' denotes a constant that does not depend on the step nor on the step size a, and which may vary from line to line. First note that the risk R C on the compact set [-1, 1J2 In particular, it is A-smooth function for some A > 0, in the sense that its gradient is A-Lipschitz continuous. Thus
R RC (18411, R (N,, + (VR (VR
Ket1 - Kit V4t1 vi
+
A
{
Kt+1 - Kic
2
V--1 1/
ie..
R² (8441,1)4)1) - R (R
< (O,RR))(HHH) ne) + (2,R v4) A [(544] -Ra)² + (ut+1 - -1) (22)
Our goal in the following computations is to derive an inequality of the form
R (841,
a(0, R (Nt, < aa0 1 v2)
+ +Ca - x²) C
which shall give us descent lemma for a small enough. To this aim, observe that, by definition of the iterates (Ke, vi) given by (12)-(13), one has
a0,R - x2)
Ke+1 Kr
Kt
(23)
VI+ (
(0. RT)
=aa(a,R )1 xx+)
(Rr
V a² a
24

---

Published as a conference paper at ICLR 2025
As consequence,
Kx41
Kr + ((++((22<(((1x))
- 1 | a(0, - x3)
V1 k3) a (0.R5)(1 - 13) -(0,R -x,]
(0
(24)
where the second incquality holds by Lemma 16 and the last bound holds since the function (x, v/ H IK 1 a (k, 11) x²)l is uniformly bounded for all a & 1. This bound has two implications.
First,
x2) (0,R5)(mea - ni) + (a10 (x)) = (0,R5) -Ka) + a(0,R k2))
< R < || Ni a R (1-1))
< Co (
(25)
where we use the fact that R* is bounded, and the bound (24). Second, since the square function Lipschitz on compact sets, we have
R$
[ -x,²²
Thus
(Ke+1 - (2) (0 -nt)
(( (1)1)(()
(26)
We also obtain analogous bounds (25) for v. namely
- v1) + a0,, (11 12) < (a
(27)
and
(V++1 - v1)² < Ca (a, ((1
(28)
Plugging the bounds (25) into Eq. (22), we obtain the desired inequality
R (KE41, #(41)
R (Ke, v.) < a0, ( x²) a(aR (1 1/) +
By choosing the step size a & 2z, this ensures that
R (Ket1, R (KE, < (0,R5) n2)
a (1,
This shows that the risk is decreasing along the PGD iterates. Next, introducing Re (x.)).0) R (R, v) and using a telescopic sum, we have, for all T 2 0,
2R 14) - Rei > R (Ko, - R (KT,,VV))
a T
(0.R5)((1-K() + [10-3)9)--)++0-------------
2
Since the left-hand side is finite, and the terms of the sum are nonnegative, we conclude that the series converges as T - cc. In particular, the generic term (0,R$)? n&?) + (02R5) n²) of the
series converges to as L + oc. Therefore, the accumulation points (Koo, (oo) satisfy
{ O 0.R (Roo, Poo) 0
0.R. (Nox; Mox) = 0
1
Inspecting identity (23), we observe that the convergence of the general term also implics K4,11 - KGg 7 O. We obtain similarly that 24+1 14 0.
25

---

Published as a conference paper at ICLR 2025
B.7 PROOF OF PROPOSITION
Recall that the risk in terms of (Ki, 1/) is given by
d/2/2n
d x,, 2 +
(K, = - ver
+ +
1C((0,33) +
VI + 2
Then the gradients of R are given by
O₂R (
X d/2n
(u-erl"
d/222x
- \ 2(1 + 2x222) erl'
/1++22772
(1 2X2)))(11111 4X2₂2)
and
Avd/2x
R,R s (K, -2-2erl
V1+22222
Therefore, the solutions of the system (14) satisfy
+ erf (c) = 0
K
d with C1 V1A2-241 (1+22*77)
The solutions of this system
(k, = (0, 0) or (K, = #1)
B.8 PROOF OF PROPOSITION 10
Since RS is smooth function, the extrema of this function on [-1, 1|2 are either critical points (admitting null derivatives) or points on the boundary of the square [-1, 1J3. Starting with critical points, the only critical point is (0, 0), and it is saddle point. Indeed, the Hessian of R at (0,0) is
HR (0,0)
V PM
d
C
2(1 + 2X222)
==
where c (T(4125241) 2(142x272 0. Then, as det(M) 11 the two eigenvalues of
HR (0, have opposite signs, (0. 0) is thus saddle point. The extrema of R must therefore be on the boundary of the square, which we examine next.
For any (16, v) (-1,1) one has, by inspecting the signs of the gradients given in the proof of Proposition 9,
R. (x (-11) and R (1,11) < R (-(10) <
This shows that the minimum of R on (R, 1), E [-1, U ((L (1, 1), and the maximum is reached both at (1, 11,-1) and (-1,1), since R is even. Using again evenness of R², we conclude that the extrema of R on the whole boundary of the square, and
is reached
thus on the whole square, are the minimizers (1, and (-1,-4 -1), and the maximizers (1, and 1-1,1).
B.9 PROOF OF PROPOSITION 12
We prove the statements of the proposition by one.
26

---

Published as a conference paper at ICLR 2025
The mapping g is local diffeomorphism around (0, 0), whose Jacobian matrix has eigen- value (0. 1) and one eigenvalue in (1. oo). Consider the Taylor expansion of the first component g(k.2) of g(., v). Since R₂ R=,200,00 (0, O, and R is smooth, letting (K, we have (O. (k, 0(11=]13). Thus,
a(d (N. (K. (1) x3)
0(K,V)1 (K,
k - alO,R (K,, 1)) -22)
v + 0(1==12)
-K- a1--00, (ne 11) x²)) ) (1 0(1x=13))
= ad, (Kr, v) + 0(1x=113)-
Proceeding similarly with the second component of g. we obtain that the Jacobian of at (0, 0) is given by
(,(0, = 12 o (0,0) = 12 + a Vi PMV 2(1 + 2x2y2)
1
where 7x((1>-2+11) 2[1+2x=77]
0. Since det(M) 1, choose small enough
so that one eigenvalue of J, (0. 0) is strictly between 0 and and the other one is strictly larger than 1. Therefore, (0, is invertible, showing that g is a local diffeomorphism around (0, 0).
The mapping g is differentiable on 1-1, 1|2, and its Jacobian is not degenerate. The mapping clearly differentiable as a composition of differentiable function. The more delicate part is to show that its Jacobian cannot be degenerate. To show this statement, observe first that, for [-1, 1]. we may write g(x) = x + ah(x), where the first component of is given by
h(,,, 1/1 = (g(ni v)1 x))
K - a0ORR= (Ki, v))(1 - x²)
-k)
/1 a² R- (K, v)) k33)
(0 v))(1 x2) + a² R (x, (1, R?
(O,R (x, () k²2)
=: =: (K.M) (K.
=:fa (R.U)
Let us prove that the gradient of h(,,, v)1 is bounded uniformly over a < The uniform boundedness clear for the gradient of f!?), which writes as composition of functions with uniformly bounded gradients for a & 1. Moving on to (())) and letting
[
+
9: ())
VI+a+t 1) B
sup (0,R* (K, (K) xx²). (k.
we observe that få is the composition of with a smooth function independent of a. In particular, suffices to show the uniform boundedness of Vg deduce the of V v We further have, by Lemma 16, and for a & 1,
1
abb Oag(a, b) /1 + ab
1
< <ab
B
and
Oog(a, b)] =
aa
a 1 2
<
+ o
2
Therefore, the gradient of h(K, v)1 is bounded uniformly over a < 1. Proceeding similarly with the gradient of h(,, v)2, obtain that the Jacobian of h(,, is uniformly bounded over & 1. Recall now that (n, v/) = I2 + (K, v). Therefore, taking a small enough, we obtain that the eigenvalues of Jg have to be bounded away from zero.
27

---

Published as a conference paper at ICLR 2025
The mapping g is injective. The computation above shows that h is B-Lipschitz continuous with / independent of a (for a small enough). In particular we can choose a such that a < 1/3. Now, let # y E [-1, be such that g(x) g(y). Then
| vll < [h(=) h(yill & aB1lz vll < |1 -yll-
This is contradiction, showing that g is injective.
B.10 PROOF OF PROPOSITION 13
Recall that -1) and (-1, are maxima of the risk R on [-1, 1]2 by Proposition 10, and that the value of the risk decreases along the iterates of PGD by Proposition 8. Thus the only possible way to converge to these points to start the dynamics from them.
The case of the point (0, 0) is more delicate. We apply the Center-Stable Manifold theorem (Shub, 1987, Theorem III.7) to g. which a local diffeomorphism around (0. 0) by Proposition 12. This guarantees the existence of local center-stable manifold Wis Ioc. which verifies the following properties. First, its codimension is equal to the number of eigenvalues of (0, 0) of magnitude larger than 1, that is, 1, by Proposition 12. Hence it has Lebesgue measure zero. Second, there exists neighborhood B of such that Nt=o g (BB) C Was- Then, let W be the set of all æ which converge to (0.0) under the gradient map g. and take x E W Then there exists T such that (ar) B for all t > T.
This means that g² (I) E N== g " (B), and thus (x) E lse- So, I E (Wes). We have just shown that
W c U (WW T30
Finally, we prove that the pre-image of sets measure zero by g" has measure zero for any T > This shall conclude the proof of the result since countable unions of sets of measure zero have measure zero. To show this, note that is injective by Proposition 12, and therefore g" injective inverse T of g" defined on the image of g", and the pre-image by g" of Wiss is exactly the image by g -T of W (intersected with the domain of definition of g g-²), Furthermore, by Proposition 12, the Jacobian of g" is invertible. This guarantees that g is differentiable by the inverse function theorem. The conclusion follows by recalling that differentiable functions map sets of measure zero to sets of measure zero
too. This allows to define
C EXPRESSION OF THE RISK BEYOND THE INVARIANT MANIFOLD
In this appendix, we provide expression of the risk R beyond the manifold M that extends the one provided in Lemma 6. This result is not needed to prove Theorem 5, and its proof is more involved that the one of Lemma 6. However, we provide it since it might be relevant to follow-up works that would study the dynamics if not initialized on the invariant manifold M. is also useful for the numerical simulations (see Appendix E).
Proposition 15. We have the closed-form expression
R (K, (1,, 1. 8² + 7
Av A/2/2K
d - -2772
1
erl A
-2r2verí
/1 + 2X272
V1+ 1-2222
V 1 2X222
2Xp erl"
Avd/2N
/d K. x2?2)
+G5PP + +731C
1+ 1+22222
VI+ 22272
(d
X222p'
41y2
Add/2k
erf' X/d/2n
+ V2
0p
erf
1 2X
2X29
(1+4227)) (1+2222)
VI+2X22
4127402
Av
+
erf'
V 2 22222)
V1+ 11272
+ (I. - 1) [s(0,)9)
812
412
+
(I. 1) 2)p
T 4
( + 212)
(1 + 2X²)
28

---

Published as a conference paper at ICLR 2025
4ML-11 1)p + 2X²)T
(Vgrort
Avd/2K + (1++2x2
Xn2p erf'
\v +
V(1+2M
21222222
V1 + 2X2g2
Proof. We first recall the notations for the five scalar products that used throughout this proof.
v= == k k
==v² - 1 k v²
p= k= V.
A first decomposition. We start back from the expression (15) obtained for the risk. By expanding in &. then expanding the square, we obtain
R(k, 1) = EE (e((()) [x] 0- Σ x7 vote(((AX) +
=B((xTo-- x x(
Σ x xT + E²
- (xTeo=xxTeer((A)
IE (x] x)))
=
+
=:R
=:fq
E b)XJ 1XX k))
f1
=1
-2
E
(XT o0 Xi ver((AX] X3 ver((AX, (i) to²
t=2
Computation of Ri.
By expanding the square,
T (x7 o²
xj -x xe]xe((Ax]]]))
= =E((XT(*))) 215 [x] xxTuxx vefee k)) E (exTeef((XX])))
These three terms are computed hereafter. First we have
E ((xT))))) = (E(XTT*))² Var(X] u*) = (0
Second,
E x 0z XI ver[(AX, 1:))
(v (d
(x:)) w² + 7,1 ( (V2 2z)
erf (N/(
(
E
d 9 Zz erf A / 5k + 2Zz ]
E
Zi
(V2
with
Z1 Zz
TTu
IT y1-
N 0,7² -
1 V "
V 1ppa ( 1))
N
0,72
V
0"
u 1
k
Z
k
k
n
P
Recall the multivariate version of Stein's lemma (Stein, 1981), which states that, when Z, Gr.....CC.. Gp are centered and jointly Gaussian, and IR" + R.
E (Z,,
Σ
29

---

Published as a conference paper at ICLR 2025
Therefore,
E Xi o- x verf(AX; k))
\v d
+ Xy22n n [[V
d + Zz
arf' A AZs) )
=T'VE
erf
2 K + AZ3
2
\ /d
/d
d
erf' (N//ee g + 223 )
= =TrVE
erf
+ AE3
+ x 2 V non
2
+ x2 nop erl"
2Vz AZ3
a/2K
11 erf
Avd/2n VI+22292
verf
11+ 2A292
Xpptmp 1+ 222 erf"
dd/2s
1++ 2X373
by using Lemma 18(1) (iii). Finally, using again Stein's lemma and Lemma 18(iv) (vi). the computation of the last term as follows:
E
(X]
eex b))²
d 0+ 222 k*)
d kT +
27
E
erf
X1
XZ3
2
d
d
d
(>/22² d
2 XZ3 )
E
8² erf²
+ -7%3)
2
07z ert2
+ E
Z2 erf
X1
AZs
d P
X// d
erf2
+ AZ3)
2
2
erl' \/ 2 n + a2s) )
d
/ d
+ -2722
2
OpE
erl Av 2
x + X23)
+ +
erl²
(N/ z AZ5)
d \
+x2+) 273 )
+ +2272pB PF
Zz orf
2 K + AZs
orf' > 2
d
=(202 + +1D)E
ert2
1/ 2
+ AZ3
d * NZ3 )
+ +4x72
2 OpE
erf
A/
K + XE2
err' >//
2
/d.
lert" 1/ n N75o ]]
/d 2
/d (// S + NZ3a A2s))
+ +222740²
IF erl
\1222
+ 273
+ E (ent')2
2
-G & (
d
5 Rr {2n2)
ld 4x72 V
/d
erf' (AV/ 2 AZs) )
OpE erf
>23
2
222rtp²
d
erf' (AV + 2E3
-2X 1
F
erf
+ AZ3
1+ 2X222
2
30

---

Published as a conference paper at ICLR 2025
\122 /
2 2x= XZ5)) +273) ))
+E
(erl')2
1/ d
10 + tg2)c(1
/ 1279)
V/2
4272 V/2 / 0p
dd
1222p'n
A / + X2s
/d d
erf' ( // XZ3
(N//
erf
2
1 + 212y2
222
d
E
(err")²
Av 2 + >Z3
1 + +22272"
d 0² + (,2))(2)/) d
22
15, 1292)
d
2272p²
4272
\vd22k
er AV d/2m
+
0p
orf
1 212g2
VI+ 2X2y2
V(1+ 41272)(1 2X2,2)
11+22272
2A²
2
Avds
erl'
+ 2X2y2
V 1+ 422g2
VI+41893
by Lemma 18(ie) (vz).
Computation of Rz.
We have
Σ E ((((Tver((XXT)) -L 1)B (Xe((X]eer((AA) k))²)
R2
Thus, using previous calculations with -==
= 1, 0 0, and x = 0, we obtain
412
R₂ = (I. 1) (30,2))+ V + 1(1 + 2X2)
Perrt (0)
= (L 11 ((0,)2) + TV TX4 + I(L + 2X
8X2
Computation of Rg. lemma, one gets
Regarding the cross- cross-product terms, by independence of the (Xe)'s and Stein's
E [ ver((AX, x) ver((AX] b)) = X verf(AX, 1) k)) (x] ver((ux]] k)) c2p2, with C = (AX, b)) 22/ / (1 + 2X by Lemma 18(1). This leads to
412
Rs
[1 1 (L 1) 2)p²
Computation of Rg. We have, again by independence and Stein's lemma,
E (( v* Xx verffx(A Xx) crx((XX k))
= XiT 0² x1 veree((AX]T]) B) XT (eTe(((AXX]]
- vex] E. (xTTer(((x31))
B(X
EEXT veTT(e Cp
2Ap
E(XT vex
v(I + 2X2)
Note that, still using Stein's lemma,
-E(X] ver((XX] k))
31

---

Published as a conference paper at ICLR 2025
d -V 2 (Aat) (k*))Ter((XxX]]) B((XT
-(((X) w- / (kt)) (0 orr((AXT))
E(1 =-B(V
( k)) Cor
Cov
Cor(xT X7
crf()x;
V B k)) ACOn Xi v, xT 1 (ex(()XX))))))))))))))))){)))(
V a erf
+ +1+27222 (x /
AV d/22
1
orf'
A d//2
2
(1 1 2X2g?
VI+ 2X2r3
where we used that A k & \dd/2k + G with G ~ 220,) in combination with Lemma 18(1) (it). Thus
R4 4A(L 1)p
V erf Add/2x
XTP erf'
Avd/2k
+
V(++ + 2
V1+2
V1+ 272x2
11+2222
All in all. Putting everything together, we obtain
Rk,,o) c
Avd/2K
dd
A/d/2k
+-² - 22y-vert
- -277
2210
VI+ 2222222 erf
/1 + +22273
V1+ 22222
21p
orF"
A/22
+ +1200 + (2// Nr x2)22
d
1 1+ 2X292
V1+ 2A2g2
d
22
4172
\v
erf' Avd/2K
+V
0p 2X22
erf
2
2X222
+1+
V1 + 2X272
1
Avde
erf'
V
V 7 (1 2X292)
V1+ 41272
(L 1) [cc0,2)) TVI+4x* (1 2A2)
812
02
4X2
(L 2)p2
(1 2X
11p (1++2X7)
(Voort
A/d/2n
X-2p
A \vd/2x
orf'
11+22822
VI+ 2X2
(I++2822
This concludes the proof.
D TECHNICAL RESULTS
This section gathers formulas that are useful in the proofs, in particular regarding expectation of functions of Gaussian random variables involving erf.
Lemma 16. For > 0.
1 uu.
V
Proof.
The argument of the absolute value is non-positive for u 2 0, hence we need to show that
1
S(u) =
V1 + u
is non-positive for u >> 0. Just note that
f(0) = 0 and f'(u) (1 -1 [1+u)3+2 & 0.
32

---

Published as a conference paper at ICLR 2025
Recall that the erf function is defined on R as
2 e- dt.
/
erf
V 0
Lemma (Properties of the erf function), We have
(uf) =
e
π
4
erf" (u)
ue -ueerf (u),
v π
[erf(u)] < |u].
Proof. The first two statements clear by usual differentiation rules. Regarding the last statement,
since erf is an odd function, it is sufficient to prove the statement for u > 0. Moreover, erf is concave on [0, oo), so we get, for > 0,
[erf(u)] = (erf(u) - erf(0)] < erf (0)u= u,
V
which concludes the proof.
Lemma 18. Let G ~ NO), For t E R,
(i) E er (t + G)]
1 erf
V1+273
1+223//
(i)) [er((t + G)] crf(
V1+2,3
(i)) Ber G)] 1
V1+272
(iv) E [(ert)) 1+87 erf'
(-
v2t
VI+s7²
(v)
(1 (++G)) = 22 -2 B|err(t-
(t+ - (t G))
(vi) E [erf(t + G)rrf Gi]
V142z2 V orf
erf'
((14419((4477)
V1+2y=
This lemma reveals the importance of choosing the erf function as the component nonlinearity: there are closed-form formulas for the expectation of erf and its derivatives applied to Gaussian random variables. Extending the results to any nonlinear, bounded, increasing, equal to 0 at O, and differentiable activation function an interesting next step.
Proof.
(i) By Lemma 17,
V
2 (1+1)²
E
eerf (t + G)]
24 dg
Ty
V2
2y² 1 2n
2gt El &
dg
for c:=
V2
(xlet)2
Hot dg
V2
(1
e
(otee)²
dg
n²?
= =/ac
2
-t2 (
27
exp
1
/ 7( + 2)
33

---

Published as a conference paper at ICLR 2025
2
cxp
V V + 27
1+ 27²
(ii) By (1),
E [erf(t + G)] =
E [[erf'( (s + G)] ds
-00
2
s?
exp
ds
/T + 27²
+ 27
142y²
2
exp (-u²) ds
V
erl
2?
(iii) By Lemma 17, and following the same steps in (i),
E [erl"(L+ (1+ G]] - 2v2
/ (1 g)e ((+g)) g)e0 e 27 dg
VTQ
2/2 6 -o(-o) /(t+ g)c
(atat)
dg
TT
2v2 e
-P(1--
(1/TT0/ (VC(N-/T 9))
2/2c (1-o -c)
VTY
4
(1(1)))
VTr(1 + 22(2)
1 272
4t
8²
VT(1 + 2//23/22
exp
1
(iv) By Lemma 17,
1
E
G)]
(erf") +g)e
-, dg
V2Ry
2/2
.2 2(1-g)²
8 252 dg
773/2
2v2
4gt e
221 dg
with r²: r²:=
343/2
2r
2v2
(0tar202 2
2p 26 dg
3/3/2
2/2
-2((1-472)
/
dg
7//3/2
2/2
-2-(1-11")
(2m11
/2T11'
3/3
4
262
exp
(T 47²
T / + 4-2
(v) We use Lemma 17 and then Stein's lemma:
Elerf(t + G)erf"(( (t G*]
-2E [(t G)erf(t ()erf'( G)]
34

---

Published as a conference paper at ICLR 2025
=tB crf(( G)erf (t G)] 2E + G)]
- [er[(1 G)erl (1 G)]
27² 2 (TE [erf' (1 + G) [er[(L+ ()erl"
Reordering terms, this gives the desired equation.
(vi) We define the function
((t) - E [erf(t + G)erf
Then, using Lemma 18(.). have
ft) E |erf (t + G) + E [erf(t + Gt)rf
2t
E (1 G)
[er[(1 G)erl' G)]
1 272
272
E [(rrf (t. + G])
1 + 2?
2t
E [erf (t + G)) - 1 1 2-,2 f(t)
1 27
We solve this differential equation by the method of variation of parameters: we have
dt (f(t)e=/(4zz)) 1 2y2 [trf (1+22)))
We use Lemmas 17 and 18(iv):
d di (1())e* (00223)))))
err'
/2t
e 21442))))
vT
(1 27
47
/I+++4 4,2
4
2f (1/(111)) ee (11(+227)
(1+
1+ 472
4
t² (1 272) + 472)
/1 44 exp
T (1+
V (1 2,2) orf'
(1+22 + 472)
As the distribution of G is symmetric and erf is an odd function, we have that f(0) E [erf(G))rff' (G) = 0. Thus integrating the above derivative, we obtain
f(1)2/1(+277)
2
/ dserf'
V 242) 1 472
V(I+ (1++2?) 212) 4
2
t
erf
VI V
27²
21²)(1 152)
Using again Lemma 17, we obtain the claimed result:
t
t
f(t)
erf'
erf
VI+ 222
VI+ 27
(1+ 27 4
E EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
Our code is available
httpgs///gitauu.ccc////eeeeeerrr s singleellocation-regessss
We use the Transformers (Wolf et al., 2020) and scikit-leam (Pedregosa al., 2011) libraries for the experiment of Section 2, and JAX (Bradbury et al., 2018) for the experiment of Section 5. All experiments run in short time (less than one hour) standard laptop.
35

---

Published as a conference paper at ICLR 2025
E.1 EXPERIMENT SECTION 2 (NLP MOTIVATIONS)
Data generation. We use settheetccall-ggeerrtee data for this experiment. To create our train set, generate sentences according the patterns
The city is [SENTIMENT ADJ] [PRONOUN] [COLOR ADJ] [ANIMAL] is [ADV] [SENTIMENT ADJ].
and
The city is [SENTIMENT ADJ] [PRONOUN] (SENTIMENT ADJ] [ANIMAL] is [ADV] [COLOR ADJ]
where ADJ stands for adjective and ADV for adverb. Note that the difference between the two patterns is that the locations of the sentiment and of the color adjectives are swapped. Each element between brackets corresponds to a word, which can take a few different values that are chosen manually. For instance, some possible sentiment adjectives are nice, clean, cute, de delight fu,, mean, dirtyy, or nasty. A possible value for some words 0, meaning that we remove the word from the sentence, which creates more variety in sentence length. By doing the Cartesian product over the possible values of each word in brackets, we generate in this way large number of examples. Then, the label associated to each example depends solely the sentiment adjective appearing in the second sentence. For instance, the words n nice, cl ean, cute, or del i ght ful are associated to label +1, while the words de delight ful, mean, and dirty are associated to label -1.
We now explain how the test sets are generated. We generate four test sets in order to assess the robustness of the model various out-of-distribution changes. The baseline test set the same sentence patterns and the same sentiment adjectives in the training set, but other words in the example (e.g.. animals, adverb) are different. In particular, given sentence cannot appear both in the train set and in the test set. Then, we generate another test set by using sentiment adjectives that are not present in the training set. We emphasize that the sentiment adjective fully determines the label, using unseen adjectives at test time makes the task significantly harder. The third test set uses the same adjectives as in the train set, but another sentence pattern, namely
Hello, how are you? Good evening, [PRONOUN] [COLOR ADJ] [ANIMAL] is [ADV] [SENTIMENT ADJ].
Finally, the fourth test set combines a different sentence pattern and unseen adjectives. The size of the datasets is given in the table below. All datasets have the same number of +1 and -1 labels.
Name
Number of examples
Train set
15552
Test set
4608
Test w. OOD tokens
3072
Test w. OOD structure
144
Test w. OOD structure+tokens
96
Table Size of the generated datasets.
Model. We recall that there exists several families of Transformer architectures, which in particular are not all best suited for sequence classification. An appropriate family is called encoder-only Transformer, and foremost example BERT (Devlin et al., 2019). We refer to Phuong & Hutter (2022) for an introductory discussion of Transformer architectures and associated algorithms. Here, we use a pretrained BERT model from the Hugging Face Transformers library (Wolf al., 2020), with the default configuration, namely bert 12 layers, the tokens have dimension d = 768, and each attention layer has 12 heads. was pretrained by masked language modeling, namely some tokens in the input are hidden, and the model learns to predict the missing tokens. We refer to Devlin et al. (2019) for details on the architecture and pretraining procedure. We do not perform any fine-tuning on the model.
-asse-uunasee. The model has 110M parameters,
36

---

Published as a conference paper at ICLR 2025
Experiment design. Our experiment consists in performing logistic regression on embeddings of [CLS] tokens in the hidden layers of the pretrained BERT model, where we recall that the [CLS] token is special token added to the beginning of each input sequence. This is particular case of the so-called linear probing, which is a common technique in the field of L.MSs interpretability. More precisely, let E denote layer index between 0 and 12, where the index 0 corresponds the input to the model (after tokenization and embedding in R²). Then, for each value of é E {o..... 12]. we train a logistic regression classifier, where, for each example, the input to the classifier is the embedding of the [CLS] token at layer & (that is, d-dimensional vector), and the label is simply the label of the sentence as described above.
Results. For l 0 (blue bar in Figure lb), the embedding of [CLS] is fixed vector that does not depend on the rest of the sequence, so the classifier has a pure-chance accuracy of 50%. However, as soon f > O, thanks to the attention mechanism, the [CLS] token contains information about the sequence. We report in Figure lb the average accuracy over & E {l...., .....,122 12) for the train set (in orange) and the test sets (in green). We observe that the information contained in the [CLS] token is actually very rich, since logistic regression achieves a perfect accuracy of 100% in the train set. In other words, the data fed to the classifier is linearly separable. We emphasize that the size of the train set is significantly larger than the ambient dimension d, so it is far from trivial that this procedure would yield a linearly-separable dataset. Therefore, obtaining linearly-separable data demonstrates that the model constructs linear representations of the input inside the [CLS] token. Moving on to the test sets, the accuracy on the baseline test set is very good (95%), which suggests some generalization abilities of the model. The accuracy on the out-of-distribution test sets degrades (between 64% and 75%), but remains largely superior to pure-chance performance. This suggests that the internal representation built by the Transformer model is some extent universal, in the sense that it is robust to the specifics of the sentence structure and of the word choice.
E.2 EXPERIMENT OF SECTION 5 (GRADIENT DESCENT RECOVERS THE ORACLE PREDICTOR)
We begin by providing additional results before giving experimental details.
PGD with initialization on the sphere and constant inverse temperature schedule. As emphasized in Section 5, the dynamics of PGD with general initialization on (3M-1)² depend on the choice of the inverse temperature schedule X₁ The experiment presented in the main text in Figure 4a is for a decreasing schedule A 1/(1+ 10-4t). We report in Figure 7 results when taking constant inverse temperature. We observe distinct patterns depending on the value of this parameter. With a large inverse temperature (Figure 7a), we observe that the dynamics in (r. 1/) always escape the neighborhood of O. Furthermore, the direction v* is almost perfectly recovered, i.c., V/ However, the value of k* is only partially recovered: the dynamics stabilize around ~ 0.3. Moreover, the excess risk plateaus at high value, while the dynamics stay far away from the manifold M. In the case of smaller inverse temperature (Figure 7b), the situation is different. We observe that some initializations lead to a convergence to the point (x, v) = (0, 0), in which case the dynamics stay far from the manifold M. In other words, there is no recovery of k- and v*. Other initializations lead to perfect recovery of Ke and v*. In all cases, the final excess risk is low. Theoretical study of these observations is left for future work.
Implementation details. The implementation of the PGD algorithm (10) requires to compute the gradient of the risk. To this aim, we use the formula for the risk given by Proposition 15. Note that all quantities appearing in this expression have explicit derivatives. The only quantity for which this is not directly clear is the function Ç. which needs to be differentiated with respect to its first variable to compute the derivative of the risk with respect x. However, recall that ((t, 72) == E [erff (t + G)]. Then, by Lemma (18),
0,5(1, = 2E [(erf erl")) GG]
t
VI+ 2? erf
erf'
v11+ + 4 + 27
VI+ 2T
Evaluating Ç itself (and not its derivative) is not required to simulate the dynamics, but is useful for reporting the value of the risk. For this, we also use the formula above, and numerical quadrature
37

---

Published as a conference paper at ICLR 2025
Delane la dhe manifoll M
20p
Excose ttek
Altannent with
10-2
10-1
2
-0.
10-9
M
30000
100000
X,000
10000
50000
100000
Step
(a) For A
0.9.
6
1.00
1.0
Dilance to the musileed M
130000 16k
Aligomenn with
0.75
0.5
1001
01.00
0.0
10-1
0.225
-05
10-4
0
3.00
-00
10
10000
10000
20000
Step
Step
(b) For A = 0.1.
Figure 7: Dynamics of PGD from random initialization on (8d-1)?, for two iteration-independenr values of Ay. Left: Excess risk function of the number of steps. Middle left: Alignment |m| kk k-1 and | = to "*| with the oracle parameters. Middle right: Trajectories of k and v in few repetitions of the experiments. Each repetition corresponds to a color, the end point of each trajectory is in blue. Right: Distance to the invariant manifold M. In all plots except the middle right ones, the experiment repeated 30 times with independent random initializations, and either 95% percentile intervals are plotted or all the curves are plotted. Parameters are d = 400, I, 10. and 7 1122.
to compute the value of
((1,92- O,C(o, 22)s.
We report the figures the value of the excess risk, i.e., the risk (k, v) - 22 To compute the distance to the manifold M, recall that is defined by
M = {(ks E s gd-1 k_ o² 0, Tk k* 0, k 0}
For point (k, u) E gd-1 x gd-1 its distance to M is therefore computed as
d v)) / (k (kT (u (KT
Parameter values.
The following table summarizes the value of the parameters in our experiments.
Name
Figure 4a
Figure 4b
Figure
Figure
Figure
400 10
400
80 10
400
400
10
10
10
1/V2
1/v2
1/v2
1/v2
1/v2
X
1/(1 + 10
0.1
2/(1 + 10-4))
0.9
0.1
a Number of steps of repetitions Batch size
4 10
4- 1-003
-
10
4 10 -3
120k
20k
200k
120k
20k
30
30
30
30
30
0.1
Table 2: Parameter values for the experiments on recovery of the oracle predictor by gradient descent.
38

---

Published as a conference paper at ICLR 2025
E.3 ADDITIONAL EXPERIMENTS
Transformer layer. The most general formulation of the Transformer layer we consider writes, for X Rexd
X concat(r,
1 L K, I LNCK) Of (L+I)XP px((+1) (L+I)XP pxd
X x
softmax
(29)
T(X)
X ReLU + 16T) + +1b.
where
concal((; E R((--1xdd adds new token the beginning of the sequence by concate- nating r E Rd to X E Rxxd This token corresponds to the [CLS] or register token (see Section for discussion and references). In all our experiments, E F. is a vector with i.i.d. Gaussian entries of variance 1/d, which is not trained;
LN denotes layer normalization, softmax denotes row-wise softmax, and E RL+1 the vector filled with 1;
the parameters are Qh. Ka, Vn, On E RRdx, W1 Rdxm, by E I W2 Rixdd and b2 E IRd.
Experiment with single-head Transformer layer single-location regression. We first consider the case of single-head attention, where H = 1 and = d. For ease of notation, we drop the subscripts h in the parameters of the attention layer. We also set O be the identity matrix. We aim at training the Transformer layer on the single-location regression task, to check that our simplified model a good description of the Transformer layer. First note that the output of the Transformer layer (29) is matrix in R (L++)) +)x1)x while the target of single-location regression is scalar. Thus, consider only the first row of TXX), corresponding to the register token, and learn a linear projection of this row to R. In other words, the Transformer layer should learn store in the register token global information about the sequence, as described in Sections 2 and 3. Overall, letting @ E IRd, our risk writes
R(Q, K, V. Wi,,, W2, b2, = B(( -(1--T(x))00)
where (X, Y) are distributed according to the single-location task as described in Section 2. We train using single-pass stochastic gradient descent (meaning that fresh samples are used at each step), for 8, 000 steps with batch size of 128 and learning rate of 0.01. The experiment is repeated 20 times with independent random initializations, and 95% percentile intervals are plotted (but are not visible when the variance is too small). Parameters K, V. W1. W2 are initialized with Gaussian entries of variance 2/ (cl. + dou)). The bias terms initialized to 0, well as the query matrix Q, following standard recommendation in the literature on signal propagation in Transformer (Yang et al., 2021; He et al., 2023; He & Hofmann, 2024). The output weights 0 are initialized with Gaussian entries of variance 1 1/d², following the mean-field regime (Chizat et al., 2019). Parameters are L. = 10, d d= p 80, 200, x2 0.01, 0.5.
Results are given in Figure 8. We observe in Figure 8a that the Transformer layer is able to solve single-layer regression. Furthermore, as shown by Figure 8b, it does so by encoding in its weights the underlying structure of the problem, namely the oracle parameters k and v", as in our simplified model (see Section 5). More precisely, in the case of our model, we showed that the two parameters k, E (IR')² converge to (k* v²). To make appear the equivalent of l and in the more complex parametrization (29), we let h be the first left singular vector of and V(I + W, W2)0/ ( W, W2) We check numerically that the weight matrix QK is nearly rank-one after training' which validates taking & as the first singular vector of in the present experiment. It also validates considering vector-valued parameters in simplified model. The role of the vector & is to select the relevant token among all input tokens, while the vector e describes how successive transformation (the value matrix of the attention layer, the MLP with skip connection,
"The ratio between its first and second singular value of the order of 10 the end of training.
39

---

Published as a conference paper at ICLR 2025
1.00
Alignnent with oracle paramoters
Excess risk
0.75
0..0
0.25
Aligument with k Aligument with
0.00
2000
4000
6000
S000
2000
4000 Step
6000
8000
Step
(a) Excess test function the number of steps.
(b) Alignment between Transformer parameters and oracle parameters and v*. We plot |k k*) and e" function of the number of steps, where the first left singular vector of K. and == V(I + W/ (V Wi W2)
Figure 8: Training full Transformer layer on single-location regression. The Transformer layer solves the task, and encodes the structure of the problem in its weights.
and the final linear projection) map this token to the output of the model. We observe that these two vectors align perfectly with k" and v², This confirms that our simplified model is good description of how the Transformer layer solves single-location regression.
Multiple-location regression. A natural extension of single-location regression is when the output depends on > tokens instead of just one. This task, which we name multiple-locatee regression, can be written as
Σ x Joy +
(30)
where J(),...., (s) are latent discrete random variables on {I,..., ......,,)) LJ,, all different, and such that, conditionally on J(1)....,J{s),
N ( (2 Ia) for e (J(1......J(I))
XJ(M)
X
Experiment with simplified predictor on multiple-location regression. accordance with the above, natural extension of the model presented in the main text is the multi-head predictor
T P kn,
X
Σ erf (AXkk)
Xun.
(31)
The hope is that each head (kn, vn) should align with of the oracle directions (ki vf). As first attempt in investigating this question, we run stochastic PGD in a setup similar to the one presented in Figure We take
the pair (J()), (2)) takes uniform values among disjoint pairs of indices in {I...., L). The directions (kj vz)) and (kg vz)) are sampled independently uniformly on the sphere, such that (k:) ) og O. Parameter values are the same
Figure 5, except that the number of steps is set to 10"h the number of repetitions is set to 20, and the inverse temperature X, is constant after 2.5 1 steps. Results are given Figure We observe (Figure 9a) that our predictor able to solve the task. However, the recovery of oracle parameters is only partial, shown in Figures 9b and 9c: each head partially aligns with the oracle parameters, but the alignment not perfect. In other words, the model is not well able to separate the signal coming from the different X P(()) This calls for additional research in understanding how attention heads differentiate from each other in order attend to various signals, and why in setup the heads not well-separater sel--separated.
40

---

Published as a conference paper at ICLR 2025
10
0.8
h-1
l=2
Alignammi wish
0.8
b-1
h=2
Excoss risk
10P
0.4
0.2
0.2
10-4
25000
50000
75000
100000
25000
30000
75006
100000
25000
50000
75000
100000
Step
Step
Step
(a) Excess test risk as a function of (b) Alignment jk kg with the ora- (c) Alignment |v" vgll with the ora- the number steps.
cle parameters.
cle parameters.
Figure 9: Training the multi-head predictor (31) on the multiple-location regression task (30). The predictor is able to reach low-risk region. The recovery of oracle parameters by the predictor is partial. In the middle plot, for each repetition and each oracle parameter K we look at the end of training which head among ki and k2 is closer to kj. and report the alignment between ki and that head along training. If the alignment were perfect, this quantity would be close to 1. The same holds for the right plot.
Experiment with multi-head Transformer layer multiple-location regression. We train multi-head Transformer layer on the multiple-location regression task (30), taking = 2. The data is generated as in the previous experiment. Parameters are in the experiment for single-head Transformer, except the dimension = d/H = 40, the number of repetitions set to 10, and the learning rate set to 0.02. Mimicking the single-head experiment, we let k2. be the first left singular vector of Ks, and 1 = (7 W W Wa)0 + W; W2) We also check numerically that all weight matrices Q₁ K, are nearly rank-one after training. Results are reported in Figure 10. similar to the previous experiment: the excess risk is low at the end of training,
The conclusions but we observe partial recovery of the oracle parameters (although the recovery is somewhat better than with the simplified predictor, especially for kt)) This suggests that simplified predictor might be first good testbed to understand the training dynamics of multi-head Transformer for this task.
1.0
10P
Elooss sbl
Alignneet: with
0.5
0.5
10-1
Alggament
0.4
a=]
0.2
2-2
10-2
0.0
Step
Step
Step
(a) Excess test risk as a function of (b) Alignment Ik 1% with the ora- (c) Alignment o "f with the oracle the number of steps.
cle parameters.
parameters.
Figure 10: Training the multi-head Transformer layer (29) on the multiple-location regression task (30). The predictor is able to reach a low-risk region. The recovery of oracle parameters by the predictor is partial. For each & E {1, 2), we let K2 be the first left singular vector of Ki. and un = V,, (I + Wi W,W W2)0/[[Vs,] (I + Wi W2) oracle parameter ki, look at the end of training which head among k and k is closer to kj and report the alignment between bj and that head along training. the alignment were perfect, this quantity would be close to The same holds for the right plot.
[]. In the middle plot, for each repetition and each
F FURTHER DISCUSSION RELATED MODELS
We begin by discussing some related works training dynamics of Transformers (Jelassi et al., 2022; Nichani et al., 2024; Wang et al., 2024), to illustrate the originality of our task and predictor. Jelassi al. (2022) study how (vision) Transformers learn spatial patterns in the data by relying positional
41

---

Published as a conference paper at ICLR 2025
encodings. This differs significantly from our task that is invariant by token permutation. Further, in their model, the argument of softmax (i.e., a matrix E Rxxx) is directly a parameter of the model. This is radically different structure from the usual attention, and from setup, where the data appear in the nonlinearity o( X k). Next, Nichani et al. (2024) explore task involving fixed latent causal graph over the positions of the tokens. Here again, positional encodings play critical role in their analysis, whereas our task is invariant under permutations of the tokens. Moreover, in Nichani al. (2024), the output is expressed providing the necessary context for this computation. In our setup, however, the output depends on token whose position varies and must be identified within the context. Closer to our approach is the recent paper by Wang et al. (2024), which also incorporates notion of token-wise sparsity: the output computed as the average of small subset of tokens, where the subset is identified by comparing the positional encodings of each token with that of a reference token. We outline two key differences with our setting. First, we do not make use of a reference token, but instead learn the latent direction k* to identify the informative token. Second, in our setting, the tokens also encode an output projection direction v* on top of k*. In other words, our task involves learning a linear regression in addition identifying the relevant token, which not the case in Wang et al. (2024).
function of the last token, with the previous tokens
Besides, we also note that our task shares similarities with multi-index models (McCullagh & Nelder, 1983) and mixtures of linear regressions (De Veaux, 1989). However, our task (Ream)) has more structured nature, involving sequence-valued inputs and incorporating a single-location pattern.
Finally, could imagine multi-layer perceptron (MLP) designed specifically for single-location regression, where the weights have a diagonal structure with respect to the sequence index, namely
MLP(X
Σ t=1
We b1) + b2.
In such a setup, the first layer could learn the projections along k" and v², while the subsequent layer could learn to map these projections the ouput Y (in somewhat similar spirit to multi-indes models). However, this architecture is far from resembling those used in practice. If we do not assume diagonal structure and instead use traditional MLPs, the number of parameters must scale at least linearly with the sequence length, which is highly suboptimal and may lead to very slow training. This highlights the efficiency of attention layers, which perform single-location regression with fixed number of learnable parameters, independent of the input length. We leave rigorous study of the learning abilities of MLPs in single- regression for future work.
42


Reviews:
[{"review_id": "PFDbB70K0L", "reviewer": "Reviewer_iaWj", "summary": "This paper introduces a new theoretical approach to understand attention mechanisms by considering a task called single-location regression. The authors consider a simplified model that's related to self-attention layer and shows that the model can achieves asymptotic Bayes optimality, while linear regressors fail. The authors also use PGD to show that the non-convex loss function still converge.", "questions": "1. What's the function of input in Figure 1 (a)? It seems that the Y label just depends on the output and the input is not related to the sentimental label?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow. For example, the step-by-step illustration of how to connect the construction to the attention mechanism in Section 3 is helpful for understanding.\n2. The choice of model is good, like using [CLS] property which is observed in empirical study in (5) is natural and reasonable, and using erf as nonlinear weight function is also reasonable. In general, the theoretical result is solid.\n3. Also contain some empirical result for showing the convergence of constructed model", "weaknesses": "1. The setting is restricted to single position token, although it focus on the sparse attention settings and it's already difficult to analyze, it's still far from the real-world case. Besides, the authors haven't done experiments on real-world experiments (like sentimental tasks as shown in Figure 1) to support some claims in the paper, this may kind of reduce the impact of the theoretical analysis. But in general it's already good as a theoretical-centric paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YQNRGbM9PS", "reviewer": "Reviewer_9Wat", "summary": "This paper studies the ability of attention mechanisms to deal with token-wise sparsity and internal linear representations. In order to demonstrate the capability, the paper proposes a simplified version of a self-attention layer to solve the single-location regression task, showing an asymptotic Bayes optimality and analyzing training dynamics.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written and easy to follow. A novel task called \"single-location regression task\" is introduced to satisfy the sparsity of the token and model real-world tasks to some extent. Despite the non-convexity and non-linearity, the paper is able to analyze the training dynamics and show the asymptotic Bayes optimality.", "weaknesses": "1. The proposed task may be over-simplified and lack generality. For instance, it assumes that the tokens other than $X_{J_0}$ have zero mean and only one token contains information.\n2. The paper shows the connection to a single self-attention layer by using the assumption that $p=1$. Although the low-rank property may come true after the training process, it is so strong to make this assumption directly.\n3. It is uncommon to use the function $erf$ to replace the softmax function. To demonstrate the feasibility of this simplification, more explanations or theoretical backups should be provided.\n4. The paper shows the asymptotic results of $k_t, v_t$, while a non-asymptotic result is needed to investigate the convergence rates of these two parameters.\n5. The initialization is limited to the specific manifold, which is better to extend to a more general one.\n6. The current experiments only validate the theoretical results on synthetic datasets. It is recommended that the authors consider adding some experiments on real datasets to test the effects.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qKgMUCNKYy", "reviewer": "Reviewer_1VPX", "summary": "This paper introduces a new task, the single location regression task, showing it is solvable by a predictor resembling an attention layer, when a linear predictor fails. The result is theoretically well grounded and of good significance given the limited theory on attention layers and their striking efficacy. This task relates to key famous NLP problems known in the existing literature, serving as a good testbed for studying Transformers through a theoretical lens. The statistical performance of this attention-like optimal predictor is shown to exceed the one of an optimal linear predictor, under minimal assumptions. (projected) Gradient descent, is demonstrated to reach the optimal solution, wiith experiments provided to support the theory.", "questions": "- Could you point out (and explain) the steps in your proofs that would break if softmax were considered instead of erf. My guess is that preserving the independence  is key (via elementwise application) but could sigmoid or any other nonlinear bounded increasing and differential elementwise activation work as well? It might be helpful to mention early on that the analysis could extend to a broader class of activations if that’s the case (modulo adjustments of the formulae in theorem 1).  \n- Linear predictors are shown to fail at solving the task and the comparison to them is much appreciated. Do you have a sense of how non linear predictors would in turn perform, thus this could show how attention layers are to be preferred to fully-connected layers for instance in such contexts (an analogue of proposition 3 in this case may not be true and probably non trivial to show, but in any case, interesting to discuss or to investigate). \n- Are the tokens assumed to be independent (line 74) or independent conditionally on J_0 (line 85)? I did not read the proofs so I couldn’t figure it out  myself but would be good to clarify. More generally, could you explain at which points of your proofs the independence is needed to help readers understand how this assumption could be relaxed for future works concerned with more realistic scenarios.\n- Minor: line 1832 \"Our code is available at [XXX]\" in the appendix. \n- Extra minor: line 469: a space is missing between \"Figure 4a\" and \"(right)\"", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is extremely well written and very easy and pleasant to read. \n- The analysis is theoretically strong and rigorous\n- More generally, this work promotes an approach that is worth being acknowledged and valued: looking into a simpler problem than the ones practitioners can face, yet relevant, and solve it completely and rigorously.\n- Additionally, the problem is well connected to practical concerns, with the authors made a convincing case for the significance of their analysis.", "weaknesses": "By decreasing order of importance\n\n- The present work’s approach is not so well connected to the existing literature in line 103: \"note that our task shares similarities with single-index models (McCulllagh & Nelder, 1983) and mixtures of linear regressions (De Veaux, 1989)\". I see the differences between those works and the present one being hightlighted in the following sentence , but the exact nature of these similarities is not clear. Could this connection be elaborated?\n- Minor: In the caption of figure 2, it may be helpful to add a note about the size of the squares, that presumably indicates the level of alignment.\n- line 234: \"We emphasize that empirically this simplification of softmax using a component-wise nonlinearity has been shown not to degrade performance\". The cited paper Wortsman et al, 2023 indeed indicates such behaviour but their experiments include a normalisation. \n- Broadly speaking, the question of whether softmax is needed in attention layers remains an open and unresolved one in the commuinity. To avoid overstating the case, consider rephrasing to reflect this ongoing debate. Note that I don’t think having used erf in your analysis diminushes the value of your work in any way.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8DVw7trcoj", "reviewer": "Reviewer_gchu", "summary": "This work examines the transformer mechanism for solving the single-location regression task within a simplified non-linear self-attention setting. It provides a Bayes optimality guarantee for the oracle network predictor and demonstrates how the training dynamics allow the network to converge to this solution using projected gradient descent.", "questions": "1. What is the specific convergence rate in Theorem 5? Since this forms a key part of the contribution in characterizing the training dynamics, a more detailed presentation in the main paper would be beneficial.\n\n2. In line 334, why is the entire sum on the order of $\\Theta(\\lambda \\sqrt{L})$?\n\n3. The PGD analysis relies on a relatively strong assumption that the initialization lies on the manifold. Could the authors discuss any possibilities for generalizing the current analysis to accommodate a broader range of initializations?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The mathematical analysis from both statistical and optimization perspectives is insightful, offering a comprehensive theoretical guarantee.\n\n2. The paper is well-written and easy to follow.\n\n3. Experiements are provided to support the theoretical findings.", "weaknesses": "My main concern lies in the motivation for studying this specific setting, characterized by token-wise sparsity and internal linear representations.  While these features are relevant to real-world NLP scenarios, the essential mechanism for transformers to succeed in this setting is, first, to identify the underlying structure (the latent sparse pattern in this work) and then to perform a task-specific operation on this structure (linear transformations here). This intrinsic mechanism has been extensively explored across various settings (e.g., [1-3]), even with more realistic softmax attention. Thus, the technical significance of studying a simplified attention model in such a specific setting remains unclear. Could the authors elaborate on the specific technical difficulties encountered in this setting? What are the key technical challenges in extending the analysis to more realistic softmax attention?\n\n\n\n[1] How Transformers Learn Causal Structure with Gradient Descent. Nichani et al., 2024\n\n[2] Vision Transformers provably learn spatial structure. Jelassi et al, 2022\n\n[3] Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot. Wang et al., 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)