High-Dimensional Analysis of Single-Layer Attention for Sparse-Token Classification
Authors: Nicholas Barnfield
Hugo Cui
Yue M. Lu
Venue: ICLR 2026 Poster (2026)
Keywords: Theory
exact asymptotics
high dimension
high-dimensional statistics
attention
Primary area: learning theory
TL;DR: We analyze an attention-based model for a classification task on sequential data with weak and sparse signal.

Abstract: When and how can an attention mechanism learn to selectively attend to informative tokens, thereby enabling detection of weak, rare, and sparsely located features? We address these questions theoretically in a sparse-token classification model in which positive samples embed a weak signal vector in a randomly chosen subset of tokens, whereas negative samples are pure noise. For a simple single-layer attention classifier, we show that in the long-sequence limit it can, in principle, achieve vanishing test error when the signal strength grows only logarithmically in the sequence length $L$, whereas linear classifiers require $\sqrt{L}$ scaling. Moving from representational power to learnability, we study training at finite $L$ in a high-dimensional regime, where sample size and embedding dimension grow proportionally. We prove that just two gradient updates suffice for the query weight vector of the attention classifier to acquire a nontrivial alignment with the hidden signal, inducing an attention map that selectively amplifies informative tokens. We further derive an exact asymptotic expression for the test error of the trained attention-based classifier, and quantify its capacity---the largest dataset size that is typically perfectly separable---thereby explaining the advantage of adaptive token selection over nonadaptive linear baselines.

Published as a conference paper at ICLR 2026
HIGH-DIMENSIONAL ANALYSIS OF SINGLE-LAYER ATTENTION FOR SPARSE-TOKEN CLASSIFICATION
Nicholas Barnfield Department of Statistics Harvard University Cambridge, MA, USA
nbarnfieldêg. harvard.
Hugo Cui
Center of Mathematical Sciences and Applications Harvard University
Cambridge, MA. USA hugoëcmsa. fas. harvard.edu
Yue M. Lu
John A. Paulson School of Engineering and Applied Sciences Harvard University
Cambridge, MA. USA yuelušseas harvard.edu
ABSTRACT
When and how can an attention mechanism learn to selectively attend infor- mative tokens, thereby enabling detection of weak, rare, and sparsely located features? We address these questions theoretically tion model in which positive samples embed weak signal vector in randomly chosen subset of tokens, whereas negative samples are pure noise. For a sim- ple single-layer attention classifier, we show that in the long-sequence limit it can, in principle, achieve vanishing test error when the signal strength grows only logarithmically in the sequence length L, whereas linear classifiers require vL scaling. Moving from representational power to learnability, at finite L in a high-dimensional regime, where sample size and embedding di- mension grow proportionally. We prove that just two gradient updates suffice for the query weight vector of the attention classifier to acquire a nontrívial align- ment with the hidden signal, inducing attention map that selectively amplifies informative tokens. We further derive an exact asymptotic expression for the test error of the trained attention-ben based classifier, and quantify capacity- the largest dataset size that is typically perfectly separablc-there explaining the ad advattgee adaptive token selection over nonadaptive linear baselines.
sparse-token classifica-
study training
Attention-based architectures Vaswani et al. 2017) have proven in recent years be a major driver of progress in a wide spectrum of learning tasks, ranging from language processing (Kenton & Toutanova] 2019 Brown et al.) 2020) computer vision (Dosovitskiy al. 202T). core strength of these models is the ability of attention layers to dynamically weigh the importance of different in- put tokens, enabling the model to selectively focus on the most relevant information. This flexibility makes transformers particularly effective at capturing subtle patterns and features within complex, high-dimensional data, even when such information is dispersed throughout the input sequence.
Despite the ubiquity of attention- based models in contemporary deep learning practice, a rigorous theoretical understanding of their working mechanism still in its early stages. large body of theoretical works has focused
understanding the benefits of attention in simple solvable models, 2023: Ahn et al.) 2023) Von Oswald et al., [023 Edelman et al. [2022) Hahn Ahn
e.g. (Geshkovski et al.
2020
Bordelon al.
2024 Bietti al. 2023
Maulen-Soto et al. 2025). with particular focus

---

Published as a conference paper at ICLR 2026
devoted single-layer architectures. Recently, line of studies has demonstrated the advantages of attention- architectures for sparse token regression tasks settings where labels depend only
on small subset of imput tokens (Oymak et al., 2023 Marion et al. 2024 Sanford et al. 2023
Wang et al., 2024] Mousavi-Hosseini et al., 2025 ZZhan et al., 2025 Ren et al. 2024). In such tasks, attention mechanisms dynamically identify and prioritize the relevant tokens, significantly enhancing learning efficiency. In contrast, fully-connected architectures require exponentially more
samples (Mousavi-Hosseini et al. 2025) or neurons (Sanford et al. 2023 as the input sequence
length grows.
In many applications, however, the sparsity of informative features is frequently compounded by additional challenges, notably the weakness and rarity of the underlying signals. For example, cancer diagnosis from computed tomography seans involves detecting lesions- that are typically subtle (weakness), appear in varying locations (sparsity), and occur infrequently (rarity). All these characteristics significantly complicate the detection problem. Motivated by scenarios of this kind, we examine statistical classification problem in which positive samples contain weak signals embedded within a small, randomly selected subset of tokens. We analyze the capacity of a single attention layer to learn to adaptively identify and enhance these sparse, weak, and potentially rare signals. Specifically, our main contributions are as follows:
In the limit of large sequence length L, we show that an attention model can detect signals that are exponentially weaker in L than those detectable by non-adaptive linear classifiers.
Moving from representational power to learnability, we study training at finite L and derive an exact characterization dodoriizatiiooooonnnn to explicit constants of the test error for the attention model after two gradient updates, followed by full optimization of the last-layer weights, in the limit of high-dimensional token embeddings with proportionally large sample size. These sharp asymptotic results quantify precisely how the test error depends on the number of samples, the sequence length, and the signal strength.
Our analysis demonstrates that merely two gradient steps suffice for the attention model to develop meaningful internal representations. Consequently, the classifier can dynamically identify and selectively focus on the relevant subset of hignal-beringg tokens effecciveel amplifying the signal-to-noise
outperform línear classifiers.
To provide complementary perspective on the advantage of attention, we characterize the capacity of the attention model, defined as the maximal dataset size that can be perfectly fit with high probability, and compare it with the corresponding capacity of linear classifiers.
RELATED WORKS
Theoretical analysis of transformers and attention models. The expressivity of attention-basee architectures has been extensively studied in recent literature. Fuu et al. (2024) established that single multi-head attention layer with fixed weights can represent a broad class of permutation- invariant functions. [Edelman et al. (2022) observed that the statistical capacity of bounded-norm attention models scales only logarithmically with sequence length, suggesting strong inductive bias toward sparse functions dependent on only a subset of input tokens.
Sparse token regresssioncllasiffcatiinn tasks. A special class of sparse functions is studied in further detail by (Sanford et al. (2023), who consider a sequence-to-sequeel quences, where outputs correspond to the average of dynamically selected subset of R << L tokens. Whereas culle-coneeeted architectures require 22(L) hidden units to represent such functions, atten- tion models only need (2R)) and can provably learn the task via gradient-based training on the population risk (Wang et al., 2024). Complementing these findings, MMousvvi-Hosser establish corresponding results demonstrating significant separations in terms of sample complex- ity. Similarly, MMarin et al. (2024): (urannhhn et al. (2025) prove that a softmax attention layer can Maaiiar learn a single-token regression up to Bayes-optimal error, whereas linear attention fails, and linear regression on flattened samples performs poorly due to its inability to adapt to dynamic sparsity. Ad- ditionally, recent work by Zhang et al. (2025) analyzes a sparse classification task where the relevant token locations are fixed across samples. Closer to work, Oymak et al. (2023) study a related classification task with the same model as the one considered here, and prove that reaches a good accuracy after three steps of gradient descent, outperforming linear regression with average- peooing..ooling.
task on length L se-
et al. (2025)
2

---

Published as a conference paper at ICLR 2026
Our current work builds upon and significantly extends this line of research along multiple fronts. On the technical level, we crucially extend the analysis of sparse token tasks to arbitrary convex losses beyond the square loss which is considered in prior works (Marion et al. 2024, Mousavi-Hosseimi eet al. 2025; Wang et al., 2024] Oymak et al. 2023). Our extension importantly includes classical loss functions such as the logistic loss, of particular relevance for classification tasks. Furthermore, while most theoretical works have focused on studying the challenges posed by signal sparsity, we further address the often concurrent hurdles of signal rarity and weakness. We demonstrate that attention mechanisms can adaptively address all three challenges by dynamically selecting informa- tive tokens and amplifying their signals. In these respects, our manuscript provides a fully rigorous and encompassing analysis of empirical risk minimization in a classification setting.
PROBLEM SETUP
Sparse token classification. We consider a binary classification task on L x d covariates, seen as sequences of L tokens embedded in d dimensions. Positive samples contain a weak signal added to a random subset of tokens; negative samples do not display the signal. The learning task consists of discriminating samples with the signal from those devoid thereof. In a similar spirit to the sparse- token regressionlclasiffcatiin problems studied in (Sanford et al. 2023 Oymak et al. 2023 Wang jet al., 2024) Marion et al.. 2024) Mousavi-Hossent
et al.) 2025). the difficulty of the task lies in the fact that the location of the signal varies from sample to classifier must dynamically detect and attend to the relevant tokens. Formally, let D (X;, y/ Letn)
any successful
be the training data where each sample X; R Lxd has rows representing token embeddings, and the labels y E [1] are such that P(y: = 1) =: æ E (0. 1). We assume that the token matrices {Xi)ve(m) are independent and drawn from one of two probability distributions. Specifically, for negative samples (namely given y/- = -1),
X = Zz
(1)
where Z, E L is matrix whose entries are i..dd standard normal random variables. In contrast, for positive samples (y:- = 1),
X; = Ou;E + Z.,
where & is a fixed signal vector with ||E|| = 1, 0 > 0 is the parameter indicating the sig- nal strength, and v; is a random binary-valued vector indicating the location of the hidden fea- tures: Uj = [lien,
[LER]] RG denotes the subset of tokens that contain the signal, and assumed have fixed cardinality |R, R E N. The law of 2; is thus supported on {o E 10,1)* Ee It = R}, and we furthermore assume its marginals p, = P(,, = 1) for j E [L] to satisfy ||p|| S CR\VI for some constant C' > 0. This assumption essentially requires that the distribution is sufficiently spread out across tokens, and is not localized on any privileged tokens- thereby making its detection particularly challenging. In particular, when the law of the non-zero elements of e is the uniform distribution with the capacity to generalize on the task must be able to adaptively identify the subset RG contain- ing the signal, if the sample is positive, in addition to learning the signal vector g. The latter point is further rendered non-trive aoiviiivial by the observation that in (2). the signal part ou,E_ is of norm ( R),, which is considerably weaker than the background noise term 2+ Ld) when d and/or L are thggeby making the signal hard to detect. Note that this scaling differs from that considered in (Marion et al. 2024) where both terms are comparable in size-a a regime corresponding more easily detectable signal in the limit of large dimension d.
all subsets of [4]. Ilell = R//L. Therefore, algorithm
Intuitively, the data distribution 22) could be interpreted as simple model of a vision task, where each token corresponds to patch of an input image (e.g. computed tomography scan), and where the location of the feature & signals the presence of certain pattern (e.g. lesion) at the correspond- ing position. This pattern is sparse (R < L), weak ([l6g{ || < ||Z]D, and potentially rare (small TT. The data distribution and task is similar in spirit to that considered in (Oymak et al., 2023). with however two important differences. While in OOymak al. 2023) the signal is present in all samples, in the current work the signal is totally absent from negative samples, posing the additional challenge of rarity. In addition, the relevant tokens R, are devoid of any noise in (Oymak et al., 2023) and contain only the clean signal. On the other hand, in (2) the weak signal ş is corrupted by the additive noise Z;, posing the challenge of signal weakness.
3

---

Published as a conference paper at ICLR 2026
1.1 TWO LINEAR CLASSIFIER BASELINES
We first introduce two simple linear classifiers that will serve as reference models, providing bench- marks against which the attention model- specifeed in the next subsection-will be evaluated.
Vectorized linear classifier. The first baseline flattens each matrix-valued input X; E 1Rdxdd into an Ld-dimensional feature vector vec(X.) = [(X?) ... [I which is then fed to linear classifier. Explicitly, the classifier is
(3) As noted in (Marion al. 2.24), the location of the signal within the vector would then be shifting from sample to sample due to the randomness of R. Ra-maning it challenging for this vectorized linear classifier to pinpoint the relevant features.
vee (X) = sign( ((w, Svee + b) where (vec()) = vec(X), w E E DE
Pooled linear classifiers. A possible remedy would be to instead average the input along its first dimension, rather than flattening it. More precisely, the classifier becomes
pool [Xo)(X) (X) = sign((w, (Xoo)(X)) + 6) where (xool(X) = 1 Exx, web
x² w EER (4)
KE|L]
While such average-pon pooaggeppooligg featurization bypasses the challenge of dynamically shifting signal positions, it introduces another complication. Specifically, after averaging, the norm of the signal term || /cll LR) can become significantly weaker compared to the background noise term ||/]l O( d/T.), especially when R is small and L is large. In other words, the averag- ing procedure effectively reduces the signal-to-noise ratio. These intuitions will be made precise in the following section by Theorem and Proposition]] II, which show that a large signal strength is needed to counteract these limitations, in order for linear classifiers to generalize.
1.2 AN ATTENTION MODEL.
Ideally, to remedy the issue of signal dilution suffered by the pooled linear classifier, a non-uniform, sample-dependent reweighting of the tokens should instead be deployed, selectively placing more weights on tokens that embed the signal. As we will discuss and formalize, such a reweighting can be readily implemented by an attention- -aase mechanism. This intuition motivates the principal model analyzed in this work: a single-layer attention- based architecture designed tackle the sparse token classification task. Specifically, we consider the model
(5) This attention model Aq, b is parameterized by two trainable weight vectors g, w E Rd and a train- able scalar bias E R. In (5). the parameter 3 represents the inverse temperature of the softmax activation. The formulation (5) is simplified attention model widely studied in theoretical contexts (see, e.g., (Oymak et al. 2023 Marion et al., 2024)), in which the representation (X) can be viewed as analogous the [CLS] token used for classification and readout in transformer archi- tectures (Kenton & Toutanova 2019). A detailed discussion connecting this simplified model with standard self-attention architectures can be found in (Marion et al., 2024 Tarzanagh et al. 2023). While we devote the theoretical study herein to the model (5). we compare empirically with more complex attention- based architectures in Appendix A
Ag.u, (X) sign (fg (X), w) + b),
with
fx)) X TT (30ft
Dynamic reweighting and signal amplification. An important feature of the model (5) is that the weight vector w acts not directly on the raw input X, but instead on the attention-batt based feature:
3(x* m)
fx(X) E
X
(6)
oiix .
KEILI ccL4)
where each token x is reweighted according to the scores 3(x* 4). Crucially, in contrast to the naive average- average-pooling (4) discussed in subsection 1. (which corresponds to the special case of q = 0g), the attention scores dynamically adapt to the input tokens. Therefore, in principle, the attention mechanism can allocate greater weight to tokens containing the signal E. thus mitigating the diminished signal-to-nn toisee ratio described following 4 Such improvement occurs when the internal attention parameter q aligns non-trivially with the signal vector &; this alignment increases the inner product (x q) and consequently enhances the attention weights (6) for the signal-bearing bearing tokens. In Section 3 we formalize and rigorously prove this intuitive mechanism.
4

---

Published as a conference paper at ICLR 2026
2 OPTIMAL TEST ERRORS IN THE LIMIT OF LONG SEQUENCES
Before analyzing how effectively the attention model 5 and the two baseline linear classifiers ( and 44 perform when trained on the sparse classification task described in Section it is instructive to first determine the conditions under which these models can, in principle, learn the task. In this section, we examine the optimal test error of the considered hypothesis classes, measuring their intrinsic ability to represent the sparse classification problem. Formally, the optimal test error for any predictor Dw (X) parametrized by some mmite-dimensional parameters W is defined as follows:
test [p] inf Exest [Dw] where [iw := Px,y low (X) # y]
The optimal test error corresponds the smallest misclassification error achievable by the classifier, provided its parameters W are selected optimally. Concretely, for the vectorized and pooled linear classifiers defined herein, W is given by (w, b) E R Ld x R and (we b) E IR x IR respectively whereas for the attention model one has W = (w, q, b) E F x Rd x I. In this section, we view 8, R as sequences depending on L, and focus on the + DO regime.
Proposition 1. Suppose that the limit SNR == lim on/VI exists. Then, the optimal test error of the pooled linear classifier satisfee
if SNR 00,
lim Ejett [Pool] L-000
(1 ()) + T SNR) if SNR (0)
(8)
min (T, 1-T) T)
if SNR
In the above display, b" = SNR
log(1/T - and Φ(.) is the cumulative distribution
SNN
function of the standard normal distribution.
We note that similar result appears in (Oymak et al. 2023) (Appendix A) for the pooled classifier, but for a different data distribution, and without a trainable bias
Theorem 1. Suppose that the limit SNR == limg-+00 vL
2R
exists. Then the optimal error of the
vectorized classifier satisfies
0 min(( ,,, 1-)) m)
if SNR = 00; if SNR = 0
lim Eiett [ vec L-700
(9)
and lim inf Ejst
[L vec
SNR E (0)
The proofs of Proposition and Theorem 1 are detailed in Appendix C Note that since the pooled classifier can be viewed as a particular realization of the vectorized classifier with tied weights, the optimal error of the former (Proposition I upper bounds the optimal error of the latter (Theorem T.. Concretely, to generalize perfectly on sparse signals = (()), both the pooled and vectorized linear classifiers require a strong signal strength 0 = ((VL). In this regime, the optimal weights of the pooled (resp. vectorized) classifier are proportional to the signal & (resp. to concatenation of & L times), and allow for vanishing test error. If the signal is weaker, namely 0 o( (()L), the model performs no better that the naive predictor that always outputs the majority label and Eesst min( т). In the case where SNR (0, xo), Theorem I shows that the optimal test eτror is bounded away from zero by a strictly positive number. In sharp contrast to the linear classifiers, the attention model can perfectly classify data with a much smaller signal strength:
Theorem 2. Consider the attention model A given in (). In the limit T, oo with R O(1), suppose that the signal strength f satisfies lim inf - 8/log > O. Then, one has Eiett [A] = 0.
The proof of Theorem can be found Appendix A direct consequence of Theorem that significantly milder signal strength of order 0 = Tog L suffices for the attention model (5) to perfectly learn the sparse token classification task task- provided employs optimal parameters q. w. 6. Similarly to the linear classifiers, perfect classification is in particular achieved for weights g, colinear to the signal ç. Similar results appear in (Oymak et al. 2023) on the optimal error in a related task, but are restricted to simpler noiseless case (Z, = 0). While Theorems [ and 2 and Proposition I paint a clear separation between the attention model and the two linear baselines in terms of representation power and oracle test errors, they leave the question of learnability largely
5

---

Published as a conference paper at ICLR 2026
open. Furthermore, this clear-cut distinction, which happens in the large-L limit, becomes less pronounced when the sequence length L is finite. Thus, a more nuanced analysis of the training at finite sample complexity and sequence length is warranted. This is the objective of the following section.
3 PRECISE ASYMPTOTIC ANALYSIS OF THE LEARNING
In what follows, we turn our attention to the study of the training of the three models on finite datasets, aiming to precisely characterize the learning behavior of the attention model (5) and the two linear classifiers 3 and (4) in this regime. Such exact characterizations become tractable in the high-dimensional embedding limit, as demonstrated by a growing body of literature on high- dimensional attention mechanisms (Rende et al. 2024] Cui et al. 2024a| Troiani et al. [025 Tiberi
et al. (025). We adopt in the remainder of this manuscript the following high-dimensional, length scaling regime:
al. 2024. Cui 2025 Erba et al.
Assumption 1 -iigh-dimensional, finite-length limit), We consider the limit of large embedding dimension d and comparably large number of samples n, namely d, + oo with fixed ratio
= x/d (1). The chosen scaling ~ d is such that the detection of the weak signal from the background Z is statistically possible (Lesieur et al. 2015), yet non-trivial. Meanwhile, the sequence length L, signal strength 0 and sparsity R. along with all other parameters. remain finite and fixed.
Training procedure. We now turn to the learning process. The attention model (5) can be trained to solve the sparse token classification task defined in subsection by performing empirical risk minimization over the dataset D = (X,, vel) formulated as follows:
i wr, b E argmin&D b), with RD w,b) 1 E e(Ug((), + b, v) 2 |v:11 (10) (X.Y)D
g,ur6b
Here, (-1,) TR loss function that convex with respect to its first argument (for example, the logistic loss {(z, v) log(1 + exp(-yz)) or the quadratic loss {(z, v) - {= y)
Notably, compared to The empirical risk (TO) also includes a ridge regularization of strength Mousavi-Hosseini al.)
prior studies sparse token tasks (Sanford et al. 2023 Wang al. 2024
Marion et al.] 2024; Oymak et al. 2023). our setting extends beyond the squared loss to general convex loss functions. A natural approach solving the non-convex optimization problem 10) is to run gradient descent on the set of trainable parameters q, w, b. In fact, as demonstrated below. just two gradient steps are sufficient for the query weights q to achieve an alignment with the signal E. This alignment enables the attention model (5) to develop internal representations capable of effectively identifying and amplifying the hidden signal. Specifically, we consider the following training procedure:
2025
1. Initialization. Consider a partition of the training data D D, U D2 into two disjoint sets of sizes no and n. = n no respectively. We assume on no/d = (1)), and aj = m//a O(1).
Initialize the weights of the attention model (5) as w)()) ()) = O.,, ((0) = 0.
2. First gradient step on b, q, w. Perform first gradient step on each of the trainable parameters
on the risk R.D (q, w, 6). using the training set D, with learning rates no. ng Twi
d
b(1) ((0) Th db
R Do (4, w. bllb:o (00))
w (1) = (0) - n V R (g. w. b) b003 1000 (w())
q 1) g00 (0) na v y RD. (g. w. b) b00) (00)) w0)0
(11)
More explicitly, we obtain the parameters
( noL
Tw Σ 0, Iz, e.Sng
((1) = ?7b Σ h,(0,0,0)
q0) - 0z
(12)
no
where 1, E L denotes a vector with all elements equal to 1, and we have also introduced the shorthand notation he(g, w. b) = ((((((g)) (g), w) + b;

---

Published as a conference paper at ICLR 2026
3. Second gradient step q. Note that after first step, q0)) remains zero. For the attention model 5 to develop non-trivial internal representation parametrized by * Od, a second gradient step
on q is thus needed:
q (2) q(t) TVV.D. (g. w. 600,, g00, 600, will)
n Lo
Σ h, w ) (I- 1217 I
Xww
(13)
4. Full training of w; Having developed a meaningful internal representation parametrized by g( the readout weight w and bias b are finally fully updated by empirical risk minimization on the retained data Dj:
ú b = argmin RD, (q( w. b).
(14)
web
The performance of the trained model A 10, is measured by its training loss and test error
Eraiin RD, (q (2) ú. 6),
Exest Px.. A AA .n (x)) +
(15)
The latter simply corresponds to the probability that the trained model misclassifies a fresh test sample (X, y) at inference time.
The primary purpose of the dataset partitioning performed step 1-splitting the data into two subsets, used respectively for steps 2-3 and step 4-is to simplify the subsequent analysis of step 4. This partitioning ensures statistical independence between the learned query weights (8)) and the dataset Dj. In Appendix [A] we note qualitatively similar performance when the data from steps 2-3 is reused for step 4. Adopting a more practical viewpoint, Do can also be viewed as pre-training dataset used to train the query weights q, which can then frozen as the model deployed on other datasets, with only the readout and bias being fine-tuned.
Similar stage- training protocols with sample splitring have previously been analyzed in the context of two-layer neural networks (Ba et al., 2022 Moniri et al. 2023 Cui et al. 2224bb Dandi ctt al. 2024 2023), demonstrating how even a single gradient step on the first-layer weights can yield meaningful internal network features. Analogously, in our setting, two gradient steps the query weights q are already sufficient for the attention model to develop informative internal repre- sentations. For transformer models, similar few-step analyses were conducted for instance in (Bietti) et al. 2024 Oymak et al. 2023), however without the final step of full empirical risk minimization. This final optimization the output weight w can be taken an analog transfer learning, thus lending to more practical insights for real training procedures.
We now in position present our main technical results: precise characterization of the test error (15) achieved by the attention model (5), trained using the four-stage procedure detailed in subsection In the following sections, we first analyze step 3 demonstrating precisely how the query weights (2)) develop an alignment with the signal &. resulting in nontrivial attention weight- ings. We then examine how this learned attention mechanism leads to an improvement in the test error (15), as compared to the baseline linear classifiers and (3) at the conclusion of step
Characterization of the attention weights after two gradient steps. The first technical result characterizes how, at the end of step (see subsection 3). the query weights q = g2)) develop a non-zero alignment with the signal vector ç. As will be discussed in subsequent subsection, this alignment allows the attention model (5) to develop internal representation adapted to the task.
Theorem (Characterization of the query weights (2)) after two gradient steps). the asymptotic limit of Assumption ! [2) and (2) { converge in probability to deterministic limits, whose
expressions are given in Appendix
Theorem precisely characterize the parameters of the attention model (5) at the conclusion of step of the training procedure described in subsection 3 The detailed proof of Theorem 3 is provided in Appendix F A direct consequence of Theorem ii that the alignment between the query weights after two gradient steps (²)) and the signal &, captured by the cosine similarity ") tends rapidly in absolute value to its maximal value of 1 as the sample complexity ag is increased, at a foo rate. This observation is formalized in the following Corollary.
7

---

Published as a conference paper at ICLR 2026
Corollary 1 (Cosine similarity). In the asymptotic limit of Assumption 7 the cosine similarity (3/1/// converges in probability to limit Sq. furthermore the right-hand side of (22) is non-zero, its absolute value admits the expansion |syl 111=11 C/ao + + The expression of the constant C is detailed in Appendix
Characterization of final test and training errors. Having described steps 1 of the training procedure 3 we now focus step 4, where given g(²), the readout weights and the bias b are fully trained on the held-out data batch Dj. We note that once q is fixed, the empirical risk minimization (14) amounts to training linear model with weights w. b on the high- -ith-diinnsional non non-linear fea- tures fo (X) (6. While the behavior of such linear classifiers is in general very well understood
Liang & Sur 2022 Montanari al., 2019 Mai et al. 2019; Loureiro et al. Mignaaco et al. 2020a). such works very often build on the assumption of simple (c.g. Gaussian mixture) data distributions. In the present case however, the features f (X) possess a highly non-trivial distribution, as they result from the non-linear at- tention mechanism. Fortunately, the softmax acts only on the low-dimensional projection y E R- of the tokens along the query weights q(², which can be handled separately. The idea of the proof, detailed Appendix F proceeds from this observation. The final results are succinctly summarized in the following theorem, while the full technical statement is deferred to Appendix
in the asymptotic limit of Assumption [ (Candès & Sur 2020
Theorem (Test and training errors after step 4). The test error and training loss associated the empirical risk minimization (14) converge in probability in the limit of Assumption 7 to deterministic limits Erest (A] and Etrain (A), whose expression are deferred for clarity to Appendix
0.50
0.200
0.45
025
0.175
Vectorizee
Pocled
020
0.150
0.125
Etroin
015
Pooled
Attention
0..000
Vectorized
Attention Pooled
0.050
0.025
Vectorized
0.0
0.000
C2:
Figure 1: Test (left) and train (middle) errors achieved by the attention model (5) and the pooled (4) and vectorized (3) classifiers, for I, = 10, R 1, 0.5, 0 = 5, À 10-5 "bb.g... 0.1, 0gg on trained with the square loss, as function of the normalized number of samples or Solid lines correspond to the theoretical characterizations of Theorem 4 Dots represent numerical experiments in dimension d = 1000. Error bars represent one standard deviation over 8 trials. (right) Training loss Etrain for the attention model (green), and the pooled (red) and vectorized (blue) linear classifiers, as a function of the sample complexity a1- = 2, R = 1, 2, 0.3. The attention model has unit norm query weight with alignment 7 0.99 with the signal &. Dots correspond to numerical simulations in dimension = 2000; error bars represent one standard deviation over 20 trials. Dashed lines: theoretical prediction of the separability thresholds, given in Conjecture
Theorem provides exact characterization--rree attained by the attention model (5). trained according to the procedure described in subsection 3 within the high-dimensional limit specified by Assumption [ The resulting expression is formu- lated in terms of a small set of scalar summary statistics, which are determined as solutions to system of self-consistte equations. While the latter still possess rather intricate form, they can considerably simplify in some simple cases, yielding valuable insights. We detail such an instance in the following, for the case of a square loss in the ridgeless limit. Let us remark that while (Oymak) et al. 2023) also provide error bounds for three-gradient-steee protocol, Theorem 4 offers tight error characterizations, exact down to explicit constants. While the same work also reports sharp characterizations (Theorem 8) for the special case of the square loss, those results are restricted to much simpler learning protocol that involves neither gradient steps nor empirical risk minimization, and that necessitates further oracle information the of relevant tokens.
down to explicit constants-of the test error
8

---

Published as a conference paper at ICLR 2026
Baseline classifiers. Having characterized the test enror and training loss of the attention model, now turn to the case of the two linear classifiers u.,b Lpool whose parameters w, b are trained on the dataset D through the empirical risk minimization
i E argmin Ro, (w,,b) with Ro, (w,b) = Σ ((U(()) w) + bzu) + 2 1001, (16) (X..)DD
web
where S E (pool, Jver and is an arbitrary strictly convex loss function. As for the attention model, a tight characterization can be reached for the associated test error and training loss, leverag- ing the observation that the distribution of the features (X), ppoal (X) are in fact simple Gaussian mixtures with respectively (x) + 1 and 2 isotropic clusters. The test error and training loss of gen- eralized linear classifiers in the high-dimensional limit of Assumption T for such data distribution
has been characterized in prior works (Mignacco et al. 2020a; summarize the corresponding results below.
Loureiro et al. 2021). We briefly
Theorem 5 (Errors for the linear classifiers). L'Luureiro et al. 2021 [ In the asymptotic limit of
Assumption 7 the test error and training loss for the pooled (resp. vectorized) linear classifier converge in probability limits Etrain (Loool and &test (test(Looo)) (resp. Errain[Lveo] and and Erest [L )
We defer the precise exposition of the expressions of Etrain Freon]Lveo), completeness, and to help readers connect and compare the proofs of Theorems 5 and we also present in the same Appendix an alternate sketch of proof using the same leave-one-out approach as that leveraged in the proof of Theorem
Etest [Lveo] to Appendix G For
Comparison of the three models. The theoretical predictions for the training and test errors from Theorem and and [ both the attention model 5 and the linear baselines (3 are compared with numerical simulations in dimension d 1000 in Fig. demonstrating excellent agreement. The figure clearly illustrates how the learned attention mechanism leads to superior test performance compared to the linear classifiers, which lack this adaptive representation capability.
To garer further quantitative insights from the technical results of Theorem and let us focus on the particular case of a quadratic loss function f(z, y) 1/2(y =)², in the limit of vanishing regularization À = 0+. In this setting, the characterizations of Theorems 4 and 5 considerably simplify, revealing further insights, which we describe in the following Corollary.
Corollary 2 (Ridgeless quadratic loss). For a quadratic loss function E(2,y) 1/2(y 2)2 and = O, the asymptotic limits Erest [A), Erest (Loos) and Erest vec
[Lvec] characterized in Theorems andd 5
tend to their a1 + 2o limits E [A] and E E [[Pool]] E voc at rate 1/a..
A number of interesting conclusions can be gameered from Corollary 2 First, all three test errors tend to their respective a - xo limit at the same 1/a, rate, as the sample complexity 01 is increased. Furthermore, the two linear classifiers L [pool LL Lvee tend to common limit E This finding somewhat echoes the intuition from Theorem I which already suggested that both models share similar oracle- and thus plausibly infinite sample complexity-behateee wonder which of the limiting test errors & (A), & [L] is lower-in particular, whether the attention model always achieves a lower error provided it given sufficient data. The answer is more nuanced, and crucially depends on the alignment sy (see Corollary I between the query weights g(2) and the signal & achieved after step 3 of the training protocol. As shown in Fig. 5 in Appendix F. Eee [A] > S [L] can hold in some settings for s, sufficiently small. In simple words, when the query weights have insufficiently aligned with the signal-e.g. as a result of insufficient data og or bad choice of the hyperparameters b the attention suffers from misaligned internal representation, and achieves a worse error than the simpler linear classifiers. For moderate and large Sy on the other hand, Eos [A] [ [L] and the attention profits from the advantage of the dynamical reweighting implemented by its internal representation.
Lastly, one may naturally
The previous subsection compared the three models in terms of their test errors. We adopt in this subsection complementary perspective, and analyze the capacity a* of the models ,u, - and L L. defined as the (normalized) maximal number of training samples that can typically be fitted by the models to vanishing training loss. More formally, let 인 E (A, Loool L L be one of the three models, and let Erain (01) designate the asymptotic training loss characterized
Capacity.
9

---

Published as a conference paper at ICLR 2026
in Theorems and in the limit of vanishing regularization À 0, for the logistic loss f(y.2) log(1 + exp( y2)) The capacity of the model ÿ is then formally defined as
a
sup 222 20 [a] = 0}
(17)
For agj. the training set is small enough that it can with high probability be perfectly sep- arated by the model and Etrain [ú] (a)
At large sample complexities a ag. such perfect classification becomes typically impossible, resulting in a positive training loss Eraiin (a) > O. The capacity of a model captures how easily it can classify samples from a given data distribution, with a higher capacity thus intuitively reflecting higher adequacy of the model to the task. An analytical expression for the capacity can be extracted from the characterizations of the training loss Etrain provided by Theorems Saareessaadde and which we report in the following Conjecture.
Conjecture 1. The capacities of the models L Lpool, L Lvec Ag(e), an admit the following expressions:
L(1 - 5²)
-c max sE(0.1].b T (b+ vL ut + x) x)
max aA
my me
cgut -00- mg #egme)
uduu
(18)
and Goooo = Rio//.. The expectation in the expression of a bears on y, Cz. Cg, Cq whose joint law is detailed in Lemmal 7 and depends in particular (g
The derivation of the expressions 18 is detailed in Appendix Because they involve some heuristic
step, state the result as conjecture. The capacity of linear classifiers has been studied in rich line of prior works, e.g. (Candès & Sur 2020; Mignacco et al., 2020a Loureiro et al. 2021), impulsed by the seminal work of (Cover) 2006), albeit analytical expressions have been to awareness reported for the data distribution considered in the present work. Such results are on the other hand scarce for attention-baset models. Conjecture Caauureecooooo contributes bridging this gap, by reporting an analytical expression for the capacity of the simple attention model considered in the present work. The theoretical predictions (18) are plotted in Fig.[, where they are overlayed upon numerical evaluations of the training loss Etrain. for the three models, revealing good agreement. In the probed setting, oxeec > a > Poot the attention model displays higher capacity than the pooled classifier, while the higher capacity of the vectorized classifier can be explained from its operating in T.-times higher dimensional space. As we discussed above, this higher capacity of the attention model intuitively hints at a better suitability to the considered data distribution.
Finally, we note that this ordering can vary depending on the parameters of the problem, and cru- cially on the alignment sy achieved by the attention model between its query weights y(2) and the signal a, as characterized in Corollary [ We discuss in Appendix how a small sy rysrullg, for instance, from insufficient pretraining data or bad choice of hyperparameters Te can result in the attention having a lower capacity than the pooled classifier, namely OA < "poot This echoes a similar observation the level of the test error made in the previous subsection, and discussed in Appendix
CONCLUSION
We study the sparse token classification task of detecting a sparse, weak, and rare signal embed- ded in sequential data. For long sequences, we rigorously establish a clear performance separation between linear and uttention-basee classifiers. showing that attention- icantly weaker signals to achieve perfect generalization. For finite sequences, we provide sharp analysis of the learning for simple attention model in a high-dimensional limit. Specifically, our study demonstrates how merely two gradient steps suffice for the attention mechanism to learn meaningful internal representations, enabling the model to dynamically identify and focus on to- kens containing the relevant signal. Moreover, we derive a sharp characterization of the resulting test error, quantifying precisely the performance gain achieved by the attention model relative the linear classifier baselines. Finally, we put these results in perspective by analyzing the capacity of the three models.
models require signif-
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We thank Odilon Duranthon, Lenka Zdeborová, Claire Boyer, Pierre Marion, and Bruno Loureiro for insightful discussions. NB acknowledges support from the Fonds de recherche du Québec Nature et technologies (FRQNT) in the form of a Master's s Training Scholarship. HC acknowledges support from the Center of Mathematical Sciences and Applications (CMSA) of Harvard University. Y.M.L is supported by a Harvard College Professorship, the Harvard FAS Dean's 's Fund for Promising Scholarship, and DARPA under grant DIAL-FP-038.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614 2023.
Luca Amaboldi, Bruno Loureiro, Ludovic Stephan, Florent Krzakala, and Lenka Zdeborova. Asymptotics of sgd in sequence-single index models and single-layer attention networks. arXiv preprint arXiv:2506.02655 2025.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivias On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310.19793, 2023.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36, 2024.
Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head trans- former dynamics. arXiv preprint arXiv:2405. 15712, 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners, 2020.
Emmanuel J Candès and Pragya Sur. The phase transition for the existence of the maximum like- lihood estimate in high-dimensional logistic regression. The Annals of Statistics, 48(1):27-42, 2020.
Thomas M Cover.
Geometrical and statistical properties systems of linear inequalities with ap- plications in pattern recognition. IEEE transactions on electronic computers, 3:326-334, 2006.
Hugo Cui. High-dimensional learning of narrow neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2025(2):02342,, 2025.
Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. A phase transition between positional and semantic learning in a solvable model of dot-product attention. arXiv preprint arXiv:2402.03900. 2024a.
Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue Lu, Lenka Zdeborová, and Bruno Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. Pro- ceedings of the 41 International Conference on Machine Learning, pp. 9662-9695, 2024b.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. arXiv preprint arXiv:2305.182270 2023.
Yatin Dandi, Luca Pesce, Hugo Cui, Florent Krzakala, Yue M Lu, and Bruno Loureiro. A ran- dom matrix theory perspective on the spectrum learned features and asymptotic generalization capabilities. arXiv preprint arXiv:2410. 18938, 2024.

---

Published as a conference paper at ICLR 2026
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbor,, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021.
Odilon Duranthon, Pierre Marion, Claire Boyer, Bruno Loureiro, and Lenka Zdeborová. Statis- tical advantage of softmax attention: insights from single-location regression. arXiv preprint (rXiv:2509.21956 2025.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference Machine Learning, pp- 5793-5831. PMLR, 2022.
Vittorio Erba, Emanuele Troiani, Luca Biggio, Antoine Maillard, and Lenka Zdeborová. Bilinear sequence regression: A model for learning from long sequences of high-dimensional tokens. arXiv preprint arXiv:2410. (tt::2410.1885555 18858, 2024.
Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. Advances in Neural Information Processing Systems, 36, 2024.
Elizabeth Gardner and Bernard Derrida. Optimal storage properties of neural network models.
Journal of Physics A: Mathematical and general, 21(1):271, 1988.
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. mathematical per- spective transformers. arXiv preprint arXiv:2312.1079) 2023.
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8: 8:156-171, 2020.
Noureddine Karoui. On the impact of predictor geometry on the performance on high -igh-dimensional ridge-regularized generalized robust regression estimators. Probability Theory and Related Fields, 170, 02 2018. doi: 10. 100007776004400006-7755599
Jacob Devlin Ming- Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. 2015.
Werner Krauth and Mare Mézard. Storage capacity of memory networks with binary couplings. Journal de Physique, 30020))3057-3066, 1989.
Thibault Lesieur, Florent Krzakala, and Lenka Zdeborová. Mmse of probabilistic low-rank ma- trix estimation: Universality with respect to the output channel. In 2015 53rd Annual Allerton Conference on Communication, Control, and Computing (Allerton), pp. 680-687. IEEE, 2015.
Tengyuan Liang and Pragya Sur. precise high- high-dimensional asymptotic theory for boosting and minimum-ll-nome interpolated classifiers. The Annals of Statistics, 1669)16669 2022.
Bruno Loureiro, Gabriele Sicuro, Cédric Gerbelot, Alessandro Pacco, Florent Krzakala, and Lenka Zdeborová. Learning gaussian mixtures with generalized linear models: Precise asymptotics in high-dimensions. Advances in Neural Information Processing Systems, 34:10144-10155 2021.
Xiaoyi Mai, Zhenyu Liao, and Romain Couillet. A large scale analysis of logistic regression: Asymptotic performance and new insights. In ICASSP 2019-2019 IEEE International Confer- ence on Acoustics, Speech and Signal Processing (ICASSP), pp. 3357-3361. IEEE, 2019.
Pierre Marion, Raphaël Berthier, Gérard Biau, and Claire Boyer. Attention layers provably solve single-location regression. arXiv preprint ar xrX1v:2410.01537, 2024.
Rodrigo Maulen-Soto, Claire Boyer, and Pierre Marion. Attention-based clustering, arXiv preprint arXiv:2505. /3/12, 2025.
Francesca Mignacco, Florent Krzakala. Yue Lu, Pierfrancesco Urbani, and Lenka Zdeborova. The role of regularization in classification of high-din high-dimensional noisy Gaussian mixture. Interna- fional Conference on Machine Learning, pp. 6883. PMLR, 2020a.
12

---

Published as a conference paper at ICLR 2026
Francesca Mignacco, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Dynamical mean-field theory for stochastic gradient descent in gaussian mixture classification. Advances in Neural Information Processing Systems, 33:9540-9995 -550,, 2020b.
Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban. theory of non-linear feature learning with one gradient step in two-layer neural networks. arXiv preprint arXiv:2310.08891, 2023.
Andrea Montanari, Feng Ruan, Youngtak Sohn, and Jun Yan. The generalization error of max- margin linear classifiers: High-dimensional asymptotics in the overparametrized regime. arXiv preprint arXiv:1911.01544 7, 2019.
Alireza Mousavi-Hosseinii Clayton Sanford, Denny Wu, and Murat A Erdogdu. When do trans- formers outperform feedforward and recurrent networks? a statistical perspective. arXiv preprint arXiv:2503. 1272, 2025.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In International Conference on Machine Learning, pp. 26724- 26768. PMLR, 2023.
Yunwei Ren, Zixuan Wang. and Jason D Lee. Learning and transferring sparse contextual bigrams with linear transformers. arXiv preprint arXiv:2410. 234338 2024.
Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. Mapping of attention mechanisms generalized potts model. Physical Review Research, 6(2):023057, 2024.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of
transformers. Advances in Neural Information Processing Systems, 36:36677-36707, 2023.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308. 16898, 2023.
Lorenzo Tiberi, Francesca Mignacco, Kazuki Irie, and Haim Sompolinsky. Dissecting the in- terplay of attention paths in a statistical mechanies theory of transformers. arXiv preprint lrXiv:2405.15926, 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. arXiv preprint arXiv:2502.000 00901, 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances Neural Informa- fion Processing Systems, 31, 2017.
Roman Vershynin. High- Dineeniimessi Probability: An Introduction with Applications in Data Sci- ence. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. International Conference on Machine Learning, pp. 35151-35174, PMLR, 2023.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. arXiv preprint arXiv:2406.06889; 2024.
Chenyang Zhang, Xuran Meng, and Yuan Cao. Transformer learns optimal variable selection in group-sparse classification. arXiv preprint arXiv:2504.08638, 2025.
13

---

Published as a conference paper at ICLR 2026
A IMPARTING TO MORE COMPLEX ARCHITECTURES
We considered so far the simple attention model in (5) and the training procedure described in sub- section Bo allow for tractable analysis. In this last section, we provide some synthetie numerical experiments evidencing the parallels between our setup and more complex attention mechanisms. For comparison, we consider three models that build upon the attention model (5) which refer to as Moo For weight matrices Wq. Wx, Wv E Rdxd let Q = X XWQ K WW V XWv. and define the attention weights anism considered in the seminal work of Vaswani et al. (017)) let M1 and M2 be the models with outputs:
softmax (QKT) /va). Akin to the classical self-attention mech-
MI: : f(X) 1
L Σ (AV),
sign (ff()), w) +b,
(19)
outut
M2: = (WW
sign (h(X), w) + b),
ourpat
where h- ReLU and (Wh, c, w, b) are learnable weights. For final comparison, also consider a multi-head, multi-layer attention model My (4 heads and 2 layers) with linear activation and final output defined analogously 19
Reused data New data
Af
Mr
Etest
Mo
0.3
M2
0.2-
Mf
M6
0.2
0.1
9.1-
2
a1
a1
Figure 2: Test errors achieved by the models {M, for = 10, = 3, = 0.5, 3, and À 10-2, trained square loss using Adam optimizer (pre-training stage) before freezing inner model weights and optimizing readout weights (fine-tuning stage). Error after epochs (left, middle) and 100 epochs (right) of pretraining are shown. Comparison between reusing pretraining data versus generating new data for finetuning (as in subsection 3 is also shown (left). Curves represent numerical experiments in dimension d 500; error bars show one standard deviation over 8 trials.
We plot in Fig. the learning curves of these different models. We employ mini-batch Adam
(Kingma & Ba 2015) instead of full-batch gradient descent, and vary the number of pretraining epochs. Fig. 2
(left) shows that using dataset D2 for the training of w, b, as we considered 3 yields the same behavior reusing the dataset Do employed in the first pretraining steps. Qualita- tively, in all probed settings the test curves for the model MM display a strong likeness to the analytic provided in Fig. This similarity in the shape of the loss curves between the more complex models and the simple model analyzed above can observed even after 100 epochs of pretraining. For instance, the double-descent phenomenon of Fig. rrmmann present. As point of contrast, when using only epochs of pretraining, Mo out- out-performs the other models (Fig 2 middle), which may be attributed to the much larger parameter spaces being optimized over by more complex models. Unsurprisingly, this observation flips with more pretraining (Fig. right).
B AUXILIARY RESULTS
Throughout this appendix, for two random variables X and Y, we write
to mean that the two random variables are equal in distribution. For example, as is used often in our derivations, given a matrix G pm with i.d. A/ 1) entries, an independent Gaussian vector g (O, Im), and another independent random vector E R", basic fact is
Gu
|ullg
14

---

Published as a conference paper at ICLR 2026
Moreover, for two (possibly random) sequences (ax) and (6,). we write
x b
Loo (a, - bn = O, where the convergence may be taken in the almost-sure or in-probability sense depending on the context.
We first present statistically equivalent representation of the feature vector (x) in the attention model defined in 66)
Let R and 2 I& be two independent random vectors with i..dd standard normal entries. Define two probability vectors
Lemma 1.
4 wof.max(3((11))))) + O)v))
and :- s0flma
(20)
We have
1x(0) tu = +13 (4) (g,8+)g + (0v, 8.15 + ||8 1|
|1
and
[x(x)) (-=-11 -1 (d) (g. s 1)9
+ |1s 1|P
Igll
where
99
P1
1
(g. 4)
the orthogonal projection onto the subspace orthogonal to g.
Proof. By the rotational invariance of the isotropic Gaussian distributions, write
Z
94
7P
(21)
||
where is an independent copy of Z. The result is straightforward after inserting the representation 21 into (6). which provides
1(x0 lu +1)
99
7p/ +OuET)
softmax (Bllallg + (300(a))))
|qll|
(4)
(g, (0,84)9 (0v, s+1E ||s 1s
114ll
In the above, we have used the facts that Py q 02 and Zs+ (d) 1s+1lg. The signal-less case (for
y 1) follows analogously.
The following result gives simplified form for the test error that is valid for any q, w Rd and and
b E IR. It will be used in the proofs Theorems 2
Lemma 2. Let
(g. w)
#1
112 ()),),
#3 Vlwl22 2¹7
The test error is
b +
(g, )p4
)
n.. EEg ( - (0v, s4)222 (g. (4,, )]
Erest = =(1 TT- Bg
ual
8
|all8+1|
where s+ and s_ are the two vectors defined in 20).. and 40-) is the cumulative distribution funcctioooonnuncioon of standard normal distribution.
Proof. By (5)
Exst = (1 (1- TTPP (X)( w) +b00 = -1) xP (IG(X)), +b w)) ==1=)000)+)))>> +1).
The result then follows from the statistical representations given by Lemma
15

---

Published as a conference paper at ICLR 2026
C PROOFS OF PROPOSITION I AND THEOREM 1
Proof of Proposition
1
Notice that the pooled classifier corresponds to setting in the attention model (6) and we have
(d)
[w [y=1}
OR(w, E)
(fo(X), w) + VL
+6
L
for 2 ~ NO 1). Absorbing the factor - -VL/ L//l by redefining the variable b, we obtain
@R(w. E)
Exest
inf BER
(P- 6) TP 2
|I0lLvL <
p
inf ,c(-1,1),bER
(1- b) + TP
<
VL
AR
inf BER
(1-m)) (1-=)IP( b) + n|
<b
VL
OR
inf (I - 7) (-b) + オΦ 6
(22)
beR
VL
Set ly 0R/VI and Denote gL (b) (1 - (-b) + @b (L) the function over which the infimum is taken in (22 For any 9L. admits the derivative
$ +
(b(b)
-1+
RE
(23)
V2TTT
We assume without loss of generality that & > 0 since the asymptotic test error shall remain the same for when & + as + co. We have that the derivative g' (b) is zero at
b 1 (cc 2/e2 log (*/1 x)),
(24)
where it switches sign from negative to positive. Therefore, the infimum in 22 is attained at 6% and
Ejest [0] (1 π z))((bb;,) (b; fr.)
1 (er. 2/2. log (5/1 (=)) + +Φ ( (le. 2/1, log (7/1 (=/1 Cs.
(1 )
(25)
Inspecting (25). by continuity of Φ we immediately that when e = one has limy Etest [ O and when L E (0, oo) we obtain the corresponding expression the statement of Theorem Notice that for f 0,
-0,
lim -2/r, log (T/1-x)= 1e
x,, π
0,
1/2
and so, again examining r).
25 follows that under this regime limms Ejstll) min(m,) 1
Proof of Theorem
we begin with simplification of the optimal test error. Writing w = (Wi,..., wT. with wy we have
Before dividing into the two separate cases of Theorem
(voc(X), w) + b = || 2 1 0
ve we { +
16

---

Published as a conference paper at ICLR 2026
for 2 2 ~ 1). Absorbing ||||| into b gives
(wk, E)
Gistlll
inf WERLd,bER
(1- (T= b) + xPP
- 0
|w| <
t=1
inf westa-l, (1- b) + xP
2 +
Ui ve E) <
inf
(1- b) TP Ov, a) < b)
(26)
aE8t-inE4,nd 1oe.. DER
The last line follows as any optimal w will be the form wy a,5 for S l < where az Z and ||||| 1.
With the representation for Eest ! rem.
given
(6), we now establish the separate results of the theo-
1. Recalling the assumption
R VI.
|ll| =
where Pj = P(U; = for j E [L],
(27)
there exists C > 0 such that ||l|| s C(R/VL) for all L 2 1. To begin, defining the
random variable 0(v, a) and the decreasing function fx(x) + b), notice that (26) is equivalent to
inf aEs4 indf bER
+
(28)
where the dependence on persists through u and the expectation taken with respect to
L - oc. To this end, consider "flat" solution a = 1/VI . 1 and notice that this gives u = E[u] = 9R/VI. Thus, we have [u]f((w)] fo(on/vI) Φ--on/vI b).
We first show that if 8 co, one has E [0]
Taking 0R/2/1, we have
Ejstt (1 10)-))0(--00) (-)0((0/221) + (--0R/2111 L 0.
Next, we show that Eiett [i]] + min(T, - T) if l" = 0. Observe that for v > 0,
Bu(fo(u]] 2 [., (
2 (a)))) - k)
2 fb(k) (1
E|u]
(29)
where the second and third inequalities above are due the monotonicity of fo and Markov's inequality respectively. Note that
R
E(u] 0(p, a) 01|2|1 < C VI
by our delocalization assumption. Setting V22 (COR/vI) 122 from 28 and 29
we have
E [+]] 2 inf (1 T) (bb) + (--0(-v1 - b) Uz).
bER
Since Φ(-VL bb(11 vz) L.X,, Φ x(-b) uniformly in b IR, we have
lim inf Eiss 2 lim inf L-+00
inf BER
(1 TT) - - 6)(1 -vr))
(30)
L-+00
= inf
(1 - (1-x) (bb + (-b-))
(31)
beR
min(m, - x).
(32)
17

---

Published as a conference paper at ICLR 2026
On the other hand,
lim sup Eeet [ < min( 1 m)
1-000
as the upper bound above can be achieved by setting the original weight vector w = Od- when
This establishes that lim Etest [g]
0. Finally, we consider the case
where f oo). Setting ū = max u, -b), we have
E[f((u)] 2 [[)( (b( = Φ
b)
(33)
where the first inequality due to the monotonicity of fo and the second comes is by Jensen's 's inequality seeing that fo(z) is convex for x: 2 - b. This provides a lower bound
&ostt (
inf
(1 π c (b) क ( b)
aES ans bER
where we remark that E|u depends on both and Defining g(b) -Elu one notices that is concave, piecewise linear, and non-increasing. As R < 0|pp| one finds that the function
g(b) = { 0||l||
b 250 0
0|pl| b, b>0 0
is a minorant for g(b) and so
test M 2 (1 T) + (o(b))
(34)
pll-01pp1), (1-x)/2,
650
(35)
b>0 0
Applying the delocalization bound on ||||| then yields the lower bound
lim inf Ejsst
[v] min
2 π
x x x >
L
where C >> 0 was such that ||pl| CR/Vr. This completes the proof for the first set of
assumptions of Theorem !
2. We now turn the the uniformity assumptions, namely when π = 1/2 and v has a uniform
distribution on its support. Setting G(a, b) be the objective function of (2 and
g(b, t) 1/20( (b) + 1/220(b t)
for 2 O, observe that (bg u]] G(a, b) where we again recall that u depends on a. Following the same minimization over in the proof of Proposition [
we see that
inf g(,,t) (-11/2
BER
2
and so
E[4(-1/))
4(-E(u/2)
Ejst 10]
inf aesa InRd.bER
G(a,
2
inf
inf
ae inad
2
ae 1ned
where the last inequality follows from Jensen's inequality as 4 is convex for 0. By monotonicity of (b)(-)) and since the choice 1/VI. 1, maximizes RR(u) we have
$ ((-(1.12)
Exest
2
where é, -R/VL. Here, one notices that the right-hand-side corresponds the optimal test error found for the pooled classifier in (25) when π = 1/2. Notably, the above is indeed an equality which is seen by evaluating G as the previously considered values (a, b).
Hence, the uniformity assumptions reduce the optimal test error for the vectorized classifier to those of the pooled classifier. One then obtains analogous result to (9).
18

---

Published as a conference paper at ICLR 2026
D PROOF OF THEOREM 2
We detail in this Appendix the proof of Theorem 2 The proof builds on the following intermediary proposition, which gives sufficient condition for vanishing test error, when the query weights q constrained in norm.
Proposition 2. Consider the attention model A (5. For > O, let TB = {x E R ||=||| <
Th we consider the optimal test error EiaalA, 7] gnggerBd,wER bER Erest (ess)Ag,,,b). restraining the minimization on to vectors of norm less than or equal to In the limit L + co, R (1), allowing the signal 0 to depend on L,, suppose that
00
lim
(36)
L-00
L
Then, the attention model A achieves optimal test error of Ejst 0.
Proof. We remind that from Lemma 2 for any w. b the test error can be expressed as
|| wil- I|s ||2 b+ {w, 9)
Erest(A9, =(1 - π )P
||
(9.
(37)
||41 (g, 84) (0v, 80), un) v0)
+ P [P w|ll ||8 [|==|=<<<> < -b (w, 4)
(38)
|111
where
8 softmax(31lgllg) and 8+ = soflmax(S(llglll + (q. (,00),
P+ La qg" / / q), and (O, 1) is independent of g AOO Tz). The probability TP bears jointly over the random variables v, g, 2. To derive an upper bound on the optimal test error, we can consider the special case q = 7f, w E. The expression of the test error then simplifies
Exest Aq,n,)] =(1 - IP (0 < + (g,8 + xP (0 < -b (g. 8+) (Ov, 8+))
(39)
Note that
Σ
2r9
222 Bre
KETL
+
VT. 22
(g, 8-)
(40)
Σ
e0r9
82,2 2
VT. 21
TE[T.]
We have introduced the random variables
Σ LELL)
e BTg. - Le
22,2
Σ
gr - LBBr 2².2
SE[L]
VL From the central limit theorem, 21, 22 converge in distribution to standard Gaussian variables. By the same token, one can rewrite
22
VL
(41)
(g. 8+ + Ov)
Pro B
00 L
A
Bre
22-2
+
VI
22
(42)
12,2
A + e
VI
introducing the random variables
A
Brg:
B
arg
(43)
(E(R)
(E(R)
Using the change of variables b = BT 8/VL allows to reach
i
Erest GreelAAre,. = (1 (1= (g.
x)P
BBT>
mP
(g + θr, 84) Br <
VL
vL
(44)
(1 x)P
22
BT
>
(45)
g4,2
VT."
0 Pro
A
Br0 -1 (B
(T + 222
BT
+ mP
VT.
VT.
<b
(46)
5re
12.2
+
+ VT.
19

---

Published as a conference paper at ICLR 2026
Let > O. We first focus on the first term, which one can bound as
P
22
BT21
> b
s P 22 BTz1 > b ( e 2 -1)) + +((12] > VT...
82,2
(47)
g2,2
e
vL 21
Let M
Veerf((1 1 c/8), and let
22 8z72 - 22 (20 (26 1))
b
erfo(1 - -c/8).
(48)
82,2
-1 -1+6 e
and
b2 = 2 (1+ 223 er)er[( - +/8).
(49)
We now fix b = max(b], b2). Let L, be such that for L1,
i 8(e"2 52.2
-1) -1)
P 22 12-33aa BT21 >b ( 12,2 (e $ S1- - +
+ 8"
21+ (1+ 3272 127+
20 (2 - -1))
(50)
P([21] > >M $ 2- 2222-220 2M) + 8'
(51)
The existence of such L1 is guaranteed by the convergence in distribution of $1. 22 joint normal Gaussian variables. Then for any L > max M, L1),
BT22
P
b
S
(52)
g2,2
6
2
+
VT.
21
Turning to the other term,
0 0e VL
A
+
B - AT + 22
BT21
P
LL
<
<P(22 BTZ1 < -b) + P([21 2 VL)
A + 2
22,2
+
(53)
2b(e
22.2
1)
8-0 -
B
+ P
A <
err 11
(5T + b/vII)
VI
VI
(54)
Let T.2 be such that for T, T,2 Tzz
P 22 - > i ( T -1)) s
$1- -1-0 ②
b
+ e
22 3272 - 2 32 02 - -1))
8'
(55)
P(|=1 > M) 2- +
(56)
8
Finally,
x2,2 2b
7 + 1) 800 - B
25 2,2
+ +1) +
5ro
Vo
P
A
vT.
<P
A <
VT.
P( Vθ).
650
e0r -1 (BT + 5/vV))
20rl -1 (3T + 8/vV))
VI
vI
VI
VI
(57)
Now, note that
300
L-000
xc = 0
L1000 xo.
(58)
L
20

---

Published as a conference paper at ICLR 2026
From Markov's inequality,
B|
P[[ > v0) -
(59)
v0
Let Ly be such that for all I 2 L3,
P([B) > v0) '
(60)
8
Finally, let introduce the shorthand
2b(e
52,2
+ +1) + 20r0 v0
vL
(61)
-d-0-1
(BT b/vT.)
VI.
VI.
Remark that h + 0 as L > co. Let La be such that for S Ls, h < Using Mill's inequality,
P(A < h) < h) S
28 log
(62)
log h
BT
The right hand side tends let L5 be such that for all L S L5 it smaller than s/s. conclusion, summarizing, for any L 2 max(M [, L1, Lz.Lo, Lg, Ls),
0 S Eiest (A, 7] S Exest A
<e.
2014888--- 2 (2n 1))erte(( -1/4) max (1002,11211)
(63)
Thus,
Ejst [A]
L 0.
(64)
This concludes the proof of Proposition 2
We now prove Theorem 2
Proof. Suppose
0
lim inf
0.
(65)
log L
There then exist C > 0, Lo. such that for all 1 Ta, 0 > C Coog L. Then, setting
1
== C3
(66)
observe that
2 Clog LVL L-000
(67)
VI
From proposition 2
0 < E [A] < E A,
L-000 0
(68)
L
E PROOF OF THEOREM 3
Output Weights and Bias.
We prove in this appendix Theorem 3 which we summarized the main text. We first give the full statement.
21

---

Published as a conference paper at ICLR 2026
Assumption 2. The loss function is of the form ((2 E" (yz) for some convex function (.). This assumption is in particular satisfied by the logistic and quadratic losses on IR x (-1, We further denote C(E) -yd_ ()) ==)== For the logistic (resp. quadratic) loss, C(e) 1/2 (resp. C(E) =)
Theorem 3 (Characterization of the query weights q( after two gradient steps). Let w( ((0) be the readout weights and bias of the attention model A (5) at the end of step of the training procedure detailed in subsection _ In the asymptotic limit of Assumption | the summary statisties || and (w² ε) converge in probability to deterministic limits, given by
(
C(e)mo(2T 1),
(69)
while
|1001|| P +
d-00c 71 = 1m
ao + (TOR/L),, (w())E) P
d=o00 72 = The C(t)
OTR
(70)
L
Similarly, let (( denote the query weights at the end of step 3. The summary statistics (2 and (q(²),)) converge in probubility to the limils
1
Eg "2" (L-1)nf ((L-1)E7+
ag
+02 (R- 27)
(13 2
(1
4 +0 (R- 2) bg
(71)
14000
and
(S, (2))
P
ng ng [[I- + 02 (2)
(72)
d
Here, Er; F2, Es, Fa are constants whose expressions given in the proof.
For this bias, the Law of Large Numbers yields
50)
Σ h,,0,,,0)
ny
Σ C(1) & C(1)) -1) 1)
710 S
since Ely] = IP(y = 1) - Ppyy = 1) 2π 1. Now, decomposing the noise Z, by
2,
s U.
E IR Lxd
U, E RdxL-1
8, E Rd,
and setting
y1
S = [81 ... 8n and
=
yn
we have
(1
Tee Σ 0,(0,0,0)x711
nL
isno
(d)
C(f)
!w
(V.Sy
Σ you
L
iSno
Sy
TORE
C(1) The
(73)
nu VL
L
where one applies the Law of Large Numbers for the last line above. Using the representation of we obtain
(w .g)
7/2
d
22

---

Published as a conference paper at ICLR 2026
since (g. NO, no) and
|Syl1
=000)a nht.
+ (=OR/L)2
P. 71
d0000
as (Sy,Sy) /no X2- Query Weights. Setting
:= (0, w(1),())) and Ap: x 11 /
have
((0)
nap E (pAawa)) noL uSno
(74)
It will become clear that we require only the first and second moments of cu conditional on yp to characterize |[qll and (g. {) for large no and d. Specifically, set
F = [cpp
E2 = -Eeeul = 1], Eg = =E[2])
FA = EE = 1].
(75)
Concretely, cu is given by
d
cx = dz
((2,
(76)
where
my = (2) xT 1/D3 + ((2)
To find the distribution of Ma. write
C(l)n (yr: X P 1),
())
w (1) + Aps
A
(77)
nnL
where w obtained from all samples except μ and is therefore independent of X Substitut- ing (77) gives the exact identity
m₂ = (ā w + (Tps Ap) (1(1)
cerm
The ||_ 1 no L
C(f)
Yu + C(c) Tee
1 (5 (z.: Srest) + (21 1),
Lo
with Syest : Ejtu Vo X3 1 (independent of X,). The above representations lends to the following
conditional distributions:
{yH = 11 ~ ((E)))(2T - - C(E)
C(
mp
aT?
{Up = ~ ((1))s -20- 1)) (
0²
C(E)
mp
L
aI?
(78)
From the above, see that marginally m, is Gaussian mixture. Knowing the distributions muly and yı facilitates the computation of Ei,..., E4. This can easily be done to machine precision example.
such via Gauss- quadrature
Returning to another piece of (74). set by := Lv. ((211) /VL - IL)U) R and decompose the Gaussian noise U,, by
U₁ [gp Val, 9N E Rdd ERd,
V, E R_2L-2
23

---

Published as a conference paper at ICLR 2026
We then decompose the feature gradient Aμ by
A₁ UNU (R (yp=1) RR2 10,=1)/ L)EET + , DEb
= 9u9n + Vivi + 0 (R 1 1w0=1} - R² -1 (1mn=1)/1]55
0 R [vm- RR 1l- (ou5" + +ggH))
-ug9 + VIV 02 ORTEEE + Oon +E9H))
for
-I)- V SOOVR-R111 2 |ل4 (m
Now, set
[c
h1
G = =
Inl E Rdxn
==
=
On
hnn
and A, == diag(x) E Bixk for x E Re k E N. Observe that
1-2
Σ VV
Σ
VAAV
F=1
V3 li V Rdxa and V has i.d. (O entries. Making final
where abusing notation
decomposition the noise:
g G
V
G.. U ER", Gu, V. RRnd-1
VT
obtain
(2)
Σ cu w((1)
1oL
usmo
3
GAGG + VANT +0² Σ Ciic
EE + +GGAnCE T + (OGAnCET)T
w)
L
L. ind. copies
(d) n
||
GA.ãs +
VAgU.
+ (l()) GAne
L
1. - 2 ind capies
+ ((0lu0t11c Ang. + (2) Σ
3
H Har
71
GAols +
VA
+0 GARe
12 L
2 ind. copies
+
0 0.71.0 c. Ange + 07.22 Σ Eunh 5
Since no
T Angs P. as no > oo, we have
9 (2)
Bg L
H (
71 - -GAgGs +
VA.U. )+0 72 GAre
0 -2 (R - (R-R2/L) E2
110
ind. capies
N
M
3
(M + N -E)
--
24

---

Published as a conference paper at ICLR 2026
Therefore,
(k. x na 15T M NN
and
ngB L
2
V MTM 2NE M N2.
By rotational invariance of the isotropic Gaussian, we may take & to be the first standard basis vector in the following derivations. We then have,
w()T
(
<
M
|
n
GA..,
+0. 72
GAae
no
L ind. copies
( GADC
(1)T
GAc. ind. copies
+0- 72
n2
71
no
(1) T Wi 9. AcI.
110
(T. 1) 72- Fr.
This gives us the alignment
(2)
[LL 1) + R2/L)E2 =7
d L
as claimed.
Finally, to compute the magnitude of g(2) all that remains is to determine M M. We have,
M
M
1
71
Σ
V
VOAcū, + n
2 0m1 AAG G....
TO
ISUjSL-1
+
0 n
2
c.
AT
VA
+ . c G GARC.
no
2 ind. copies
Examining each term separately, note that by repeated application of the Law of Large Numbers we obtain the following:
1-1
3 7
Σ
A.V. VYACE, -
0 A.V. VA.U, +
1
.7. Σ CAA.
11
mg
18W5
iF
1
1 Ad V. Act + (( Ac0.)))
x (L 1)/?
(i. Ac0.)²
+
(1
n0
n
x
(L 1)
(L- 1) Ef+ +
Es
a
2 n
2 -y112 c. G
Any T, GAeg. + A
VA
x 0,
ng
L ind. copies
and
-0-3 c. A GAre (d)
0223
|
9 91
n1
71
EA
x 02
(R - R²²
25

---

Published as a conference paper at ICLR 2026
0.9
0.8
0.7
0.6
(09) A
(w
[|
0.0
0.5
1.0
1.5
2.0
ao
Figure 3: Cosine similarity between the signal & and the query weights q(2) (blue) and readout weights w()) (red) after step 3 of the training 3 for L 10, R 3, 0.2, 0 6, n 0.5, and logistic loss L, as a function of the normalized number of samples ag- Solid lines: theoretical prediction of Theorem 3 Dots: numerical experiments in dimension d 1000. Error bars represent one standard deviation over 10 trials.
Putting all the terms together, we obtain
M M x (I- (1- (I.-
Eg a
+ +0 (R R2/L)) Ea
(0-11)EF
F
and so
(|
n7 L
(I. 1m??
(I.- F?
E3
2 (R - R2/T.) E1
a
1/2
+2N 1) .22
E -
N2
where we recall that
N N 02 72 (R- 12 LL) E2.
This completes the precise characterization of the magnitude and g-alignment of the query vector (2) where the definitions for the relevant constants Eg...., E1 are found in (75).
LARGE og BEHAVIOR
To conclude this appendix, (e)"))) ( 91 of the attention weights w, after one or two gradient step with the signal vector in the limit of large sample complexity oo >> As summarized in Corol- lary [ in the main text, the cosine similarities rapidly approach in absolute value as the sample complexity ag is increased. We give here the full technical statement.
discuss the asymptotic behavior of the cosine similarities
Corollary (Large ag asymptotics). Let w()), (()) be the readout weights and query weights of the attention model A (5) at the end of step 3 of the training procedure detailed in subsection 3 In the converge in probability deterministic limits Sw, Sy from Theorem We further assume that (E. (2)) + 72
asymptotic limit of Assumption ! the cosine similarities (w(") (o("),))/nn ()) (g .S/ (0
When then further taking the limit og + co, these limits admit the following asymptotic expansions
1
Sw
20d (TRR))
+0
Opp
(79)
1
1
(80)
200
2
(1-
G + (1- + (L- (011-))(G23))
+L
R²
(G
(L- 1)= +(1 (1 T)G + (R R2
),)))
(81)
+
(82)
00
26

---

Published as a conference paper at ICLR 2026
We denoted
G =f= =< - -1) T R2-02 .1).
G = (C(())(22 - (83)
2L2
The sign of By is on the other hand given by that of
- [[L 1) 02 R(1 R/t)] + -1--)=)
(84)
Proof. The proof of Corollary followw straightforwardly from ag expansion of the expres- sions of Theorem
Corollary [ establishes how the weights of the attention model recover the signal vector & when provided with sufficient data, rate of 1/ao. The sign given by intricate but explicit condition 84 on all the parameters in the problem e, т. 0, R, I., Tbb, yuu, and can in certain cases be negative signaling that the query vector y detrimentally anti-aligns with the signal Ç. In order to avoid such scenario, the condition (84) can offer some guideline for choosing the hyperparameters Th. Tw l. For example, for the logistic loss (y,, = log(1 + exp( yz)), when æ < 1/2 (resp. π > 1/2). choosing ?/ sufficiently large (resp. negative) ensures $q > O, namely that the query weights g g(2) properly align with when on grows.
F PROOF OF THEOREM 4
In this Appendix, we detail the proof Theorem which summarized in the main text. We now present the full technical statement.
Theorem 4 (Test errors after step 4). Let denote the query weights after step 3 of the training procedure 3 and ú. b be the minimizers of the empirical risk (14) at step 4. We denote y ig, E) In the asymptotic limit of Assumption] 7 to
the associated test evror Evest (15) converges in probability
0 (9,8-)/1 +b
i (Ov, 84) (g. 8
Ezest [A] (1 - T -g,
+ t..gg
$
|a|[8_11
M3 0/||s+ll
(85)
with μ3 [v2 (u7 27 -mi² .The description the joint law of the finite-
dimensional random variables g, s s_ E R² is given Lemmal The scalar statistics b, pii #2, are defined as the unique solutions of the following variational problem:
Pl, 12, b = argmin dhp Me, b) + l0gg
Me]
1
7
Ug
(86)
Y
ue
In the above display,
A 2
bgg Mg, b) := Begsere [6(2" + Cylla + cEpe + b. v)]
+
(87)
2
where Cz2 Cus Cg are scalar random variables whose joint law is detailed in Lemma and N( 1). Finally,
2" (2* C2U2)
(2
2 -Ec.
Ec
+ X.
(88)
Xx
Q1x
1 1+ (x-)
We used the shorthand 2 prox (x+exx((++4pppp touue (c.uz). converges in probability the minimizer of the right hand side of (86).
Finally, the training loss Evrain
Leveraging the equivalence between the attention model with zero query weights ao (or, equiv-
alently, vanishing softmax inverse temperature 8 0) with the pooled classifier L pool w.b
similar
characterization for the latter directly be deduced, as summarized in Theorem
27

---

Published as a conference paper at ICLR 2026
Corollary 3 (Test enror and training loss of L The training loss and test error if the pooled
linear classifier L trained on the empirical minimization (16) converge i probability to limits Erroin Erest [L [ooo]] whose expressions fan read from Theorem 4 if one sets /3 =
Remark 1 (Length generalization), Note that the characterization of the test error in Theorem A readity generalizes to the case where there exists a distribution shift between the training data and the testing data, when the model is tested on samples with different length Luest + L and sparsity Reest * R. The characterization (85) can be adapted to this case by using Luest Reest in the definition of the joint law of 9; 8 8- in Lemma | unchanged. Fig.5) shows the + oo error achieved by the attention model trained on sequences with the square loss, and tested on different Luest Reest, and shows that the model is capable of length generalization.
with the definitions of b, 21,2.3 otherwise
F.1 NOTATIONS AND ASSUMPTIONS
We take the following definition from Karoui 2018
Definition 1. Lei
X= (X,,
Y= (Ynu): TT(u))nENN,uEU
(89)
be two families of nonnegative random variables, where Un a possibly n-dependent parameter set. We write Xn OL (Y,)
sup [l)X,,u)) O( sup (u)l)( wEt
wetn
where "O" refers to the classical big O-notation. That is, for two deterministic sequences (an), (bn), we say an = Ob,)) if there exists some C' such that a S Cb2 all sufficiently large.
We make the following assumptions the loss function & (with the first argument denoted z):
(A1) E is non-negative.
(A2) L is convex in its first argument.
(A3) f E Cd in its first argument.
(A4) has bounded second-fourth derivatives.
(A5) L is coercive, i.e.,
lim ((2--1) + ((2;1) oc.
Remark 2. The above assumptions are satisfied for many natural choices of loss functions such as the quadratic loss, Huber loss, and logistic loss.
Remark 3. Having a bounded second derivative immediately implies the existence of a quadratic majorant of l since for any z R. a second-orde Taylor expansion yields
{(2) - ((0)+ + ( ()- (
((0)) (())
2
F.2 EMPIRICAL RISK MINIMIZATION
In what follows, we study the following learning problem:
min 1 Σ w.b 741
E((fe, w) + b; w.
|wll
(90)
where {(2; y) is a loss function that is convex with respect to z. Let w* be the optimal weight vector and 6* be the optimal bias for (90). Our goal is to characterize the following quantities:
#1
(q. w"),
P2
(E, w")
P.. w"
(91)
and b*, where p.ls denotes the projection onto the space orthogonal to 9 and &. Having b*, P1s P2s and will provide for full characterization of the test error due Lemma
28

---

Published as a conference paper at ICLR 2026
Remark 4. Recall we have assumed that || and at loss of generality we also take |1q 1, absorbing |[qll into B. Moreover, as a reminder, ? = (E, 9). It is easy to check that for y # =l.
Py.e" 19
5]
1
7
Hg
and that
?
[gg
||wl1 = (ug
mel
+ || w|l2
1
me)
for weight vector with
My = (q, w))
Me = (5, w)
From Lemmal I
we can rewrite the feature vectors {fi}
Ji = Cq.i9 + Ce + C2 Py. g. where {cr,es Cg,s C2.i)ign ignn are scalar random variables that are independent of the isotropic Gaussian vectors {2s hig. We write the joint law of Cg.is Cg,s, C2,s as
y1 if yı The exact specification of the joint distributions given in Lemma 1 Specifically,
[ (e: P Cc, C2),
Cgas Ce Czft
= -1
7 71 11++1 |
|| 1|=0
P+(cg
Cq = 8+)
cE = (Ou, 8+) +
C₂ = ||
V1 72
v1 72
(92)
Cy = (-, ) 7 7 1
|| 20 VI
ce, Cc)
CE
= 1==1__<|
(93)
VI-1"
With this new decomposition of the feature vectors, the empirical risk minimization of 90 splits into (i) a three scalar variable problem of Py (e and 6, governing the q. plane and bias, and (ii) a (d - -2)-ddmmessoonll sub-problem determining the orthogonal component to span(g: E). The next display formalizes this sequential optimization problem:
> (Mg Mel
1 7
1
[gg me
min Φa(P's Me, b
2
where
DalMy, Mg, b) = min wERd
Σ E((c2,22,,)) + Cq.lty + cgall + bs) + 11111 #E[ma
(94)
and {ziisnn is a collection of (d 2) dimeenssio isotropic, normal random vectors.
Henceforth, our goal is to characterize the asymptotic limit of o(g Mg. b) and v² || where x= denotes the optimal solution to (94). Since x² is stationary point, must have
x=
1
(((2,,2()x") + Cu.li + cg.tte b; (c),, VE[ML]
Σ
Thus,
22 = 11
Σ
E((c,,22)xx)) + Cq.l4 + cgsll + b; b: b;
(95)
πr
(E(m)]
In the following, we will denote
Cqjilly Celllee + b,
l(uu ((u+E)) Ei) == f(u e.)
(96)
to elicit parallels between our derivations and those present in Karoui 2018). For simplicity of notation, we further write
je C2.12
29

---

Published as a conference paper at ICLR 2026
F.3 LEAVE-ONE OUT: DETERMINISTIC ANALYSIS
The key probabilistic structure in our problem is that different feature vectors are independent. This naturally prompts to consider leave-one-out analysis. We first need introduce some notation. From this point forward z is in Rd d Let
2(ff.))) + + X ||]|
1 = min (x) = min
x arg min F2(()) (z)
(E(R))
$ 22 = min (z) = min 7h
Σ j7i
50S + + 2 |1 z2 arg min f2
denote the optimal values and the optimizing solutions of the original optimization problem and its leave-one-out version, respectively. Going forward, we will often omit the d-dependence of these quantities to alleviate the notation.
F.3.1 LEAVE-ONE-OUT ANALYSIS
A key step in the following consists in constructing a close approximation ã; of x², with simpler distributional properties. To that end, we introduce the surrogate optimization problem:
L Layz + min Fa,(x),
H == arg min (z)
x
where
1 (e((F. (x=))) xx 11 11, 2(d)
(z)
(97)
n1
and
H 1 (e(tnaa) + + +
the (leave-one-out) Hessian matrix. Heuristically, this surrogate problem may be viewed as quadratic approximation of x- in the vicinity of o.. It is straightforward to verify the following lemma.
Lemma 3. Let (x; denote the Moreau envelope of (x(x), i.e.,
M.(x:n) = min (x()) + (x 2)²
2
and let
(x: 2)2
Proxx(x;() == arg min (2()) +
24
be the corresponding proximal operator. Then holds that
F: = (f.. [x) + = Prox, (F.As))))))
(98)
where PLyy == (.. + and
7/
1
fo H H fc.
(99)
n1
Moreover,
(100)
and
$ 9 Milli. (.)))
Th
Remark 5. Let Prox(c:7) 7). It is often convenient to recall the following identity:
((x) + I-C 0. 7
(101)
30

---

Published as a conference paper at ICLR 2026
F.3.2 ON THE BOUNDEDNESS OF E!
key technical difference with the closely related analysis of Karoui (2018) lies in the assumption made therein that L' is bounded. We would like the present results to hold for the quadratic loss in particular, which does not satisfy this assumption. The following lemma bridges this gap by showing how the optimizer of the inner problem using loss coincides, with large probability, with that of modified loss with bounded first derivative.
Definition 2. Given I >
O, define the clipped loss Letipes v)) R - IR as follows:
1. lelip Ci.andd and convex
2. (elip() = ( [-1.1]
3. Letting M = SUPsE(-1.1) (z)l,, we require that 11 S 2M
4. we further require that kelip L.
The construction given for felip in definition 2 can be achieved in the following manner. Consider the "bump function"
(1)) exp
(
if (0,
0.
else
and, fixing O, define
R - 10, by
0,
2 1 (1,1+
l 1/--1/- +0
n(=)= =
Jo ((t))dt
2 +
Note that E C with bounded derivatives of all orders. Now, consider the left and right linear extensions of E (=))
(())) I),
which allow us to define fellip as the piecewise function
I- (z),
5 - - 2 (-I-
(=)(-=))(()+
Celip =
{(2),
3 [-
(1 (1-(())(((s)) + (=))++()), (2);,
(I,I+)
>2I+4
The prescribed properties definition
2 are then easily verified from basic calculus.
Lemma 4 (Clipped loss derivative). Recall that
x² arg min
Σ (l((f,))) + ca)+ + 1|10
(102)
n1
For given 8 (0, 1), let
dRa = 2 B(((c:)),,
2n1
(107 + m&)
I I: (1 (1+
2
+ + p n2 +
+ bb| L82
1 1- 72
(103)
where 6 ~ Cgllg + ceflg + b. Define
dhpip arg min
Σ
lelip.:( + (x) + ||ar].
(104)
x
se(mi]
Then, with probability at least 1 - 6,
Fellp
(105)
Four times differentiable with continuous and bounded derivatives.
31

---

Published as a conference paper at ICLR 2026
Proof. The strategy consists in controlling the supremum supie|m,] ! clip of the residuals r; ri
C_MZ
Since by construction ||Caiplee < CpolyTog(m))
apply the results of
Karoui 2018 to the clipped problem, showing that
|r: <
ProXalip (( ((r fi) + E 7el + 8(1) 2 |gi Ii x 8(1) + leel
(106)
using the contractivity of the proximal operator. We used the shorthand gi = (aam
Note that from
Karoui
2018). )(1) sup, Ir! Proxelip (roxelip((221p... ((eaip fi) + ei, vill
OLe
polyLog(n1) VBi
From the identity Fatip.i(=2rip vi) S (e one can bound
| kel? 11
Σ j
Eotip (()l()))))))))))))))))))))))))) 1 9A² + 5(22)
(107)
Ani
with 8(2) = =O1, [oolyTor(a)) v"1
Using the identity 1v++ x - 1 S |x), we have ||
R + [8(2). Summarizing,
sup Ir; cip < (sup (.)l)) 15(2)) + 5(1) + sup |sl-
(108)
For n large enough, from Markov's inequality.
8
PMA 1
2|(8()
11 - 6
(109)
We now need to control the term sup, lex]- From 992 for any given and fixed pe, May b and remem- bering ||s S 1, one can bound
| S a 11 + 20.1 + ce
(110)
with
a. = Ipgle
b.
M4 /1 x2
+
= |6| + LL22
(111)
VI
We remind that all entries of g E R², alongside with 20, are normal-distributed. Then, for any h, using an union bound
P sup e: 2
o + 12h Ce < Σ P
aelllsi11 + b, b.1 2 Va² b2h
(112)
*E[ma]
Examining more closely the summand IP(ue |1 ll1 + bel
belz0.ol 2 h1. one has
(aellgell be/zaxl 2 aa + 12h
Σ
Jass19:11 + ... be 2 02 12h
.E(-L+1)611
(113)
from coarse union bound, remarking that the left hand side appears in the right hand side sum Now that one has ridden of the absolute value, observe that each term in the summand is distributed (0, / ag 12). Thus,
- $62
Pr suple + + co
< 2 m1
(114)
In particular, for h = /2 log n1
Pr sup | 2 V 02
12 /2 222goog ce
2 1+1
(115)
/2 log n
Let us again suppose n large enough that this probability is smaller than 5/0.
32

---

Published as a conference paper at ICLR 2026
Thus, for nı large enough, the probability of the complementary event of A =(81) < 1 n {
< 1) (sup, laal Va² + b2 22llloo n + ce) is bounded as P
v 6/2. Now, for any
t 12 2 + V -h + 2 log + Cr.
- 8(2)
P Plsup ,r x
4) IP
sup Igal
(116)
+ |
5(1)
sup; |4
8
$ sup [g:l
+
(117)
9 + |
2
t
(v b2 V2V22llgg n2 + +.+)
sup l
+
(118)
ni +
a 1/2 2 log 1I +
P
|
(119)
+
(E(ns]
nie
+
(120)
where the last line follows by Mill's inequality. In particular,
P
sup
clip
(1 + R) v
2 log
2n1
+ 1
a
02
log m + c 5.
(121)
8
The last step of the proof comes from the simple observation that with probability least 1 8, for all i E [mi], (r,)) = (Fi): and so under this event we have
Σ ((r)))
Σ (rii..((ri).
(122)
D1 vE(m))
m1
EE(MI
Therefore, x² satisfies the stationarity condition for the clipped problem. By uniqueness of the minimizer Elipp we have
x²
Jelip
(123)
in this event.
A consequence of Lemma 4 is that assume, without loss of generality up to an event of probability that the first derivative E' is bounded. More precisely,
|| (polyLog(n1)).
(124)
This enables in particular the borrowing of number of results from Karoui (2018), where such an assumption is leveraged. Henceforth, we worth under the (1 - 8'-pbobbbility event where x" diipp and work strictly with the clipped loss felip- however we omit notation and simply write for simplicity.
F.3.3 CONCENTRATION RESULTS
We first introduce and recall several quantities of importance in this section. For i, E [nı], write
(f., x3) +G, 8,, = - PAN
(,, +
The following lemma establishes that the introduced surrogate estimator 2, constitutes a good ap- proximation of the full minimizer x" as well as further concentration results.
Lemma 5 (Approximation by surrogate estimator). We have, for any k.
polyLog(n1)
sup Ix" [E(mg]
3, OLi
and
sup ||
Fzll
OLi
(125)
n1
ve [E(n]]
/ni
33

---

Published as a conference paper at ICLR 2026
Moreover,
Var(Il>"13)) = (polyLog(m))
(126)
Furthermore, at the level of the residuals, one has the bounds
polyLog(m))
sup =E/mi]
[r P.| = Ozn
(127)
Proof. The proof follows directly from Lemmas C.2, Theorem C.6 and Proposition C.7 of Karoui
2018 The statement on the residuals corresponds Theorem 2.2 of the same work.
The lemma thus shows that the squared norm ||~*|| concentrates. We denote in the following by 22 == E [[]] its limiting value. The statement on the residual can be further complemented by the following lemma, which covers the off-diagonal terms.
Lemma 6. We further have
polyLog(mm)) 11
fra)= On
(128)
where we write T3 (fzr 2.) + Ez.
Proof. Denoting Fi E (R(n 1)xd the matrix with rows (as we need to control
[polyLo(mm))
(polyLog(n1)
Cr, F2,)²² F1 =2)1 = POL.
OLi
ni
n1
jy
(129)
using Lemma to control ||x- T:ll-
Lemma 7 (On %). We have
1
sup hy 22,x =OL se|m]]
where
x= tr|H 11
(130)
[T11
nn
We called
= 1
Σ
GOALE +
(131)
jE
the full Hessian.
Proof. This follows from Corollary D.7 and Lemma E.4 in Karoui 2018
Lemma (h" concentrates). We have
Var [ O -ppolyTog(aa))
(132)
Proof. Appealing to the Efron-Stein lemma, we have
far[
<
E ((="))FF((att)))
(133)
sE/ma]
The summand can be controlled as
E (F*
F(((),)))]
2B
(Fi(e)) -Fulot))
2
[rs))
(134)
nn
34

---

Published as a conference paper at ICLR 2026
We first control the second term.
E [((r)) (4:(0)² + 11 [r31) ni
(135)
As we will later show in Remark 6 the moments of r; are indeed bounded, making the right hand- side OpooyIoog(m))/n) Note that the current result is not used reach Remark 6 so there is no circular argument. We finally examine the first term. By the mean value theorem,
x 2
F*(x*) - (((c)))
Σ 1#1
(136)
T11
2
where rj belongs to the (unordered) interval (rj, A,))) We now show that both terms in the scalar product are small. First, we will use the fact that the first term is close to Fx((og,), which is by definition of o vanishing. More precisely,
+
(G(()
(137)
11
2
181
2
1
booby - Farall 2 -
-
(138)
11
jæi
(olyTog(m))
O
(139)
n
Since it (1isfgad)
b8g - T S lr, TyAul = O polyLog(n1)
(140)
Vni
from Theorem 2.2 of Karoui (0018). From Theorem 2.2. and Lemma C.2 of Karoui (2018),
we
further have |x" 7. = Or, (polsLog(ns)//ng
Therefore, from Cauchy-Schward
[F*) - F,((o = OLn (polyTog(1))
(141)
71
Putting everything together,
EE (1*("") -R,(0t,))) =0 polyyLg((1))
(142)
ni
and
Var[ø*] polyLog(n1)
(143)
m1
from the Efron-Stein inequality, concluding the proof.
L
Lemma 9 (x concentrates). Recall = 1/m2 ur|H-
The following concentration result holds:
Var[x] = polyLog(n1) )
(144)
Proof. From the Efron-Stein lemma,
Varixl <
Σ
B[(x x]
(145)
sE(m1]
where X/+ 1/11 tr| HH We recall
1
H
((SW)) + Ala,
(146)
T1
35

---

Published as a conference paper at ICLR 2026
Let us further decompose
E [xx x1.19 s 2B [[x 80) 21E [xt: x,
(147)
defining
X1
1 tr(H; 1
H.
1 Σ
(FG + +aa-
(148)
n1
n1 iti
We first focus
E
[(x {.
1
X
Xi
tr
H
( (H:- H
(149)
71
HH fs
1 (r(r))
HH
=
(r,)) e" (F,,,))
+
(150)
11
n1
111
(
- fga))
5j H H
1 (ri))
Ji TTH-HH111 fi
(151)
n1
n1
n2
n1
where we used the mean value theorem and rty E Fjs). Thus,
1
pooylooooyLoog(n)
Ix
Xil S (8, e) + Oz
(152)
711
X2m
introduce the vectors o, R with elements
6, (rg Fz.))
(153)
f H
2 = (F)
(154)
The latter can be controlled as || = O ( /oolooyoyg()) )) while from Lemma OLe (polylog(mi)//mm)) Thus, the Cauchy-Schwar inequality implies
7 ||6||
(polyLog(m))
Ix Xel O
(155)
n1
We now examine
[(x - ) From a Taylor expansion,
H
H fo
X:
X/x
(nrala - )):
(156)
71
H H 13
(3) B3))2
(157)
2n
71
for some 8, E (F,,s R)) From Lemma C.4 of Karoui 2018 F Fuuu = Or (polyt.ogtas)//ml]
from which it follows that the second term is OL, (poly (olyLog(m1)) lm.)). ). The objective is now to approxi-
mate H, in the first term by the fç- - independent Hessian H.s to unravel all statistical dependencies
on f.. The correction is
H (H H )
2 (Frw ()(1)
(158)
M
n
1
sup jfi
|H, -17--1111
(159)
An
1l
< ||( 11o0 sup
1F
|H - H.rll-
(160)
A
jfi
36

---

Published as a conference paper at ICLR 2026
But
|| -1,1l =
Σ
-xo),, s sup ( (tg)(( - Frie) ||aم
(161)
11 2/4
jti
where E is the empirical covariance of the features, excluding the th. Putting everything to- gether yields
1 Σ jjt
53 H. (H, - -H, 1
polyLog(m1) )
=0
(162)
n2
Thus, going back to the original objective,
27
polyLog(n1)
E [ xx] = E
(ls) (FWW)F5
0
11
m
nf
jæi
(163)
Leveraging the closed-form expression of f.. Fy(ss the first term be written
fJ H
E
(Fi.)) Σ
(FrWf
2
H.
H n1
E Li
Σ f+i
(33) (F))
1 [9"]]
(164)
using Minkovski's inequality; g 1), 1) in the expression above. Note that, introducing the vector h E - with elements h, ((3) (fT "
f//m
H -1 F R. s An |
(165)
81
jfi
But
[|,, O2 /nopolyLoggn )). and
F, H
1
||b||| S /n1lle(3)
sup jz{
< |1
1oo
sup O2, (polyLog(m))
71
(166)
Thus,
2
1J H H.,²
)
E
(Fr.))
1 n1 it
Σ
(F polyLog(m))
n2
nf
(167)
To complete the proof, need control of EE [F) which is provided by the proof of Lemma where we established that (f(r.,) = O1. (1).
F.3.4 LIMITING RESIDUAL DISTRIBUTIONS
It now remains to ascertain the law of F, which we describe in the following lemma. Lemma 10 (Limiting distribution of (x)) The leave-one-out residual admit the simple representa- tion
polyLogg(1))
FAN = = + C2 / + OL2
(168)
T1
with 7 ~ (0 1) independently from ei, C2,4-
37

---

Published as a conference paper at ICLR 2026
Proof. We have
Prg - 6 = (R, |لتلحم||
(169)
and Z = 1 (f., NN 1). Furthermore, from the proof of proposition C.7 of Karoui
2018),
polyLog(m1)
|,a =1xx11² +L2
(170)
n1
polyTog(m1)
(ollyoog(1))
=1² O
+L2
(171)
V
71
polyLog(n1)
=2²+00 +0z
(172)
Therefore,
polyLog(n1)
polyLog(n1)
|k / 1+ Lz Vn1
= v + OL2
(173)
vn1
using the inequality VI++ -1| s |ar| in the last step. Finally,
- -11) =0 =) (olyLog(mi)
E
2
(112,11l -))² E
(174)
n2
in other words
polyLog(m1)
Z (1 -v) Or
(175)
Lemma 11 (Limiting distribution of Ft,)) Setting XE == E[x]] we have
7,,4 = Prox(e; + C,, (2, + +OL, polyLog(m1)
(176)
Proof. Let us introduce the shorthands 8, = -Ai -; Cu.iV and 8x = 7/ 22, [x]. From Lemma 3
Fi, Prox(e 22; [ [x))
(177)
Prox(E: + v2 222624202 + 8pi 22 E [xl + 6x) Prox(e: - Cc,
(178)
(F)
1+ (xx 8, c (F)
8x.
(179)
using the two-variable mean value theorem, and eliciting the derivatives of the proximal function. x on the line between the points (Frz e, C2 ö,, E [x] + 5x) and (Fry & + C, e2, [x]). From Lemma
8, Ol2 (polyLog(m1)//mm)) For the second term
((())
E
8x
$ =110111000211
(180)
1 + cg (F)
But
polyyog(n1) TT1
15x1 s bn: - c2.Xl + x [x]] Or
(181)
One thus reaches
(())
(polyLog(nn) n
1 e 8x = Ol,
(182)
Putting everything together, one thus reaches that
polyLog(n,)
F.1 Prox(E: + C2 cz e2,
E [x)) OL
(183)
n1
38

---

Published as a conference paper at ICLR 2026
Remark 6 (Second moment of rz). The second moment E [r3] of the responses is O(1), for any
i E [mi].
Proof. Fix any i E [n1]. The moment E [r.]] can be controlled as
E [r]
s 2E [r: f:)² + 2E
(184)
(oolyTog(m))
s 2E
Prox((; + 12 E2, [x)2 + 0
(185)
polyLog(mi) (1).
$ 4E 2 1 +
(186)
n1
F.3.5 COMPUTING THE EXPECTATIONS
Self-consistent equation V
Lemma 12. The expected squared norm 1/ = E [||x*13] satisfies
V X Bzy.e..c [e' (Prox(e; c.vgZi (xx)) + c, Prox(e; + c.vgZeexE)) +0 Vni
(187)
where i a random variable distributed as i, given y = yo-
Proof.
Using the stationarity condition,
1 Σ C(r.).)
Ax
(188)
11 [E(Ps]
Thus,
- 1 Σ
(r(()r)(ra -s)]
(189)
sE(ma]
Since
((rr)rr -4) - 4)1 |1 11 [ra]] A
(190)
From Schhhy-Schwartz's inequality and Lemma
[[]] Ioc + + (1el + (ra)] In All < [[I + (lx + 1ra)))('// o (oolyLo((m))
/ni
(191)
The boundedness of the first expectation follows from Remark 6 and the existence of the second moment of follows from the proof of Lemma Thus
1 Σ ((rr)(r -c)=
1
Σ C(F)() c.) Oz, (polyTog(mi)
(192)
12 sE[m]
m
sElm]
We now appeal to Lemma 11 to elicit the second term. Let 6; Fi.s Pi. using the shorthand
= Prox(e + ce, Md). Then,
(((F)) e.) - (p:)(pp a)) = (0" + +3, - es) + E' (pi))
(193)
< 1 Ma
83 + 2 cz, 12)] + ||02
(194)
Using
inequality, and the fact that 5, Or (polytor(m11)//m from Lemma 11
the term in square brackets is Or,, (polyLog(n:)/ //TT) Thus,
(F.)) 4.) - (((()) - =s)l=Orx Oi polyLog(n1)
(195)
/mn
39

---

Published as a conference paper at ICLR 2026
and
((r)(r-- -
(196)
T1
(E(m]
& (Prox(es + (, + tz 2,,1, - Ei)
(197)
(E(ms]
polyToo(mm))
+ OL,
(198)
Taking expectation,
|
v22 A Ez,, [e' (Prox(cs + c.UZ; + +c,g) Prox(G + c. [x)] + +0
(199)
which completes the proof.
Remark 7. Note that alternatively, v may be expressed
Prox (e + C,V2;
Proxx (e (xx - cvVEZ
E
+0
AXE
(200)
by applying 1001 to Lemma 12
Self-consistent equation for X -
Lemma 13. Recall x = 1/1, -r(H- and XE = F[x] We huve
(Prox (e 23
1
+
(ooyLog(mi)
AxE + E
(201)
1 ("rox + (p
Proof.
By the construction of the Hessian matrix H. we have
H (m) +
It follows that
Σ COA
n1
Applying the matrix inversion lemma then gives us
(r.)
Σ
+ +-x=
71
(
a
where
Xi n
2, A 22
and
H, - E es GODL +
m1
We note that X is close to 1/4s tr A, To formalize this intuition, introduce
(oolyLog(m1)
X1+ 2, H 1/4 |r| Or.
(202)
/nn
40

---

Published as a conference paper at ICLR 2026
the last equality following from Lemma G.3 of Karoui (2018). But
1 -2 H
1&= Ep =
( H.)
(203)
1 (1)
Σ
(fr)(r) FA
(204)
xs
j
1 (polyLog(1))) ) sup rr TyAel = OLI
(oolyLog(zz)
(205)
X
jti
vnL
The derivation mirrors the steps of Lemma 9 and the last bound follows from Theorem 2.2 of Karoui 2018). Thus,
1 |r' Oz
polyLog(m1)
X:
(206)
70
Vni
Now in trace form, we approximate 1/41 tr [ back by 1/4, tr(H, This can be done along the exact same lines the previous approximation, finally yielding
tr|[i; Or. polyLog((1))
8: m
(207)
mm
We now show that Ys is close to
X:
18 xl
tH,
tr
(H H.)
(208)
n1
(r:))l
tr 2 H "1
(209)
polyLog(m1)
<
e |i H '11.f Ol.
110
(210)
n1
n1 X2
Furthermore, we can also approximate (r:) x (F.). More precisely,
- ( = Ol (1100)) polyLog(m1)
(211)
n1
Thus,
(r)))
(polyTogg(m)) ) TT1
+OL (F)
(212)
1 (Ti)
Observe that further
(F))
(F.F))
(xI( + ExD
(213)
1 1+000
1+
(A.)
polyLog(n1)
< 1lx
(214)
Vnn
using the concentration of x. see Lemma 9 and that 0 S C2,s S 1. Summarizing.
(r:)c,, (
( + O
polyLog(m))
(215)
C+
n11
Finally, let ō, Fr. ps using the shorthand p: ProX(Ei Cz cz 22
[x).
One can control
E 1 1+ (F)
11(03)
(216)
[x]
1 1+ (p.)),
using X 1/2. From Lemma
11
5; Oz, (polyLag((n)) 1/rrr).
Putting all intermediary results
together, and taking the expectation, it holds that
e" (Proo (c.PEZ:c(x)) 1 e" (c:
1 - +0 (olyLog(m1)
+ AXE
(217)
a
m
proving the lemma.
41

---

Published as a conference paper at ICLR 2026
F.3.6 LAST STEPS
We begin by defining the constants and X as solutions of the following self-consistent equations:
1 R
12"
C.vZ
(218)
(2" u)2xx
E
+ Ax
(219)
(=
where
= Prox (e + Cg c 7z c
and take for granted that and x exist uniquely. further assume the regularity conditions for map (Mg: Mg; b) - (v,
Assumption 3. The map (Pgs Me, b) + x) is continuous and (VE, xE) - (v.
as 71 -> oc, where the convergence holds uniformly over (Heg Mg, b) any compact set.
Define the asymptotic inner objective function by
dA(Mg: Mg: b) = . 10(="' + +e + A
where we recall that e Maca + MelE + cc NO 1) independent of cq. cg. and Cz. Let
G: min (e(pg Mg. M
Ma da ue: b) + A 1ug
Msl Y
?
1
My
2
1
ug
(220)
and
A lmg galMy bc. = @A(Ma +
1
Mg
GA min 9A (Pgs Mg, b), Pg.le.bb
mel
(221)
that G denotes our original optimization problem and GA is the surrogate problem where the random dy-dependent function & has been replaced by Φ In establishing the result of Theorem remains establish the asymptotic equivalence between G and GA. We begin with the follow- ing brief result which establishes the sufficiency in considering minimization of and over compact set in R
Lemma 14. Let u (Hg. Mg, b) and set
0 = arg min ga(u),
arg min 9A (v). ecRs
veRs
For 8 E (0, 1), there exists compact set V = V(5) C TR3 not depending on (equivalently on
n.),, such that
0 1 EV
for all d E N, with probability exceeding 1 1-8. 8.
Proof. If a function h IR + IR is coercive, in the sense that
lim h(x) +00, ||
then h has bounded level sets
levx(c) (xr (x : S cf forc R.
To show that AA is coercive note that |f + oo, but (1 M4) remains bounded, then necessarily
|b] + oo and (A5) implies - 0o. If indeed [1(gg M&)ll - oo, then due to the quadratic regularization term
1
Qu (e)
uel
1
Pg mg
7
42

---

Published as a conference paper at ICLR 2026
have
lim 9A(Mg: Hg, = oc.
since 2 O, and so 94 is indeed coercive. Moreover, the map v + (v, x) is continuous by Assump- tion [3] and 6 is continuous in by continuity of and the proximal operator. It then follows that 9A is continuous in v and so, having established coercivity, its level sets are closed and bounded, hence compact. From similar observations and reasoning, we see that ga(v) is also continuous. Since {(E;, 3/) J221 are sub-Gaussian and has quadratic majorant by Remark {[(Es; vi) are sub-exponential and so by Bernstein's 's Inequality (Vershynin, 2018 Theorem 2.8.1). for any k sufficently large,
(
P
Σ
((Gu)) E(((g)) > n = 2exp (-Ck
n1
SE(mi]
where C >
0 is an absolute constant. Therefore, taking K En, 21 2exp (-cx S ó, by union bound, one can ensure that for
large so that
Ω =
(leny) - S $
n121
RE(ma]
F(2,)) 2 8-8. In the remainder of the proof, work the event Ωg.. Letting Amin >> 0 denote the smallest eigenvalue the positive definite matrix in Q(1gg we have
AAmin
(e (E,a))l S ga(vz) S [[(((e,))] =:
30
where write the optimal solution by 2g (Mg.d) aa b2). The first incquality above is due to ga(u) 2 Q(pq Ms) whereas the second follows from ga("z) I ga(0). Thus, find that the first two components of oj are bounded uniformly (independent of d) in particular
2,0
(1 (Ugga:HEea)) S V Admin
Bo
Moreover, as
ga(u) 2 E[f(e, 8)] - K + 00
|b| for fixed (Mip me) by (A5), there exists B,
independent of d- such that
inf ((nol)))lSR 1B|.
ga(Mg: Me. b) 2 30 +
Hence, on Ωg the minimizer vß of ga lies in the set
u = ((ug. ug, b): (1 se)ll < Bo: |b| S B1}
which is compact by continuity of Ja and importantly does not depend on d. Therefore, taking 3 max 3o + 1), having previously established the level-compactness of ga. we have that
v:= U lov (3)
is compact and contains i and vi for all E N.
The compactness yielded by the above lemma is an important fact that will be carried in the subse- quent results. Notably, we remark that all preliminaries that have been established hereunto involv- ing (b,,) errors terms for some sequence (b,,) hold uniformly over the above defined set To see why, simply recall the meaning of writing O(b,) is to infer the existence of an n-independent constant C such that
Gn S -bn
for n sufficiently large. Revisiting our previous results, one can check that, given a sequence (v) parameterized by e V, the map C(v), namely the map from the parameter to the order- defining constant, is continuous. This turns out to be a simple consequences the continuous the loss f. Therefore, sup,ey C(v) < co, and, stated, all previous results hold uniformly over V.
43

---

Published as a conference paper at ICLR 2026
Lemma 15 (Uniform convergence to @A). We have
sup weV
[Bd(v) &a(1) -
as d + oc
Proof. Let
S, : Prox(e+ c.
noting that the dependence on ni z comes through deterministic n1-dependent quantities XE and vg. Recall that by Assumption (vg, XE) - (v. the proximal operator, applying the continuous mapping theorem together with Slutsky's 's theorem yields convergence of
uniformly over V. By continuity of
Note that this convergence holds uniformly over V as the proximal operator is non-expansive (i.e. Lipschitz). For some f; lying between F3 and ri and i, between ñ; and zn, a Taylor expansion yields
1
Bloa(v)]
Σ
Bl((=nniid)) + (fi - F.) BB ((()))) - -=%aa]
11
(E(m1)
X + V
1
V
= =E(((=h, :y)l + polyLog(n1)
+ 2 U
where in the second equality used [1' oc OppolyLog(na and applied the upper bound on Ira - Pal from Lemma 5 and Lemma 11 bound
i 2n Now, for M 0, decomposing
E[((=,, w)] = E [( (eten,
e(2,, (2) (elet,
we have that
((= u) [e(e,,
s3 (e(et,,o)MM + E [e(=" + +
uniformly over V by the Dominated Convergence Theorem. Uniform convergence of (VE, Xx) (N, x) yields uniform boundedness in L of (e(=", :1))m.21 since l has bounded second derivative. Namely.
sup sup E[[(2,, u) < oo EN
which provides uniform integrability of (E(=", : 1))m1111 That is for arbitrary e > 0, there exists
M > 0 for which
sup
[( (( (e(=,, (e(e, <E
wEY
as 11 - oo and so, uniformly over V, one has
E [((2n v)] + E [l(2" + +e;
Lastly, by Assumption 3 2 - Av2 /2 uniformly over V, which yiclds the result.
Lemma 16 (Uniform convergence to Eő(v)). We have
sup (o(u) - E[((v)] P0 vey
d
44

---

Published as a conference paper at ICLR 2026
Proof. We include the parametrization of in x2(e), (U), and other quantities where the param- eters v = (f4g. Mg. b) were previously fixed and hence omitted in the notation. Note that continuous differentiability of the map t (Mg. Mg. b) + E(( 2) + Cq.iMy + Cg + b; yz) carries to the map + =²(v) because strong convexity from the regularizer A/2|2211 ensures unique minimizer and the Implicit Function Theorem provides that the minimizer depends smoothly on v. Thus, the map + (1) is uniformly bounded over V as the set is compact. Then, observe that
|= (0; vEV V
sup
1
Σ 1E
((G((): O(polyLoggnn))
by compactness of V and since supisn, le.l OppolyLog(n:)) by the proof of Lemma Again, invoking compactness of V and Lemma have that
sup o
V c
O(polyLognn)))
V
T1
:E|ma
since by the Implicit Function theorem, (v) (pp) and have that ||e" OpolyLoog(n,)> ). Putting these results together, we have that da is Lipschitz on V with a poly- logarithmic constant which we denote by La. Namely,
16(u) <(w)) S Lalle w|| ||e w|l O(polyLog(m1))
for w E V. Lipschitzness of Ee follows by linearity of the expectation and thus the centered
process Z = Φb Ep is 2ip-Liisch We finish the proof with covering-net argument. Fix e > 0, set
E
4La
By compactness of V, let v(1), ...., (Na)) be points in V such that
c NN
(v (m))
where Bs, (s,(m()) (v(m) denotes a ball of radius 8a centered at ((()) A standard volume argument shows that we may take Na (pooyLogg(m)) as La (ppoyyogg(m) Using the variance bound of Lemma & Chebyshev's inequality yields
E /
= OppolyTo((m))/m)
for m E [Na]. A union bound then provides
P max [Z((ml)) >
2 O(polyLog(mi)/mm)
m.N. N
Na (polyyoo((n)). By construction of the cover (n(1) ,..., exists y(m)) such that
(e()).....((N))) ((N)), for any V. there
IZ(z) (m(m)) 2
Hence,
sup
(o()) E[o(v)] >
<P
max
Z(u(m))] - + 2
as + oo which concludes the proof.
The following result marks the grand conclusion of the section and completes the proof of Theorem
Lemma 17. We have
|G - GAl P,
(222)
as d + oc.
Without loss of generality may assume 1/ is a closed sphere of radius and by (Vershynin) 2018 Corollary 4.2.13), Na (2ró 1)..
45

---

Published as a conference paper at ICLR 2026
Proof. Let v² and Vn be the respective minimizers of g and gA. hiding the d-dependence for nota- tional ease. Setting
A =sup (o(v) o(xlll VEV
have
G G - G-
- ga(xa) 1 g(va) - gA(()) SA.
By symmetry, we obtain
G - A
and so the result follows by the triangle inequality in applying Lemma 13 and Lemma 16
G PROOF OF THEOREM 5
Appendix F details the asymptotic characterization of the learning of the attention model in the asymptotic limit of Assumption [ We now expound the related characterization for the linear
classifier baselines [Pool and w.b (B), summarized in the main text in Theorem 5 The first part
w.b
of the latter for the pooled classifier pooo was already covered in Corollary in Appendix F as
it coincides with special case of Theorem for the attention model.
We consequently turn to analyzing the learning of the linear classifier acting the vectorized inputs Formally, let us consider the empirical risk minimization
vec (3). described in subsection I.1 beb problem
> [| argminø(b) Σ E((Fi, w) (u((us), w) + b, /
w", D" argmin wERLd,BER iE|ml
(223)
BER
where we denote fi = vec(Z; Zi) the flattened background term of the inputs, and p(e)) (v))(u,ET) We denote by Pv the law of over {0, 1), and recall that p. 02) by definition. Note also that in these notations, y
8,,0, is a function of v.
Proposition 3 (Test eгror and train loss of the linear classifier on vectorized inputs). The test error and train loss of the linear classifier acting on the vectorized inputs, described in subsection T.] trained with the empirical risk minimization (23 1o
concentrate in the asymptotic limit of Assumption
Eraiin min Ey,e' 18(2" + +b b.y ,222
b
$ (--10 6z (-b- --m(mu))
Eyest (1 )
+ Euzor
For any E R& noted m(v) mn .... + Ur..L. where the summary statistics v, X, {m hele), b" are given by the set of self confisonnistent equations:
m4 X Byn [0(22 + b², (62uh + m&k 2)
1
2²
X By [e'(2" + ())=
L Ey,
(" + b².y)
+
ax
1 (-=""("* + b", y)x
where v' ~p. y y=1 80,0n, (0, and
bt arg min E 10(=" + b², v)] + A 22
We employed the shorthand xt proxxe( w)) (vz + m(v')).
Note that the data distribution formally coincides with a Gaussian mixture with 2 + isotropic clusters, and the analysis of logistic regression on such data is covered in Loureiro et al. (202T). In
46

---

Published as a conference paper at ICLR 2026
this appendix, we rather give more concise derivation in the specific setting considered, leveraging once more the leave-one-out approach. The following derivation closely follows the steps of the proof of Theorem [ detailed in Appendix[F] For the sake of conciseness, mal sketch of the derivation. Before doing so, let us observe that the equations (224) are amenable to being massaged into form closer to that of of et al. (202T); MMgnaaco al. (2020a).
only provide an infor-
Remark 8. The system of self- equations 224 can also be written as
1
L a
X
E
1
1--proxy( m) (vz + m
X
(224)
X
x
+ X
0²
m& X + x
mg
E [[= uzz-m(o')ik)
my
(225)
x
10 + お
mi
a
6² = x2 B [(= - -m(')) =
2²
(226)
x)2
Proof. We begin by noting that the derivative of the proximal operator reads
0prox,e(./e (w)
(227)
0
1 78th (prox (w))
Therefore,
(= + b) 1 ("((= * + b", ,,, v)
X = E
(228)
and the last equation of (224 can thus be written as
L
x x
(229)
Let us now focus on the first equation of (224). We have
0 =mm + (C -b,,))(( + m E [e' + + y))²
(230)
8² -Amx-
(= + b,
E[(22 - v2 m()')u(] + m&B
(231)
1 xl" (2" + b,,y)
X
Thus,
m&
m,-
(232)
x X
Finally, starting from the second equation of 224
0 ==>=²_ B [(= - v2 -2- -z2-m()))
[[2
22 m(ro))(uz + mm(')))
(233)
X
x
=xv²_
1
[[=- - 12 -vz
m(
1
Σ
m 1²E
(" b" b",y)
(234)
x
g2
1 + (3 + br,y)
k=1
1 L 1
E [(x" - vz 15-mmo --===2up m(
(235)
x2 a + X
(2' + b,, y)) 1 x+ + b".
m + v?
(236)
02 x))
Thus
L + Φे Σ m
(237)
(A + x)2
47

---

Published as a conference paper at ICLR 2026
0.06
0.40
0.05
Etrain
Eost
0.03
0.30
0.02
0.01
0.25
0.5
1.0
1.5
2.0
0.5
1.0
1.5
2.0
Figure 4:
Train loss (left) and test (right) of the linear classifier acting on the vectorized
outputs, as discussed in subsection [f the main text, for L = 3, = 2, 0 = 2, π = 0.5, = 0.01. Solid lines: theoretical characterization of Proposition 3 Dots numerical simulations in dimension d = 1000. Error bars indicate one standard deviation over trials.
Sketch of the derivation
For a given b, let us introduce
① argmin Σ
(((f), w) + (u(wz), + +,,35)
||w||
jE(m)
P argmin
Σ
(((fs, + (u(w;)) w² + b, ) + 2 |100
Y1
J
+ 2 min
[e((f) + (u(v;), + + 55(w= -w--u w w Hi w
where the Hessian is defined as
H 1 (" w) + (1(15), + bb,,) ((;)))( +(u)))) + ALa
Then it holds that
(fi +i++u(un u()s)))),"") ((fi
where
M()) 1 st u fa u(v)) H ((v.) + +2f ((ms) tr(l1-1]
x
n
We used that (Ię(t)))(=;)) 1I. I(t.) < I|
Probabilistic analysis In similar fashion to the proof of Theorem one can show that the parameter x satisfies self-consistently
L Ez,y.
("(2",8) +
ax
(1+1"(2",))) ()xx
where =² prox (00xx((..) (vz + m(u)). with m(v) = (u(v), w²). Using the stationarity condition
"º -
e'(1s, w²) + (M(U)), w²) + b. U,) + m(v,))
n2 JE(m]
allows to reach
3/
E Ey.z [e'(2" + b, x)*)
48

---

Published as a conference paper at ICLR 2026
and
m(v) - -lyyy +b +.,0)((())(())) + m(v), 2- /ll(t))|2-m(o)//0u)
yy ((2=' +; m(v)
where 2 ~ (N, 1). Finally
((b) = E (((=", v)] + 3
This completes the derivation, but there is one further simplification. Let us introduce the unit vectors fer E Rd4, where the + 1 to (k + 1)d-th elements of eg correspond to 0C, with all components otherwise zero. Note that all these vectors orthogonal to each other. Then one can write
u(v)
Σ UnCk. kE|L
Then, simply, one has
of
and
m())=
Σ KE[L]
ukm
with my : (w0, px). This simplifies the equation for m(v), and yields the characterization of Proposition
The theoretical predictions for the test and train errors of Proposition 3 are displayed in Fig. and show good agreement with numerical experiments performed in dimension d 1000.
H PROOF OF COROLLARY
The full technical statement of Theorem presented in Appendix F and of Theorem presented in Appendix G provide tight asymptotic characterization of learning errors in terms of small set of summary statistics, characterized turn as the solutions a set self-consistent equations. For the case of the square loss E(y, =) 1/2(y =)². in the limit of vanishing regularization À 0+. these equations considerably simplify, making possible to reach closed-forn expressions for the test eтror in particular. These expressions
summarized in Corollary 2 in the main text. In this appendix, we provide the full technical statement. For ease of presentation, we break the statement into three proposition which we derive in succession, respectively for the attention model Aqu (6).
the pooled linear classifier pool and the vectorized linear classifier w.b
vec ur
H.1 ATTENTION MODEL
Proposition 4. From Theorem in the asymptotic limit of Assumption 7 the test error of the attention model converges in probability to a limit Ereat [A]. For the quadratic loss function Ely.2) 2) 1/2(y 2)?, this quantity admits well-defned limit in the limit X 0. This limit admits the
49

---

Published as a conference paper at ICLR 2026
expansion:
Etstt [A]
&est [A]
(238)
1000 (9.aa)
(ō6 (g.5_)6p1 (8 + [g u1 jous/ug
E
"g |a_
+
1
π E
(239)
1T
11 s-1
(9v...................... 20nge 10.0. 704
[|]
E
(240)
a1
V 2T
66
(0u, S0)0p2
(g. 8))ppp + (8 + (00, s+)p" (g., (t) )
(241)
1/3 |8+1
+
(242)
The limiting error is
000 + (g. s 8-)
8:00 [A]
=(1
T g.
(243)
11 ||5
joc
(Ov, s 11
(9, 84
N 1.
(244)
113 ||++|l
We introduced
11:00
(100-1100,
6u12 5b
8J + 0T (100--
(245)
800
800
where
[ea]
E [cgeg]
E [cg]
E LyC
7
E [cgeel
E
E [ce]
SICE
(246)
E [cg]
E [cg]
2T
E [cg0c="]
E [e
E
.8T
IE
[cec]]
8J
E
yoec Lyc
(247)
E
E
[c2]
[22]
I
Finally, we denoted ōur = 1/23 (1/20² + p uz 811 7p 8112 - ypo 8pm) We remind that the joint law of ce. Cg, Cy is given in Lemma D and y is defined in Theorem 3
u 6us 5us/
Sketch of the derivation- In what follows, we consider the case of quadratic loss
1 E(:; v) マ (ye - 1)2²
In our problem, ((x)) ((x; + (.)))) We have
(x)x) = x; + V and (x)x)= 11
Moreover, for this case, the proximal operator assumes compact, fosedd-form expression
Prox{(x:7) = 1 1+ (y: es).
x 7 +
These closed-form expressions allow us to greatly simplify the self-consistent equations appearing in Theorem 4 Specifically, we can rewrite (88
E
+ Xx-
(248)
01
1
x2
50

---

Published as a conference paper at ICLR 2026
and
E (x(6-r)
(1+2x)
2²
A ++ E
(1+02x)²
Let be the unique solution to (248) check that
In the ridgeless limit (with X - 0+), is straightforward to
lim
Xx
1,
for an < 1.
X
a1
and
X
Xridgeless
for a1 > 1,
X-
where Xridgeless is the unique solution to
2
1+
We focus on the latter an >> case in the following. In the ridgeless limit, the fixed point equation for further simplifies
c-x(u-e
E
11+02
D2
E
Then, the function ((2g (g; assumes the simple form
(y c)
olug, Me, b)
(249)
Requiring that the gradients with respect to plg ug. leads to the following characterization for the minimizers f1, M2,
[4
I(a1)
A22
((1))
(250)
with
c
E
E
E
1+c2 x
1+c2
1+02x
I+c2 X
I(a1)=
E
Cg
E
E
J(az)
E
(251)
l+
X
1+c2x
++
E
E
E
I+cix
I+cEx
I+c+x
I+c2x]
Note that M(a1) ) the Gram matrix of the random variables (Cgs cg, 1) for the inner product (a, b) E [ <xx], and is thus invertible since the random variables linearly independent
Large ar behavior In this limit.
We now study in further detail the regime of large sample complexity an
> 1.
x
+
(252)
a,E
a1
while
1
(253)
E 2
a1
Note that the limit 2/2 01- ( implies that for large sample complexity, the readout weights lie in the span of g, We denote の0 c + cgpu3 + Jos with
(254)
100
51

---

Published as a conference paper at ICLR 2026
where
E
[
E
[cg]
[3/c4
7
E
[g
E
[cg]
J
E
CCC
(255)
E
[cg]
E [CE]
1
2x
The corresponding residual test error is then simply given by adapting
85
to obtain
1000
+
(g.
8-
Etst
Etest -(1 - TT ) Bg
(256)
13
88
1000
(0v. 8.
(gs 8 84/199
+ n 9.54.8-
(257)
11 | ||
with n3² = 3/ 2 + (2g9)² 2yp (119)3]²
We now turn to ascertaining the
leading correction. We introduce
En
6pm
12
+
81l2
(258)
b
8000
a1
a1
with
p
0/12
(19)-1
5J +8
(259)
8b
600
where we denote
1
E
c2
E
E
E
yJ
81
E
C
E
E
[ege]]
8.J
E
yosc [yc²]
(260)
[[]]
E
[co]
E
E
E
[
E
Finally, let denote 8143 (1/222
1 5μ1 + 1120 8p12
7/4 6102
7112 811)
u sus /
Then, the following asymptotic correction holds: Erest
&eest
(261)
500 +00 (00
(bb 8- - (i + (g. 8- yo13000 /O133/18
"
e
+
(1
π T)
(262)
01
2m
M3 |
600
(Re. 30 (0. 1
E
M3 [x]
(-68- (00.e4 1a2g (a. (6 (0 4 2 +to (
a1
"g 11
(263)
+0
(264)
H.2 POOLED CLASSIFIER
Proposition S. From Theorem 4 in the asymptotic limit of Assumption 7 the test error the the pooled
classifier model converges in probability to limit Erest [ p. For the quadratic loss function Eyy 2) = 1/2(y - 2)², this quantity admits a well bell-defined limit in the limit + O. This limit admits the expansion:
21-1-x22
Etsst [
Pool
eno
pool - (1 -T]
E
2T 1 TX7
D²
2
2T T)
)2
(265)
2r- 1
2T T -T)
u2
1
+
(266)
2/2T
2T (17) TT)
a1
52

---

Published as a conference paper at ICLR 2026
The limiting error is [L
T 2TT 1 - x22 2x - n
(-2--111 2 (1 - x)
2x x2 mT
Egest!
+ T
(267)
We denoted the signal-to-noise ratio X' = BR//I.
We remind that the pooled classifier corresponds setting the softmax inverse temperature in the attention model to zero, namely / = 0. In this limit, the joint distribution of the parameters 8 S-, C22 Cg2 Cy detailed in (92) simplify to
Sketch of derivation
11 I' 84 = 8_
(<
1
r)):
8,,1 I.
N
02:
L(1 - 72)
17
c. VL
(268)
Then, the limiting summary statistics ui uu, joo are given by 1e = 2/z² and
T
1
TX
T/)
(
(269)
1
800
2T -
i..
112
2T - TT)
(270)
VL
1 21 - T))
2T n Tx22 xt)
j00
(271)
1 + x2² (1- x)
The residual error then reads
b
Xux -b-Xxu /VE
(b
E00 =(1 =))
+ (1 (1=
(272)
112 /V.
11 /VI.
2n
1 - T22 (1 (1-T) TT)
(-2m 2Tx - T))
(_2T
1 + T22 1T)
= +)
+ +
(273)
2 - T))
We used the identity
a + bg }
a
Ey
= -Eg.o' (1-a-botcg2ul =
(274)
VP + c
Finally observe that I = a/1 + 00 _J = + aa
As a consequence,
111
(/)
11
(275)
6
80e
and
2 =
1+88- - -26(2T 1) +
(420)2
(1 + T22) + 26 - 1) 112
-1//T
(276)
a1
L
VL
I 1+ 22<2 1) + (W
1. (1+ TT&2) + 2b 1)
P3 11-12 72 +
(277)
201
VI 11-11 22 x
(278)
follows that the leading order correction to the test error reads
Erest -
(279)
1 2u x2 x) }
c
2xx11 m)
2TT - 1 1- 1 - x)
x²
(1 (1-77) T))
(280)
2TT
2 - T)
(2
-
2
e
2m
1 + xx2 - T)
u2
1
+ π
to
(281)
2v2m 2m
2T (1 -T)
01
53

---

Published as a conference paper at ICLR 2026
Fax: attenten
m, md
Sap. classifers
0.16
0.22
0.14
L 4, Lust
0.20
0.12
0.18
0.10
0.16
0.06
0.14
0.04
0.12
0.02
0.0 0.2
0.4
0.0
0.2
1.0
Y
Figure 5: (left) Residual enror Ejst in the on + 0c limit function of the alignment ? (1,E) between the attention query weights and the signal vector, for the attention model (blue) and the linear classifiers (dashed red), trained with the quadratic loss at vanishing regularization. = 5. R = 1,0 0 = 3, = 0.75. Solid lines correspond to the theoretical characterizations (256) and (272). Dots correspond to numerical simulations in dimension d 100, and large number of samples = 10', averaged over 10 trials, with enror bars representing one standard deviation.
(right) Residual error of the attention model for different training lengths L, R = L and test-time
lengths Leest Riest Lest and 0 =
Comparison with the attention model We contrast in Fig. the residual errors Exs achieved in the limit of large sample complexity on >> 1 by the attention-based and linear classifiers. As we detail in Appendix G the vectorized and pooled linear classifiers share identical residual test errors. Interestingly, for small alignment between the attention query weights q and the signal vector & the attention model performs worse than the linear classifiers, as the discrepancy between q, E can cause the model to spuriously privilege tokens devoid of signal.
H.3 VECTORIZED CLASSIFIER
Proposition 6. From Theorem in the asymptotic limit of Assumption 7 the test error of the vectorized classifier model converges in probability to a limit vorr. For the quadratic loss function E z) = 1/2(y 2)², this quantity admits well well-defined limit in the limit X + 0. This limit admits the expansion:
1-x
1
60
Etest [L vec]
= [vee] +
(1 7)
(282)
v2n
11xx9
(x +(600 (1600))-2(2-2222 1ybx0)
+
(283)
14x22
01
a1
with
2T 1)
2T 1 -T(1 x)
(2/0022
80
(284)
I+ [ x)x
1+ T X)
The limiting error is
2T 1 Tx2 E [L vec (1 T
2T 1 TX2 Tt)
(285)
2 - T))
2Tx (1 - TT)
We denoted the signal-to-noise ratio X @R/VT..
54

---

Published as a conference paper at ICLR 2026
Sketch of derivation
For the quadratie loss and vanishing regularization, the fixed point equa-
tions of Proposition
3 simplify to
0'p(1 - 6*)
m- 1 p1 + (T. 1)p)
(286)
2² + 0pp + 02 1)p 11
LL - b²)² (1 + 0p + (1-1))))
(287)
1 02ppp11+ (L 1)p)
2x 1)b² + (b*)²x
(288)
x = L
(289)
where p = aR L
and
R(R- - π
6- = =1 (2T 2) A 8 Lp²
p= 0= I(L 1) p
(290)
We used shorthand : compactly rewritten, introducing the X' introduced in Theorem [ setting, X' admits the compact expression
+ + p(1 + (L. - 1)) These expression are amenable being more
We remind that in the current
X 2R VL
(291)
The self-consistent equations then simplify to
6" = 1 (2 2)(1 + T2²2)
(292)
+ T( x)
1 (1 -b)) b²)
(293)
R 1 TX²²
+ 2222
n x2 b*)
x² - bb)² (1 + xX2)2 2x(2x - 1)b* + (b*)²xx
2 = X
(294)
+ x2²
01 - X' = = > L regime- - Following a similar derivation as the ones detailed in the previous subsections, the test error found to admit the large an residual
Eest = T (-x-
600 y00
+ +11
(295)
400
with
(100)2 11 -20 1222211-000
2TT - 1)
(296)
(1 + x222
1 7(1 1 - T -)
(2t 2) + TX2)
2T - 1 111 x)
00 = 1+
(297)
1 + (1 x)x2
1 1+ - x
and the asymptotic expansion
e --x
1000 + X'000
e - - )² b00
Exest Ege
π
2m
2(p00
(1 T)
V2T
2(100)3
(298)
4
2 - (1 14
+
o
(299)
On
an
Remark (Comparison with the pooled model). Note that the residual error Eoe can be explicitly expressed as
Exest (1) (2x 2)(1 + +T
2TT - 1 7x22 T) (1-)
+ лΦ
(300)
2xx(1 x)
2 (1-x) x)
55

---

Published as a conference paper at ICLR 2026
This incidentally corresponds the residual error achieved by the pooled classifier trained with ridgeless quadratic loss (272). since for the considered data distribution X' = on//I. We also furthermore have similar correspondence at the level of the summary statistics, namely b Jo0 - ,00 leading order corrections are related by simple factor L:
" |VT,, where i "0" (VL are defined for the pooled model in 270 Furthermore, the
veest,vector
Eno
1. + o(1)
(301)
Etest.pool
Eoc Cest
Note that a consequence of Remark 9 is that in the a - oc limit, for ridgeless regression with a quadractic loss, the pooled and vectorized models converge to the same solution, in the sense that the weights of the vectorized model correspond to that of the pooled model stacked L times. Both models furthermore yield the same limiting test error. Let us also comment that Arnaboldi et al. (2025) also observe a similar speed up between related flattened and pooled models learning from sequential data, in related task, in terms of weak recovery time. The result of Remark instead bears on the coefficient of the leading asymptotic correction in terms of sample complexity.
Remark 10. We note that the joint limil 0 L + oo, = 4/L = 0(1), = O(1) can also be setting
analyzed, and is simply given by equations 292
(302)
b b
Study of the o1 do residual error We now examine the behaviour of the residual error & with the signal-to-ne noise ratio X. We first examine the case X' - c0. In this limit,
6 = = +1 ol)))
(2 + 0
(303)
X²2
X²
The residual error then decays to zero as
e x²
poc test
(304)
X
In the opposite limit of small signal X 0,
800
2 1 + o(1),
(v00)3 (1- +(x22)
(305)
Then
Eest
X-0
min TT.
(306)
These limiting errors stand in coherence with Theorem
I
I DERIVATION OF CONJECTURE
As corollary of Theorem 4 we derive in the appendix the capacity of the three considered models, namely the largest number of samples that can typically be perfectly classified, up vanishing training error. The corresponding separability threshold was characterized in the seminal work of
Cover 2006), and revisited in many later works, e.g. Gardner & Derridal (1988); Krauth & Mézard (1989 Candès & Sur (2020); MMgnacco et al.] (2020a). Note that at the level of the representations vee J,(),, Spont( (-), the capacity intuitively reflects how well the representations separate positive
and negative samples in feature space, with a larger separability thresholds signaling more markedly separated classes.
(Separability threshold). Consider the empirical risk mininization problem 14 for the attention model, or the related problem for the linear classifier models, with logistic loss f(z: y) = log(1 + e-v and vanishing regularization X 0+. As stated in Theorem 4 the training loss converges in probability in the considered asymptotic limit to limit Evainn We define the separability threshold a² of the model as
Definition
a² ssup {a 0 Errain 0).
(307)
56

---

Published as a conference paper at ICLR 2026
3.4
---- ai
3.2 %6
apol
1
3.2
110
3.0
24"
22x109
2.8
0.0
0.2
0.4
0.6
0.8
Y
Figure 6: Separability thresholds for the attention model (green), and the pooled (red) and vectorized (blue) linear classifiers, as given in Conjectures 2 and 2 function of the sequence length L (left) and the attention query/signal cosine similard larit 7 (right). Left: 0 2, π
0.3, 7 = and R 1 is kept fixed while L is increased. Right: 0 2, 0.3, 2, R 1 and ry varied.
closed- form characterization the separability threshold a* can be heuristically derived from Theorem 4 for each of the three models. We first provide the characterization for the vectorized classifier.
Conjecture 2 (Separability threshold for the vectorized classifier). The separability threshold for the vectorized classifier is equal to
L(1 s2)
Ok
vec
max
(308)
eE e
( (b X's u) + (1 - т T) (u 6)] dd
We have used the shorthand X OR/VI.
Proof. First note that the following identity follows from Proposition 3 and most conveniently seen from the rewriting of Remark 8
2² ==m"==BBee
a + + +,,))
(309)
for any given We assume that the loss function is of the form {(2,y) = f(yz). and satisfies lim;-+00 {(2) = O, while being convex. We assume i to be decreasing. with a monotonically in- creasing and negative derivative satisfying lim (2) = 0-. We denote k lim (z)x), which we assume to finite. Note that all those assumptions are satisfied in particular by the logis- tic loss function. We again assumed all token locations are symmetric, leading to solution m, = m for all A [L]. Introducing the cosine similarity
vImfov [0, 1]. and the normalized quantities 7 x/v, b t/u, and introducing the random variable u = ('( * + b)y)
2
Z [u
(310)
But 2 8y 8y.1 Rm + Xyu 0 by definition of the proximal operator, and 2" = yē- b. while 2 ~ No, 1). Furthermore, u E (-x, 0) Thus
= 2 /
T ( (2)
Rm + xu
i (-)(2) + + Xu
xu
E 1
+ (1 - (1--))
udu
-x
(311)
(~)(2)
i (-) + b yu
= ②
=
b Xs yu
+ +(1 (1-r)0
udu.
(312)
57

---

Published as a conference paper at ICLR 2026
Following Mignacco et al. (2020b) aim determine the necessary conditions such that there exists solution satisfying / = 0o, = 0o - which should hold for a solution achieving zero training loss. We conjecture the following limit
lim IB [u²] = 2 7:0-000
( (-b - Xs - u) (1 T) (b u)] udu
(313)
(T (b Xs u) + 7) (u b)] du
(314)
where we remind that Φ' is simply standard Gaussian density. Then, necessary condition for the existence solution with y 7 = oc the existence of an s E [o, 1] so that
L(1
(315)
O (b + Ns + u) + (1 T (u 6)] u²du
Note that the pooled classifier can be mapped special case of the vectorized classifier, formally evaluating the expression for the vectorized classifier for L, R - 1,0 - OR/VL. Leveraging this connection yields the following conjecture.
Conjecture (Separability threshold for the pooled classifier). The separability threshold for the pooled classifier equal
(1 s2)
apooll
(316)
+E(O.1].b
I (6 + Xs u) (1 π T (u b)] u²du
L
Finally, similar characterization be conjectured from Theorem for the attention model.
Conjecture 4 (Separability threshold for the attention model). The separability threshold for the pooled classifier is equal to
a
max My m2g. b
(317)
2o cx f 4' (cuttu(tt,mmtemmm)) uudu
(Stutulbteem
E Egcc,
Proof. The derivation proceeds in close likeness to that for the vectorized classifiers. First observe that from Theorem 4 the following identity holds:
2
2 [=2 + Cally + Cellg b, y))
(318)
Introducing the normalized quantities 7 = x/u, = 6/u, Mig = Mofu, mg = ne/u, and introducing the random variable u = =(2" + Cglly + Cg/l + b)y), this identity can be compactly rewritten as
1 = 2 [c
(319)
But -C 22 Xy/ 0 by definition of the proximal operator, and =" - b Callg CeHis while 2 ~ (0, 1). Thus
c y(b m. c
E 2 -
E
2
udu
(320)
Co
Again, following Mignacco et al. (2020b) we aim to determine the necessary conditions on such that there exists a solution satisfying и numerical experiments, as illustrated in Fig. 7
00,7 7
oc. This assumption is further motivated by
where v, are observed to diverge as À + 0. Then,
y c. ccms)
lim 7 Z [cfu2] 700000
Cg
udu
C y(b
C. cmmg)
c
u 2d
(321)
which concludes the derivation.
58

---

Published as a conference paper at ICLR 2026
10a
10
10
10a
10²
10
10a
10t:
Figure 7: Parameters 7, involved in the derivation of Conjecture I for the attention model trained with the logistic loss, as function of the regularization A. The curves correspond to numerical experiments dimension 1000, averaged over 3 trials. The problem parameters are L 2, R = 1,0 = 2, π 0.3.
The theoretical prediction of Conjectures 2 and contrasted with numerical experiments in Fig. revealing good agreement with the point where the training error defined as the fraction of misclassified training samples ceases to be zero. Note interestingly that the separability thresholds avec,puol for the vectorized and pooled classifiers are related by factor L. The latter can be ratio- nalized by the fact that the vectorized classifiers operates in R.d while the pooled classifier acts the smaller space IRd. Moreover, observe that while the threshold a for the attention model lies for large query/signal alignment 7 above "pool it becomes smaller for small values of y (see Fig (6 right). This temptingly suggests the intuitive interpretation that when the internal representation of the attention is misaligned with the signal, the attention model displays a smaller capacity than the simple pooled linear classifier. This conclusion echoes similar observation the level of the residual errors, see the discussion of Fig 5 and its discussion in Appendix F
59


Reviews:
[{"review_id": "RsG04K5G5G", "reviewer": "Reviewer_sdvV", "summary": "This paper introduces a sparse-token classification task, where positive samples demonstrate a weak signal in a random sparse subset of tokens, while negative tokens are i.i.d. Gaussian noise. The authors prove that by learning to recover the sparse subset, a single-layer of attention with a single learnable query token can achieve vanishing test error with the magnitude of signal only logarithmic in sequence length. On the other hand, the authors consider two baselines, fully vectorized linear classifiers and average-pooled linear classifiers, where they demonstrate that the signal magnitude requires to be at least of order $\\sqrt{L}$ for sequence length $L$ for vanishing test error.\n\nThe authors additionally study the effect of two gradient descent updates on the attention model, which allows the query token to achieve non-trivial alignment with the signal vector. They further characterize the exact limiting train and test errors of the above algorithms, as well as their capacity to overfit.", "questions": "* Isn't average pooling a specific case of the fully vectorized classifier? Would that imply that in Section 2, the optimal error of the pooled classifier is lower bounded by that of the fully vectorized?\n\n* Minor typos:\n    * Line 103: negative signals -> negative samples\n    * Line 122: I believe the authors meant $CR/\\sqrt{L}$, but it looks like $R$ is the exponent of $C$ in that expression.\n    * Line 458: the capacity if linear classifiers -> of linear classifiers", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "While the ability of attention to recover a sparse subset of important tokens has been the subject of many recent theoretical works, this paper presents a strong case by studying a fundamental task in classification. A particularly interesting aspect of the paper is on the technical side to accommodate attention in high-dimensional asymptotic analyses of classification error.", "weaknesses": "* The problem introduced does not really need to use the full power of attention. This setting can be seen as having a fixed query vector $\\mathbb{1}\\_d$ of all ones and a trainable query projection $\\mathrm{diag}(q)$. Recurrent/state-space models could also efficiently achieve vanishing test error by going through the input sequence and memorizing tokens that have a high correlation with $q$. It would be interesting, perhaps for future work, to extend this setting to the case where the number of queries in the attention mechanism would also grow, which would enable attention to outperform even more baselines.\n\n* I think some additional details in the main text could provide valuable intuition to the readers. For example, how do the optimal weights for attention/vectorized/pooled models look like in Section 2? I also thing some details might be missing on the loss, e.g. Corollary 1 could not hold if the loss is constant, since $q$ would be left at $0$.", "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": "l535M0kfxV", "reviewer": "Reviewer_CDaj", "summary": "This paper theoretically proves that a single-layer attention model can detect signals exponentially weaker ($\\theta = O(\\log L)$) than linear classifiers ($\\theta = \\Omega(\\sqrt{L})$) in sparse-token classification tasks where weak signals are randomly scattered across sequences. \n\nIn a high-dimensional asymptotic regime ($d, n \\to \\infty$ with $n/d$ fixed), the authors show that just two gradient steps enable the attention mechanism to align with the hidden signal and provide exact characterizations of the resulting test error and model capacity.", "questions": "Please refer to W1-W6.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The paper tackles a meaningful problem--detecting weak, rare, and sparse signals (e.g., lesions in scans)-- where attention’s adaptive token selection offers clear benefits over fixed pooling.\n\nIt establishes an exponential gap in representational power, showing attention detects signals at $\\theta = O(\\log L)$ vs. linear classifiers needing $\\theta = \\Omega(\\sqrt{L})$, where $L$ is the sequence length. The analysis provides exact asymptotic characterizations of test error and shows that just two gradient steps suffice for meaningful query-signal alignment.\n\nExperiments (Figures 1, 3) align closely with theoretical predictions for test/train errors, confirming the asymptotic analysis in moderately high-dimensional regimes.", "weaknesses": "W1.  The 4-stage training protocol (Section 3, lines 270–291) with data splitting $\\mathcal{D} = \\mathcal{D}_0 \\cup \\mathcal{D}_1$, only 2 gradient steps on $q$, then full ERM on $w, b$ does not reflect how transformers are actually trained.  \n\nW2.  The attention model (Eq. 5) is single-layer, single-head with no key/value matrices, feedforward layers, layer norms, or residual connections.  \n\nW3. Theorem 2 proves attention can detect signals at $\\theta = O(\\log L)$ by showing existence of optimal parameters $(q^* = \\tau\\xi, w^* = \\xi)$, but provides no guarantee that gradient descent actually finds these parameters. \n\n\nW4. Lemma 4 defines the clipping threshold $I$ using $R^2 = (2/\\lambda) \\mathbb{E}[\\ell(\\epsilon; y)]$, but this expectation depends on the solution $w^\\star$ being bounded. The proof that $x^\\star = x^\\star_{\\text{clip}}$ holds with probability $1-\\delta$ relies on quantities that themselves depend on the solution’s properties, creating potential circularity that needs explicit resolution.\n\n\nW5. Conjecture 1  provides the main capacity characterization via expressions in Eq. (15), but the derivation involves an unverified “heuristic step” (line 2356) in taking the $\\gamma, \\nu \\to \\infty$ limit. \n\n\nW6. The high-dimensional limit takes $d, n \\to \\infty$ with $\\alpha = n/d$ fixed while keeping $L, \\theta, R$ finite (Assumption 1). This is opposite to practical transformer settings where sequence length $L$ often far exceeds embedding dimension $d$.", "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": "j0qJhlXwQL", "reviewer": "Reviewer_4vxC", "summary": "This paper presents theoretical results regarding the learnability and representational capacity of a single attention layer on the task of sparse token classification. Compared to two linear model baselines, the attention model is shown to detect signals exponentially weaker in the limit of large sequence length, and is proved to require only two gradient steps for learning meaning representations that align with the hidden signal. Further, it provides a precise characterization of the test errors achieved by the attention model, highlighting a significant gap of the representation power of linear and the attention models. Overall, it presents solid theoretical evidence for the superior performance of the attention model in the sparse token classification task.", "questions": "Please see the weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. From my knowledge, it is a novel theory work that studies attention in the task of sparse token classification. It shows a clear advantage of the attention model in terms of representation capacity and learnability in this task.\n2. It extends existing literature beyond square loss to general convex losses.\n3. This paper is well-written and clearly presented.", "weaknesses": "1. The attention model considered is rather simplified, without multi-heads and separate K, V states. While this benefits theoretical analysis, there exists a clear gap between this simplified model and real practice.\n2. While a minor weakness, the capacity result remains a conjecture rather than full proof.\n3. Some experiments, even synthetic ones, could further strengthen the work and provide more evidence for the real-world applicability of the presented results.", "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": "tHwEM4Ph2h", "reviewer": "Reviewer_sbqr", "summary": "This work presents a theoretical analysis of a sparse-token classification task, in the high dimensional scaling limit of number of samples and embedding dimension, under the information theoretic limit of $n/d = \\Theta(1)$. For a simple attention model, the precise asymptotics of the test error under an optimization procedure is shown, for a given finite sequence length $L$, as well as a fine-grained analysis of how small the signal-to-noise ratio can be. To complement this, a lower bound is given for linear feedforward models, matching intuitions from previous works.", "questions": "- Finite sample/dimensional guarantees are given - what are the bottlenecks to extending the analysis to this setting?\n- A natural question is whether or not the learned attention model length generalizes to longer length, as shown in the regression setting of Wang et al. (2024). I suspect that it does up to some reasonable upper bound - could the authors comment on this?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The results are very rigorous, and hold under the high-dimensional limit of $n,d$ for a given sequence length, and follow upon a line of works analyzing the fundamental task for attending to sparse token signals.\n- The analysis shows that two gradient steps suffice for learning the hidden representation of the query matrix, which matches intuitions for weak recovery of the signal from other high dimensional settings such as multi-index models and such.\n- Lower bounds are given for linear models, in the sense that the signal strength requires $\\sqrt{L}$ scaling compared to $\\log L$ in the attention model.", "weaknesses": "See questions below.", "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)