Understanding Softmax Attention Layers:\\ Exact Mean-Field Analysis on a Toy Problem
Authors: Elvis Dohmatob
Venue: NeurIPS 2025 poster ()
Keywords: softmax self-attention; theory of transformers; statistical physics; optimization dynamics
Primary area: theory
TL;DR: We prove that a softmax self-attention layer trained via GD can solved the so-called single-locator regression problem

Abstract: Self-attention has emerged as a fundamental component driving the success of modern transformer architectures, which power large language models and various applications. However, a theoretical understanding of how such models actually work is still under active development. The recent work of (Marion et al., 2025) introduced the so-called "single-location regression" problem, which can provably be solved by a simplified self-attention layer but not by linear models, thereby demonstrating a striking functional separation. A rigorous analysis of self-attention with softmax for this problem is challenging due to the coupled nature of the model. In the present work, we use ideas from the classical random energy model in statistical physics to analyze softmax self-attention on the single-location problem. Our analysis yields exact analytic expressions for the population risk in terms of the overlaps between the learned model parameters and those of an oracle. Moreover, we derive a detailed description of the gradient descent dynamics for these overlaps and prove that, under broad conditions, the dynamics converge to the unique oracle attractor. Our work not only advances our understanding of self-attention but also provides key theoretical ideas that are likely to find use in further analyses of even more complex transformer architectures.

Understanding Softmax Attention Layers: Exact Mean-Field Analysis on a Toy Problem
Elvis Dohmatob 1,2.3
:Concordia University MMila-Quebec Al Institute Meta elvis dohmattob@concordia.coo ca
Abstract
Self-attention has emerged as a fundamental component driving the success of modern transformer architectures which power large language models (ChatGPT, Llama, etc.) and various other types of systems. However, a theoretical understand- ing of how such models actually work is still under active development. The recent work of (Marion et al., 2025) introduced the so-called "single-location regression" problem, which can provably be solved by simplified self-attention layer but not by linear models, thereby demonstrating a striking functional separation. A rigorous analysis of self-attention with softmax for this problem is challenging due to the coupled nature of the model. In the present work, we use ideas from the clas- sical random energy model in statistical physics to analyze softmax self-attention on the single-location problem. Our analysis yields exact analytic expressions for the population risk in terms of the overlaps between the learned model parameters and those of an oracle. Moreover, we derive detailed description of the gradient descent dynamics for these overlaps and prove that, under broad conditions, the dynamics converge to the unique oracle attractor. Our work not only advances the understanding of self-attention but also provides key thcoretical ideas that are likely to find use in further analyses of even more complex transformer architectures.
1 Introduction
Understanding the theoretical foundations of transformer layers [Bahdanau et al., 2015] (also see (Schmidhuber, 1992]), particularly self-attention (SA) [Vaswani et al., 2017], remains critical and largely unresolved challenge in machine learning. SA stands as a comnerstnee of modern large language models (ChatGPT, Llama, Gemini, Mistral, Deepseek, Cluade, etc.), driving their unprecedented success across diverse tasks. Unlike classical layers such as feedforward, convolutional, or recurrent architectures, SA enables capabilities that these traditional mechanisms cannot replicate, including efficient capture of long-range dependencies and context-aware representations in single forward pass. While classical layers benefit from decades of rigorous analysis and well-established theory, the inner workings of SA remain poorly understood, with only sparse body of research attempting to unravel its operational principles. Moreover, the training dynamics of SA, or how evolves during optimization, present additional complexities that are crucial for developing more robust, efficient, and interpretable models. A deeper theoretical understanding of SA and its learning behavior is therefore essential not only to explain the empirical successes of transformers but also to unlock their full potential and inform future advancements in neural network design.
Recently, Marion et al. [2025] have considered simplified version of SA and showed that unlike linear models, it can solve a so-called single-location regression problem: the teacher model f. sees an incoming input X = (#1,..., made of the embedding vectors 3y in Rd for sequence of tokens (c.g words), and must correctly locate secret block zt. at a random secret index l. [L]. Refer to Figure This token index is special in that except for additive Gaussian noise,
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Token Features X (#I,..., E Ie
Jf.
y o= 0.)
Figure I: The Single-Location Regression Problem. Each ze E Rd corresponds a token embedding. The embedding for the secret token at index l. contains signal aligned with hidden vector u, The embeddings for tokens at all other indices are pure Gaussian noise. The label y is computed using only this secret token, and all other tokens are ignored. Optionally, we also introduce link function a to capture non-linear problems (a was taken to be the identity function in [Marion et al., 2025]).
it is aligned with an unknown unit-vector u. which can be thought of as encoding the position of lu; all the other blocks TEgt, are Gaussian noise, Once this block is identified, the model must then approximate the output given by = (xr) == x._ v. where v, is another unit-vector perpendicular to u,. Thus, presumably. the model must somehow figure out the directions u, and u. in order to solve this problem. This problem captures some aspects of the sparse parity problem, except it is considerably simpler; for example, it does suffer from the the well-known exponential query complexity lower-bound which characterizes the latter problem. Marion al. [2025] considered a simplified transformer model
Pointwise Self-Attention: g(X;
Σ peaf v, with Pe 0(xor u), 2=1
(1)
for some point-wiss activation function @ and inverse temperature parameter A. The parameters of this student model are pair of unit-vectors (u, v).
Going beyond [Marion et al., 2025] which considered pointwise/separable SA (1), we consider softmax SA which better reflects what actually used in transformers. Our student model is then
Az M
SoftMax Self-Attention: f(X; u. v) Σ (x) v), with P Ete Am E=1
In our extension, we also include possibly nonlinear link function (known the learner) used to compute the labels using the the embedding of of a token at secret index a azx v.) (instead of J/ v.) for the true labels, and a (x v) (instead of ai in (1)) for the values in our attention model (2). This link function should not be confused with the softmax layer which is always present in the separate extensions of (1).
setting we consider in our work. These
Importantly, our theoretical analysis is valid for all L up to a limit which is super-pol poper-polynomill in the dimension d, i.e., log L (d). In contrast, the analysis of Marion al. [2025] is only limited to sub-linear number of blocks = o(d).
Main Contributions. Our contributions can be summarized as follows:
Exact Analytic Formulae for the Risk. In appropriate asymptotic scaling regime for d and L (refer to (9)), we obtain precise analytic expressions for the population risk of our softmax self-attention model (2) (Proposition 1 and Proposition 2). Our approach uses ideas from the elassical analysis of the Random Energy Model (REM) [Derrida, 1981, Lucibello and Mézard, 2024] to handle the softmax, which maps the Gibbs distribution induced by the disorder in corresponding REM. In order to incorporate the nonlinearity a, we extend a recent result of [Zavatone-Veee Veth and Pehlevan, 2025]. See Proposition 12, Proposition 13, and their corollaries (Appendix C).
Optimization Dynamics. We study the optimization landscape of projected gradient-descent on the population risk relative a manifold corresponding to spherical constraints the model parameters. We classify the stationary points and show that for a large variety of link functions, the induced dynamics always has the optimal model parameters as attractor (Propositions 10, and 9).
No Need for Special Initialization. Proposition we focus on the linear link function (t) = and remove critical initialization assumption which made in [Marion et al., 2025]. Indeed some

---

of the main results about optimization in the aforementioned paper assumed that initialization be selected from a peculiar manifold, which effectively assumes some knowledge of the teacher oracle parameters (u,; v.), which is not feasible feasible requirement in practice.
2 Related Work
The self-attention mechanism, introduced by Vaswani et al. [2017] drives much of modern deep learning, notably in natural language processing and computer vision. By adeptly capturing long- range dependencies, it eclipsed recurrent neural networks many tasks. Empirical breakthroughs like BERT [Devlin et al., 2019] and Vision Transformers [Dosovitskiy et al., 2021] expanded its reach, sparking theoretical exploration of its mechanics.
sparse literature of studies have unpacked aspects self-attention's 's expressive power, complexity, and optimization. Schlag et al. [2021] crafted a minimal attention model without sans positional encoding or normalization, to study its core behavior, while Cui et [2024] explored solvable dot-product attention model, identifying phase transition between positional and semantic learning driven by positional encoding. Another common simplification of attention, is the so-called linear attention mechanism, where the softmax layer is removed altogether. Such models have been intensively studied in the setting of in-context learning to derive theoretical insights on the internal workings of transformers [Ahn et al., 2023, Von Oswald al., 2023, Zhang al., 2024, Lu et al., 2024].
The work which is most related to ours [Marion et al., 2025] which considers simplified attention model (1) with the softmax layer replaced by a pointwise function, and study the generalization profile and the optimization landscape induced by such models. In contrast, we consider the more difficult (and practically relevant) case of softmax attention and provide a complete theoretical picture.
Let us note that Dong al. [2021] showed that except if MLP layers and skip-connections also used, single-layer SA transformers have a strong inductive bias to converge to rank-1 matrices, a simplicity bias which would limit the applicability complex problems. Fortunately, the single- location problem studied here and in the reference work [Marion et al., 2025] is just simple enough to be captured by single-layer SA transformer without need of MLP layers or skip-connections.
Preliminaries
3.1 Problem Setup
Data Distribution. Let d and L be positive integers, > 0 and E (0, 1) be real numbers, and consider the following data distribution: P on [T.] x RRddd x TR given by (fx, X, ~ P iff
(Secret Token Index) l. ~ Unii((L)),
(3)
(Secret Token Features) zt, ~ NN(cvdu, La), with V1-22
(4)
(Dummy Features) (x))s sdd
N Ia).
(5)
(All Features) X (II,,.., E R
(6)
(Label Noise) & ~~ ~0,00, e), independent of X,
(7)
(Label) y (z] +
(8)
We will denote the marginal distribution of the features X by Px. Here, U., E S4-1, where S the (d - diddieensional unit-sphere in Rd, while o R - IR is a link function. The unit-vectors u, and are fixed but unknown to the learner. The constant ratio (SNR) of the problem. As argued in [Marion et al., 2025] in the limit y - 0+ the dummy features vanish, and the problem reduces the usual Gaussian linear signal noise model, which is solvable via simple linear regression with large enough sample from P. The situation graphically illustrated in Figure 1.
(0, 1) controls the signal-to-noiss.
The Single-Locator Regression Problem and Softmax Self-Attention. As already mentioned in the introduction, the task is to approximate the true label function +> ( v.). Of course, neither the sample-dependent index b,, nor the unit-vectons and v. are known to the learner. For pair
3

---

of unit-vectors (2, u) E -/-1. consider the following simplified softmax self-attention (SA) model ff introduced in (2). The inverse-tempee temperature À > 0 controls the sharpness of f(x; u, v). We will scaling À = 3vd, with fixed 3 > 0. Within this class of models, the one with parameters (u.; v.) will be referred to as the teacher / oracle model and denoted f-. We shall see in Proposition 2 that this oracle model can indeed approximate the true label function the feature noise parameter 7 is not too large. On the other hand, if the learnable parameter vector is close to the oracle version ue, then the softmax will concentrate its mass around the right index L., allowing the model to select the value a (z, v) from all the other values a (æz v). Then, if v is itself close to the oracle version Us, the output of the model will approximate the true labels o (x, v.). Thus, goal is to learn the oracle parameter U.),
impose the following inverse-temperann
We work in the following asymptotic regime where d and L are large but L is exponential in d, i.e d. oc, log(L)/d a E (0, oo).
In our theory, taking the limit a + 0+ will correspond to the extreme case where L is at most sub-exponential in d (e.g., polynomial, even constant as in [Marion et al., 2025]).
Risk / Test-Error. We will be interested in the average L2-squared error of the parametrized model defined in (2), relative to the data distribution P. i.
R(u, v) == E ((X u,
u, v) y) -
(10)
which measure how well the model solves the single-location task. The offset e2 corresponds the irreducible error of the Bayes model Bages X + Ely X). due to the label noise.
x(x):
o(x): =x²
oix: max(x
2.0
Alignment v. Risk Rtu,
1.5
1.0
0.0
200
400
600
800
1000
200
400
600
800
1000
0
200
400
600
800
1000
Iterations
Iterations
Iterations
Figure 2: Illustrating the optimization dynamics for various for different choices of link function a. For this experiment, we use input- unputdiimension d = 100, L = 20 blocks, (normalized) inverse- temperature 3 = 1, 7 = 1/v2, and label-noise level c = 0.1. The Riemannian gradient-descent 0.01. The population risk R is replaced by empirical scheme is used (29) with step-size version R 11 (J=,, u. r) y;)² where (X1, v1), )....)....:: (X.,, yn) is iid sample of size n = 1000 from the data distribution P. The final risk R(ug, 2/. shown is evaluated on an independent test sample of size 10000. Broken lines correspond
our theoretical predictions (Proposition 1). Notice the perfect agreement between experiment and our theory. The oscillations in the curves for the ReLU (3rd sub-plot) are reminiscent of the non-smoothness this link function.
4
A Mean-Field Approximation
4.1 Main Idea: Equivalence to the Random Energy Model
Fix the parameters (u, v) E Sh_1 of a softmax self-attention model defined in (2). For random data point (L., X, 3) P. one can express the output of convex combination like so
s(X;uvv= Σ malaTT) = (x + (1 pr. (2r.)]((X;u v), with
(11)
f((;uu o(oT,o), f2(X;u,)) = Σ
0 (xT v),
(12)
Z-r.
tft.
p::=
Yl E [L],
Z-t, = Σ
@Vdao us Z VVdeT. + - (13)
ttf.

---

We have isolated the contribution of the secret token index L, E [L]. The other terms (captured in f2) is linked to a d-dimensional random energy model (REM) [Derrida, 1981] with = end configurations with random energy levels E, = vdz, u drawn iid from (0, d). It turns out that the mean-field description for such a system completely captured follows:
log Z_e.
+ 1 =
(14)
3d
3
Ô log -_e.
+
(15)
d
03
03
where Ø = 0( 3) = a + fr--B-ss s(())
(16)
r -=r(a, = min(B,Boric) = [3,
if B < Berits
(17)
(Berits
if B 2 Berits
Berse =- V2a, s(r) := 12/2.
(18)
Condensation. The value Terit 1/ E (0, oo) is the so-called condensation temperature. For temperatures T 1/B less than this value (i.e for 3 > Perit). the system freezes; only handful of equilibrium configurations carry the maximum energy, which is of order E' ;= maX Ee ~ v.v Finally, in the limit a + 0+ in which L now at most sub-exponentially large in the parameters Serit and r vanish and the the system becomes permanently frozen / condensed, for all values of B. Our main results for optimization (Section 5) will focus on this regime.
Remark 1. The above is mean-field description in the sense that it ignores corrections of order 1 1/d which could cause statistical fluctuations. Such corrections would require more advanced treatment of the REM, as done in [Bovier and Kurkova, ]. example.
4.2 Mean-Field Representation of Models and Their Risk
For theoretical analysis, the softmax in the parametrized model class (2) troublesome because all the blocks are interacting (via the partition function 2). This is precisely the reason Marion et al. [2025] decided to forgo pointwise / de-correelttd self-attention (I) instead.
Overlaps. It turns out that in the high-dimensional limit (9), (2) admits a simple description in terms of the "thermodynamie" quantities r, Φ. to introduced earlier in (17) and (14), and the overlaps Ju, U, C, 7: P E [-1, [ 1] defined by
:: u U.. v:= U.. 17 := u v.: Ç: --uuu.. Un, p:= p:=uTa TTv
(19)
In particular, μ and v capture the alignment between self-attention transformer f given in (2) with parameters (u, v) and the oracle parameters (m.,v.).
Now, the μ overlap parameter will enter the picture via the following function p: [-1, 1] + (0, 1] which will play a crucial role in our analysis
p(u) 1 e- (cu-v
(20)
This is effectively the high-dimensional limit of the probability that f correctly locates the secret Px.
index l, E L is random data point X (#1,...,,L))
Remark 2. Notice a sharp phase-transitiee for < bb/c, the probability p(u) is exponentially close to zero (because the input dimension is large), while for μ. w/c, is exponentially close to 1.
Auxiliary Functions. Given any p, E [-,, define an auxiliary function 0.,5.0
-R - RR by
(t). == olc + 7t) (prvd).
(21)
In particular, we write (t) == (t) = a (?t). The so-called "dual kernel" associated with d denoted o. 1- IR, will play crucial role our results, and is defined by
ā (v) := (G)),,
(22)
where (G1, G2) ~ Nr, bi-variate Gaussian with unit variance and correlation coefficient v. For example, for the linear link function (t) := we get (t)
(v) = 22v.
5

---

We shall also need the auxiliary functions H1, H2 : [-1, + R defined by H2(u, p)
Hi (v, (. Biosss.
(23)
for (G1, G2) ~ N2. Note the implicit dependence of the H functions on feature noise level y and the thermodynamic parameter introduced in (17).
Also define simplified versions h1, h2 [-1,1]] + R of the Hu's corresponding to setting < = 0, i.
(v, 5) = (.
(24)
Remark 3. Note that H2(u, 0, = Hi 0, p) and Hx 0) (v) k 1, Also note that ifr = 0, then Hx (u, P) = (u, {), and the H1 (u S. p) doesn't 't vary with p.
Our mean-field analysis will need the following technical condition on the link function o. Condition 1. (A) The link function o is square integrable w.r.f (N 1), and (B) a is positive- homogeneous, meaning that there exists m >> 0 such that o(u)) = u² (() for all > 0, t E R.
Examples include: linear link function a (t) == t: sign link function (tt) := sign(t); ReLU (t) ==
max(t, 0); quadratic link function (t) := t²: power link function (t) := t (with m > 0); etc.
The following one of main results. Proposition 1. Suppose Condition 1 prevails. Then, for any model parameters (u, v) E S and for a random data point (e,,X, ~ P, it holds in the limit (9) that
f(X; u, v) = (a v) + (1 - (pr(pp
(25)
where p = p(u) is as defined in (20), and μ. v, TB. C, and are as defined in (19). Furthermore, the population risk of the model is given by R(u, v) = R(y, v, (,P), where - 2pH22v,S.P) pH
R(p, ((,,v,,,1 : :)
(1,
(26)
Proof Sketch. The backbone of the proof (provided in the Corollary 3 of the Appendix) uses ideas from the classical analysis of the REM [Derrida, 1981, Lucibello and Mézard, 2024] to establish (25). In particular, extend recent result of Zavatone- Veth and Pehlevan [2025] for our purposes. Once formula (25) is established, the formula for the risk is matter of Gaussian integration.
In order to apply Proposition 1, one needs to compute the auxiliary functions H1 and H, defined in (23). This is done in Proposition 11 of Appendix B. The said formulae can then be readily exploited to get explicit expressions for the surrogate risk R appearing in Proposition 1.
Linear Link Function. Consider the special case of the identity link function (tt) t. In this case, thanks again to Proposition 11, we know that H₂ functions appearing in Proposition are:
3. (v) -"²,, H1 H (p = ā. (v) + (a (a-b)²,
H2(1,,,,) = (=(,)) (a 6)6, (27)
with =cVV and b:= prvdd. We obtain the following corollary.
Corollary 1. Under the conditions of Proposition 1, and in the limit (9), it holds that f(X;u, u, v) =
pad. + (1 p)pr Moreover, the population risk is given by R(u, = R(y, v, S,P), where R(p, v, S. (peÇ + (1-p))r)2a 2p -2p 1).
Geometric Insight from Corollary 1. For any fixed u E S4-1, the restriction of the population risk + R(u, v) on the set {u E Sa-1 | u v =0= 0} behaves like quadratic well I ~ (pv" u. ). y² ppu ux 11? with deepest point u()) given by
v(u) (n2 (p) La pdu. u.) -u.pp x U.
In the above calculation, we have used the Sherman- Further, if cu" T u, 2 (1 & (i.e u is within a spherical cap around u.), then
formula and the fact that u, O.
p p= 1/(1 + (ca T u. -))d 1/(1 + - for large d,
and so v(u) = v./p ~ v.. It is then easy to see that any such (u, ru)) minimizes the population risk R. Thus, we only need to get an ((1) alignment of the parameter with the oracle counterpart u..

---

4.3 Bayes-optimality of the Oracle Model: A Sharp Taase-Transition
To be sure we are actually in business, we must ensure that the best possible choice of the parameters
(u, u) Sa-1 for the parametrized family (2), namely the oracle parameters (u,, Us),, ), does indeed achieve zero risk R
As the next result shows, turns out that Condition 2 is necessary and
sufficient condition for this purpose.
Proposition 2. The oracle parameter V.) are indeed optimal for the risk functional R restricted the parametrized family (2), ie R(u,, (nf(w.)) R(u, v).
Moreover, we have the following sharp phase-transition (recall that c:= 11-12) 72))
(A) Ifc (1 + 5ys for some 8 (0,1), then in the limit (9), it holds that R(u,, v.) - O. That is, the oracle model with parameters (Mx, v.) is Bayes-optimal.
(B) If c S (1 8)bb. then in the limit (9), it holds that R(u,, = 211), more precisely, R(u.,Va) (1) 0. That is, learning is not possible!
Proof. Indeed, by definition (u,, v.) is optimal for the population risk functional R. Moreover, thanks to Proposition I, we have R(u.; v.) = R(1, 1, O, 0), with
= -12,0),,,, -2000000000000000--0000000000000
H1 (1, (
2p (10, = (p- 1)2o, (1);
where p = p(1) = 1/(1 + (c Finally, since ã, (1) > 0, observe that RHS in the above
display vanishes for large d iff tb and the result follows.
Motivated by the above result, we shall need the following condition in the sequel.
Condition 2 (Realizabilitt). 2 (1 + (2(1)), i.e 2 +b + 8) for some > O, where v is as defined in (14) and we recall that /1 72. For most of our analysis, weaker condition e > v.
work under the
For example, the condition is always satisfied for 0 corresponding to sequence length L which cub-exponential (e.g., polynomial) in the dimension d, because we have yi - 0+ in this case and v (1 8)c trivially, for any 8 and sufficiently large d.
5 Learning and Optimization
From Proposition 2, we know that under Condition 2, the oracle model parameters (u, v.) solve the single-locator problem. But, can numerical optimization actually find it? As mentioned in the introduction, unlike the case of pointwise attention considered in [Marion et al., 2025], the analysis is complicated by the softmax (2) which characterizes genuine attention layers [Vaswani et al., 2017] used in practice.
5.1 Preliminaries
Projected Gradient- SeaiiieetDeeeeentGrriieeet--ee
For simplicity analysis, we shall consider learner who has infinite samples, and therefore can directly optimize the population risk R over the parametrized family (2). We shall study the dynamics of the following projected gradient descent
(PGD) scheme
uy+1 Us sP.l V, R(u, v) ())((.))) (afl,,))m(mprpp) (u).)))
ūk+1 == - sP'.. R.R(u, (((0))))) (28)
Uk+1 = u++1/ [[u++1)) I|l, Uk+1 = Ux+1/ 1/
(29)
where 0 is the step-size and P,l Ia uu." is the orthogonal projector onto the tangent space
to unit-sphere Sa-1 at the point u. The projection step on the second line ensures that the iterates (ux+1, 1k+1) remain on the the manifold S For sufficiently small step size s, the dynamies (29)
are captured by projected gradient- flow (PGF) on the manifold $3_,
given by
ú(t) = PPl R(u, (,)) ((.)))((0)))000))
é(t) - -P R.R(u, (((()))))()))))))))
(30)

---

The dynamics (29) and (30) induce a corresponding evolution equation for the order parameters (p,1,, v4 17 (. which will be our main object of study.
Some of our results will concern the following submanifold M C $
M = {(u, v) E $2-1 v- u-uu. u, u1, -0,=0
(31)
introduced by [Marion et al., 2025]. It is clear that M contains the oracle parameters (u., U.). On this submanifold the dynamics reduce to just μ and v, a two-dimensional system.
Condition (Sub-exponential block length). In this section, we shall shall work in the frozen regime r O, corresponding the case where Borit - 0+, that is the number of blocks L is sub-exponential in the input-dimension d, ie log L o(d). This condition is very mild, and subsumes the regime L o(d) considered MMaiion al., 2025] as special case.
Since p u only enters the picture via the product pr, in this regime the effect p is permanently lost; this variable disappears from the picture under the above condition.
5.2 Analysis of Optimization Dynamics: Arbitrary Initialization
Proposition 3. For any u, E Sa-1, define μ == u.u,, = IT Vo. == u Vos S w Uo.
= u v. If p = 0, then in the limit (9), we have the following:
w T V(u, v) = Ti (M,, v, Syw Un, w T V,,RR, v) = Ta U, S)w u. + To(u, u, Aww Vos (32)
uniformly on w E Sa-1, where the functions Ti, T2, T2 [-1, + IR are as given in Appendix E.
Thus, (asymptotically) the gradients of the risk R are trapped in the 2-dimensional subspace of R spanned by the oracle parameters (u., v.).
Corollary 2. For sufficiently small step-size in the iteration scheme (29), the equations of motion for overlaps (u.V, m. () are given by the following 4 d-dimensional gradient-flow
ix(t) Fi (1(()))
v(() ((),
(33)
(t)) =
((() = (())(0),(0))(()))
(34)
Stronger Results in the Case of Linear Link Function. Even the greatly simplified case of point-wise non-softmax attention (1) considered in the reference work [Marion et al., 2025), their analysis was further restricted to the setting where the optimization scheme (29) is initialized on the sub-manifold M. This is problematic because by definition, choosing point on M presupposes knowledge of the oracle parameter (w), Our next result closes this gap and shows that this sub-manine mandmanifold indeed contains all the stationary points, and is therefore eventually attained by the dynamics (29), irrespective of the initialization.
Proposition 4. For the linear link function I= the sub-manifold M contains all the stationary points of the 4-dimensional dynamics (30). In fact, the only stationary points are (11, 0), of which (1, 1, 0, 0) (corresponding to the oracle model parameters (wa,v,) is the only stable one (more precisely, it is attractor/sink); the others are saddles and sources.
5.3 Dynamics on the Submanifold M
We now consider the dynamics (29) the regime where it has entered the submanifold M defined in (31) (e.g., via initialization). and show drastic simplification of the picture.
Proposition 5. The 4-dimensional dynamics (29) fixes the submanifold M. That is, once the dynamics (29) enters M, remains there.
Proposition 6. In the limit of vanishing step-size, the equations of motion of the overlap variables μ and v, induced (30) are given by the following 2- 2-dimensional gradient-flow:
((t) = (P((t), b(t) = (F(t),
(35)
where the scalar fields F F2 [-1, 112 + R are defined in (71) and (72) respectively.
8

---

Stationary Points. Let E = ((,.)) Ç [-1, (( = (2 0} be the stationary
points of the 2-dimensional dynamies dynamics (35), Consider the following subsets of []1,
E1 = {(=1, Eg = {(=1, 2) v (-1,1), d,(x) = =],
(36)
E3 = (11 p = -
(37)
- p(p)= F4 == [(4,)) E (-1,1)² - (v))==
(38)
Proposition 7. E = Es U2 U Eg U E4 is partitioning of the set of stationary points of the
2-dimensional dynamics (35).
Proof. The E of stationary points correspond to pairs (H,, for which Ff(,, = (2, = 0. By definition this translates to
0 == (1 1 (01-)2)A(0)) A())) (1- (1-x))) p2)A(()))(((()))),0(00)
(v(),,, 0 0= -
But A(1) (1 p(u))cBd and p(1) are always positive (since p(1) (0, 1)) and they can be canceled out in the above equations to give
(1 - (1-)))p(())0_(0(0
(2)) =0 (1-12)0;(1)) =.
By manipulating the equations, it is easy to that the all the solutions be organized into the sets E1. E2, Es, and Eд respectively. Each Eg vanishes exactly one of the two factors in either of the two equations above (four possibilities in total).
A complete classification of the stationary points is provided in Appendix A.
5.4 Convergence A Stationary Point
Our analysis would be incomplete without showing that the PGD dynamics (29) actually converges stationary point. The following result is proved in Appendix G.
Proposition 8. Under some smoothness conditions on the link function (made explicit in the appendix), the following holds: If (ug, vo) E M, then the PGD dynamics (29) converges to a stutionary point of the population risk functional R.
Taken together with with Proposition 4, Proposition 5, Proposition 7, alongside the results in Appendix A on the classification of stationary points, we infer the following:
In the case of the identity link functions (1) population risk function R converges to the group truth model parameters (v., v.), irrespec- tive of the initialization (un, un). This result is much stronger that [Marion al., 2025] which required (ug, vo) E .M, even though the latter considered the much simpler case of pointwise self-attention (1).
(projected) gradient descent on the
In the general case, we have the same convergence as above, provided the initialization (wo, vo) is on the manifold M.
Concluding Remarks
We present end-to-end theories of softmux self-attention and an interesting task, the single-locutor regression problem proposed in [Marion et al., 2025]. Building on the pointwise/non-sootmaa analysis of Marion et al. [2025], we give closed-form formulas for the population risk and gradient flow, pinpointing why attention solves tasks that defeat linear models. Our proof unites tools from statistical physics, mean-field theory, and Riemannian optimization, while eliminating the delicate initialization assumptions of previous work. Thus a ubiquitous yet opaque layer becomes one whose behavior can now be predicted, analyzed, and engineered.
Limitations and Future Directions. nature next step would be to extend our proposed theory to (1) empirical risk minimization (i.e snnte-sample analysis) and (2) multi-locator regression problem, where instead of a single secret token index [Marion et al., 2025], several secret indices must be recovered. The latter would encompass tasks such as the well-known sparse parity problem, opening the door to richer combinatorial analyses of attention. With some extra effort, the core ideas technical ideas developed in our work would directly extend to the this setting.

---

References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA. USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Anton Bovier and Irina Kurkova. Derrida's s generalised random energy models 1 models with finitely many hierarchies. Annales de I' I.I. Probabilités et statistiques, 40(4), 2004.
Hugo Cui, Jaron Kent-Dobias, Florent Krzakala, and Lenka Zdeborová. phase transition between positional and semantic learning in solvable model of dot-produee attention. Physical Review Letters, 132(13), 2024.
Bernard Derrida. Random-energy model: An exactly solvable model of disordered systems. Phys. Rev. B, 24,17 1981.
Jacob Devlin, Ming-
Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Association for Computational Linguistics. Association for Computational Linguistics, 2019.
Yihe Dong, Jean- Cann-Bapiiste Cordonnier, and Andreas Loukas. Attention not you need: pure attention loses rank doubly exponentially with depth. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2793-2803, PMLR, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Yue M. Lu, Mary I. Letey, Jacob A. -evatooee-Vthh Anindita Maiti, and Cengiz. Pehlevan. Asymptotic theory of in-context learning by linear attention. arXiv e-prints, 2024.
Carlo Lucibello and Marc Mézard. Exponential capacity of dense associative memories. Phys. Rev. Lett., 132, Feb 2024.
Pierre Marion, Raphaël Berthier, Gérard Biau, and Claire Boyer. Attention layers provably solve single-location regression. In The Thirteenth International Conference on Learning Representa- tions, 2025.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research. PMLR, 2021.
Jürgen Schmidhuber. Learning to control weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1), 1992.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz 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, editors, Advances Neural Information Processing Systems. volume 30. Curran Associates, Inc., 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordv- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research. PMLR, 2023.
Jacob A Zavatone-Vi Veth and Cengiz Pehlevan. wadsonya-watoon kernel smoothing as random energy model. Journal of Statistical Mechanics: Theory and Experiment, 2025. jan 2025.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context.
Journal of Machine Learning Research, 25(49), 2024.
10

---

Appendix
Contents
Classification of Stationary Points: Sinks, Sources, Saddles
11
B Mise: Additional Theoretical Results
12
Relevant Statistical Physics for Machine Learning
12
C.I High-dimensional Analysis of Abstract Nadaraya-Waid Watson Estimator (Local Learning)
12
C.2 A Novel Extension of the Nadaraya- Watson Estimator..
Wadsraya-Watson
14
D High-Dimensional Representation of Estimator and Its Risk
16
D.1 Case of Linear Link Function: Proof of Corollary 1
16
D.2 Proof of Proposition
16
Equations of Motion
16
E.1
Auxiliary Functions
16
E.2 Proof of Proposition 3 (High Biagh-Dimessonal Representation of Gradient of Population Risk)
17
E.3 Proof of Corollary 2
19
Proofs of Classification Theorems
19
F.1
Jacobian Matrices
19
F2 Proof of Proposition 9
20
F3 Proof of Proposition 10
20
F.4 Proof of Proposition 4
22
G Proof of Proposition 8 (Convergence of PGD to Stationary Point)
23
G.I Step A Descent Lemma
23
G.2 Step 2: Convergence to Stationary Point
25
A Classification of Stationary Points: Sinks, Sources, Saddles
We start with the following general result which shows that oracle model with parameters (u,, ux) corresponding to μ == of nonlinear link function a.
stable stable stationary point of the dynamics (35), even in the case
Proposition 9. The point (1, 1) (corresponding to the oracle model parameters (Mx, v.)) is always an attracorrsinkk of the 2-dimensional dynamics (35).
As with all the other results in this section, the proof of the above result is provided in Section F.

---

We now give complete classification of the stationary points E = Uj for various choices of the link function o.
Proposition 10. We have the following classification of the stationary points of the 2-dimensional dynamics (35) induced by different choice of the link function a.
(A) Linear link function. For (()) := I, the stationary points are: (+1, of which (1,1) is a sink (stable), (1, is a source (unstable), and (-1, =1) are saddles (unstable).
(B) Quadratic link function. For ott) := e?, the stationary points are: (+I, (=1,0), and (e, 0), where dp is the thermodynamic parameter introduced in (14). Moreover, (1, +1) sinks (stable); (1,,11), (-1,0), and (2), 0) are saddles (unstable); (1,0) 0) is a source (unstable).
Note that because of the evenness of this link function, the stable stationary points (1, both correspond to the oracle parameters (We,Us) v.).
(C) ReLU link function. Consider the link function (1) == max 0), and suppose y E [1/2, 1).
Then, the dynamics has 4 stationary points: (11, of which (1, 1) is sink (stable), (1, is saddle (unstable), and (=1,-1 are degenerate.
Thus, for all these link functions, the iteration scheme (29) with sufficiently small step-size is guaranteed to converge to the oracle parameters V.).
B Misc: Additional Theoretical Results
Proposition 11. For any v. p. E [-1,1], set := ccvdd. h= h= b b= bu max(b, 0), and q == (h) + (yp(h). For different choices of the link function o. the function dual function J.,. and the H2 in Proposition 1 have the following closed form expressions for any v, V.S.P [-1, 1].
(A) Linear link function. (t) == t, then oz(v)= = v2, and - Hi(v.S,p)=
+ (a (a--b) b)?, H2(2, = o,,(v) - (a (a- b)b.
(39)
(B) Quadratic link function. (f == t then o, (v) 11-1-1111211 and
Hi = -,) at 2a 2 + (21) - -26²(a² + +72) +64
(40)
H2(v,S,p) = -,(1)) (a +,²) +04
(41)
(C) ReLU link function. Ifa := max 0), then (v) = (VI - 2² + virccos(=u)) the well-known "arc-cosine" kernel, and
TI(,,cp) = 0v()) + 22 (2, + 1h) VT V - -2qbo + +bg
(42)
VI+
Hzlu, = (v( + (1) v- - (q+ -
7
ybo + +b
(43)
VI+
where ip is the standard Gaussian pdf. 4 is the corresponding cdf, and v2(-,vv) is the cdf of the standard bi-variate Gaussian with unit variance and correlation coefficient v.
The proof for the case linear and quadratic link functions reveals that the results can be readily generalized to general powers by making use of hyper-geometric functions of type 2 F1.
C Relevant Statistical Physics for Machine Learning
C.1 High-dimensional Analysis of Abstract Nadaraya-Watee Watson Estimator (Local Learning)
Fix unit-vectors u, v E R.N. Consider a size-n random energy model (REM) [Derrida, 1981] in N dimensions, with energy levels Ei...., En, where E. = VNx, u, i. independent energy levels from N (0 (O,N), and consider the sum
g (x // V/VN), with Pi = SE,
2 ==
3E
Z
12

---

for some link function a E ()), which doesn't depend on N. We seek deterministic equivalent in the limit
log nl N
N, n2 + 0o,
+ a.
(44)
The following result is an adaptation of the main result of Zavatone- Veth and Pehlevan [2025] to the case of Gaussian covariates.
Proposition 12. In the limit (44), holds a.s that
log Z
v/ - a(r);
(45)
3N
with
= uTT, p==u -= 0/B, 0= Br /2/2,
(46)
r = min(B, Berit)
Berit
/2a.
(47)
Proof. We use Laplace method of integration coupled with the recipe of Lucibello and Mézard [2024].
Zavatone-VV Veth and Pehlevan [2025]. Thus, we write
+ I] colt (q) dtdg
(48)
cd(t.))N dtdg
where we have introduced overlaps
1= -.//VN==////a
q=x/u/VN.
The potential has an energetic part 3t and an entropic part a - stt, q). In order to determine the function s, use large-deviation methods. Consider the random bi-variate Gaussian random vector 2 [ ut x v), with covariance matrix NE, where
- I):
Σ x-'= A --
with p= TT k: 0².
(49)
p 1
The (normalized) Log-MGF of w is given log Ee
eT NE2/2 = T. Σ6/2 ((0), for any E R
(50)
We take to be the Legendre transform of i.e
(cc) sup eT (()) c. Σ 1 12 (2/(21)) ptg/n + 92 (2x)
CE
(t pq) (2x) + (1-p²)g (2x) = (t - (2n) + 92/22
for any c= (1, 4). The condensation region the corresponds to s(c) 2 a, i.e C Ec/2 2 a. which is ellipsoid c-space. Consider the re-parametrization in polar coordinates
pq = T - p² cos
r sin 6.
(51)
That is, rsin 0 and = / p² cos 0 + pr sin O. This gives s(r, 0) = r22/2, and our potential takes the form
(r, = (t(r 0) + s(r,0) 0) a + Br. (V1- cos psin 0) -2/2
(52)
and we must now maximize w.r.t (r, 8) outside the condensation region. Maximizing ( 0) w.r.t gives
arctan --
This in turn gives cos 0 = /1 p² and sin 0 = p. Plugging into @ gives
+ = a Br - r2/2,
pr.
Moreover, the condensation region is now given by S s(r, 0) = /2,r V2a. Outside of this region (i.e for /2a),, maximizing 4 w.r.t gives B.
Putting things together, we get the following final form for our potential
Φ = (a, 3) = a + Br - r2/2, with r = min(3,80 Bera), Berit = V2a. It is standard result for the REM that (8N) -1 log Z + 6/B. Finally, plugging everything into the RHS of (48) then gives ~ (e (pr) ~ ( No (pr), as claimed.
D
13

---

For our purposes, we would like to compute sums of the form 9 Σg (2) v) and not Esprt(a (x.T(r/T//// v/vN) as in Proposition 12. To address this, can't simply consider new func- tion by (1) == (lVN) to write / = E, (x( (NTTTv//VVV v/vN) and then apply Proposition 12. The issues that hN is not fixed in (2, 1)) but itself varies with the dimension N which is tending to infinity. In the special case when is positve-homogeneuus, we can effectively factor out this dimension-deperr dependence and correctly apply Proposition 12. Viz,
Corollary 3. Let (pi); be the Gibbs distribution from Proposition 12. If E (2 1)) is positively-homogeeeoaw then in the limit (44), it holds that
Σ =)orr\))
(53)
Proof. Indeed, positive-homogeneitty means that there exists 0 such that (ut) u a() for all 0 and IR. In particular, we have (zx 0/ v/vN) axx um/Nm/2 i.e. a v) Nm/22 a (n] v/V The result then follows from Proposition 12:
Σ (I N
pio (x. v/vN) N ((rrVN).
C.2 A Novel Extension of the Kedaraya-Waason Estimator
We now extend Proposition the case of multivariate link functions. A slight modification of our arguments from Gaussian to spherical data immediately gives non-trivial extension of the main result in [aaaaoone-V Veth and Pehlevan, 2025].
Let u, ......, E RV be unit-vectors, and let Pj == u v be the cosine of the angle that makes with each l/.. Let .1...., Zn be iid from (0, IN). and consider the following generalized NW estimator
BVNIIN T
VI/NN at 0 Ve/VN), with Z = BVNNTu
(54)
Z
where the F is an (N, 1x)) function and & 1 is fixed integer. We seek deterministic equivalent for in the limit (44).
Proposition 13, In the limit (44), it holds that - F(pri...... prk) a.s.
Proof. Once again, we will follow the line of thought of Lucibello and Mézard [2024]. The Laplace method let's us write:
g S ... colt.al...... .x) F(q..... ..., (s) dt dg dgx
(55)
SS cll..n...... dl dg ........... dga
with the definition of overlaps
1 N/VN, @ <- x. vn/VN,
Ok + x.. V./VN.
(56)
Now, we need to compute the log-MGF of the Gaussian random vector (x, u, xI. v1.... 27 vx). Its covariance matrix is NΣ, where E E (++1))x(++1 is the covariance matrix of 2/ /N, given by
Pi
V1
1 P P
Σ
with p
e Rx1
and V ==
E R
(57)
VV
PR
We thus compute the normalized log-MGF of
(()) log e
T NE6/2 Ea/2, for any E *=1
14

---

The Legendre transform of is of course given by
((c) sup
é" c ((é) = c/22, for any (1, (1......... 91) E R
cEa
Recall that the condensation region in c-space is then the exterior of the set by s(c) S a, i. the ellipsoid 3 C R g++1 given by
Σ c/2
Now, for any fixed value r²/2 of scc). the sought- potential ø has energetic part 3t and an entropic part a s(c) = a -2, i. has the form
Φ = 3t a a 12/2.
We now maximize & subject to the constraint s(c) = /2 - This is linear maxi- mization problem (since c e1) with quadratic constraint. The method of Lagrange multipliers gives Σe,, i. is proportional to the first row of E which is the vector (1, (l,....... Pk). Plugging the constraint T -T--- == r² gives rEee = (r, pr...... prk). The potential then takes the form -=--===--
Φ = + Br - 12/2, t= c e
9 Mrj. Vj.
(58)
The condensation region E' is then 2 22a Berit- In this region, maximizing ø w.r.t gives r = B. Thus, we must have r = min(3, (eitt). Combining with (55), we deduce that
g S.... e ce F( Fpri prr) dt dq .q. eoN F(pr pPk) colt.gi .......- dt dqi ... dgp
EP.
-(pri,...,,,r
as claimed.
We have the following generalization of Corollary 4.
Corollary 4. Let F be in Proposition 13. in addition F is positve-hommoeeeouss w.r. each imput (the orders of the homogeneity are allowed to be different), then in the limit (44), it holds a.s that
EpF((sT vi xf
xTO) vk)
(59)
Recall that being positiv--homogenety order
R$ means that
h
F(u1tı..... uctx) = F(t tx) II
for any .,,
.t E IR and U1 ......, up 0.
Proof of Corollary 4. Indeed, taking uj = 1/ and t, = x.7 Uj gives
F(x Te........ Un...... o 2&)N E, m;/2 F( 01 /u///N V. ... of / v//N)).
Thanks to Proposition 13, we know that RHS ~ F(rp .IPp,....,TP... rpu). We deduce that
F(IJun.... "u NE, / ((pp.........
claimed.
15

---

D High-Dimensional Representation of Estimator and Its Risk
D.1 Case of Linear Link Function: Proof of Corollary 1
We start with an instructive self-contained proof of Corollary 1. Following the decomposition (11), we only need to estimate f2 = -efe, qe v, where qe == Rx /Z-e,, Et == vdz; u,
2 Eefee eB We can decompose x. u = pa + v pize, where the zς are (0, 1), and
independent of Xu. This gives
J2 = Σ gro u + v p2 Σ Ext.
qezc- Egt.
Conditioned on Xu, the second sum is a centered Gaussian distribution with variance equal to (1 - Eese. oi p² $ 1. For the first sum, observe that
Σ q EpC.
d
10 Zee. 2Z
1
Z-L. E, =
3E
06 03
(60)
vd
d EFC.
03
and the result follows. The step "-+" is a classical result in the analysis of the REM, while the last step "AgØ follows from the definition of & and in (14) and (17). We deduce that f2 = pr and so f(X;u, v) = p u + (1 - pr as claimed. The claimed formula for the population risk R(u, v) is then obtained by plugging the previous formula for f into definition (10), and then computing basic Gaussian integrals.
D.2 Proof of Proposition 1
Indeed, following the decomposition (11), we know that ~ po pz (1-p)) and we only need to estimate f2 Eet,, (2, v), where == ex /Z- Ee == vdx Z == Eese. e.E, and p p(u u.) is defined in formula (20). Corollary 3 with N = and L gives
f2 ~ (prrdd).
We deduce that f f pa + (1 (1- (pr d) as claimed.
The claimed formula for the risk R(u,, v) is then a matter of direct Gaussian integration.
E
Equations of Motion
E.1
Auxiliary Functions
Define auxiliary functions A [-1,1] R and B, T,,,, T2, T3, Fi, F2, F4 [-1,1]3 R. and
F3 : [11,1]
IR by
A(u) p(n)(1 - (1)(11-(4))33d) 0,
(61)
B(n,, = (12v)) T,, A(m)B(,,,,C),
(62)
(63)
T2(,, = 122
(64)
T,(,, =-
- 2h2(1,
(65)
F1 = -(1 - (1-,2)Ti(H,,,,))
(66)
Fe(u, U,C = (VTT,,,,, - (1 - (2)Ta(p,, U,)
(67)
(,, == (1,T)(
(68)
F, ==
(69)
where (v, = (v, Ç, 0), and the H₃ are as defined in (23). particular, when 17 = (= 0 (as
on the sub-manifold M)) the above formulae drastically reduce to
B(u, v, 0) = 2(p(u)o (1) - -,(())) (H, v. = A v,0), Fi (i v, 0) = (H( = -(1 11--p -11-12)1 (2,
(70)
(71)
T2(,, v4 0) == (V), Fe(u, v, 0) Ed (u, == - (2)TE(H,,
-11-v² (1-11)22
(72)
16

---

E.2
Proof of Proposition 3 (High-Dimensional Representation of Gradient of Population Risk)
Recall R(u, u) = [[ u, )) where 8(X; v) == f(X; - e.). Differentiating w.r.t and gives
Valluu,)) VX
- (Xl v)]
(73)
We already know that
6X;u,u0 = pa (x). v) + (1 - ()) - ((( v.).
We now need to estimate (x u, 2) and V. (X; u, v). Setting À == Bvd, one computes
exx; o' (x) (d[T
Vos-
o (z)
(xTu)e u)sz = po (x) ujze. + +(1 (1-13 p)
Σ
Z
Z
tH
uxe u
1
V of
Vu
o (x. v)
Z
dat "
du
o (wt v)
Σ
e
dal 2
7
2
7
1
AxTT u
2x, a
ol0 1) - f
Σ
Ft
Z
Z
8
xx27 n
exoT a
= p. (o(z),)) v) f)ze. + (1 (1-p) p)
Σ
o ox v) - f 1.(1 (1-p) p) E
Z-e It.
Z-c.
where p - p(p)= = 1/(1 + e (cH- () and ;== u u, as usual.
Now, for any w Sa-1, Corollary 3 and Corollary 4 (with N = d and n = L 1) give
f=
(74)
8 = p(o(tt, - -c(x) (prVal)
= po -0
(75)
(x -f f ~ (1- (1)))
b) (1- ())
= - (p) (p))Ve)) = (1 - p)o u),
(76)
cx u
Σ Ext,
(x, u)x w = o
(77)
Z_e
c u
Zie E w ~ rv w.
(78)
Eph.
exTT u
o' (IzTu)); ~ o' d' v d o ddu w.
(79)
Z
EtC.
From the above display, we deduce that
w v po u (1
w u a1= N = =PVo(o(rd,0) - -o(rv T w + Vno(prva)
=pVioncolz v) w + (pr()) Tw.
17

---

Likewise, we have
w Vuf
= (0(2_ - w o(prv w (po + (1 p)o(pr vajjruT w
(1 p)d
= =p-(o(z o(pr w p (o(m w) - o (pr w.
w vuf
= (o(2) u) o d)) ) (o( u) - o
p(1 - P)d
= (cuu w ru + + T o)(( w).
Putting things together, we have shown that
V vof w -V. - o(prvd))) w voo(prvd) w,
(80)
Vuf w
p(1 p)d = ons.p(2 vouu, w ru ru-ww))))))))))) w)va 70 ( (o)(2
(81)
8 p(o(x]] -(prv)) (o(z) v.) (prVa)).
(82)
Gradient w.n.t Parameter. Proceeding from (81) and (82), we compute
Vul((u,)) = E 8 w v 2p(1 p)1
p(1 - p}
(cu, w ru w) V dpo v) - -ox (0taop=TTee)eelll (o)))ancoo w)
+TB[(pon,c.ol (v) o2,0,(zz v.))o,,sill (on(=z))) oo)(=TT) w)l
= (cuu, w - ru ( - H22u,,
+>2Vd((T v. (1 2p)w 2 (),,,p H2(1,
Since 0 by assumption, the Remark 3 tells that all the H3 functiosn no longer depend on and we get
p.
1
p(1 p)d V VuR(u, T w B(u, v, doV w, where U(U,,, = (1,,)))((1,)) h2(u, ()) usual. Recalling the definition
(83)
A(r) = p(1((1 plu))(Bdd p())(1 - p(u))cdVad
we can write the above as
V. u(, t = (w (.))(())(() () = (i u, S)u, w.
(84)
This proves that V. V. v) = Ti (1, v C)u. as claimed.
Gradient w.r.t v Parameter. Using (80) and (82), one computes
w V, R(u, v)
2 Elg V.fl
~ El(por,c,0=> -)) (o..ooolT v.))
ww w Voo ww)
n2w Bloxs.olz (V Voo v)] pB (
pw T u. V.o,,,o T b)] - (, (
Now, using the well know formulae g (1/2)Vg² and gVh Vgh hVg. we have
2E(onsc.n= VV v.).,... [w]
(, V,H,,,,,.
El T v.) V.VV Tw]] w)] - - To) Vo TT,.) v.)]
-0,Ho(u, (,p) +
+ (v.
-Bloxc,o u) V V 100 v.)].
18

---

Observe that V V.0g.u.pe v.) = - (.o(pryd) and so
1 V,,,y,,,, T e.)] = Blonss,
(85)
Putting everything (once again under the assumption that = 0), we get
V u( w = pV.(phn (1, - 2h2(v, ())" w
=p((@ch((,S)) 20, 5, p))u, - 20,h2(u, [u.]"
(Ts(mn, S/u, + T2(,,, V, ()v.)"
We conclude that R(,, To(m, Uu, + Te(M,, v, S) as claimed.
E.3 Proof of Corollary 2
From (30), the dynamics for (14, V, 17, < is given by
ju(t) in() - P.T -..P.TT u(1) V R(u((t), u(t),
i(1) - v, e,'U( i(() = -- Uu m(2) V. V. o()),
vi({) T,Tú(
P.T V. R(l(((t), x()),
((t) = u. é(t) - u. Po(e) P.,
R(u((),, ott).
On the other hand, computes
P. V. R = Ti (u. pmm); P.V. T3 (u. - (v) T2 (v. vv),
" P.L V V.R R = Ti 11.(1 112) = = Fi, " P. V. T3 (1 (2) T2(u =: =FF
P.. V.R = T: F3. v. V = -T3C + T2 (1 v²) =: -F2,
and the result follows.
F Proofs of Classification Theorems
For ease notation, we shall use the following shorthand
(1)) p(u)) B(,, v) = B(,,v,0); Te(U, = Te(m,,, A(p) = Fi. (M, for k = 1,2,4; Fa(u,)) = F,(,,
(86)
(87)
F1 Jacobian Matrices
Classical theory of dynamical systems tells that the classification of the different stationary points (p, v) E E can be done by studying the signs of the real parts of the eigenvalues of the Jacobian matrices
( = (.,Ti(m.)) T,T2(1, (,TT(H, 0,T2(1,D)
(88)
One can use these Jacobian matrices to classify the stationary points sources (J(u, v) only has positive eigenvalues), sinks (J(u, only has negative eigenvalues), and saddles (J(1, v) has positive and negative eigenvalue).
Let us now compute the entries of each J J(4,, v). From the definition of Ti and T2. we have
0,Tiiu,,)) = 2(c3d)² ((2p 1)p(1 p)B -
(89)
= 2(c3d) p(1 - p) ((2p - (112p-1)BBa p(1 (,TT(u,v) = -2p((1 - p)cBd d'((v,, 0,T2(H, = -2p(1 - p)cBd dv),
(90)
(91)
(92)
0,To(u,P) = (v)
(93)
19

---

Thus, for any stationary point (1,D) E E, we have
02F1(H,)) -11 (1 + 2T 2H,
(f, Ei U E2, if (1, E Fs U E., ((1 -=) - (2) <0,
( v)ju
0,F((H,)) (1 - -1))) =0,
0,F2((D)) (1 =2)0,,T((4.v)==0
0,P2(1,)) -(1 () D)R,
= - -11- v) = =2(1 1-l1-22 ((()))((
(f, E BEUB E3: if (1, E Eg U..
We deduce the following Lemma which will be crucial in the proofs.
Lemma 1. Recall the definition of A(μ) and (p, v) from (86). For any stationary point (μ, v) E E E U E Ez Es of the dynamics (35), the Jacobian matrix (88) has the following form.
2A 2A( v)
If (U, Ev, then V(
2A(p) ((, D)
(f, E Ez, then J(u,v)
2(1- (1)-)))())
-2(1 /² A())(())
If (f, Es, then (H,
-4p(1)o5(000)
2(1 - 12 A(12200
0
If (u.v) E Ea, then J(,
2(1- (1=)))((())(((1)
F22 Proof of Proposition 9
Thanks to Lemma 1, the Jcobian matrix is given by
1(1, = ((1,)) 0
0
-4p(1)0;
Note that by definition the A and B functions in from (86) that
A(1) 1) = p(1)(1 - (()))-- 2(p(1) - 1)(= (1) (1)<0 <0
because p(1) 1) and (1) > On the other hand, -4p(1)0'
(1) is negative since (1)
Σn21 ne 0. We deduce that J(1, 1) has only negative eigenvalues, and the result follows.
F3 Proof of Proposition 10
We now consider a few important choices for the link function a,, and provide a complete classification of all the stationary points of the induced (j,, -dynamicss. The picture will of course depend on the underlying link function a.
Linear Link Function.
Consider the case (tt) t. It is clear that
(v( 72, d'((v) y, d"(v)= = 0.
is then easy to see that Eg E3 E, 0, and so the set of stationary points
E Up E E₁ (8-(111) 1), (-1,
(1,
Observe that for any stationary point (H, v) E E, has (i,, = 2T (H, (,, 0, (, = 2T2 0,Fi(H, = F,, 0.
20

---

Now, Ti(=1, = (p(=1) - v)A, T2(+1,)) = -1)p(=1), with A == p(1)(1 p(1))oBd > 0.
The Jacobian matrices at each stationary point is then given by
-2(p(-1) - +1)+1)47
1(1 =
(94)
2y 2
J(-1,1) 2(0( - 1) - -
0
(95)
-2
-1) 2(x(1) + 1)A4
(96)
0
27²x(1)]
1(1 2(p(1) - 1)A4
0
-2y2p(1)
(97)
(98)
Since A > 0. p(n) 1 > 0 and p(n) 1 < 0 for all Jr, deduce that
J each have one negative and positive eigenvalue: these are saddles.
(1,, has both eigenvalues positive: this is a source.
1), has both eigenvalues negative: this is a sink.
This proves Proposition 10(A).
Quadratic Link Function. Here (t) = t and thanks to Proposition 11, 11, 01,) = -111+
+ 2v²).
Consequently, we have
o, (+1) = 27 > 0, d²(v)= = d)(()= = 47
(99)
- 0} = It is then clear that E2 = ((#1, WE -3,
((11,0)}
We now consider E3 = (1, (0))
= [ 0}.
Now, 0 p1))o,(1) (+1) iff
p(u) o, = 1,
(1)
which is impossible. Thus, Eg = 0. Finally, we compute Eq. By definition,
E₁ [(L,)) E ()² d²(()) p())= (1) [v]1 = ((1,0) (())),()) (0) = =0.
Now, 0- (1)) - (0) iff
p(u) (0) 1
i.c., p == V.
(1) 2'
which is feasible iff vp < 1.
Therefore, the stationary points are E = E, U2 U Es, with E1 = [1), Eg = and E, {(tb, 0)}. Recall A(u) and B(H, v) from (86). We have the following classification
((=1,0)) ((H1))))
#1) at each such stationary point, the Jacobian is given by
J(H, = (2A(w) B(u, v)μ
[AA w)
0 -16p(u)>
4p(1));
Note, that for any μ [-1, it holds that
B(u, +1)
(+1) 2(p(m)
(1),(1) < <,
since (-1) o > 0. Thus, (ign(A())B(r)) v))) = - sign(u) for all (p, v) E
+1),, Thus 11) each have one negative and positive eigenvalues and (11, +1) each have only positive eigenvalues. We conclude that (1) are saddles while, (1. +1) are sinks (stable stationary points). It should comfort the reader to know that +1) are equivalent representations of the oracle (i.e. Bayes-optimal) parameters (u., v.)) because
due to the evenness of the quadratic link function, replacing v, by -P, doesn' 't change the oracle model.
21

---

(2A(4))((4 O) 0
2p(4));((0) In particular, we see that J (-)1,0) has one negative and one positive eigenvalue, while (1,0) has only positive eigenvalues. We conclude that (-1,0) 0) is a saddle, while (1, 0) is a source.
(=1.0): Here, J(,,))
(v, 0): This is a stationary point only if 0
+ <1 In that case, we have
J(b,,) 0) 2 - (1,)
0 =
2(1 2 A(1)
0
(0
4-21
Thus, J(bb,) 0) has one negative and positive eigenvalue. We conclude that (v), 0) is saddle. This proves Proposition 10(B)
ReLU Link Function. Consider the case where (1) = (1)+, thanks to Proposition 11, we have
22 (V1 2² + varccos(-v))
ov(x): == 2x
One then readily computes
(1) = 72
12 2T
0y(-1) = =0, d²(1)
arccos( --V),
2
and so
Fgg = [(+1,1) [ (u) = 0) = ((=1 u) | V (-1, 1), arccos(-u) = -0} =
E4 ((,N) (-1,, p(1)) (-1)) (v( 0}
{(,,2) (--1,1),, -((()) arccasccos---- = 0} 0,
E3 = (()--1) (-1,1), -1)(-())) = =0}
6
= (- 0())) -)) 0.
Thus, the stationary points are E E = ((#1,=11 Now, for any (H.V) E E1, the Jacobian of the dynamics is
Ju, v) B(, v)
0
B(A 4)
A4p())o((V)))
0
-2p(u) vaccos( -v)
More explicitly, we have
-1, = (2( v) 0
0
J(H, 1) = (A( v)
(100)
0
0
-pp(1)
Since B(u, 1) : 21)==(() --, ongoing constraints), we deduce that each of J(=1, -1) has one negative and one zero eigenvalue; while J[-1, 1) has one positive and one negative eigenvalue; [1, 1) has only negative eigenvalues. We conclude (1, 1) is a sink (stable stationary point). (-1, 1) is a saddle.
221 po (1) < 0 (because (,( > 0 under the
This completes the proof of Proposition 10C.
F.4 Proof of Proposition 4
For linear link function o (1) = l, we have for any n. u, E [-1,
h ((v.))= a² 12( ,²v, with cVvd.
(1 () - 20, 2a = 2dcVd, 0ch2(u, () 7
We can thus simplify the auxiliary functions F2, E2 Fs Fa appearing in Section E.1 like so:
Fi - (1-1,,
(101)
(1, = (VT33 - y2,
(102)
Fs(H,,,7.5) () -mHII,
(103)
Pa(m, () -(1 -c2)Tss CUT2,
(104)
V, Telp, () -2p12, Ti(,, ((,,)
(105)
B(N, U. -2phh (1. h2((u))) 2(p (92 + a2) - -22v),
(106)
Ta () = ((())0((1 (h (1, - (2, 5)] = 2p(1)a- = 2p(u)csvd
(107)
22

---

Now, the stationary points (n, v, 77 () of the equations of motion given in Proposition 2 are defined by
the equations Fi = F2 F3 F4 O. We will prove that any such point must verify 1) Ç 0.
We prove that Ç = 0. Suppose on the contrary that ( * 0. Then, the equations F2 = F4 = 0 give
= /C and c² = 1 - (,,2)2/C²² where C' == cv > py². This implies
B/2 = p- (72 + -²) - ~ー = p. (02 + d² /C) = = (12(1 1/0) +c) 0.
Thus, we must have B o O, and so the equation Fi = 0 gives u T u. = = +1. Now, μ = Tu u, =1 implies u = _u,, and so Ç == v iv' u tp 0. By reductio ad absurdum, we must conclude Ç (=0 0.
Note that with r = 0, the equation F2 0 now gives = =1. This means v = Iv, and so 77 == D. =u v =p- O. We conclude that every stationary has 17 Ç 0.
We now show that the only stationary points are (+1,+1) +1). Now, plugging dj = Ç = 0, the stationary point must satisfy
= (u, = - (111111/(piii ,²)TT(H, 0) = - -1=--=1111(ppppppp (1))))))())))((()00000_ - -1)
-22(()(1 - -2) -
- -
0
= (11- ----" (01)))=111ppppp))) 1²)2p(())²
Because p(,))?² > 0. second equation tells us that we must have y +1. Plugging this into the
first equation and dividing through by the factor 2A(μ) > gives (1 - 11 (p(() F 1) 0. But
p(u) = 1 / 0 because p(u) E (0, 1) for all μ. and we conclude that μ +1. This shows that the stationary points are (+1, 0, 0) as claimed.
The classification of the stationary points then follows from Proposition 10(A).
G Proof of Proposition 8 (Convergence of PGD to Stationary Point)
G.1 Step 1: A Descent Lemma
We shall need the following regularity assumption for the link function a. Assumption The link function a is C2 on IR with I|o'll ||o"|loo < cc. The case of ReLU
activation needs special treatment (not provided here).
One can show that on 53 the R functional
R(u, v) == p² 11 (p - 2pHH2 S. (1,00. with μ : Uns == u. Tee,, (:= ===v" vJu..
Uo
p = v u,
(108)
p = (1 lew -3ddd
is L-smooth on $2_1 for some finite O. Now, consider the following canonical the extension of R from $2_1 to all of Rd
r(u, v) = R(u/Iuu], w/|l=]]
Note that V Vur(u, v) = that for any 6 E (0, 1). the functional r is Lg-smooth with L3 == : 4LL(11 - 5)2, on the tube
v) for all (u, v) E $2_ S (8) C R22d Then, one can show
2 = {(u, v) E R2d - 5 < ||u|]. ||//| < 11+6% 67.
This means that
r(u',0') < v(,,)) + (rru,, v), A) + L ||
(109)
2
for any u, u.v,a u', o' (5), with A (ul u v) E Rdd Furthermore, because r is radially symmetric, we know that Vwr(u, v) 1 and v,r u) 1 v for all non-zero u, e Rd.
Now, define gr- = (ax, bz), where ax by E ddd are defined by
ag = Var(us, vx) = P V.R(u, bk = Vor(us, uk) P.l. V..((u Ue).
23

---

In (109) above, taking u = Uk. = Uk. u' ux++ = uy sat (as in (28)), and v == 0x+1
Uk sgk. we get
(1++1,0++11 r(ux, uk) >(Vr(ux, (k): uk):
19x,11
= =r(ux; 0i)) -111xx122= + L6
(110)
|1
2
=r(u, v&) (s - L
We shall now control the deviation of (uk+1, Uk+1) from (ūx+1: 1 ). By definition, have
||u4+1 ux+1ll == |1(
1)ūx+1l1 = [ || 1
|u++1|
(111)
1
||o++1 uk+1ll := |1(
1))x+11l
0 1/.
|
Now, because uk+1 Uk
s9u. with
|| and
ut 1 Jux-
have
s²
u
v s² Iu 1 < 11+
9uu 11?
2
where the last step uses the elementary inequality v1 + S Uh++1 we have
a/2 for a 2 Similarly, for
22
1 S |0++11l| = / +1++s2l99
s211 2 S 11+
11 11
2
On the other hand, (28) and (29) give
Uk+1 uk+1
1)ū++1, Wk+1
0k+1 =
1)8++1
|
|0++1l
We deduce that
g²
|lu+1 ux+1/l [ || 1/5
||ab11
2
(112)
z²
|1++1 - ūk+1l
||
1 - 1/5
|
Using this in (109) above gives with u us+1, Ui 21'
= uk+1 u of = uk-1, so that
Ck = (us+1 Ms+1, Uk+1 (k+1) gives
r(u441, S r(0x+1, 1w) + (r(#x+1: Uo-1) La
11Ga
(113)
5 L6
= r Th 1) +
8 Igull
where we have used the fact that Vr(ūx- (x+1) ) 1. Ga- Combining with (110) above gives
8 Ls
r(ue+1, 1k+1) - r(ux, uk) - (s - (s=22L_/2)|oon
|1
8
Now, one can show that (u, v) 110 Ms < do uniformly on ( Thus, if < 8 < 1/Ms,
then s s''||=]1* = (s ||gr ))² |g I 32 "| 12 F, and we
r(ux+1, Ub41) S r(ux, vk) (s- s²LLL -02 Ls/B\ll2x1
r(us, Ub) (s 156-522 L58>La/8)llaa
(114)
$ R(uk, v4) sL8 lg S rrux((u 2k) 4 $L $Lol19a11,
16
provided the stepsize is sufficiently small in the sense that
O <<s< < min 1/Ms). Noting that R on S S (0), we get the following descent lemma.
Lemma 2. If the step size is sufficiently small in the sense that 0 < min(1/(2Ls), Mo), then
sL6 1| with R2 = -R(us vx),
Ra+1 I R4
(115)
4
where recall that gs (ak, bb), ak = P.l. V. Puuuu vk). by P V R(uK, vk).

---

G.2 Step 2: Convergence to Stationary Point
The above inequality can be rewritten sLsllgx11 R2 Rk+1, and summing both sides gives
s LL8
K-1
11
- Ro Rx S Ro Rmin < 00,
h=
where Rmin == meS, k O. Now, one computes
R(,, v)
O. We deduce that || ll + 0, and so - 0 in the limit
ūx+1
|k++1
uxl
ux|l
|
Uk
s0k
l0x+
8
|
uxll
=
[
Up
0
()1
|
+ s|
|ux+11l II
82
||ag 1|2 s s||ax:]|
because | || 2
2 |axll + s s|lag] because s||ax|| S
3
+|
2
Analogously, we get |16+1 2/ S (3/2)s|) [bxl)>|bbl|l. Combining gives
9 2 219l O.
|1u+1
u + || 12
4
We deduce that the PGD iterates (ug, vn) given in (29) form Cauchy sequence in $2-1. Due to completeness of $2_1, this sequence has limit (u Use) E $2 We now show that (# Uso) is stationary point of the risk functional R.
For simplicity of presentation, we focus on the case of linear link function (()) =
Thanks to Proposition 3, if (uo. vn) E M,, then (uk, 0k) E M for all k and one computes,
(T u VV R(us, vk) = (T u Tiu.. = (u. - Ti, Teen T22 =(u- 1xR4) T2,
(116)
by (I u v. R( vk)
where T3 = Ty(se, Pk, 0) for j = 1,2, as defined in Appendix E.I, with μi := ug u, and P% == Ug v. Note that we have used the fact that T' (pk, 1kk 0) thanks to the equation (107). We deduce that
I1xx 2| = ||ax|| 11? + ||bıl1 = (1 2 T? + (1 -1- (1-1
Thus, the limit point (M Uso) is such that
(1 2 )Ti (Moc Moc 0) 0, (1 - 12 )T2 (p Voc 0)
This is precisely the characterization of stationary points risk functional R established in Proposition 7. This concludes the proof of Proposition 8.
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes].
Justification: The main contributions outlined in the introduction clearly linked to the appropriate sections and adequately support each claim.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes].
Justification:
Yes we discuss limitations in the last section of the main paper.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
26

---

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

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [NA].
Justification:
This is purely a theoretical work. Question doesn't 't apply.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //niss... wttpcccppppppeeeeeeeeeecccccceet htteggddde/eooeeeniiiittss /CCo///oodissssseeeeeeeee ssionsoiccy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: We provide all the experimental details in dedicated sections in the appendix (supplemental).
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No].
Justification: This is pure theoretical work.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
28

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative етrоr rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes].
Justification:
with single CPU laptop, and took less than 30 minutes in total. We provide all the experimental details in dedicated sections in the appendix (supplemental).
This is pure theoretical work. Experiments
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics htttp:://eeurpgee
Answer:[Yen Yes].
Justification:
We have read the code of ethics carefully and we conform with it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer.[NA].
Justification: Our work proposes theoretical analysis of self-attention, the main main building block of waanformmr-basee models which power LLMs like ChatGPT, Llama, DeepSeek, etc. Building rigorous understanding transformers has the potential help understand how they fail (c.g hallucinate), or drastically perform well on tasks previously though to be unsolvable with MI/AI.
Guidelines:
29

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negatíve societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA].
Justification: This a purely theoretical work. No new gadgets (models, architecture, datasets, etc.) are being proposed here.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The original creators of all models, datasets and algorithms used in this work are properly credited, with citations in the manuscript.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
30

---

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

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [No]
Justification: LLMs are not part of the core, method development of this research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: /ttpa/:///euiipee ct/ cc//oonfeeeecces/2222555 for what should or should not be described.
32


Reviews:
[{"review_id": "sjKJfpk0jm", "reviewer": "Reviewer_Mp31", "summary": "This paper explores a theoretical understanding of softmax self-attention, a key component of transformer models. It investigates how softmax self-attention functions in the context of the single-location regression problem, where the goal is to locate a specific token in a sequence. By using concepts from statistical physics, the authors provide exact analytic expressions for the population risk in softmax self-attention models, and they analyze the optimization dynamics under gradient descent. They also prove that under broad conditions, the dynamics converge to the optimal solution.", "questions": "- The authors mention that this work removes critical initialization assumptions that previous works (e.g., Marion et al.) relied on. Could you elaborate on how this will impact the practical training of transformer models, especially when working with models that involve multiple layers or more complex structures?\n\n- While the toy problem offers theoretical clarity, how do you envision the insights from this work being extended to more complex transformer models or real-world datasets? Are there any assumptions that would need to be modified for those contexts?\n\n\n- Can the optimization dynamics described here be generalized to other attention mechanisms (e.g., linear attention or variations in the attention head)?", "limitations": "None", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "#### Strengths\n\n- The paper offers a clear and precise mathematical framework, drawing from statistical physics (Random Energy Model) and mean-field theory. The rigorous treatment of softmax self-attention is a valuable contribution to understanding this crucial model component.\n\n- Impactful for Transformer Models: By improving our understanding of self-attention, this work has the potential to influence the design and optimization of transformers, which are foundational to modern AI systems like LLMs.\n\n- The use of well-structured propositions and corollaries, with complete proofs in the appendices, makes the paper both rigorous and accessible to those familiar with the theory.\n\n#### Weaknesses\n\n- The work is quite theoretical, with few experimental results provided. While this is understandable given the focus of the paper, empirical validation would have added strength to the claims.\n\n- The focus is on a toy problem (single-location regression), which is useful for theoretical insights but may not directly translate to larger, more complex tasks that real-world transformers handle.", "was_revised": true, "final_justification": "All of the concerns are addressed, and I will keep my score.", "rebuttal": "## Response to Reviewer Mp31\n\n### Impact and Extensions\n**Potential impact on practical training of transformers.**\nThe optimization dynamics of transformer-based models is still not well-understood, and for good reason: an attention layer is much more complex than an CNN or dense / full-connected layer, say. This is mostly due to the softmax non-linearity, which renders the system highly non-separable. The same thing which makes transformers highly performant (compared to traditional layers like dense layers and CNN layers), makes their analysis nightmarish.\n\nOur work contributes to a very limited (but growing) pool of theoretical works which attempt to explain the mysteries of optimization of and learning in transformer models. Our results classification of stationary points, and on the robustness of the initialization provides non-trivial much desired insights into the robustness learning into our simple transformer model. transformer models. \n\n\n**Extension to more complex transformer.** Yes, we strongly believe that the key technical ideas in our work can be leveraged to attack the case of more complex transformer architectures (multiple layers, multiple heads, etc.). In particular, with some additional work it should be possible to develop versions of **Propositions 11** and **12** (the main technical ingredients in our work) which are adapted to such settings. This is left for feature work alongside the other things mentioned under the **Concluding Remarks** section of the manuscript.\n\n**Linear attention.** Note that, the case of linear attention is drastically easier to handle, and has already been done in the literature (e.g. see [Yue et al. 2024], [Yedi et al. 2025])\n\n### Additional References\n- [Yedi et a. (2025)] \"Training Dynamics of In-Context Learning in Linear Attention\""}, {"review_id": "IMtxcHdKbG", "reviewer": "Reviewer_t3fK", "summary": "This paper closely follows the work of Marion et al. (2025) and focuses on the single-location regression problem. The authors provide a comprehensive and formal definition of the problem and extend the analysis to include softmax-based attention mechanisms. They derive an analytic expression for the risk and use gradient-based analysis to demonstrate that the optimal relation can be achieved through learning.", "questions": "The theoretical derivations presented in the paper are impressive and technically well-executed. However, many of the results feel like intermediate steps that ultimately lead nowhere concrete. It is unclear how these derivations contribute to a deeper understanding of model behavior or guide practical model design. Strengthening the connection between the theoretical results and their implications would greatly enhance the impact of the work.", "limitations": "yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The paper provides a comprehensive and rigorous theoretical walkthrough.\n\nWeaknesses:\n- The motivation for this work is unclear. In Marion et al. (2025), the analysis is grounded in a meaningful setup where one word decisively determines the output, making the theoretical insights directly relevant to real-world model behavior. In contrast, the current paper lacks a clear justification for why the theoretical results are useful or what practical implications they have. As a result, it is difficult to understand what is gained from this analysis or how it informs model design or learning dynamics in more realistic settings.\n- Too many typos. Like here on top in the title, an unintended backslash (\\) appears, and in the paper on line 241 the phrase “shall shall” is repeated.", "was_revised": true, "final_justification": "The authors clarified their motivation and contributions in the rebuttal. However, the work appears to be a follow-up with mostly incremental theoretical extensions over prior results. The presentation remains difficult to follow, and the lack of empirical validation limits the perceived impact. While I recognize that others may see theoretical value here, I remain unconvinced and lean toward a borderline reject.", "rebuttal": "## Response to Reviewer t3fK\nWe thank the reviewer for their feedback. We shall correct the typos mentioned by the reviewer. Below, we provide a detailed bullet-point response to the main comments of the reviewer.\n\n>*The motivation for this work is unclear. In Marion et al. (2025), the analysis is grounded in a meaningful setup where one word decisively determines the output, making the theoretical insights directly relevant to real-world model behavior. In contrast, the current paper lacks a clear justification for why the theoretical results are useful or what practical implications they have.*\n>\n>*As a result, it is difficult to understand what is gained from this analysis or how it informs model design or learning dynamics in more realistic settings.*\n\n\nAs in [Marion et al. 2024], the mission of our work is to contribute to the theoretical understanding of learning in transformers and optimization dynamics. Our work considers actual **softmax** self-attention layers (unlike the reference work which considers pointwise attention), and considerably extends the range and scope of the findings (global convergence to teacher parameters, etc.).\n\nThe setup of our work and our main contributions are well spelt out clearly in the paper see **Introduction** and **Concluding Remarks** sections. Let us repeat a summary thereof here:\n\n\n**[Marion et al. 2024]** considers solving the so-called single-locator regression problem (described in the introduction of our paper; see **Figure 1**) via *student model* which is simplified self-attention layer where the softmax is replaced by a **pointwise activation** function $\\theta$, namely $f(x;u,v) := \\sum_{\\ell=1}^L p_\\ell x_\\ell^\\top v$ for any input $x \\in \\mathbb R^d$, with $p_\\ell := \\theta(\\lambda x_\\ell^\\top u)$, as in Eqn (1). Here $L$ is the sequence length and $d$ is the input dimension. They show that this simplified model can be trained to recover the oracle / teacher parameters $(u_*,v_*)$. They also study the optimization dynamics and show that convergence to $(u_*,v_*)$ is ensured assuming gradient descent (GD) is initialized on a certain manifold $\\mathcal M$ formed by parameters in a peculiar configuration. This is local convergence.\n\n**Our work** extends [Marion et al; 2024] in many directions:\n\n- **1. Softmax attention and general link function.** Keeping student model same as in [Marion et al. 2024], we change the student softmax attention (which is what is actually used in real-world transformers!): $f(x;u,v) := \\sum_{l=1}^L p_\\ell \\sigma(x_\\ell^\\top v)$, where $(p_1,\\ldots, p_L)$ is now the softmax of $(x_1^\\top u,\\ldots,x_L^\\top u)$ at a temperature $\\lambda \\asymp \\sqrt d$, and $\\sigma$ is a generic *link function* which can be linear or non-linear (e.g quadratic, ReLU, etc.). Note that these are two separate extensions of [Marion et al. 2024]: the softmax attention and (2) the general link function $\\sigma$. Importantly, softmax attention maps more faithfully to what real world transformers (as noted by the other reviewers, e.g. **Reviewer mLDi**).\n\n- **2. Realizability and local convergence.** For this considerably more complicated (but realistic) student model, we show  (**Proposition 2**) that optimal parameters are in fact the teacher parameters $(u_*,v_*)$, and that GD on square loss finds it (**Propositions 5, 6, 8, 9**). For general non-linear link function $\\sigma$, our convergence is local, as in [Marion al. 2024]\n\n- **3. Global convergence.** In the special case of linear link function $\\sigma(t) \\equiv t$ (considered in [Marion et al. 2024]), we show in **Proposition 4** that for our softmax-attention student model, GD converges to $(u_*,v_*)$ irrespective of the initialization.\n\n- **4. Very long sequences.** The results of [Marion et al. 2024] were established on the very restrictive condition of short sequence lengths: $L=o(1)$, i.e $L/d \\to 0$. In contrast, our results for a vast range of sequence lengths $L$, upt o exponential in $d$, i.e $\\log L=O(d)$. Such long sequences capture the situation in real world NLP / LLM workloads (a desirable property of our theory, also noted by **Reviewer mLDi**), and marks an exponential increase compared to [Marion et al. 2024]. See **Remark 4** of our manuscript.\n\n### Reconsideration of the reviewer's score\n>We hope the above points provide the reviewer with a better understanding and appreciation of our contributions, and they increase the score accordingly. Currently, the reviewer scores our paper with a $1$ (**Strong Reject**), which is the **lowest possible score**, only reserved for papers with \"well-known results or unaddressed ethical considerations\" and **dismissive** and **very unfair**, which is not the case of our paper which makes a broad range of contributions to a fundamental problem. Thank you in advance for your understanding and consideration. We are happy to answer any questions the reviewer might have regarding our work.\n\n>Regarding the numerous typos which have obstructed the reviewer's understanding and appreciation of our work, it has been very difficult to keep track of everything as a **sole author**. I sincerely apologize. These typos will be carefully corrected. Thanks in advance for your understanding and consideration."}, {"review_id": "7fZbtHpQkz", "reviewer": "Reviewer_CnGc", "summary": "This paper focuses on the single-locater problem (a special regression problem) to investigate the learning dynamics of self-attention with (projected) gradient flow. The most important contributions of this paper are the closed-form descriptions for the population loss and the gradient flow dynamics, which employs techniques from statistical physics and mean-field theory.", "questions": "1. Could the authors give how the population risk explicitly depends on $t$?\n\n2. Could the authors explain why the stationary points would be eventually obtained by the dynamics?\n\n3. Please refer to the questions in the Weaknesses part.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n1. The results are generally novel and insightful to me, showing how techniques of statistical physics can be interestingly applied to solve problems regarding the learning dynamics of self-attention, though the problem being studied is a toy one. This might inspire future works to try to apply more analytical tools from physics to solve a wide range of machine learning and deep learning problems to further unveil the mystery of their success. \n\n2. The statement of the problem setting and the corresponding motivation are clear, which might be generalized to further study the mechanisms and advantages of self-attention over other models.\n\n## Weaknesses\n\n1. Unfortunately, I believe that the current manuscript is hard for readers to understand, due to many typos and insufficient discussion and explanation of concepts/definitions/notations/techniques involved. I would strongly recommend that the authors improve the readability of their manuscript so that the results can be more easily appreciated, and currently I struggled to understand the details. In particular,\n\n   - line 135-136, the authors directly present the mean-field description for the system, without any explanation on why it has to be these quantities. I would suggest the authors clearly explain the motivation for these definitions and why/how they are essential for the analysis of the system, as many readers in this venue in fact do not have a statistical physics background and it seems like that these definitions are coming from nowhere.\n\n   - line 157 \"index $\\ell_{\\*}\\in L$ is a random datapoint\", an index cannot be a datapoint; \"Lemma xyz\", where is this lemma? \n\n   - There are too many different definitions, notations, and propositions spreading out the paper, and it is unclear why and how we have to have these definitions, e.g., what is the exact definition of $\\simeq$ in the context of this manuscript? The readers must scroll back and forth constantly to find them when they suddenly appear in some propositions and corollaries, e.g., it is unclear why we have to define the auxiliary functions in Eq. (21) to Eq.(23).\n\n   - The explanation for techniques is far from sufficient, e.g., line 721 how can we directly write Eq. (48)? line 755 how do we have Eq. (55)? line 787-790 what are these auxiliary functions for?\n\n2. The population loss indeed explicitly depends on the overlaps defined in the main paper, but we still do not know how the population loss explicitly depends on $t$. As the considered problem is a toy one, I think it would be better to derive an exact solution for the population loss.\n\n3. It is not clear why the stationary points would be eventually obtained by the dynamics (line 265-266), as we only have the formulation of the flow dynamics and there is no discussion on convergence.", "was_revised": true, "final_justification": "The authors have addressed my concerns on the convergence to stationary points.\n\nOn the other hand, as noted by me and other reviewers, the paper is hard for readers to follow due to the presentation issue and many typos as discussed in my original reivew and replies. Hence, I recommend a borderline score.", "rebuttal": "## Response to Reviewer CnGc\nWe thank the reviewer for their feedback and insightful comments.\n\n### Main technical concerns\n>*The population loss indeed explicitly depends on the overlaps defined in the main paper, but we still do not know how the population loss explicitly depends on. As the considered problem is a toy one, I think it would be better to derive an exact solution for the population loss.*\n\nWe’re not sure we understand this comment from the reviewer. The population loss $R(u,v)$ is given derived explicitly in **Proposition 1**. The formula given is valid for all unit-vectors $u$ and $v$ (the student parameters), irrespective of whether these depend on time $t$ or not. We’re happy to follow up on this if the reviewer could clarify the question.\n\n>*It is not clear why the stationary points would be eventually obtained by the dynamics (line 265-266), as we only have the formulation of the flow dynamics and there is no discussion on convergence.*\n\nIndeed, strictly speaking our statement on **Line 265 -- 266** is referring to continuous time dynamics Eqn (30) and not the discrete-time dynamics (29). The former is established in **Proposition 4**. As long as we are not worried about speed of convergence (not a concern of our analysis), the conclusions of the proposition carry over to the discrete-time dynamics (29) as long as the step-size is sufficiently small. We shall modify **Line 265 -- 266** to make this clearer.\n\n>*The readers must scroll back and forth constantly to find them when they suddenly appear in some propositions and corollaries, e.g., it is unclear why we have to define the auxiliary functions in Eq. (21) to Eq.(23).*\n\nThe auxiliary functions defined in Eqn (21) -- (23) appear explicitly in some of our results (e.g., **Proposition 1**). Without those functions, the statement of those results would be rather cumbersome and unwieldy, making the manuscript even more difficult to read.\n\n>*line 787-790 what are these auxiliary functions for?*\n\nThe $T_k$ functions are used to streamline the proofs our our results on classification of stationary points of the gradient-flow dynamics (**Propositions 4** and **7**). See **Appendix F**. We shall add a sentence to make this more explicit.\n\n### Comments regarding statistical physics\n>*line 135-136, the authors directly present the mean-field description for the system, without any explanation on why it has to be these quantities. I would suggest the authors clearly explain the motivation for these definitions and why/how they are essential for the analysis of the system, as many readers in this venue in fact do not have a statistical physics background and it seems like that these definitions are coming from nowhere.*\n\nEqn (14) -- (18) Define statistical physical constants on which the phase transitions and mean-field approximations in our theory depend. We'll add a paragraph explaining the philosophy in words to ease understanding by a broader audience. However, note that one doesn't need to know any statistical physics to follow the actual statements of the results (propositions, etc.).\n\n>*The explanation for techniques is far from sufficient, e.g., line 721 how can we directly write Eq. (48)? line 755 how do we have Eq. (55)?*\n\nEqns (48) and (55) are the starting points for the so-called Laplace method which is a standard procedure for approximating sums by integrals and extremizing them [Lucibello and Mézard (2024)]. I'll add a section in the appendix to detail on this technique.\n\n### Notation and typos\n- The standard notation $a(n) \\simeq b(n)$ means $a(n)/b(n) \\to 1$ for large $n$. For example, $n + \\sqrt n \\simeq n$, etc. For random quantities, the notation means $a(n)/b(n) \\to 1$ in probability. We shall add a small notations paragraph in the introduction.\n- **Line 157.** \"index $\\ell_* \\in L$ is a random data point $X=(x_1,\\ldots,x_L)$\" should be \"index $\\ell_* \\in [L]$ in a random data point $X=(x_1,\\ldots,x_L)$\"\n- **Line 157.** \"Lemma xyz\" is a typo; indeed there is no Lemma xyz. Let us just proof the fact here in-line (this will be added to the manuscript): the probability of locating correctly picking the right index $\\ell_*$ is $p_{\\ell_*} = e^{\\lambda x_{\\ell_*}^\\top u}/Z$, where $Z := \\sum_\\ell e^{\\lambda x_\\ell^\\top u}$. From Eqn (13), we can further write\n$$\n\\begin{split}\np_{\\ell_*} &= \\frac{1}{1+e^{-\\beta\\sqrt d x_{\\ell_*}^\\top u}Z_{-\\ell_*}} = \\frac{1}{1+e^{-\\left(\\frac{x_{\\ell_*}^\\top u}{\\sqrt d} - \\frac{\\log Z_{-\\ell_*}}{\\beta d}\\right)\\beta d}}\\\\\n&\\simeq \\frac{1}{1+e^{(cu^\\top u_*-\\psi)\\beta d}}= \\frac{1}{1+e^{(c\\mu-\\psi)\\beta d}},\n\\end{split}\n$$where we have have used Eqn (14) to write $(\\beta d)^{-1}\\log Z_{-\\ell_*} \\simeq \\psi$ (this the first place where statistical  physics plays a role in our theory), and also $x_{\\ell_*}^\\top u/\\sqrt d = cu_*^\\top u + \\mathcal N(0,\\gamma^2/d) \\simeq cu_*^\\top u$ thanks to Eqn (4)."}, {"review_id": "DieK8VC7zK", "reviewer": "Reviewer_mLDi", "summary": "This paper develops an exact mean-field theory for a single-layer softmax self-attention network on the *single-location regression task*. Building on ideas from the random-energy model, they derive a closed-form expression for the population risk in the high-dimensional limit where the sequence length grows as $\\log L = \\mathcal{O}(d)$. They uncover a sharp phase transition that delineates learnable and non-learnable regions as a function of the feature-noise level $\\gamma$ and the length parameter $\\alpha$. Finally, for linear link functions they prove that Riemannian gradient descent converges from arbitrary initialization to the teacher parameters, which are shown to be the unique stable fixed point of the dynamics.", "questions": "- **Task correspondence.** Which concrete NLP tasks (e.g., keyword spotting, pointer networks, span extraction) map most closely to the single-location regression problem?\n- **Finite-sample effects.** The optimisation study uses population risk. How do you expect finite-data noise or mini-batch updates to alter the dynamics? Could replica or dynamical mean-field tools quantify these effects?\n- **Sphere constraint.** How crucial is the unit-sphere constraint on $u$ and $v$ ?  Could the analysis extend to unconstrained parameters with $l^{2}$ regularisation, and what obstacles would arise?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n- **Realistic modelling of attention**. Deriving the population risk for a softmax layer, rather than a decorrelated surrogate in previous study, makes the analysis markedly more linked to real Transformers.\n- **Scales to long contexts**. The regime $\\log L = \\mathcal{O}(d)$ captures the much longer sequences found in modern NLP workloads, which is more realistic than earlier work restricted to $L=o(d)$.\n- **Explicit learnability boundary**: The closed-form condition that separates solvable from unsolvable regions offers valuable theoretical guidance.\n- **Strong optimisation guarantee**: For $\\sigma(t)=t$, the proof that the teacher parameters form the sole attractor, without any special initialization, significantly strengthens prior results.\n\n## Weakness\n\n- **independence and Gaussianity assumption**. Real token embeddings are correlated and far from Gaussian. Although the paper moves the theory closer to practice, lifting this assumption remains an important next step.\n- **Dynamics only under Condition 3.** Convergence is analysed in the frozen regime ($r=0$). A more detailed discussion of why dropping Condition 3 is technically hard—and what partial results might still be feasible—would be helpful.", "was_revised": true, "final_justification": "My questions have been resolved through the rebuttal. I will maintain my current rating of “accept.”", "rebuttal": "## Response to Reviewer mLDi\nWe thank the reviewer for their insightful feedback and questions.\n\n>***Task correspondence.** Which concrete NLP tasks (e.g., keyword spotting, pointer networks, span extraction) map most closely to the single-location regression problem?*\n\nAll the examples mentioned by the reviewer are captured by the theoretical setup considered in our work. Let us also mention: in-context learning, topic understanding, sentiment analysis, etc. as they all have some kind of , token-wise sparsity and internal linear representation.\n\n>***Finite-sample effects.** The optimisation study uses population risk. How do you expect finite-data noise or mini-batch updates to alter the dynamics? Could replica or dynamical mean-field tools quantify these effects?*\n\nThe reviewer's comment is spot on. Indeed, as mentioned in the **Concluding Remarks**, a future extension of our work is to study finite-sample (S)GD. Pushing our mean-field approach further is definitely a good candidate for such an analysis, even though the results will be considerably weaker (no surprise). What will be crucial in such an analysis is the relative scaling of the sample size $n$, the sequence length $L$, and the input dimension $d$. This is ongoing work. \n\n>***Sphere constraint.** How crucial is the unit-sphere constraint on $u$ and $v$? Could the analysis extend to unconstrained parameters with $L_2$ regularisation, and what obstacles would arise?*\n\nThe unit-norm constraint on the student parameters $u$ and $v$ plays a number of roles.\n\n- The single-locator problem is such that what really matters is the alignment between the student parameters $(u,v)$ with the teacher parameters $(u_*,v_*)$, and not the norm. The temperature parameter $\\lambda$ is of order $\\sqrt d$, so that $\\|u\\|$ (for example), doesn't relly matter as long as it is of order $\\Theta(1)$.\n\n- The main ingredients for our analysis, namely **Propositions 11** and **12** rely on $\\|u\\|=1$ so that the random energy levels $E_i = \\lambda x_i^\\top u$ are of order $\\sqrt d$ as in the classical random energy model [Derrida 1982], accorss the entire optimization trajectory. Without this constraint one would have to rescale the temperature and keep tract of it. Unfornately, $L_2$-regularization can't accomplish any of this."}]

Author rebuttal:

Meta-review: This paper presents a rigorous theoretical analysis of a single-layer softmax self-attention network. By leveraging tools from statistical physics, specifically the random energy model, the authors derive an exact mean-field theory for the "single-location regression task." This allows them to obtain a closed-form expression for the population risk and precisely characterize the learning dynamics of gradient descent. A key contribution is proving the global convergence to the optimal solution for linear link functions without requiring special initialization unlike prior work. The analysis also extends to long-sequence regimes, making it more relevant to modern NLP applications.

The application of statistical physics methods to derive an exact mean-field theory for a softmax attention layer is a novel and significant contribution that goes beyond the simplified pointwise activations studied in previous work. The mathematical treatment is rigorous and provides some insights into the model's behavior. Moreover, the paper establishes a clear phase transition for learnability and a global convergence guarantee for the linear case, which brings a valuable insight to the community. 

Although the reviewers acknowledged these contributions, they raised concerns on its presentation quality including specialized notation, several typos, and insufficient explanations to enhance intuition for non-expert readers. I think the authors should address these issued before it will be published and the paper would benefit from a careful revision to improve its accessibility. However, its core technical contributions are novel and significant in the community. Therefore, I recommend acceptance.
Decision: Accept (poster)