Statistical Advantage of Softmax Attention: Insights from Single-Location Regression
Authors: O Duranthon
Pierre Marion
Claire Boyer
Bruno Loureiro
Lenka Zdeborová
Venue: ICLR 2026 Poster (2026)
Keywords: attention
softmax attention
deep learning theory
sparse information retrieval
high-dimensional limit
high-dimensional statistics
Primary area: learning theory
TL;DR: We demonstrate advantages of softmax-based attention over alternatives in a high-dimensional information retrieval regression task.

Abstract: Large language models rely on attention mechanisms with a softmax activation. Yet the dominance of softmax over alternatives (e.g., component-wise or linear) remains poorly understood, and many theoretical works have focused on the easier-to-analyze linearized attention. In this work, we address this gap through a principled study of the single-location regression task, where the output depends on a linear transformation of a single input token at a random location. Building on ideas from statistical physics, we develop an analysis of attention-based predictors in the high-dimensional limit, where generalization performance is captured by a small set of order parameters. At the population level, we show that softmax achieves the Bayes risk, whereas linear attention fundamentally falls short. We then examine other activation functions to identify which properties are necessary for optimal performance. Finally, we analyze the finite-sample regime: we provide an asymptotic characterization of the test error and show that, while softmax is no longer Bayes-optimal, it consistently outperforms linear attention. We discuss the connection with optimization by gradient-based algorithms.

Published as a conference paper at ICLR 2026
STATISTICAL ADVANTAGE OF SOFTMAX ATTENTION: INSIGHTS FROM SINGLE -LOEE-ON REGRESSION
O. Duranthon & L. Zdeborová
P. Marion
Statistical physics of computation laboratory École polytechnique fédérale de Lausanne (EPFL)
Inria, École Normale Supéricure, PSL Research University
C. Boyer
B. Loureiro
Laboratoire de Mathématiques d' O'ssy Université Paris Saclay and Institut Universitaire de France
Departement d' Informatique
École Normale Supérieure, PSL & CNRS
ABSTRACT
Large language models rely on attention mechanisms with softmax activation. Yet the dominance of softmax over alternatives (e.g., component-wise or lin- ear) remains poorly understood, and many theoretical works have focused the linearized attention. In this work, we address this gap through a principled study of the single-location regression task, where the output depends on a linear transformation of a single input token at a random location. Building ideas from statistical physics, we develop an analysis of attention-base predic- tors in the high-dimensional limit, where generalization performance is captured by a small set of order parameters. At the population level, we show that soft- max achieves the Bayes risk, whereas linear attention fundamentally falls short. We then examine other activation functions to identify which properties are nec- essary for optimal performance. Finally, we analyze the finite-sample regime: we provide an asymptotic characterization of the test error and show that, while soft- max is no longer Bayes-optimal, it consistently outperforms linear attention. We discuss the connection with optimization by gradient-base algorithms.
INTRODUCTION
Large language models (LLMs) have recently reshaped natural language processing, enabling appli- cations ranging from conversational agents and code generation knowledge-intensive reasoning. At the heart of these models lies the Transformer architecture (Vaswani et al., 2017), where the use of softmax in attention layers has proven remarkably effective. Despite its dominance, it has compu- tational drawbacks due to its quadratic complexity in the sequence length, while being theoretically arduous study due to the softmax normalization that couples tokens complex manner.
For these reasons, numerous alternatives have been proposed. Notably, kernelized attention approx- imates the softmax function with kernel feature maps (Wang et al., 2020; Luo et al., 2021; Choro- manski et al., 2021; Qin et al., 2022), achieving linear complexity in sequence length. In parallel, state-space models (SSMs, Peng et al., 2023; Poli et al., 2023; Gu Dao, 2024) introduce linear re- current dynamics with gating to tackle long-context information retrieval. On the other hand, while less studied empirically, linear attention, which is the element-wise linearization of softmax around the origin, has been extensively studied in theoretical works (e.g. Ahn et al., 2023; von Oswald et al., 2023; Bai et al., 2023; Mahankali et al.; Zhang et al., 2024; Lu et al., 2025; Zhang et al., 2025).
Despite substantial research alternative attention mechanisms, softmax attention remains the dominant architecture in large-scale language models. A leading hypothesis for this empirical suc- cess is its superior performance on retrieval tasks, view supported by growing body of work (Arora et al., 2024; Aksenov el al., 2024; Chou et al., 2024; Hsieh et al., 2024; Shen et al., 2024; Wang et al., 2025). While we defer detailed discussion Section 2, we briefly summarize the rep- resentative study of Shen et al. (2024). These authors characterize optimal scaling laws for language
Part of this work was done while the author was at the Institute of Mathematies, EPFL.

---

Published as a conference paper at ICLR 2026
100
1.0
Pucrage Archilectuure
60
0.8
Assuneyy
HORN2
Narmel and metri
6052
CSR
0.4
ppl
App
SCROLLS LS weg
Archilecture
20
0.2
LLMMA
TNL
HGR N2
0.0
coe2
70M 160M
1M
70M 163M
410M
13
3B
7B
Number paramalers
Number of paramater ataar
(a) Retrieval tasks Hacyste-in-a-aayttack, Kamradt, 2023; Machlab & Battle, 2024). Colors correspond to various context lengths (longer context makes the task harder). The "average" curve is the average over the three tested context lengths.
(b) Linguistic proficiency tasks (accuracy on Com- Sense Reasoning and SCROLLS, perplexity WIKITEXT-2 and LAMBADA datasets). All metrics are normalized so that higher is better, and 1 (resp. 0) corresponds to the best (resp. worse) performance over all tested configurations.
Figure (Shen et al., 2024, Tables 4 and 11): Comparison of softmax Transformer (LLaMA, bolded line) with kernelized attention (TNL, cos2) and state-space models (HCRN2), function of the model size, and for various tasks (retrieval tasks on the left, linguistic proficiency on the right). All architectures have similar performance for the linguistic proficiency tasks, whereas in retrieval tasks the softmax attention systematically outperforms alternatives.
models equipped with kernelized attention and state-space models (SSMs), and then train softmax- attention, kernelized-attention, and SSM-based models with matched parameter budgets ranging from 70M to 7B. Their downstream evaluation, reproduced for convenience in Figure 1, reveals striking pattern: although optimally-scaled SSMs and kernelized-attene models achieve com- petitive performance on linguistic proficiency benchmarks, they consistently underperform softmax attention on retrieval tasks.
Yet the underlying reasons for the advantage softmax in retrieval tasks remain poorly understood. This gap limits our ability to reason about the tradeoffs and fundamental constraints of different model architectures.
Contributions and organization. In this paper, we take a principled approach to the question, with the following contributions. Code reproduce our simulations is provided.
We propose a mathematical formalization of tasks where the output depends on a single input token as the single-location regression model (Section 3). This model subsumes previous theoretical studies (Marion al., 2025), notably by introducing random sequence lengths, and gives formalization of information retrieval tasks.
By combining analytical and numerical results, we provide analysis of the performance of various attention layers, in particular softmax and linear, in increasingly involved settings (ap- proximation, statistical, computational). The analysis remains tractable despite the presence of the softmax nonlinearity by leveraging ideas from the sequence multi-index models (Cui al., 2024; Arnaboldi al., 2025; Cui, 2025; Troiani al., 2025), where properly scaled random variables concentrate in the high-dimensional limit and the learning behavior can then be characterized by small set of so-called order parameters.
We prove a gap in approximation performance between linear and softmax attention (Section 4), with the latter reaching Bayes error. We argue that this analysis captures the performance efficiently attainable with one-pass stochastic gradient descent. This gap arises from both the exponential nonlinearity and the normalization in softmax, illustrated by our comparison with kernelized and element-wise attention (Fig. 2).
We characterize the regularized empirical risk minimizer (ERM), in the high- figh-dimensional limit with sample complexity linearly proportional the dimension, the solution to self- consistent equations (Section 5). By solving these equations, we show that the advantage of softmax over linear attention still holds, and we benchmark against the optimal test risk fi-
2

---

Published as a conference paper at ICLR 2026
nite sample complexity, which we will call Bayes-optimal performance. We show numerically that the predicted ERM is achieved by gradient- based optimization algorithms (see Fig. 3).
2 FURTHER RELATED WORK
Synthetic information retrieval tasks. Notable tasks related to single-location include Needle- in-a-Haystack (NIAH) (Kamradt, 2023; Machlab & Battle, 2024), where the goal is to retrieve fact ("needle") given in sentence inserted within a large block of unrelated text ("haystack").. An abstraction of NIAH the Associative Recall (AR) (Graves al., 2014; Ba et al., 2016) task. In this task, the input contains a sequence of bigrams representing key-value pairs from random dictionary followed by a query token. For example, the correct output for the input A 2 8 C A - B? is 8. This task has been extensively studied through the lens of information retrieval circuits in Transform- ers, specifically induction heads (Olsson et al., 2022). AR extended in the Multi- Multi-Query Associated Recall (MQAR) task (Arora et al., 2024), featuring several queries for each input sequence.
Alternatives to softmax attention. Two main alternative directions feature a linear complex- ity in the sequence length. First, in kerneiized attention (often referred to linear attention in the literature), the original attention function Att(Q, K) Au(Q, K) = ()))){(KT, where is a kernel function (see, among others, Wang et al., 2020; Choromanski et al., 2021; Qin al., 2022). More recently state-space models (SSMs) were specif- ically introduced to handle information retrieval in long contexts (Fu et al., 2023; Peng et al., 2023; Poli et al., 2023; Gu & Dao, 2024). They consist in alternating linear recurrent neural networks (which can be seen as Id-convolutions along the sequence time) with nonlinear operations applied in parallel over each element of the sequence (sometimes referred to as gating). We refer to Arora et al. (2024) for a common mathematical framework encompassing many variants of SSMs.
(ottmax(QK" T /vd) is replaced by
While SSMs outperform attention -based models in the specific synthetic task of AR thanks to their capacity to handle long contexts, this success is brittle since they lag behind Transformers on the slightly more involved MQAR task (Arora et al., 2024; Aksenov et al., 2024; Chou et al., 2024; Wang et al., 2025). Arora et al. (2024) propose a theoretical explanation based on the lack of expressivity of convolution-gatiin models. Transformers also outperform SSMs and kernelized attention in retrieval tasks based on linguistic data (Hsich et al., 2024; Shen et al., 2024). In this work, we take a principled approach to understanding the henefits of softmax, by identifying an information retrieval framework amenable to theoretical study, and going beyond expressivity results statistical and computational advantages.
Methodology for sequence multi-index models. While our data model is inspired by Marion et al. (2025), it departs significantly by generalizing to variable sequence length, introducing a generic weighting mechanism to encode single location, and renormalizing to obtain proper high- dimensional limits, thereby allowing theoretical analysis of the softmax nonlinearity. The case of diverging signal strength was concurrently studied in Dohmatob (2025). The data model belongs among the sequence multi-index models as introduced in Cui (2025). Special cases of single-index models we studied in Cui al. (2024); Arnaboldi et al. (2025). Our model corresponds to a sequence two-index case, for which the Bayes-optimal estimator lated classification problem where outputs only depend on a few tokens was concurrently studied in Barnfield et al. (2026) with analysis of steps-wise gradient descent training.
studied in Troiani et al. (2025). A
Last, the concurent work Dragutinović et al. (2025) also discusses advantages of softmax over linear attention, but on a different model (in context classification) and proof techniques.
3
TASK AND DATA MODEL
3.1 OVERVIEW OF THE SINGLE-LOCATION REGRESSION TASK
We consider a sequence regression task where the input is a sequence X E LRD of L tokens, each of dimension D. The length L of the sequence is allowed to vary, remaining upper bounded by L > O. Each sequence is labeled by a scalar y E R, with the particularity that it depends only a single input token. We model this single-location dependency by setting hidden index e* E {1....,L} selecting the relevant token of the input sequence X. We emphasize that the latent index ex is samplelconexx--eeeeedenn, hence its recovery is akin to toy in-context learning task.
3

---

Published as a conference paper at ICLR 2026
Without additional structure, there is little hope retrieve the relevant token. We explore two ways to add such structure, which both involve learning a hidden direction k E RD to recover the relevant token. In spiked single-locatin regression (spiked-SLR), spike is introduced in the direction of k' at the relevant token x A closely related scheme, referred to as max-SI.R (for maximal-correlatin SLR), consists in setting e" the index of the token with the largest scalar product with k".
Our goal is to theoretically study the learning properties of such tasks. For this purpose, we introduce probabilistic data model that amenable analysis in the high-dimensional limit. This framework also unifies the two variants under consideration (spiked-SLR and max-SLR), as described next.
3.2 PROBARILISTIC DATA MODEL FOR SLR
Denote by N(w, V) Gaussian law centered at w with covariance V. and by VV w. V) its density evaluated at æ. The probabilistic data model begins by drawing two hidden directions
k ID) e" ~ ~,,
Then, each sample (X,y) is drawn as follows. The length I of the sequence is drawn from some discrete law Pt over {I....,L} L}. such as uniform or truncated Poisson, while the index e" of the relevant token is taken, conditionally on L, uniformly at random over {I...., LJ. The label is then (potentially random) function of the projection of the relevant token along v*, that is,
1 X o" + +EE
y
where & is independent standard Gaussian noise and A O. It remains discuss the law of the sequence X. A flexible framework consists in taking X as a reweighted Gaussian distribution. More precisely, conditionally on L and kt, the density of X at a point RLxD is given by
L. P(x|L,e k²) =g xx) IIIII 110;0, where x"
-nnk" ERL vD
(3)
Here, 9r fl..... L} x RRL R is a weight function, indexed by the signal strength v 2 0. Its goal is to give more weight to sequences with a large projection Xc- " k". Both the examples from Section 3.1 fall into this framework for specific values of gu-
SPR corresponds to (e,e,, x) eVPx--tt Factorizing the density of w. shows that this definition is equivalent to shifting the mean of X.- by vuk", while the other tokens are centered, as studied by Marion et (2025).
Maximum-correlation SLR (max-SLR) is a special case of the sequence multi-index model of Cui (2025); Troiani al. (2025); Arnaboldi et al. (2025), obtained with (e, x) Le"X, Σt exxx. To gain intuition, note that by Bayes' rule, this model is equivalent to first drawing the tokens as independent standard Gaussians, then picking randomly with logits proportional to the scalar product of the tokens with k", that is, P(c" = i|X) "",, /E,e",, e"X,k" When + c0, corresponds to the index of the token with the largest scalar product with k"
Other weight functions can be considered as long as they are invariant by label permutation, are prop- erly normalized so that fgo dx P(x: T., 6", k*) = 1, and at zero signal are uniform i.e. go(e, x) = 1.
3.3 LEARNING WITH ATTENTION
Since the input sequence may have arbitrary length, we adopt the formalism of working with (RD)). N the set of sequences in R that are eventually zero. We consider the class of estimators F, = (ff.k.u (I. Byy + (R.))))((,)) where for certain activation function or R + E and two
vectors k U E TR. D. the function fa,k.v is defined by
_ == L=_xU _X VD
1
fx), - ==
(4)
We focus on four cases for the activation function a. Note that we always assume that a returns on the vanishing part of the input sequence, so we only define its value on the non-vanishing part.
4

---

Published as a conference paper at ICLR 2026
Table 1: Terminology and notations for the risks. the text, we use the terms "risk" and "error" in an interchangeable manner.
&ayyes En (k, v)
Bayes risk (in population)
(Bo
Baycs-optimal risk (empirical)
population risk of the attention over (R²)² population risk parameterized over IR' R
C(k, v)
regularized empirical risk (loss)
&g (k,,,
test risk
Es
minimal population risk
(a()
minimal test risk
Softmax activation corresponds the standard choice in current large language models. For an input of length L, the softmax writes for l (1,..., L.) as a (x)) = ext /E==
Linear activation writes a (x)e = + Xe- We add a constant term to break the symmetry around X = 0, as further discussed in Section 4.2. This corresponds to linearizing the softmax around up to rescaling. Notice that the activation with the constant term is strictly more expressive than the identity (x) = X. which can be retrieved by taking large Xe and small že at constant Xe#e.
For element-wise sigmoidal non-lincarity, we focus on o(x)e
+ erf(c xe)
v retxx dxe- 20" that varies from 2, with E R. a learnable bias. While for concision we omit this parameter in the mathematical presentation below, is learned in the numerical exper- iments. Here also add constant term 1 to break the symmetry around X = O.
Finally, we investigate softplus kernelized attention o(x)= = x(xx)// Ev p(xe) in the particular case of the kernel (ax) softplus(2>) = log(1 + c*).
A limitation of our setting, compared to practical models, is the absence learnable query vectors, which not necessary here due to the structure of the task. Still, we note that the output of a standard attention layer for a so-called [CLS] query token exactly corresponds to our model (see Marion et al. (2025) for details). The absence of query vector also means that, in our setting, softmax attention is special case of kernelized attention by taking = exp. Our analysis shall still give insightful results on the choice of nonlinearity by comparing softmax the softplus kernel.
COMPARISON OF SOFTMAX WITH ALTERNATIVES IN POPULATION RISK
In this section, we compare the expressivity of the softmax attention with alternatives on our task, by theoretically assessing the Bayes risk and the minimal population risk for various activation functions. We show that the softmax reaches the Bayes risk. We combine analytical arguments with numerical evidence to support the claim that this approach also captures the performance that efficiently achievable by running one-pass SGD, which directly minimizes the population risk from random initialization. All proofs of the results in this section are deferred to Appendix A. A summary of the terminology and the notations for the risks is given in Table left.
Remark 4.1. We consider in this section the case A in the output channel (2). Because of independence of the output noise with the other random variables, results would be identical for a posilive A, up to increasing all errors by 2 corresponding the irreducible noise.
4.1 BAYES RISK AND OPTIMALITY OF SOFTMAX
In what follows, we assess the performances of the different architectures with respect to the Bayes risk (or Bayes error) Bayes = [y - E(y] k., })}}., i.e. the best achievable risk given the task defined in Section 3.2.
Proposition 4.1. Let L ~ Pt, and, conditionally on L, e Unif((L..... and x ~ N((0 Then the Bayes risk is given by
Eayee EL Et (x
(5)
Two particular cases are of interest. First, if the function gu is identically equal to 1, then the position of the relevant token uniformly distributed over 1,..., and independent of the tokens X. In this case, the Bayes error equals 1 E, 1/LL > 0. This reflects the irreducible noise stemming from the randomness of the informative token. Conversely, if (e {) = 1 . the first token always carries the information. We then return standard noiseless regression, where the Bayes error is null.
5

---

Published as a conference paper at ICLR 2026
We turn attention to the best reachable theoretical risk over the class of estimators Fo when is the softmax function. Surprisingly, this estimator reaches the Bayes error as shown next.
Proposition 4.2. Assume that, for all L 2 1, (c,c)) fi...., L} and X R24.
gu(e, g
(6)
gu(c,,x)
for some constant Cu 0. Then, for any k", ,²" RRd.
min fa. eaottmax
Ex fn (x))())))) Bayee
Note that the requirement On defining the distribution of X holds in particular for the spiked-SI.R and the max-SLR. In consequence, the softmax architectures Fsottmax reach the Bayes error both settings, and furthermore the proof reveals that the minimum is reached for & = c, and t v", corresponding to recovery of the hidden directions. In statistical physics terminology, the softmax attention is said to satisfy the Nishimori condition, detailed in the proof. This is to be contrasted with the performance of other activation functions, which is characterized next.
4.2 EXPRESSION OF THE POPULATION RISK FOR ARBITRARY ACTIVATION FUNCTIONS
We now characterize the theoretical population risk of the attention & (k: v) BB Sa,kk (X)))].
Proposition 4.3. The population risk of the attention (k, v) + (k, v) can reparametrized as function of the following 7 variables, referred to order parameters:
mak- D
TT k-"
TTTT
Mkk-
kT
mock 0
ks*
D
D
D
Qkkk D
1
k T
1 Gor D
Yok
1 D)
k
(8)
The first two order parameters, mkk and Myr quantify the recovery of the hidden directions, while 9kk and qve are the squared norm of the parameters &o and v. Finally, the cross-correlation
terms Mive,, Mak", and quk nuisance parameters. We let E, R² - IR be the reparametrized risk depending on the order parameters.
Optimization dynamics and manifold assumption. In the following, consider the case where these three nuisance parameters are null, that is, we restrain our analysis to the manifold
M = (k, v) E (RP)², mike = m = qok = 0}.
The intuition behind this restriction is that, in the SLR, the signals k* and of associated to the keys and the values are drawn independent: ; and to obtain good performances, the keys and the values should not mix and should focus 1= on 8* separately. This simplification is supported by the following observations pertaining to the landscape of the population risk & In practice, we do not have access minimizers of the risk, but can instead optimize parameters and by one-pass SGD on E., which corresponds on average to running gradient descent (GD) on the reparametrized risk E., In the high-dimensional limit D + oc, random Gaussian initialization of the parameters lands on M because the cros-correlations vanish. Then, the manifold is invariant by GD, in the sense that iterates initialized on the manifold remain on it (see Appendix A.2 for details on this statement and the following). At finite D, random initialization lands in neighborhood of M. In this case, numerically check that M is stable in the sense that parameters stay close throughout GD. For linear attention, we also show that risk has only one (local) minimizer on M, which also local minimizer of the risk over the whole parameter space, while for softmax attention we know that global minimizer is M. Taken together, these facts support that GD &, and thus one-pass SGD on &, converges to minimizer the risk on M. A formal proof is delicate, as the landscape features other local minima outside of the manifold, as shown in Appendix A.2.6. This is left for future work.
With small abuse of notation, let E, (mis+, Mov+ Rgk, Ryv) be the reparametrized risk stricted to M, with R2g 9ы* -ma-- and R², quo - m the orthogonal components. It turns that this object is amenable insightful theoretical analysis. We first write explicitly.

---

Published as a conference paper at ICLR 2026
Proposition 4.4. Let L PL and, conditionally L, <~ Unif(f1,..... L}), x IL) and & I~)(O Then, for mike; Myv+ REK. Rv E R..,
Eo (mER, M Rkks Roo) R. X. (c, 2m o (mpx. RRAE)s + (m + R (o X + RackE) R ommax + RKAE))
(9)
A consequence is the characterization of the activation functions that allow easy learning, in the sense that the gradients at initialization with respect to the recovery order parameters Misk" and mov- are nonzero. If this quantity vanishes, moving away from the initial condition harder and requires more iterations and thus samples (Ben Arous et al., 2021). At initialization, mgg+ and m concentrate around 0 while R&E and Ryn concentrate around 1. At the first order, the gradient at initialization thus equals (E, 0, 1). The next result characterizes when this quantity vanishes. In particular, a (x)
and o exf(x) impede learning, which is why we add a constant bias. Corollary 4.1. the activation function a has a non-vanishing mean, i.e. if 1 N + 0 for all E, then
(omer Eor Om...
(0.0))
(10)
If the activation function a is symmetric, iee if [o(-))) = -o(x)) -o)))) all & and all IR', then (Omas En, 3m (0, 0, 1,) (0,
(11)
We next state that on M the risk of the linear attention admits unique minimum, which together with Corollary 4.1 guarantees that GD initialized on M will converge to it.
Corollary 4.2. The population risk &o of the linear attention o(x)e = Xe admits a unique minimizer over M. It is moreover a (local) minimizer over the whole space (TR)²
4.3 COMPARISON BETWEEN LINEAR AND SOFTMAX ATTENTIONS
Letting E,, miuM &, = mingg & be the minimum of the population risk, we now compare the minimal population risk Eim for linear attention (xx) 1 + Xe to the one for softmax Esoftmax Bayes encompassing asymptotic regimes of strong signals v - oo or long sequences L + 30.
Corollary 4.3. Consider the spiked SLR model with input sequences of deterministic length L. The minimal risks attained over the manifold M satisfy
L+ (+ Ejin = I 1 L-19 while
1
Egoftmax e CCL++O(V) as V - +00 (12)
with CL > 0 a constant that depends on L.. Consider then the max-SLR model at v - +o0 where gu(c, L1 e=arg max, with deterministic L. The softmax attention is well-specified while the linear attention is not: the minimal risks attained over M satisfy
log L
Em = 11-0
as L - while
Egoftmax = 0.
(13)
The first part of the corollary shows that on the spiked-SIR for strong signal v - +00 the risk vanishes for both attentions. However, the softmax model has better dependence on v, thereby establishing its superiority over linear attention in this setting. Turning to the dependency on the sequence length L, for the max-SLR, the second part Corollary 4.3 shows a clear separation in the performance of the linear and the softmax attentions: as the sequence length increases the error of the linear attention converges to 1, which the error of the trivial null predictor, while the softmax attention reaches perfect prediction for all L.
The following corollary illustrates the impact of varying sequence lengths on the linear attention.
Corollary 4.4. Consider the max-SLR model at too where go (e, x) Lilemarg max xe: Let f(L) (x--)( max Xt. Then, the minimal risk attained over M by linear attention satisfies
1 (ELF(L)) 21 E,
1 1+ (f)
Em
(14)
B
Consider then SLR at v = 0 (i.e. gu x) = 1j. We obtain that Ein = 1 - 1/E,, and Esaftmax = 1 (1((1/T) and consequently Elin Exoftmaxx with equality when Var T. 0.

---

Published as a conference paper at ICLR 2026
The first part of the result reveals that the variance in the sequence length L hurts linear attention, since the risk for some distribution of L is always worse than the risk associated to the mean value of this distribution. This illustrates fundamental limitation of linear attention, arising from its poor normalization properties. In Fig. 2. we indeed observe that the performance gap between softmax and linear attention widens when I. ~ Unif(1, 2) compared to I. = 2, everything else being fixed. We characterize this phenomenon in the second part of Corollary 4.4, for the case of null signal O, where the position of the relevant token independent from the law of the tokens. Here again softmax performs better than linear attention whenever L admits some variance.
L-2
L U.il((1,233))
10
liwarr
a= c=1+ef
Spiled-SS.R
mnoemdiied aulgles d=sottmax
0.6
Bayne risk
0.2
0.0
5
1.0
1.5
10
15
0
5
10
15
0.8
maa-SLa
0.6
0.4
0.2
0.0
10
Figure 2: Minimal population risk E. over Fo for different attention activations a (colors), com- pared to the Bayes risk Bayes (5) (black), for the tasks spiked-SL. (top) and max-SI. (bot- tom). Softmax is the only one achieving the Bayes risk. The markers on the lines are for readability only. Population risks are computed via numerical optimization of (9). In all cases, we found that REa = Rvv = 0 was optimal, i.c. k exactly aligns with h- and t with v².
4.4 OTHER ACTIVATION FUNCTIONS
We now turn to two non- non-linear activation functions, element-wise erf and normalized softplus (see Section 3.3). Contrarily to linear attention, the expression (9) cannot be analytically minimized. Instead, we resort to numerical optimization of this expression, and report results in Fig. 2. Over- all, we observe that the population risks for these two activation functions are between linear and softmax. In particular, on the max-SLR model the two activation functions are not well-specified. Importantly, the element-wise function suffers from variable sequence lengths while the normalized softplus does not, as can be seen by comparing results for L = 2 and L Unif (1, 2, 2.3. This highlights the importance for the activation function to perform a normalization operation involving all the tokens. Furthermore, the gap between normalized softplus and softmax widens for larger I.. This can be expected since the softplus does not grow fast enough at +00 to dominate the noise stemming from the irrelevant tokens. This shows that, for kernelized attention, the kernel has to be well tuned to reach good performances, as is commonly known (see, e.g., Aksenov et al., 2024).
5 PERFORMANCE OF THE ATTENTION AT FINITE SAMPLE COMPLEXITY
So far, our analysis has focused on the properties of the minimizer of the population risk. In practice, however, we only have access to a finite dataset D ((X,,- y,) R ((()) : μa E [N]]. and must therefore rely on an empirical estimator. In this section, we consider the regularized empirical risk (or training loss) defined for some regularization strengths rs 2 and r, 2 by
C(k,v) Σ (nn (X,))
rk Σ
k
To
2
(15)
2
8

---

Published as a conference paper at ICLR 2026
Our goal is to assess the performance of an empirical risk minimizer (ERM) (k. E arg min C. Our approach here is thus statistical in nature: it quantifies the impact of the finite sample size on the performance of attention. We defer computational questions (i.e. how to practically minimize L) to the end of the section. The performance of an estimator (k: 6) is evaluated through the test risk (or test error)
(k, E [v [[v (x) 0*]
(16)
We investigate the high-dimensional proportional limit where N, - oo at finite N/D + a = ((1). In this part we take a fixed w.r.t N and D; letting a - oc one would recover the population risk discussed in Section 4. For simplicity, we consider the noiseless task A = 0 in the following. A summary of the terminology and the notations for the empirical risks is given in Table 1 right.
While in Section 4 the benchmark was the Bayes risk, at finite sample complexity it is natural to consider instead the best achievable test risk conditionally on a finite batch of data D, also known as the Bayes-optimal risk (or Bayes-optimal error). In the proportional high- sigh-dimensional limit, the Bayes-optimal error for the SLR task can be computed by extending the results of Troiani et al. (2025) to random sequence lengths I, ~ Pr. and arbitrary Iu- Although this result is of independent interest, for conciseness we refer the reader to Appendix C for the details. The outcome is that the Bayes-optimal error presents a rich phenomenology, with hard phase where the best-known first-order method fails to achieve the iioommttionthhoretiice performance, see Section B for a full discussion.
The performance (16) of an ERM is a random quantity that depends on the draw of k", 2² and of the dataset. Our main result in this section is that, in the proportional high-dimensional limit and under a concentration assumption known as the replica symmetry, see e.g. Vilucchio et al. (2025), this random variable converges a deterministic quantity, which can be fully characterized in terms of a few real-valued variables that follow a self-consistent equation. An analogous result for a class of sequence multi-index models was derived in Cui (2025). However, the (2025) did not provide any specific insight about the behavior of the generality of the result in Cui regression model studied here. We instead provide a numerical evaluation of this characterization for our setting.
Result 5.1 (Test risk for attention- bbased predictors). In the proportional high-dimensional limit N, - oo with a = N/D- = (()), under the replica symmetry assumption, the (non-rigorowe but standard) replica method predicts that the test risk (16) of a global minimizer (k, i) of the empirical risk (15) converges to a deterministic quantity
(k, 8) Es(a)
=
min
ET (o o())) w)²
(17)
(yn m. Io V.leS
where 7 =mxx== + V - m2s E R = m.y + - máZ and, for E > 1, We = sa
wi
The first expectation is over L ~ PL and the second conditionally on L over E, ~ IL) and y ~ N(O, 1) independentl. Finally, the set c IR. is the set of fixed points of the following iterative self-consistent equations
m 1+1 (rx + V/'='
d = (Fx + 02) (()h) + Gi)
veti (Fx + 1/2)-
(18)
my
[+1 (ro + 1
41 = (r.+ V.) (()t))? + +
v+41 = (r. + V/--1
(19)
m m
(V/)'** (xxxl mmi - -
= DEL.E,..ux' Dr (1,x')
(v:) (yzl - -m((Vs) Cov(z1))
(20)
(11))=2 (x)))))* (V.) -m)})
%
aE Eer.ivex 9o (1,x')
(21)
0.
(V2)) (V.)-1
(VV) Cov(xe)
LL
&E, go x")
(V.)) Cov(zz)
(22)
where we define the potential vonnt over (R its extremizers and its covariances as:
1.
Vont (x,2) = 2 -0x0 2)² +
Σ log Ixeen V1) + Σ log (ze) V.)
(23)
x', e = arg max vious (xo E (1=²)²
(24)
X.#

---

Published as a conference paper at ICLR 2026
Cov (xv) ( Mown (X' (11
E IR
Cov (2x) ( (x(,))) E ER. (25)
The derivation of this result is given in Appendix C. We simplified it assuming diagonal order parameters, which is analogous to the manifold assumption made in Section 4. Taking the limit + oC we checked that we recover the expression of the population risk of Proposition 4.4, and in particular aim (a) corresponds to the value of E, derived in Section 4.3 for the population risk. While our derivation based on the non-rigorous replica method, we expect that a formal proof could be established under the so-called replicon condition along the lines of recent progress in proof techniques from Vilucchio et al. (2025) that is able to deal with minimization of intrinsically non-convex objectives for single-index models. Extending this proof to multi-index models, such as (15), is technical challenge left for future work.
As in the population result of Proposition 4.3, the high-dimensional analysis shows that the risk depends only on a few order parameters. To give some intuition, they can be interpreted in the following way. my and my are the alignments of l and e with h- and v", respectively. Qx and Qu are the squared norms of k and v. V and V,, are related to the local curvature of the empirical risk L: they encode the fluctuations due to the finite sampling ratio and when a - oo we have VR, V, - O. Last, in the self-consistent equations, the potential touu plays the role of an effective low- www-dimesiional empirical risk. Although the resulting expressions are cumbersome, it is important to emphasize that these quantities are deterministic and independent of the diverging dimensions N and D, in contrast to the original risk (15). Thus the result provides an implementable formula for the performance of ERMs. In practice, the minimum over the set S is computed by running the fixed point method from several initializations, typically random (so-called uninformed) and informed ones corresponding to (partial) alignment of and with k- and v².
spiked-S1 -L,, p=1
Spiked-SLRR w
max-SLR, y +.00
lineat
0.8
1:00
soft max
0.6
tin
Bayesroptimal
0.2
2
0 Figure 3: Minimal test risk of the attention (linear vs. softmax) across different tasks and signal strengths v, for L. = 3. Linear attention is shown in red and softmax in blue. Solid lines indicate (ar) at finite a (Result 5.1), while markers represent the test risk of an ERM obtained via local optimization method with vND 10a. The regularizations F& and To are tuned by grid search to minimize the test risk, as detailed in Appendix D. Dotted and dashed lines correspond the value of E, in the infinite-a limit (see fosed-formed formulas in Proposition 4.1 for softmax and Appendix A.2.7 for linear). The Bayes-optimal risk is shown in black (see Section for a discussion on its discontinuity). Appendices D-E include more experimental details and results.
0
We conclude by few computational remarks. Because the empirical risk (15) is non-convex, it is not guaranteed a priori that optimization algorithms can find global minima. To assess the effect of non-convexity, we rely on numerical simulations shown in Fig. 3. More precisely, the prediction of Result 5.1 (here, uninformed and informed initializations give the same result) to the outcome of running local optimization algorithm (specifically. a quasi-Newton method) on the risk (15). starting from a random initialization, for finite but large N and D. We first note that the agreement between both is excellent. This suggests that potential bad local minima in the optimization landscape are avoided, at least for the appropriate regularization strength. Analyzing this landscape is an interesting open question. Furthermore, the softmax attention has lower error than the linear attention across all the tested hyperparameters, which shows that the benefits of softmax extend beyond the population risk analysis of Section 4 to statistical and computational advantages the empirical risk. We finally note that in this case the softmax is no longer Bayes- optimal but the gap to the Bayes-optimal risk closes as a grows, as expected from our analysis in Section 4.
compare
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We would like to thank Hugo Cui for the insightful discussions. This work was supported by the French government, managed by the National Research Agency (ANR), under the France 2030 program with the reference "ANR-23-IACL-0008 and the Choose France CNRS AI Rising Talents program. P.M. is supported by Google PhD Fellowship. This work was done in part while P.M. and L.Z. were visiting the Simons Institute for the Theory of Computing. This research was supported by the Swiss National Science Foundation (SNSF) under grant number 212049 (SMArtNet).
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 45614 45650. Curran Associates, Inc., 2023.
Yaroslav Aksenov, Nikita Balagansky, Sofia Lo Cicero Vaina, Boris Shaposhnikov, Alexey Gor- batovski, and Daniil Gavrilov. Linear transformers with learnable kernel functions are better in- context models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume Long Papers). Association for Computational Linguistics, 2024.
Luca Arnaooddi, Bruno Loureiro, Ludovic Stephan, Florent Krzakala, and Lenka Zdeborová. Asymptotics of SGD in sequence-single index models and single-layer attention networks. Ad- vances in Neural Information Processing Systems, 2025. urXiv:2506.0263.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language mod- els. In The Twelfth International Conference on Learning Representations, 2024.
Benjamin Aubin, Florent Krzakala, Yue M. Lu, and Lenka Zdeborová. Generalization error high- dimensional perceptrons: Approaching Bayes error with convex optimization. In Advances in Neural Information Processing Systems, 2020.
Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past. Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36:57125-57211 2023.
Afonso S. Bandeira, Ahmed El Alaoui, Samuel B. Hopkins, Tselil Schramm, Alexander S. Wein, and Ilias Zadik. The franz-parisi criterion and computational trade-offs in high dimensional statis- tics. In Advances in Neural Information Processing Systems, 2022. 2205.09727.
Nicholas Barnfield, Hugo Cui, and Yue M Lu. High-dimensional analysis of single-layer attention for sparse-token classification. In The 14th International Conference on Learning Representa- tions, 2026.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research, 22 (106):1-51,202)
Francesco Caltagirone, Mare Lelarge, and Léo Miolane. Recovering asymmetric communities in the stochastic block model. In 54th Annual Allerton Conference on Communication, Control and Computing, 2016. 110.v:1600
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane. Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with per- formers. In 9th International Conference on Learning Representations, 2021.
11

---

Published as a conference paper at ICLR 2026
Yuhong Chou, Man Yao, Kexin Wang, Yuqi Pan, Rui-Jie Zhu, Jibin Wu, Yiran Zhong, Yu Qiao, Bo Xu, and Guoqi Li. Metala: Unified optimal linear approximation to softmax attention map. In Advances in Neural Information Processing Systems, volume 37. Curran Associates, Inc., 2024.
Hugo Cui. High-dimensional learning narrow neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2025(2):034022, 2025.
Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. A phase transition between positional and semantic learning in a solvable model of dot- dot-product attention. Advances in Neural Information Processing Systems, 37:36342-366 -633499, 2024.
Mohamad Dia, Nicolas Macris, Florent Krzakala, Thibault Lesieur, Lenka Zdeborová, et al. Mutual information for symmetric rank-one matrix estimation: proof of the replica formula. Advances in Neural Information Processing Systems, 29, 2016. arxiv:1606.0414.
Elvis Dohmatob. Understanding softmax attention layers: Exact mean-field analysis on a toy prob- lem. In Advances in Neural Information Processing Systems, 2025.
Sara Dragutinović, Andrew M. Saxe, and Aaditya K. Singh. Softmax 2 linear: Transformers may learn to classify in-context by kernel gradient descent. 2025. 11X1v:2510.10425.
Daniel Y. Fu, Tri Dao, Khaled Kamal Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, 2023.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural Turing machines. arXiv: 2014.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's s the real context size of your long-context language models? In First Conference on Language Modeling, 2024.
Gregory Kamradt. Needle a haystack pressure testing LLMs. Github, 2023. URL https: NewdeeInnAaayeeeeeeeeeeee
//github. com/ gkamradt
Florent Krzakala, Mare Mézard, Francois Sausset. Yifan Sun, and Lenka Zdeborová. Probabilis- tic reconstruction in compressed sensing: algorithms, phase diagrams, and threshold achieving matrices. Journal of Statistical Mechanics: Theory and Experiment, 2012.
Thibault Lesieur, Florent Krzakala, and Lenka Zdeborová. Constrained low-rank matrix estimation: Phase transitions, approximate message passing and applications. Journal of Statistical Mechan- ics: Theory and Experiment, 2017, 2017. arxiv:1701.00830
Yue M Lu, Mary Letey, Jacob A Zavatone- Veth, Anindita Maiti, and Cengiz Pehlevan. Asymp- totic theory of in-context learning by linear attention. Proceedings of the National Academy of Sciences, 122(29):e2502599912 22, 2025.
Shengjie Luo, Shanda Li, Tianle Cai, He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang, and Tie- Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. 2021.
Advances in Neural Information Processing Systems,
Daniel Machlab and Rick Battle. LLM in-context recall is prompt dependent. arXiv:2404.08865, 2024.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. ICLR 2024. arXiv:2307.03576.
Pierre Marion, Raphaël Berthier, Gérard Biau, and Claire Boyer. Attention layers provably solve single-location regression. The Thirteenth International Conference on Learning Representa- fions, 2025.
12

---

Published as a conference paper at ICLR 2026
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kermion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. html.
https//ttassforneccccciiss wwwp:///2222.ooootecceeeaaaaaaad
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung. Leon Derezynski, Xingjian Du, Matteo Grella, Kranthi Gv, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Jiaju Lin, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Johan Wind, Stanislaw Woćniak, Zhenyuan Zhang, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: Reinventing RNNs for the transformer era. In Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2023.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Re. Hyena hierarchy: Towards larger convolutional language models. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research. PMLR, 2023.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong, cosformer: Rethinking softmax in attention. In The Tenth International Conference on Learning Representations, 2022.
Xuyang Shen, Dong Li, Ruitao Leng, Zhen Qin, Weigao Sun, and Yiran Zhong. Scaling laws for linear complexity language models. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2024.
Emanuele Troiani, Hugo Cui, Yatin Dandi, Florent Krzakala, and Lenka Zdeborová. Fundamental limits of learning sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholds. In Proceedings of the 42th International Conference Ma- chine Learning, 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, pp. 6000-6010. Curran Associates, Inc., 2017.
Matteo Vilucchio, Yatin Dandi, Cedric Gerbelot, and Florent Krzakala. Asymptotics of non-convex generalized linear models in high-dimensions: A proof of the replica formula. 25X2:2250..,0003, 2025.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett (eds.), Pro- ceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 35151-35174. PMLR, 2023.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv:2006.0477 04768, 2020.
Xinyu Wang, Linrui Ma, Jerry Huang, Peng Lu, Prasanna Parthasarathi, Xiao-Wen Chang, Box- ing Chen, and Yufei Cui. Resona: Improving context copying in linear recurrence models with retrieval. In Second Conference on Language Modeling. 2025.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in- context. Journal of Machine Learning Rexearch, 25(49):1-55, 2024.
Yedi Zhang, Aaditya K Singh, Peter E Latham, and Andrew Saxe. Training dynamics of in-context learning in linear attention. ICML 2025, 2025. arXiv:2501.16265.
13

---

Published as a conference paper at ICLR 2026
A POPULATION RISK
In this part we prove the expressions for the Bayes risk and the optimal risk of the attention layer, on population loss, stated in Proposition 4.1 and 4.3 and 4.4.
A.1 PROOF OF PROPOSITION 4.1: BAYES RISK
We consider the Bayes risk Bayes- We assume the estimator exactly knows k* and v². &Byee is not trivially null because the position e" of the relevant token is not known.
The best estimator Bayye minimizing the square error on given sample X E R Lx is the posterior mean y given L, ht, V² and X:
Bayes ,()L,, = Σ P(eLL, K" .)
1
x. 0.
(26)
VD
The conditional distribution of given L,, k and reads
g(Exx*)
P(c|L,k X) P(exx)
(27)
Ex-1 o(c'xx"
Conditionally on L, let J/ Xk" E LR and x" to E R be the projections of X onto the relevant directions. They are distributed according to uncorrelated standard Gaussians with law (N( It).
We express the empirical means over samples μ as expectations over one sample. The risk of OBayss is:
Rayes Er, on E., Ex, x(U
D)yyes)
(28)
Er,,., Ex-N((d,een) T. Egule xx')(x' "x"("
1
Dmayes) )²
(29)
-2: Σ P(elx)) E P(exx)))))
=1+ 1,+EB 7, Σ gol6" xx")
(30)
=1+ E, L
Σ gu(e",
-2P(E + Σ P(ex))))
(31)
=1 =1 - E,
1
Σ E.
(x(E"x))²
(32)
T.
This gives the expression stated in result 4.1.
A.2 RISK OF THE ATTENTION LAYER ON POPULATION LOSS
We consider the risk of the trained attention. We work on population loss. We recall that RD and v E Rd are the weights of the attention. They can be described by the following scalar order parameters (or summary statistics, or sufficient statistics):
MEk-
k. T
mey-
U T
Mkke
D T T
mykk
D" v TT
(33)
D
D
Okk
TT k
Jve
1
qek
k
(34)
D
D
Setting
( Tok
Tek Ook
mek-
882
+ ment y02 +
(35)
with q positive 2 x 2 matrix, ls and can be expressed as
h mag- ke u + (22/2) + (g2/2) )oo
(36)
14

---

Published as a conference paper at ICLR 2026
muk k" me v" + (g) (0112) 2000
(37)
with h- E TRD and vt E RDD two vectors orthogonal to k", o" and between themselves. We
introduce the shorthands
Rui (g
Rk = (q
Rov = (7) you
(38)
RikK Rgve and R. related to the magnitude of the components in le and v that orthogonal to k" and v" and bring no information. In the case where there is no overlap between v and & or k", and overlap between k and v", i.e. when mkke myrk quk 0, one can simply express the Rs as R2g = qkk m Re = Qov m and R&v = O. Then Res = Rov = 0 means that k is perfectly aligned with k" and with v².
The loss depends on & and v only via their projections onto the tokens X. Conditionally on L, we introduce the projections
1
x
Xk*
X VD
(39)
VD
= Xk
{=
vD" Xx²
(40)
VD
By central limit theorem have x (O(0 (t), 2* For conciseness we introduce the projections of k and V:
N(O, ~ N((,,L)) N(O,
b 1 VD
Xk mak x" + mkk 2n + RALE + RkOS E R
(41)
a X mox x + mon= R + RooS E B- VD
(42)
We introduce
(makk, Mgge, Myk' Myv= Raks Rkes Rey) IR? the set of parameters over which the loss is effectively minimized. Then the risk reads
Eg (k; v) = (9)
(43)
(e) = EL, . Be,,x -
(44)
Ei,,e or Bx-NNo.LeD). (u(""x"")6 xx)(( 0) Be i ex 9 (x) (z - a a(b))² 2
(45)
(46)
and the minimal risk is
Eg = min (k, (k)))
(47)
min &, (9) OER?
(48)
A.2.1 PROOF OF PROPOSITION 4.2: THE SOFTMAX IS BAYES-OPTIMAL.
We recall that the Bayes risk and the optimal risk of the attention estimators are
(x(,,)))
Eayye = E, Σ E. (e,
(49)
E, < min ExExx [1 2m (r, X)) x)e + m (e, () x) (mm
mike Mygg
(50)
where for E. have bouer-bound obtained by restraining the min of eq. (48) to Mok' Mike,, Rgv REE, R = 0,0, 0, 0, O. By definition the Bayes risk we have Eayes < Eg.
We search for the optimal a. We show that it has to match the Bayes-optimal estimator of We recall that, from the previous part A.1, the posterior distribution given X = vD X is
g(c,x)
P(elx)
(51)
El [L
15

---

Published as a conference paper at ICLR 2026
If, for certain m ammik= X)e P(clx), then
L
1
Ex,EE, (x x) ammux= x) = Ex,Ex I
E gele, E P(elxx)
(52)
1
gle,
= =EylEx I Σ
(53)
Ee g(
= =Eg gu(E, (m)xn x)
(54)
One can notice the similarity with the derivation of the Bayes risk in part A.1. We proved the equality
Ey gu(s, (m x)+ = E,,,, go(e, (mas= x) (m)xxx)
(55)
In statistical physics this equality is called the Nishimori condition. is satisfied when o X). matches the posterior distribution of The optimal mov- is then mye- 1 and obtain
E 1 EE gu x)o x).
(56)
1LE Σ g(c,x)²
(57)
g((,x)
Consequently E,, &Bayes
remains to show which a can satisfy the Nishimori condition. We assume that there a constant o such that for all L, all e, E [} and all X IR
gr(c, cxu(x)-xe)
(58)
gu(e',
Then
ammix= x). P(((x)
(59)
g(cxx)
(60)
Ee g(Exx))
(61)
Ee eooxe
that is to say is softmax with inverse temperature mak- Cu.
A.2.2 INVARIANT MANIFOLD
In the main part sec. 4 we introduced the manifold
M = (kk v) E PR2),, mkr- m.k- Quk =0]-
We show that it is invariant by GD. We consider the space of the order parameters, parameterized by We introduce the manifold
θ (M Muk", mgge mak= Rev. PEKS Rgv) E IR?
M = (0 R², mbu- = Mak- Rk = 0}.
(62)
At random initialization we have that = (0, 0, 0, 0, 1, E M. We consider the dynamics given by the gradient flow over the loss &, defined in eq. (44):
= (
(63)
We show that M is invariant under this dynamics. (mko* Muke Rev) on M reads
The gradient in the three directions
Omnnr E, = 2E, xe ihan &c ge ,x) (a o(b) 2 Vo(b)z*
(64)
Omny E. EE, x 4.5 (e x*) (a a(b) 2.) x"
(65)
ORo- E_ = 22, x is gele", x") (a (bb) - -2--) = (b(b) + Vo()))
(66)
At mgg- = -Myk- = R = we have a = m x* + R and = m X + RA& is independent
of n" and Ç. Then, because of the parity with respect x" and obtain that the gradient vanishes.
16

---

Published as a conference paper at ICLR 2026
A.2.3 STABILITY OF THE MANIFOLD
We show numerically that the manifold
= {θ E IR',, mkk-
Moke = RRv = 0}
is stable. We simulate the gradient flow
-= -
(67)
starting from perturbed random initial condition (mas., May' Mak-; Myy" RRks Ri Rey)
0, 0, 1, 0, 1) + with 1 Unif([-f]. +-7'² some small noise. We consider the value of reached after convergence, for several independent realizations of 7, for all the configurations considered in the main part of the article. On Figure 4 we observe that all the trajectories converge to point that belongs to M, up the numerical errors due to the integration. We additionally observe that R&s Rev Ra
spiked-SLK SLR
spiked-SLR
spiked-SLR
max-SLR
-SLR Unif( (1,2,3)
max-SLR
L (1,233)
lloear
0.0
-0.1
softmax
2100 1185
borualized softplus
0.1
0.0
0.1
B
0.0
-0.1
10
10
0
10
0
10
0
10
10
Figure 4: Values of the order parameters mike Myk+ May" Rek, Re and Rew obtained after convergence of the gradient flow. The mean, and min are taken over the independent runs. The initial noise is Tj 0.1 and there are at least twenty independent runs.
A.2.4 ORTHOGONAL COMPONENTS ON THE MANIFOLD
We show that the minimization of E., on the previously defined manifold, leads to a vanishing orthogonal component R.v 0. We can show that the second orthogonal component R& is null if we assume that the attention is linear.
We recall that the manifold gradient of the loss with respect to Rov:
-mpo' 2 Rooo and b -Ma= -RAE..
We compute the
d On.
2R ss go xx) (a (aTa(b) ((b) T
(68)
2 E.
(b) (b()
(69)
>0
where we used the parity with respect to Ç and assumed that the supports of and b overlap. As consequence R.. = 0 is the minimizer of &
17

---

Published as a conference paper at ICLR 2026
We turn to the second orthogonal component Rgk- In full generality RR:- = 0 cannot be straight- forwardly derived by considering only the gradient dR. Eo because the gradient may not be monotonous. Instead we assume that o = +1+1 X. Then the gradient
Oru Ea 2,,,, xx gult", xx) (a (1+ z/-- aae
(70)
= -2E,, Be.,BBeet gu(E" xx" ) (1 + b)
(71)
2R
(72)
and consequently Rx = i is the minimum of the loss.
A.2.5 LINEAR ATTENTION, STABILITY OF THE FIXED POINT
We consider the case of the linear attention o(x) + x. For this activation function we can state more precise result about the stability of the manifold. We show that the fixed point obtained by minimizing the loss on the manifold is stable, i.c. a local minimum in the whole space of the order parameters R²
At the fixed point obtained by minimizing the loss on the manifold we have mek" mkke Rev; Rgs; Riee 0,0, 0, O. We compute the Hessian of & with respect to Myke Mkoe Rov, REE, Rg at this point. For this we expand & to the second order with respect to these five parameters. We have
E. = xx 9u(e', E.
(73)
(2t- -miok- X Mov- 2" + RAOE Rvo ammss x - + RAKE + RkvS))
2m 2(Mmy. MKk + mmke myx- )
M "u(e"xx")xx
(74)
+ ELEex x gu(e",x") [mov T mak mx)=x"))) L(R + R2xr)
+ m mby T2)2 + min ((1 + MKK x x²)
+ 2 RRe, + R2a)2*T TT= + m2x- (Reo + Rx xx
+ 2m Mky mpx- TT 2²' + mak-x") + 2m Mike Makk (1 + make x²) 2²2² x"1
The Hessian is diagonal and positive with respect to Røv: RaK and RAV; and consequently Ron, Ras, Rive = 0,0,0 is well stable. The Hessian the two remaining directions myrk and
mkk- reads
2 mest E. = E ((1 + mak x") xx)²²
(75)
make
En = Em x²²
(76)
02 O m.
E, 2 [-x-- + mmo- T 21 + mak- xx x" may- (1 + max- T=
x*1
(77)
where we used the shorthand E for E. xx s² (o(cexx)) The trace of the Hessian is positive. The determinant of the Hessian is
det = m E ((1 + MAkk x*) E 22*
(78)
(E - + Mop- z" (1 + mAk x") Tx)Tx" + m (1 + mak- ==a=
We evaluate this quantity at mee-, Mask the minimizers of the loss on the manifold, for all the configurations considered in the article. We find that det > 0 and consequently the Hessian the fixed point is positive, that is to say the fixed point on the manifold is a local minimizer on IR?
A.2.6 OTHER POSSIBLE MINIMA
We numerically show that one does not restrict the space of the parameters to M, E, can admit several minima. Yet, among them, the minimum reached on M is the best.
We perform local minimization of E. starting from the initial condition (MaR= MAr= M Mor- RER, Rgv. R.,) = 1,, 0,0,), 0). The intuition is that the at- tention confuse and mix k* and v" For instance for identity activation function a(x) = x.
18

---

Published as a conference paper at ICLR 2026
and play symmetric roles, that we partially broke by adding a constant bias. We call mismatched minimum such a minimum where Mku * O, myk- * and m = 0, may. = 0, and set _ its risk, by opposition with the well matched minimum on .M where M&s- 0, , # 0 and mkk" = 0, Mok- = 0. On Figure 5 we show that at large enough sequence length = 10 and at the models, the attention always has well matched and mismatched minimum, and that the mismatched minimum has a larger population error 2. > Eg.
spiked-SLR
max-SLR
1.0
0.8
g = lineaz
0.6
1 erf
normalized softplus
0.4
= soflmax
mismatched minimum
0.2
well matched minimum
0.0
0
10 15
10
15
V
V
Figure 5: Minimal population risk E., reached at the well matched minimum, and population risk of the mismatched minimum 2,, for different attention activations (colors), for the two tasks spiked- SLR and wax- at L = 10. The markers on the lines are for readability only. Population risks are computed by numerical optimization of &, from random initialization (well matched minimum) from the mismatched initialization described in section A.2.6 (mismatched minimum).
A.2.7 MINIMIZER FOR THE LINEAR ATTENTION
On the manifold M the loss can be analytically minimized. We state here the expression of the minimizer. The loss is
ε, Ee, x² &x gulE", x²) (z-- - (m 2" + RRoS) m + RoKE))²
2m 2m mikk- ME Ex xx)xx
(79)
ELEEx golc,,x)) [m (1 mak- x*) (1 m
+ LIR + Lmn The mma- R [x*]]
At the minimizer R.,. 0 and REA = O. We perform the minimization over the remaining variables and obtain that the loss is
(E(1 mAR- x²)
ELEX;- x
Eg E(1 + MAA- X*) Maa- X*)
mak- Ex* - EX
(80)
where we used the shorthand for E, Ex-xx Io (o x²)) We evaluate these quantities for the two models, the spiked- and the max-SLR. We set fL,, v) Ee-xx go(c", x xx)xx",, which for the max- SLR has no simpler expression in general. We compute:
EX VD
E1Tx*
Ex X"
spiked SLR
VV
EL +
max SLR
Ef(L.v)
EL
The optimal risk for the spiked-SLR finally is
EL vBLL - 1) (EL)2 + VEL 1)
E,-
(81)
while for the -aR the optimal risk is
+ (Ef(L, v))a
E,
(82)
EL
19

---

Published as a conference paper at ICLR 2026
A.2.8 PROOF OF COROLLARY 4.3: ASYMPTOTIC CONVERGENCE RATES
Spiked-SLR large signal - oc and constant length derived in part A.2.7. We take the limit - 00:
The loss of the linear attention is
E, = L -+
L
+ 0v- (1)
(83)
+
L lv
The loss of the softmax attention is derived in part 1. For the spiked-SLR it reads
1
gle, x)2
E, 1 Ex
(84)
Ee g(e',
evix:
II [xx~Noon
= =1- E ((
(85)
E
evexe
1>1
where we isolated the first index by symmetry and made a change of variable on X1. At large v we approximate the softmax by a hardmax and
I. II
E, ~ 1- - E xI-N(EE
[Ex~N((o 8x Smaxess X
(86)
1>1
[ (0. e -ax-maxis x0)+0mmm)()) 1>1
(87)
(
(88)
Max-SLR at oo for growing lengths L. We take the limit (e x) Lô max, x: and f(L,v) = (0,,)) max xf. We assume constant length, i.e. Py is delta. The loss the linear attention derived in part A.2.7:
E, = 1 (EJ(L, v))²
(89)
EL
For large L we have the asymptotic f(L.v) L0000 (Vlog L), which gives
log
E, = 1 - O
(90)
L
The loss the softmax attention derived in part 1.2.1. For the max-SLR it reads
1
g(e,x)2
E, = - E, Σ
(91)
Σ g(e',x)
E T, Σ gle,x)
1
(92)
0
(93)
The softmax attention reaches exact recovery for any L, and in particular for large L, the limit been taken after the limit oc.
B BAYES-OPTIMAL ERROR AT FINITE SAMPLE COMPLEXITY @
In this section we state the expression for the Bayes-optimal test error in the case where samples are limited, i.e. when & N/D is finite, in the high-dimensional limit N, + co. The derivation of this expression is given in appendix C. We analyze behaviour of the BO performances and show that the SLR task presents hard phase where best algorithmic performances cannot reach best information-theeortiee. performances. This hard phase is not present
xc.
20

---

Published as a conference paper at ICLR 2026
The Bayes-optimal performances can be expressed in function of two low-dimensional order param- eters (or summary statistics, or sufficient statisties), my E R and my E R. Writing kBO E RP and Uno E R the BO estimators of k and v², my and my are defined
1 mk D T angle(KBo: k*)²,
my D "BOO" angle(UBo. u*)²
(94)
To state our asymptotic characterization result we introduce the following partition functions, for all and for B, y E R, A, R, V IRt and 7, R4:
Zx A) = EA-NN(o.)) B +BR
Zo A) = Bu-.No..Ij +
(95)
Zone (y, R,,, / dxdz I
1 Eglle, gole, x) (yl II N(xx (BNW,, (96)
where
e 23 (2-
Pout (y122.)
V
the output channel that corresponds to additive Gaussian noise. We set h2 effective distri- bution on e E {i.....4 L} defined as
halc, 7, R) - - dx dx x) II N(x.,,
(97)
Last we set the function xlogx: x + x log(z).
Result B.1 (Bayes-optimal risk). We consider the nigh-dimensional limit N, D - 0o, with a =
(1)). Let L ~ PL and, conditionally on L. < ~ NO, 1), NN Iz) and ( NO, Iz). Fix mn and m. so they satisfy the following fixed-point condition:
my; m. arg max OBO (m4; myy))
(98)
m.
1 (mgm + m, + 1 =ExxlogZZa (Viuss, mix
ΦBo (BO m.)
max ma ER
(99)
2nx
EL [ dyxlogxZout (y, Vmng, mys v Vmn. m
ExlogxZo (Vinns mm)
a
ΦBo is the free entropy of the problem given my and m.. Then the Bayes-optimal (BO) test error on inferring is:
1 EBO E, T,
Σ (e, m mx)
Σ hule, mWN, mmn) (100)
Our above result describes the information-theoreti (IT) best performance. In general it may correspond the algorithmic best performance. In fact, the SLR task presents an interesting phenomenology: in whole region of the parameters a, L of the model, there is gap between the IT and algorithmic achievable best performances. Such a discrepancy is related the existence of several maxima in the free entropy ΦBO eq. (99) and has already been studied for other models in numerous previous works (Krzakala et al., 2012; Dia et al., 2016; Lesieur et al., 2017) and we refer to Bandeira et al. (2022) for more rigorous treatment. The free entropy ΦBO plays a central role and it related to the likeliool the model; higher free entropy gives lower risk.
We can show that áBO admits one two maxima. We introduce some notation to distinguish them. At given v, L we call Galo the algorithmic threshold (or spinodal) the smallest a such that for all > Oaalg ΦBO has unique maximum. In the region < Galg where OBO admits two maxima, we set ΦisO the maximum whose basin of attraction includes a neighborhood of mı, m,, = 1, 1 and call it informed maximun. We set Φ the maximum of ΦBO whose basin of attraction includes a neighborhood of mg, m. 0, and call it uninformed maximum; it describes the algorithmic best performances, in the sense that an algorithm locally optimizing ÓBO and starting without informa- tion the solution will reach @ In general @ # digo- We call arr the IT' threshold the such
21

---

Published as a conference paper at ICLR 2026
that PiO OoO We have orr I Galg and we compute that for the SLR arr 1 for a noise-less output channel
To summarize, in the region arr maximum describes the IT best performances but cannot be reached from my m, = 0, by local algorithms, that only reach local maximum with higher risk. We depict this hard phase Figure 6 for the spiked- and max-SLR.
dalg the free entropy ΦBO has two maxima; the global
TpikddSLLL 8
spiked-SLR, v
max-SRR, L max--LR. 4,
1.0
0.8
0.6
0.4
0.2
1.0
0.8
0.6
informed
0.4
0.2
0.0
0.8
0.6
0.4
0.2
0.0
Figure 6: BO test error ÊBо. overlap m. with e" and overlap my with k* for the spiked-SLR and the max-SLR models, for different signals v and sequence lengths L, in the noise-less case A = 0. The values given by the eq. (100). The algorithmic best performances curve while for arr < a < Oaly the IT best performances are given by the "informed" curve. For arr and Olalg the two curves equal.
given by the "uninformed"
Considering Figure 6 we can state that the extent of the hard phase Oalg orr decreases with the signal и and grows with the sequence length L. At finite the upper limit of the hard phase Oralg is finite and consequently on population loss since we consider a noise-less setting v² can be exactly inferred finite a.
+oo the hard phase is not sensible. Notice that
A possible explanation of why such hard phase exists is that it is induced by the sparsity of the SLR. If one flatten each sequence into a single token X of dimension LD one would have setting similar to compress sensing where y X ō* where the regression vector e" E LD has proportion (L 1)/L null entries. Compress sensing has been shown to present a similar hard phase (Krzakala al., 2012). More generally sparsity induces hard phase, shown for sparse PCA (Dia et al., 2016) or strongly unbalanced binary SBM (Caltagirone et al., 2016).
An interesting particular case is the 0-SLR when 0 and the tokens are iid Gaussians bringing no information alone. As shown on Figure is still possible to exactly recover e" and perform better than random. Indeed one compare the train label X, v for each token If v has some correlations with v* then (y X,T 2²² likely be smaller for = c" This gives some information on the probable relevant token. Then one can refine by focusing more on the right X,. At inference time cannot use the label to guess e* and the prediction is simply Σ, Xfu.
22

---

Published as a conference paper at ICLR 2026
C ASYMPTOTIC CHARACTERIZATION AT FINITE SAMPLING RATIO: REPLICA
In this part we derive the asymptotic characterization of the Bayes-optimal performances and the performances of the trained attention, at finite sampling ratio a. We jointly treat the two cases, the BO and the attention, and the two models, spiked-SLR and max-SLR, taking variable sequence length. The derivation done for general output channel Pour such that for each sample the label is distributed with density P [y|= ). The additive Gaussian noise channel considered in the main part then corresponds to
222 22
Pout (ul=)=)
(101)
22T
Moreover the derivation is done for any convex loss function with strictly convex regularization.
C.I INTRODUCTION OF THE PARTITION FUNCTION
We consider data X, made of train samples, indexed by μ = L, ,,...,,N .N, together with N' = p& test samples, with o O, indexed by μ N 1,.... N N'. They are distributed according to the generative model defined in the main part 3.1. We will take the limit p - 0 at the end of the derivation,
information can be extracted from the test samples X XN+1S+SN+W in unsupervised way by the Bayes-optimal estimator.
The expected test square error of the attention- based estimator is
N+N Σ N=N+1
(a) = Ex N
(nu
(X,))
(102)
where
fo, = ( Xxh
VD Xuu
(103)
k.v = arg min L(ke v)
(104)
k,u
C(k, v) Σ So,, + E 7ks) + r. Σ ((1))
(105)
with f(x,,') = (x - the square loss and 7(x) = f the l2 regularization. Notice that
the following derivation is done in full generality for any convex l and strictly convex y. We will specialize to ridge regression at the end of the derivation.
The expected test square error the BO estimator is
1 N+N'
N+N' Σ
Ebo Ex, N
Σ (Un-
Ex
E
2y
BO
(i)B))
(106)
N
N=
where gBO the optimal estimator in terms of expected square error, that is, for test sample u' > N:
G = - Bu P(y (.s.))
(107)
Notice that, since owO is sampled from the posterior distribution, we have the simplification Ex,y (BBo)2 Ex,y Un DBO We can expand the posterior distribution over the test labels
P(yW2N/XX, x f /d((e,k ol) P(y
dedkdu P(ywNN e, x| BuSN)
(108)
(109)
23

---

Published as a conference paper at ICLR 2026
N+N' II dP Iu (En) x
x { S dPP(k)dPP(t)
P P( yuu 1 / (X
(110)
with Ps= = P. -N ID) the prior distribution on k and and, with an abuse of notation, P. = Unif(
the prior on Ey- We used Bayes rule and P(X k, v) & gu (k(V))) to rewrite P(c, x)x); and used the independence of the samples X pu Uy conditional cps and v to factorize the expression.
We can express the two errors (a) and Σbo under the same formalism thanks to partition function and a free entropy over and v:
N
z(X,u) f (PP(()) / N dP.(cc)) dn (sm. 1 Make PluullX Sp. k) 1=1
(111)
N+N'
II
do Iu (x. 1 xxxx P(uulx Cus K, v) E (sontor be)+ -bn)²)
ENNYY
N+1
B log 2(X v)
(112)
N
where, depending on the estimator, the densities are taken according to
(k)
Po(u) Iu PlyalX Eny (n, e e-rr E. 7(vs) e (e
v)
P(gul ens kns 616 Sak (X,)) P(u X,x Eue V)
attention e -rr E. ((k))) BÒ
(k)
Po(u)
Iv P'yulXyu pu Eui
with 3 co, for the attention k and concentrate onto the minimizer of the empirical loss. We introduced tilting variables s and & to access the expected test errors, obtained according to
Ô 1
(a) lim
80 dt p3 01a=0,100
(113)
0 1
ΣBσ = Os p
Φ $1,=0,
(114)
C.2 REPLICA, FREE ENTROPY
We compute the free entropy in the high-dimensional limit N, D -> We use the replica trick:
Ex,y log 7(X,y) y) = 0. (Ex 7X, x)) [)==
(115)
We introduce n replica, indexed by = 1, .,...,n, n. and count the expectation over the data X,y an additional replica indexed by use the superscript to denote the ground truth when it is clearer to do so. We use the shorthand O sy
0 corresponding to the ground truth. We equivalently
D + ty 0n)² for the observables. This gives
Ex, tn a Ex.y / II aP (P
N+N n
II
II dP.(C)) (ch
1 JD
X
(116)
N
Pa
P XX Ey ke, vn)
II
dØp P( X ch ha va)
a=1
#>N aml
N+N
xe Ex-N
II dig (1* (19) 0=0
II
II aPL()) a=u
(h xxxa"
(117)
N
II dy
II a=0
So ei ha va) (wna X a u0,
II P(ynd Xxs Eoo v*) II dg P(ojt X 6n ha v on
u>N
where PR. Pe, pe.ge, ia equals Pk. Pu.gu. Pif = 0 and Pr, Po, gu. P otherwise. We now introduce the two projections
1 _X,k² E BRW)
1 _x,00 E ((()
x
(118)
vD
vD
24

---

Published as a conference paper at ICLR 2026
thanks to Dirac deltas. We integrate over Gaussian X. We pack the replica into vectors of size + 1.
This gives
Ex l" x ExaN/ /
n II (P))
N4- n II
II a=0
g. (ex) (119)
II dy II
Xio ( II dy, P(y 9 II do Xp: x pa (Vul
a=0
a>N
II 100) ix (xa, J Nankk") tiž (=2. t Xnnvn)
/
II (PP(L( II T
II (P"
dx (()))))
(120)
N
II dyp
pa (uulk Gio W
II dy P( (r
II
do P(o ju Xpr c
a=D
1GN
a=1
N+N' TÍN)
mik
mat
II
II
(Xid) (Rad); ;0,
meg.
2e
mkv
0
m24
Q
m..-
Q.
where used that (K*)TK>/D> (u+)T D /D 1, uu and introduced the overlaps Mkk+ Mok Mko' M E IR" and Qux 2 2n E Ruxn defined for a, b by
1 3/ (k*) ka (mok~) )a = 1 (k*) (k*)TTo, 0", (m )a 11 [1'))A²
1 ka (mvv-) )a
1
(mxx=)a )a- DD
(T")TO"
0
D
(121)
Qut D (K
( - (o9)T (g D (T
(122)
We introduce these overlaps via new Dirac deltas. We leverage the remmica-symmetric assumption: we assume that there are makk. , myk" Mike , Movr Okk Beer Gov: Quki Qke Que E IR such that for all and b
(M )a = mMk- (me )a = muk- (M )a = mkv. (m )a = Myv-
(123)
(Qug)a Peadofb + Qikda=b (Quo)as akob Qkvda=b (Q quodapb + Quvda=b (124)
We pack these values into matrices
m (mak+ mog" ) _2x2 q
(
g2x2
Q
Quk Qky Que Quv
E 2
(125)
(lee Quo
We factorize the replica by introducing random Gaussian variables ç ~ (0, (2), & ~ N' It) and ( ~ ((0,4z). After standard computation (see c.g. Aubin et al. (2020) we obtain that the free entropy can be expressed as an extremum over the order parameters = (m, q. Q) and their conjugates è = (m e 2x & E 2 Q E 2=22))
max (0, θ) 6,6
(126)
6 θ) Tr m T m + (Trqq - +T=QQ)
(127)
2a
1 E. Tep ( 11121/25, s, mi T log Zku (g1/2 q + Q)
a
EeEe, [ dy Zout (y.mq 1/2
((=) 1 qm) log Zout (u.q12 1/ (57 (ft). .Q -q)
(5-
m
(5))
+ pF
Re,s Ee,s
/ dy Tome (v. by.mg-11 1/2
(5T)
)) m T qm)) log Ziut (v.01/2 (5
(5T) ( - 1)
(T)
25

---

Published as a conference paper at ICLR 2026
with the partition functions
4ho AB) = =P(())))(())> (+²) ("+"""(++++)) +BT (.())
(()))))))))))) (B) J (PP(())PP(u)) -(()) (())) (B)
(128)
(129)
Zour w. V) / dx'd=" / /PP(C) xx) out (012%-) II ((*) :w7, v)
Jz
(130)
Zout (ys W,, V) = / dxdz dalle) ( (()))ar
(131)
RL
Fooutus (),,,VV= V)= e / dxd=
f (c
(132)
and the observables O syly 0) +3(y 0)3
This expression of the free entropy can be simplified assuming that the order parameters m, q, Q are diagonal. We numerically extremized & for the full order parameters from uninformed and informed initializations and we observed that the cross-terms mar" myk'+ -ev, Qka and their conjugates go to This simplification is similar to the one discussed in the main part for the population loss, though there is no direct relation with gradient descent and that we have the additional order parameter Q. We simplify the notations setting
m = -akk"
Qkkk*
Qx Quk
(133)
m, makk
qu quv
Qu Que
(134)
and similarly for the conjugates. (m2, my. my my, qh: qk: go, qu. Qx, Qk. Qu Qu) ) 12 The free entropy is now
The set of order parameters becomes
-1 (mam+ + m.m.) + (ā:9k QxQx + Gogo Q.Q+)
1
(())
(135)
2a
1
E.g Zi
rink
rikk
log Z (VoLs. 9k + Qx) + a
my s, m log Zo (Vōos, Q.)
a
"
+ E2 Bg,s
dy Zout
y:
mik
C,
m
my
C,
m
log Zout (v, VoL, Qk - qs Vol, Qu - (p))
V
4x.
V
4v
m
mlg
m
+ EE,5
dy Lout
V,
C,
1
log Zont (VacE, 2k. Vars, Qu qo)
JR
V
qk
Tu
with 5 NN00 1) and the partition functions
A)
dP (k") +
(B, = dP)((k) -WAN++BR
(136)
A)B [ (PP((0")-2A0000-11
Zn(BAA) [ d -AUU+BB
(137)
Zon 7, R, V) / f dx'd2" ddP.(e)) xx)) Powe (y) III (Xxi V)
(138)
I Zon (3,7, = f dxdz / da.lee) II Wo ()N
(139)
L x.6) II Noa: (2; V)
2o R,w, V) = ddje [aje
dxde aPP(00 /
(140)
26

---

Published as a conference paper at ICLR 2026
Once & is extremized over the order parameters the test errors are computed as
/ dy (oo
(a = 3 F... Be,c
-oout O, T -==
(141)
Eg Er Bg,s / dyy(
Lut
(142)
C.3 SPECIALIZATION THE BO
We consider the BO setting, where 2 ZR. Ze Zy and Zout Zout A main simplification, known as Nishimori condition, is given by the fact the BO estimator is sampling the posterior distri- bution: have my 9k: m, = 9k. Q+ = m. qus my do and Qu 1. We set the function xlogx: x - x (z) The resulting free entropy is
1 Xl
(eeest/(Viiisi (Viss, + - Exxlogx (Vmos; m))
(mgm + m.) 2a
a
a
(143)
+ E, Bg.s dyxlogxZon (V, vm,5,1 - m&, vmm,, - m.))
JR
pB., Re,s / dy Zon (y, vm mik /mni, - m.,) log Zout (V 1 Mik Vm,C, m.y)
The extremality condition of over the order parameters is obtained by setting its gradient to zero.
We take the limit p = 0. It gives the following fixed-point equations:
m: = E. (0
(144)
my = E. 1 (0
(145)
my = BEC / dy Dorl Vout) (V~
(146)
m B., dy Zonk Zou (V
(147)
We explicit the fixed-point equations. We set
hy R) / dx II ( R)
(148)
the effective distribution
We assume that Paut is the identity channel. Then
mg
(149)
I ++ me
mo
my + my
(150)
(x- )m
my = aE, Eg,s / dy Σ hule, /mx5, - m²)
2(1 m(-m)
(151)
2T (1 mm
(y- mmg
T, Σ Σ dn vmna. mp)
214 mev
/25 my
Ammik
m. = DE, / dy Σ Ehules huuesymmss vmae, mp)
2[1 m11
(152)
/
m,,)
27

---

Published as a conference paper at ICLR 2026
(u- Gm
1
Σ
(1, Vmk,, mx) 1 vmm
201 mm!
L
me
V me))
At the fixed-point the Bayes-optimal error is given by
Εво - M,
1 L
Σ hule, vm 1 - mx)
Σ
hule, /m 1 - m2)²
(153)
C.3.1 SPECIALIZATION TO SPIKED-SLR
The spiked-SLR is defined by taking go (e(ex) = evix--5 We can integrate over X and explicitly compute hy:
he 7, Rt)
(154)
Thanks to the invariance by permutation of the tokens we isolate the index We change the variables & + E1 + viing and Çi + vI moS1 + ymay to obtain expressions that are casier to compute numerically. This gives
mi
mp
m& 1 mi
my + my
(155)
/
dy e v²²
ry = L.gs
2T
1 Eis1 -Vmx -VmA- (9 m +1[ /m
(156)
my
a
EL Be,s
/
dy
e 1²²
(V1 m 21
my
1+ Li>1
++Esee +mm /om E1) 19 m
mm)
(T-mmul- 1---m----v mmmu--mmuts)"
(157)
ev eyvmat:
ΣBO 1 mod
2
Eg
(158)
1 Etsr evimecti--1))
C.3.2 SPECIALIZATION TO THE MAX-SLR MODEL
The max-SLR model defined by taking gu(e, Leux, /(E/ exxi), X cannot be integrated out. We consider the limit v = +o0 to simplify the expression of hyi
(1 + X X
hu Y, R) = L dxx R)
II life
(159)
V2R
Thanks to the invariance by permutation of the tokens we isolate the indices / and / 2. We can still change variables Ç. /1 m.Ce ymay obtain expressions that easier to compute numerically. This gives
mp
my
my = 1+ mp
my =
(160)
+ mp
vu = AEyREs /
dy
c b/²² [(0 (1,1, Vm=5,1 - m2)) -211)---( (m)())²
(161)
2t
(x- (mpcg)
+ 2(L 1))g, (1, /mmE, - mx) (2, /mL6, -
me)e
2d mp)
(1, Vmnt,1 mele 3 )
Σ
holl, /mnss Je
mmu mp)
1>1
+ +aELEES /
dy
トル² [L 1)(0,, (2, Vmre, -mm))= --------maaaaa m )
.2T
28

---

Published as a conference paper at ICLR 2026
Car + + -,,, ( m,, 1,, VmRE, male
mee
12,
- -mm) -// -lVl-mm.v----
mmok)
Σ
ho(1, /mm.. m
19- mpp 211 mm)
1/2
my aBrBec /
dy
--tx tu? (v1 - myy mm=
ho VTN, ma)e mm..V
- my m
(1, Vm 1 male AlVi-m.p-vmmG1
hull, Vmks, 1 mx)e
(v /mm
2(1-mp)
(162)
1 εσσ =1 m, Th vmxs, 1 m
(163)
C.4 SPECIALIZATION THE ATTENTION
We consider the attention case, where the distributions Pr. P. and P(y) X,,, k, v) do not match the distributions of the model. this case is more convenient to work with the variables
V Q+ q
VA = =Q= + tE, V Qu 1 tu V. = Q go.
(164)
We perform the changes St Sh k m/// Vāk. So So u f Er &v -m / +
m/m V and Çi + // qo - m²/ / + m/ Vo. We deal with the limit & + 00 by rescaling the parameters according to me + Bing, my + Brivss qu + Bug. Go p V + VV.
V. VV. and V& 3 Vk. V. VV....
We introduce the effective joint distribution of the data, for E ER E R
x R and
1 E" = 7 x²) Pom (ulz..) II
(165)
The free entropy is then
1
1
1 Vi + VoVe Vilun +
B (00)) - - (mame + m,) +
3 V.D. + t. - Vv
2a
3
E. / (k*) log / (oL())
+ aB / dP.((v)) log / evucl))
+ a3
(166)
1 Re,s /
/ dxdz (c0Mou((xsss)) +
(2.x)))
with the potentials
(k)) = (k)
2 Vik² + (mikk" + / T
(167)
(w) = ( Vo02 (mo"" + VQss)o
(168)
Vout (x, 2; = -Tl(y, (2) + Σ Nog Vx)) + Σ log log.N(cs (Zzw), V.)
(169)
7= 7= m,X" + Vq+ mig. w m
v qu m3s
(170)
E {0, 1} controls whether the loss or the observable are active or not. We introduce the extremizers of these potentials:
k' = arg max th (k)
e' arg max (v)
(171)
29

---

Published as a conference paper at ICLR 2026
X arg max yout (x, 2; 1),
x",²" arg max Yout (x 2; Xi=
(172)
We introduce the covariances under these potentials around there maxima, with V the Hessian.
Cov(k: - (b2))())))
Cov(u) (V2 (V2))((')) 1)) Cov (21) = =- ( (ooloooul
(173)
Cov (xi) = ((V (x (x:)==((x())
(174)
The extremality condition of over the order parameters is obtained by setting its gradient to zero. We take the limit p = 0. It gives the following fixed-point equations:
/ dP.(v")
mR E. dPP(K")K>
my = B. 2
(175)
to = E3 / dPo(v")(u)))
1s,
(k*) (W)2
(176)
J
V₁. E, R
dPk Cov(k)
V. E, j
dP, Cov(v)
(177)
and
ma aEr.Ee xx
1
Σ
xixi
mk
Cov(xi)
(178)
V.
Vi
mn aEsBes dP''v.,,,xx22)
22)
1
2/
my
Cov(ze:)
(179)
Vg
Ve
92 Egs /
x
(2)
1
Σ
7
X1
-mxii -Vax m2
(180)
V2
④. = &Er, / dp"
1
Σ
r.
m.²i - v 4o - -m-mm) 2
(181)
V?
V/ EL al.
Vi EEE dP"
Cor(xa)
(182)
V22
V. F. aL V,
BELRES [at' 1
Σ Cov(z)
(183)
V
At the fixed-point the test error is given by
Ea = ELEg, / -x
(184)
C.4.1 SPECTALIZATIONS TO SPIKED- AND MAX-SLR MODELS
We consider attention trained with le regularization (xx) 1 The first fixed-poing equations can be explicited
m&
m.
(185)
Tk + Vx
rv + V.
m + 4x
m + g.
4k
9o
(186)
(r+ + V2
(r. + V.)
V/
V,
(187)
+ Vj.
To + V.
We can simplify the rest of the equations by using the permutation invariance w.r.t. the tokens to fix 1. For the two models have that, for o c", 2/ is Gaussian under Pa (y,, e", x²) and
30

---

Published as a conference paper at ICLR 2026
integrate out 2f. We assume that Pout is the identity channel so z: y. In the potential dout we have wi mm,yrr1 + / -m. for all 4. The joint distribution over data reduces to
ma + 44 m2si ~ (0,
(",, P NN(; 1) (1,x)) x") II Ni
(188)
Taking y ~ (O 1) and x ~N(00, Iz), the fixed-point equations and the error are
my a,E,, B.,,,EE,...x 90 x²)
Σ
Xixi
m
Cov(xi)
(189)
Vx
Vx
m. aLg 9u (1,
yz
m.
Cor(z1)
(190)
Ve
Ve
L
Q. = aB
Σ
Xi - mxxi
- V 9k - mmff
(191)
V2
4. aB, Be, (o(1,xx))
(+1- myy V 40 -mgGG)
Σ i>1
vihci)
(192)
V?
V2 Er. aL aBr.e.cox.
Cov(xa)
(193)
Vx
V. EEL
EL al aEL x")
V? Σ Cov(z)
1
(194)
V (a) Eg, (1, (x(x"
(195)
For the spiked-SLF SLR model we additionally have that xf is Gaussian under (y (y)x"); m&xi + V qh m (VVumg 9h) and X" can be integrated out. In the potential toott we have n Vumgdr,1 varui for all I. We have the simplifications
m4 EE, Vk
(Vixi vmK)
(196)
ás aIE, Be..B V2
(xi - Vema v
Σ 1>1
(xl V
(197)
31

---

Published as a conference paper at ICLR 2026
D DETAILS ON THE NUMERICS
D.1 TRAINING OF THE ATTENTION
The attention eq. (4) is trained using LBFGS on the full batch until convergence. The parameters and h are initialized randomly according to standard Gaussian (00 ID).
In Fig. 3 the regularizations rk and rv are tuned by grid search to minimize the test error, over 10.3.133, 10}² for the linear attention and {0.03, 0.3, 1 for the softmax. These values were determined according to Fig.
For most of the parameters of the data model we explored the training converges to the global minimum of the training loss eq. (15), the sense that the test risk of trained attention is equal to the test risk predicted by our asymptotic characterization.
In a few cases, in particular at low signal and low regularization, the training does not converge towards the global minimum of the loss, which results in discrepancy between the simulated per- formances and the predicted ones. In these cases we train the attention starting from an informed initialization x² and k k*. Note that this initialization does not necessary correspond to minimum of the loss. The performances of the trained attention then better matches our predictions, as depicted in Fig. 7 for the softmax the max-SLR. We checked that the achieved training loss is well smaller than the one starting from a random initialization.
D.2 COMPUTATION OF THE ASYMPTOTIC CHARACTERIZATION
D.2.1 BO
We compute the BO performance stated in result B.1 by iterating the fixed point equations given by the condition Vêno (m,, m.)) 0. These equations detailed in appendices C.3.1 and C.3.2 for the two models. The uninformed initialization corresponds to my. m, 0.1,0.1 0.1 while the informed initialization corresponds mg, 1 10 10-2 The expectation Gaussian random variables is computed over 10t and 10t Monte-Carlo samples respectively for the spiked-SLR and the max-SLR models.
the
D.2.2 ATTENTION
We compute the performance of the attention stated in result 5.1 by iterating the given equations. We use Steffensen's 's method to speed up the convergence. For the spiked-SI. a simplification of the equations is given in Appendix C.4.1. The expectation is computed over 10a Monte-Carke samples. For each sample at each iteration the extremizers of Pout are computed using a quasi-Newton optimization scheme. The optimization started at the extremizers computed at the previous step.
the Gaussian random variables
Notice that vout is not convex and may admit several maxima. In practice, for most of the values of 7, W, Vi and V,, oout admits a unique maximum. When it is not the case (in particular at low signal and low regularization) one should compare several different initializations of the optimization al- gorithm to find the global maximum. We tried on a few cases; it appeared that the final predictions do not change by a quantity greater than the fluctuations due to the randomness.
The minimum over the set S of fixed points is computed by running the iterations of result 5.1 from several initializations. We considered informed initialization (mz, my Ues qos Vk, V.) = for small s, partially informed initializations (ma, Ma, De, Gur Vs, Vo) 1,,1,1,1,1,1) and uninformed initializations (m2 my: qs qu Vk, V.) = (0,0, 0, formed few different runs, adding small randomness to the initial condition. For all the values of the parameters, the obtained fixed point did not depend on the choice of the initialization, up to some small numerical fluctuations. These fluctuations in which case we select the run which reaches the highest free entropy eq. (166), i.e. the lowest train loss, among those that converged before certain amount of iterations.
1). We per-
larger small regularizations and low signals,
32

---

Published as a conference paper at ICLR 2026
E ADDITIONAL FIGURES
We provide an additional figure for the result 5.1 about the test risk of the attention at finite a. Fig. gives the test risk E, for the same configurations and rv. It justifies the ranges of the grid search for the linear and the softmax attention in Fig. 3, in the sense that E, seems reach its minimal value for regularizations close to the ones considered in Fig. 7. Moreover it shows the excellent agreement between our theory result 5,1 and the simulations. the assumption
Fig. for several additional regularizations тк
We observe discrepancy for the softmax at small regularization for the spiked-SLA max-SLR at y +X. For the spiked-SIR it may be caused by the replica-symmest being false or the numerical errors the resolution
the fixed point equations. For the max-SLR Tk Te 0.03 observe that initializing the simulations at the informed point k", v v² leads to an agreement with our theory. This shows that in this case the local optimization cannot recover the global minimum of the loss from random initialization.
Spiked-SLR, y=
spiked-SLR,
max-SLR, +:00
==
m=r==3
N=rm1
Ein
0.75
==r==0
0.50
0.25
0.00
1.25
+=r==
random initiallantion Informed initisliatiinn
1.00
Farlnwe (c)
Bayes-oplimal
0.25
0.00
2
2 4
a
Figure 7: Minimal test risk E, (a) of the attention, linear (top) and softmax (bottom), across differ- ent tasks and signal strengths v, for = 3. Solid lines indicate (a)) (Result 5.1), while markers represent the test risk of an ERM approximated via a local optimization method with v VND 10a and averaged over instances. Random initialization means that the attention is initialized at dom & and v while informed initialization means is initialized at k = k." u = v² Dashed lines correspond the value of E,, in the infinite-a: limit (see foosed-ormed formulas Proposition 4.1 for softmax and Appendix A.2.7 for linear). The Bayes-optimal risk ε (ax) is shown in black (see Section B for a discussion its discontinuity). Appendix D includes more experimental details.
ACKNOWLEDGMENNTS
used large language models for helping the writing and discussing some references and related work.
33


Reviews:
[{"review_id": "t4pAyCgg6M", "reviewer": "Reviewer_tjje", "summary": "In this paper, the authors do a study to show why softmax-based attention mechanism outperforms alternative attention variants (such as linear attention). To do so, they use a simplified “single-location regression” task in a theoretical framework. They show that in the infinite-data population limit, a softmax attention model achieves the Bayes-optimal prediction error, which cannot be reached by linear attention models. Furthermore, they also show, that even in the finite case, where no model can reach the Bayes risk, softmax attention still outperforms linear attention.", "questions": "Slightly off the topic, but recently, there has been an interest in softmax-1 activation [A, B, C], which is a version of softmax designed to prevent attending in the first (few tokens). Considering the similarities between it and the standard softmax, I wonder if the findings would hold there.\n\n[A] Miller, Attention is off by one, 2023\n\n[B]  Bondarenko et al., Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing, NeurIPS, 2023\n\n[C] Kaul et al., From Attention to Activation: Unravelling the Enigmas of Large Language Models, ICLR 2025", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the paper has the following strengths:\n\n1) Principled Task Formalization - the paper introduces a clear mathematical framework for studying attention mechanisms by formalizing a single-location regression model, in which the output depends on a single informative token in the input sequence. In this way, this setup generalized earlier theoretical studies by allowing random sequence lengths and mirrorying “needle-in-a-haystack” retrieval tasks.\n\n2) Rigorous Theoretical Analysis and Key Insight - The authors develop a tractable high-dimensional analysis of attention layers using techniques from sequence mulit-index models that the softmax nonlinearity can be handles with a small set of parameters. What is quite exciting is that they show that compared to linear attention, only the softmax attention can achieve the Bayes-optimal error, while the linear attention is suboptimal.\n\n3) Comprehensive Evaluation of Performance - the study also examines the finite-sample regime (the real-world case) adding practical value to the theoretical analysis in (2). They show that the softmax model still holds an advantage compared to linear attention in the finite-sample setting.", "weaknesses": "I think the paper can be further improved in these parts:\n\n1) Limited Scope of the Task and Model - the analysis in limited to the simplified single-location retrieval scenario. While this abstraction (which corresponds to using a fixed query token like [CLS]) is mathematically convenient and appropriate for the task, it does not cover cases where multiple tokens or more complex query mechanisms are involved. Real-world Transformers often deal with interactions among many relevant tokens and use learned queries, so I am not sure how the softmax advantage might play out in these scenarios.\n\n2) Gap to Real-World Validation -  The paper’s findings are supported primarily by theory and synthetic experiments. All experiments are on simplified tasks defined by the authors (e.g. variants of the single-location regression or synthetic retrieval scenarios), which raises the question of how well the insights transfer to practical large-scale NLP tasks. As a result, the work stops short of empirically confirming that softmax’s advantage in more complex or realistic conditions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rzV9UOB96F", "reviewer": "Reviewer_Uy5X", "summary": "This paper theoretically explains why softmax attention outperforms linear attention by analyzing a single-location regression task  (Marion et al., 2025), where output depends on one hidden token, proving softmax achieves Bayes-optimal risk while linear attention fundamentally fails. \n\nBuilding on sequence multi-index models (Cui et al. 2024, Troiani et al. 2025), the authors extend these techniques to variable-length sequences and characterize both population and finite-sample risks via replica methods, showing that exponential nonlinearity and normalization are both critical. \n\nExperiments validate theoretical predictions with agreement between predicted and observed test error curves.", "questions": "Please see Weaknesses.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The authors provide theoretical results establishing provable performance gaps between softmax and linear attention ( by analyzing a single-location regression task ) at the population level (Propositions 4.1-4.4, Corollaries 4.2-4.3).\n\n\nThe paper successfully extends sequence multi-index model techniques to handle the challenging softmax nonlinearity through careful renormalization and order parameter analysis. \n\nThe authors provide strong empirical validation. Code was shared but I didn't check.  The experiments provide quantitative agreement between theoretical predictions and gradient-based optimization experiments across multiple settings (different signal strengths ν, sequence lengths L, regularizations).", "weaknesses": "- No optimization theory/landscape analysis is provided, and this is indeed a limitation as the paper proves softmax's empirical risk minimizer has superior statistical properties but provides no characterization of the optimization landscape or guarantees connecting estimation to optimization. The empirical risk (15) is non-convex with multiple local minima (Appendix A.2.6), yet the paper offers no analysis of landscape geometry, convergence guarantees for gradient descent, or the estimation-optimization gap. \n\n- The main finite-sample characterization relies on the replica symmetry assumption from statistical physics, which lacks rigorous justification. This should be stated more prominently in the theorem statement itself, not buried in discussion. \n\n\n- Expectation notation: $E[\\cdot]$ vs. $\\mathbb{E}[\\cdot]$ inconsistently used  \n- Risk terminology: switches between \"test error,\" \"test risk,\" \"Bayes risk,\" \"Bayes error,\" and \"Bayes-optimal risk\". \n\n- \"Symmetric activation function\" in Corollary 4.1 needs precise definition.  \n\n\n- Equation (3) notation confusion: The notation $\\chi^* = (1/\\sqrt{D}) x \\otimes k^*$ uses $\\otimes$ which typically denotes tensor product, but context suggests matrix multiplication. Should be clarified or corrected.  \n\n- Proposition 4.2, condition (6): States \"for all $L > 0$\" but $L$ is sequence length (positive integer), should be \"for all $L \\ge 1$\" or \"for all $L \\in \\mathbb{N}$\".  \n\n- Asymptotic notation in eq. (12): The notation $E_{\\text{lin}} = (L/(L-1))(1/\\nu) + o_{\\nu\\to\\infty}(1)$ places subscript on little-o, which is non-standard. Should be $E_{\\text{lin}} = (L/(L-1))(1/\\nu) + o(1)$ as $\\nu \\to \\infty$.  \n\n- High-dimensional limit not precisely stated: The relationship between $N$, $D$, and $\\alpha = N/D$ in the limit is ambiguous. Is it $D \\to \\infty$ then $N \\to \\infty$? Or jointly? The order and dependence should be explicit early in the paper.  \n\n- Manifold assumption justification: The manifold $\\mathcal{M}$ is introduced abruptly on page 5. While Section A.2.2 shows it's invariant under gradient descent, more intuitive geometric/statistical interpretation would help readers understand why this restriction is natural.\n\n\n-  Equations (18)–(25)  present complex fixed-point equations with no intuitive explanation. What do these order parameters represent physically? Why these specific equations?  \n\n- Figure 1 misleading markers: Caption states \"markers on the lines are for readability only.\" If markers don't represent actual data points, they should be removed as they're potentially misleading.  \n\n- Figure 2 incomplete experimental details: Caption mentions regularizations \"tuned by grid search\" but doesn't specify the grid values in the experiments section. \n\n\n- Some figure legends are too small to read comfortably (especially Figure 3).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "elY0l3pGcm", "reviewer": "Reviewer_hz39", "summary": "This paper studies the theoretical foundations of softmax attention by introducing a principled single-location regression (SLR) model. The authors derive analytical and asymptotic results using statistical physics tools (notably order parameters and replica analysis) to show that softmax attention achieves Bayes-optimal performance in high-dimensional limits, while linear attention and other alternatives (e.g., kernelized or element-wise nonlinearities) fundamentally fall short. The paper further provides a finite-sample characterization of test risk, confirming the statistical and computational advantages of softmax. Overall, it offers a clean theoretical framework that bridges information retrieval toy models and high-dimensional generalization analysis.", "questions": "Can the proposed Single-Location Regression framework be extended to multi-location or multi-head attention, where multiple tokens jointly determine the output? Would the Bayes-optimality of softmax still hold in those cases?\n\nThe analysis assumes Gaussian and independent token embeddings. How sensitive are the results to these assumptions? Would correlated or structured embeddings affect the theoretical conclusions?\n\n\nHave the authors tested whether the predicted statistical gap between softmax and linear attention appears in small-scale Transformer experiments or synthetic retrieval tasks (e.g., Needle-in-a-Haystack)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Clear and original theoretical contribution:\nThe paper introduces a new analytical framework — the Single-Location Regression (SLR) model — to study the statistical behavior of attention mechanisms. This formulation unifies previous “needle-in-a-haystack”-type setups under a mathematically tractable regime, enabling a principled comparison between softmax and linear attention.\n\nMethodological depth:\nThe work combines sequence multi-index models with replica-based high-dimensional analysis, bringing together tools from statistical physics and modern learning theory. This cross-disciplinary approach is technically sophisticated and extends recent progress in the theoretical understanding of attention networks.", "weaknesses": "Idealized task setting:\nThe SLR model assumes the label depends on a single token’s linear transformation, which, while analytically convenient, is far from the multi-head, multi-layer structure of real Transformers. Hence, the practical relevance is limited.\n\nLimited empirical grounding:\nThe validation is restricted to synthetic setups, without experiments on realistic retrieval or sequence tasks. It remains unclear whether the predicted statistical advantage of softmax can be observed in real neural models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "80ptLRimSU", "reviewer": "Reviewer_J4sf", "summary": "The paper studies why softmax attention often outperforms linear attention by analyzing a stylized but mathematically tractable setting called Single-Location Regression (SLR). In this setup, only one randomly chosen token in a sequence carries useful information for predicting the target. The authors derive population-risk expressions showing that softmax attention achieves the Bayes risk (i.e., the minimum possible expected loss), while linear attention is inherently suboptimal. They further extend the analysis to finite-sample training using replica methods and confirm that gradient-based optimization empirically follows the predicted risk curves. The work provides a clear theoretical explanation for the statistical advantage of softmax normalization.", "questions": "Please refer to the weakness section", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Provides a principled and analytically clean setup (SLR) that isolates the retrieval aspect of attention mechanisms.\n\n- Derives explicit population-risk expressions demonstrating when and why softmax outperforms linear attention.\n\n- Extends the analysis to the finite-sample regime using replica theory, connecting asymptotic analysis with realistic training outcomes.\n\n- Numerical simulations show strong alignment with theoretical predictions, suggesting the model captures key phenomena of attention mechanisms.\n\n- Offers valuable theoretical insight into the long-standing question of softmax’s statistical advantage, which is of interest to both theoretical and applied communities.", "weaknesses": "1. Limited task scope. The analysis focuses solely on the single-location regression (SLR) task, which captures a narrow form of retrieval behavior and does not generalize to multi-token dependencies or compositional reasoning.\n\n2. Restricted model comparison. The study only contrasts softmax and linear attention, omitting other relevant variants such as kernelized softmax approximations and state-space models (SSMs), which prevents a broader understanding of where softmax’s advantage truly lies.\n\n3. Reliance on idealized assumptions. The theoretical results depend on Gaussian i.i.d. token embeddings and replica-symmetry assumptions; robustness to more realistic distributions and correlated features is not explored.\n\n4. Absence of realistic experiments. All evaluations are conducted on synthetic data, with no validation on real-world or language-based datasets, leaving uncertain whether the observed statistical gap translates into practical performance gains.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)