Learning to Recall with Transformers Beyond Orthogonal Embeddings
Authors: Nuri Mert Vural
Alberto Bietti
Mahdi Soltanolkotabi
Denny Wu
Venue: ICLR 2026 Poster (2026)
Keywords: transformers
associative memories
factual recall
storage capacity
training dynamics
Primary area: learning theory

Abstract: Modern large language models (LLMs) excel at tasks that require storing and retrieving knowledge, such as factual recall and question answering. Transformers are central to this capability because they can encode information during training and retrieve it at inference. Existing theoretical analyses typically study transformers under idealized assumptions such as infinite data or orthogonal embeddings. In realistic settings, however, models are trained on finite datasets with non-orthogonal (random) embeddings. We address this gap by analyzing a single-layer transformer with random embeddings trained with (empirical) gradient descent on a simple token-retrieval task, where the model must identify an informative token within a length-$L$ sequence and learn a one-to-one mapping from tokens to labels. Our analysis tracks the ``early phase'' of gradient descent and yields explicit formulas for the model’s storage capacity---revealing a multiplicative dependence between sample size $N$, embedding dimension $d$, and sequence length $L$. We validate these scalings numerically and further complement them with a lower bound for the underlying statistical problem, demonstrating that this multiplicative scaling is intrinsic under non-orthogonal embeddings. Code to reproduce all experiments is publicly available.

Published as a conference paper at ICLR 2026
LEARNING TO RECALL WITH TRANSFORMERS BEYOND ORTHOGONAL EMBEDDINGS
Nuri Mert V Alberto Bietti,, Mahdi Soltanolkotabi, Denny Wu
University Toronto, Veector Institute, Flatiron Institute 'New York University, University of Southern California
mural@cs.toronto. edu, cestancl@uscc..m dennyvu@nyu.edu
ABSTRACT
Modern large language models (LLMs) excel at tasks that require storing and re- trieving knowledge, such as factual recall and question answering. Transformers are central to this capability because they can encode information during training and retrieve it at inference. Existing theoretical analyses typically study trans- formers under idealized assumptions such as infinite data or orthogonal embed- dings. In realistic settings, however, models are trained finite datasets with non- orthogonal (random) embeddings. We address this gap by analyzing a single-layer transformer with random embeddings trained with (empirical) gradient descent on a simple token-retrieval task, where the model must identify an informative token within a length-L sequence and learn one-to-one mapping from tokens to la- bels. Our analysis tracks the "carly phase" of gradient descent and yields explicit formulas for the model's s storage capacity-revealing between sample size N, embedding dimension d, and sequence length L. We validate these scalings numerically and further complement them with a lower bound for the underlying statistical problem, demonstrating that this multiplica- tive scaling is intrinsic under non-orthogonal embeddings. Code to reproduce all experiments is publicly available.'
multiplicative dependence
INTRODUCTION
Large language models (LLMs) routinely answer knowledge questions with little external context, indicating that substantial factual information is stored in parameters and can be retrieved by suitable prompts (Petroni et al., 2019; Jiang et al., 2020; Roberts et al.. 2020). A deeper theoretical understanding of how such parametric memories are learned and accessed is increasingly important: can guide scaling choices (e.g., trading off memory capacity against compute budgets, Carlini et al., 2022; Allen-Zhu & Li, 2024) and clarify failure modes (e.g., hallucination, Zucchet et al., 2025; Huang et al., 2025). Motivated by empirical results documenting the prevalence of parametric factual recall and its scaling with model size (Allen-Zhu & Li, 2024; Morris et al., 2025), recent theoretical works have begun to analyze the capacity and learning dynamics of transformers on controlled factual-recall tasks (Cabannes et al., 2024a; Nichani et al., 2025).
Many theoretical studies of transformer optimization work in oopulation-yynamin settings and adopt simplifying assumptions, such as treating token embeddings orthogonal one-hot vec- tors (see, e.g.. Tian et al. 2023b; Chen et al. 2024; Ghosal et al. 2024). While these choices do not always reflect practical applications, they make the mathematics, particularly gradient calculations, more tractable, and population analyses of this kind do not characterize the statistical or compu- tational complexity of gradient-based learning. In factual-recall setups, strictly orthogonal embed- dings are known not to be capacity-optimas whereas random or non con-orthogonal embeddings (i.c., superposition) enable near-optimal factual storage (Nichani al., 2025). At the same time, remov- ing the orthogonality assumption introduces token interference that leads to intricate optimization
Work done while interning the Flatiron Institute.
Code available at https: // //github. com/ wwwm///aiiineettaaaaleeeraoiiiie to- torcellll secall-experiments.

---

Published as a conference paper at ICLR 2026
behavior (e.g., oscillatory trajectories Cabannes et al., 2024b), and in practice, superposition-basee memory-efficient solutions can also be more difficult to train (Elhage et al., 2022), which highlight a fundamental trade-off between optimization and statistical efficiency and optimal storage capacity.
Motivated by the above gaps, we aim to address the following question.
Can we characterize the optimization and sample complexity of a transformer with non-orthgonal embeddings trained by gradient descent in the learning of factual recall task?
1.1 OUR CONTRIBUTIONS
In this paper, we analyze gradient-based learning of a single-layer transformer with an atten- tion+MLP block and random embeddings on a synthetic task inspired by Nichani et al. (2025); the model must retrieve an informative token from context containing many noisy tokens via attention, then map it to the correct label via factual recall. To mitigate the complex optimization dynamics arising from non-orthogonal embeddings, we follow Bietti al. (2023); Oymak et al. (2023) and consider a simplified training regime involving only a few gradient steps with finite samples on the attention and value matrices. This perspective effectively zooms in on the "early phase" dynamics of gradient descent, a common focus in the feature-learning literature (Ba et al., 2022; Damian et al., 2022: Dandi al., 2023; Vural Erdogdu, 2024; Wang et al., 2025).
Our analysis provides a fine- tine-griied characterization of how vocabulary size V, sample size N, em- bedding dimension d, sequence length L. and MLP width m interact to permit successful gradient- based learning of the recall mechanism. Our main result states that
The success of learning depends on (V, N, d, L, m) in a multiplicatiie manner: learning becomes easier as (N, d, m) increase, which reflects the benefits of more data, higher-dimensioaa. (and thus more orthogonal) embeddings, and larger MLP width; whereas learning becomes harder as (V, L) increase: that is, the task becomes more difficult with a larger vocabulary or longer sequences. This multiplicative relation is visualized in Figure 1a, where we examine how the parameter size m x depends on the vocabulary size V for different sequence lengths L. The full phase diagram corresponding to this relation, which formalizes Figure la, is shown in Figure lb.
Consequently, while optimal capacity and sample complexity can be achieved jointly for short sequences, successful learning long sequences requires either larger embedding dimension (thus sacrificing capacity) or larger sample sizes (worsening statistical complexity).
The multiplicative rate above formalizes the "tradeoff" intuition that smaller embeddinn dimension d which increases superposition and thereby improves storage capacity simultaneously yields a harder learning problem, as reflected in the required sample size. We complement this with a statistical lower bound showing that the trade-off is inherent for any estimator that accesses only gradient information from the initialized transformer. Finally, although our theory is derived for specific three-step training algorithm, empirically observe qualitatively similar multiplicative scaling when the transformer is optimized by gradient descent to low empirical risk.
1.2 RELATED WORK
Learning dynamies of transformers. A growing line of work analyzes how transformers acquire specific behaviors from gradient-based training. Much of this literature imposes population-level assumptions and orthogonal/one-hot embeddings to make gradients tractable, often on discrete syn- thetic tasks (Li et al., 2023; Bietti et al., 2023; Tian et al., 2023a; Nichani et al., 2024; Chen et al., 2024; Ghosal et al., 2024: Chen et al., 2025; Wang et al., 2025). Several works study few-step train- ing regimes lens on the "early phase" of feature learning (Bietti et al., 2023; Wang et al., 2025). Beyond discrete settings, related analyses investigate attention learning for continuous inputs and sparse-signal retrieval (Oymak et al., 2023; Marion et al., 2025; Duranthon et al., 2026). A comple- mentary thread focuses on the emergence of in-context learning and induction mechanisms: single- and two-layer attention trained on linear-regression Markov data provably implements gradient- descent-likk updates and generalized induction heads (Von Oswald et al., 2023; Zhang et al., 2024; Chen et al., 2024; Nichani et al., 2024). These results typically rely on simplified settings and do not address storage capacity. In contrast, our work analyzes finite-sample training with non-orthogonal embeddings in an attention+ architecture with a particular focus on factual recall.
2

---

Published as a conference paper at ICLR 2026
width (m)
2
ylus
Mean Bias md &
Gradient Noise mn mgt 47'
Vis
x
LL
parameter size (md)
yazs
200
2.132.133
1.40
MLP Noise
md &. V yimee
+
Stoge
2
Stoge
No Learting
(nnd 45,
20
2/1
2
2
vocah size (V)
vocab size (V)
(a)
(b)
Figure I: (a) Empirical scaling of the parameter size required for GD-trained one-layer transformer to leam factual recall, where use m x of (sce Section 4.2 for details). For small L, the trained model achieves the optimal capacity V x md (purple line). As the sequence length L increases, the scaling changes, suggesting a multiplicative rate (blue and red lines). (b) Phase diagram for theoretical scaling of the parameter count given in Corollary 2. Each region corresponds to a regime where particular noise term in Theorem is dominant. The parameter-size condition (mad) in each region is given in Corollary 2.
Associative memories and storage capacity. Classical associative memories (Hopfield-type models) study recall of vector patterns and established foundational capacity results (Hopfield, 1982; Amit et al., 1985; McEliece et al., 1988; Krotov & Hopfield, 2016; Demircigil et al., 2017; Ram- sauer et al., 2020; Schlag et al., 2021). Recent works adapt associative-menne formers, modeling inner weights as superpositions of outer products and deriving sealing laws and optimization behaviors (Bietti et al., 2023; Cabannes et al., 2024a;b). In factual recall specifically, random (othoorthogonal) embeddings enable near patarr-pramtter-coe
viewpoints to trans-
storage, whereas strictly orthog- onal embeddings are not capacity-optima (Nichani et al., 2025). Various empirical works have stud- ied the mechanisms and scaling behaviors of LLMs in factual association tasks (Petroni et al., 2019; Jiang al., 2020; Geva al., 2020; Allen-Zhu & Li, 2024). We provide theoretical analysis of such mechanisms and quantify how vocabulary size, sequence length, embedding dimension, and MI.P width jointly govern learning efficiency. Our work operates in setting similar to (Nichani et al., 2025) but allows finite samples and explicitly considers gradient descent dynamics. Our result similar the sample results in (Oymak et al., 2023), where the required sample size grows with the dimensionality and sparsity level of informative tokens, while we allow non-orthogonal embeddings and show optimal capacity as in (Nichani et al., 2025) under certain conditions.
PROBLEM SETTING
Our goal is to understand the capacity of transformers trained on finite data with non-orthogonal embeddings, in setting where the relevant information is hidden in potentially large sequence of non-informative noisy tokens. The attention operation should then identify the relevant token, while the subsequent linear or MLP block can then recall the correct label via associative mem- ory mechanism. This is similar to the factual recall task studied by Nichani et al. (2025), with simplifications that make the analysis more tractable, as detailed below.
Notation. a denotes the softmax function. ll := (1.....1) 1)T E RYY is the V-dimensional all-ones vector; e; is the one-hot vector with 1 in the i-th position (dimension understood from context). We use Z (resp. S) to mean ">" (resp. "<") up to polylogarithmic factors in V: fv Z gv < fv 2 poly (log V V)gvv and fv Z gv
fv I poly V)gv. for some fixed polynomial. Lastly, Il-lla denotes the Euclidean norm for vectors and the operator (spectral) norm for matrices.
Problem setup. Let the input/output tokens take values from finite alphabet [V] = {1,
For notational convenience, we represent the alphabet by the one-hot vocabulary V = {e) ev
Each example in the data consists of length-L L input sequence X [x) [xI...., Φt] E v and label
V generated as follows:
Input tokens are sampled independently and uniformly: (ai,...., xx) Unif()44..
Informative position is random index E ~~ Unif((Z.)) independent of X.
3

---

Published as a conference paper at ICLR 2026
Ground-truth function is permutation matrix II, [0,1) Labels are generated as the permuted informative token. = I1, while the remaining tokens are non-informative.
The goal is to identify the correct token position & and learn the target function (permutation) II..
Transformer architecture. We consider basic transformer block which first maps input tokens into a d-dimensional embedding space where d < V. The embedding layer is parameterized by RaxV RRdxV Rd x Rd, where
i Zout, Strig: ZEOS)
The input tokens are embedded by the columns of the matrix Zi E RR
Output tokens are associated with unembedding vectors, which are collected Zont E RRdVV
zig is trigger vector that marks the informative token.
zBOs is the special embedding vector that marks the end-of-sequence
Given the embedding parameters, define the self-attention head, parameterized by the key-query matrix WKQ E Rdxdd which operates on the embedded sequence of inputs Zi X E RdxL:
attn attn(X; WKK)) == Xo (Zuigee + X) WKQZEOS)
(2.1)
The trigger embedding Strig is used to "mark" the informative token with a special direction, mim- icking the behavior of previous transformer layers that may learn to flag particular tokens by adding to its residual stream" (note that the number of trainable parameters inside softmax can be re- duced d by collapsing #KQ into vector). We consider two different learning models: an Attention-only model and a width-m, two-layer neural network model Attention-MLP. defined as:
0 ZTT Vattn(X; (aka), a Zout V
Attention only
V V. WKQ)=
(2.2)
X; WKO)), Attention-MLP
where V Vedx for the Artention-only and V VR Wi. Rmmx for the Attention model. Note that compared with Attention-only model, the Artention- model contains an additional set of trainable parameters and nonlinear activation function & before the value matrix. Similar to the constructions Nichani al. (2025), the MLP allows using smaller embedding dimension d while keeping the capacity large by increasing
For the AttentionMMLA we keep Win fixed its random initialization. The trainable parameters for both of our models are (V, WKQ). We use cross-entropy loss to train our model:
C((V, WKQ), (X,p)) = - -Y pt log p.
Training algorithm. Following Oymak al. (2023), we consider a 3-step gradient- based algorithm with dataset p:)),N with sample size of N. We initialize our parameters (() 0.
(0 KQ = O and use the learning rates 7 > O:
v = (0) N En VVC((V(0) w (X.,
(2.3)
(1(1) KQ = (0)0) - NV EN VKke C((V(), w w (Xi,p.))
(2.4)
KQ
((2 ( 7 EN (VC((V(),
KQ): (X..
(2.5)
Network prediction and storage. Given our model and training method, we use argmax decoding inference and define the test accuracy as Accuracy = IP (x, .p) [p- x)) where pred(X) = arg max Py (X; (2) w jElV]
where p X:; (2). W (()) is the network output defined in (2.2). In what follows, we characterize conditions under which the model stores the informative tokens asymptotically, i.e., Accuracy + as V - 00, in terms of the relevant parameters (V, N, d, L, m)).
"The "trigger" terminology borrowed from (Bietti al., 2023), where special previous token "triggers" a retrieval operation in the context of induction heads. Our setup resembles learning only the "induction head" layer assuming the first "previous token head" layer is already in place. The triggers often appear to be single directions in interpretability literature, e.g. the "X in opposite of X" feature in Kamath et al. (2025).
4

---

Published as a conference paper at ICLR 2026
3 MAIN RESULTS
We first present our general theorem on learnability via gradient descent, and then specialize into different regimes to derive more interpretable scaling behaviors in Section 4. We provide proof sketch in Section 5, and defer the full proof to Appendix C.
3.1
TECHNICAL ASSUMPTIONS
We first state generic assumptions that apply to both the Attention-only and Attention-M.P models. Assumption 1.
Parameter range: Let I, = V for E (0, SVV log V) - and V 9((1)
Learning rate: We use sufficiently small learning rate foror the initial step (2.3), and sufficiently large learning rate "I for the remaining steps (2.4)-(2.5) that satisfy Assumption 4.
Embeddings: Let Zin out E RRAVV be independent Gaussian matrices, and let Ztrig: ZBOS R independent Gaussian vectors, all with i.i.d. entries distributed as ( 1/a).
We assume (0, 1) since in many practical pretraining setups, the context length smaller than the vocabulary size, and the condition L < V simplifies several terms in the proofs. The lower bound N 2 V log V is required so that each element from the alphabet of size V is seen at least once with high probability. The learning rates follow prior analyses (Oymak et al., 2023; Nichani et al., 2024): small r ensures that the network's predictions remain close to uniform after the first step, whereas a large is needed to push the attention scores and predictions toward one-hot vectors.
In addition to the above assumptions, we require the transformer model to have sufficient capacity to reach perfect test accuracy. Such conditions characterized by Nichani et al. (2025). For the Artention-only model, we have the following condition (see Nichani et al., 2025, Theorem 3).
Assumption (Attention-only). For the Attention-only model, we require de Z V.
With a nonlinear MI.P layer, a smaller embedding dimension can suffice if the width large enough. Hence for Attention-ML.P we require the following condition.
Assumption 3 (Attention-MLPP) For the Attention-MLA model, we assume that
Polynomial activation: Φ R + I satisfies d')))'((0) (0), (0(0) +
MLP width: md. 2 V and d 2 V where k. denotes the smallest nonzero Hermite mode of 6. i.e., k. = min(k > 0 E2-NNo.nlo(() [))== + 0} where by is the tt Hermite polynomial.
Initialization: Win E R.xd are fixed with entries i.i.d. distributed N 1).
The nonlinear MLP layer allows to compensate for the embedding dimension and go beyond the d² 2 V lower bound required by the Attention-only model (Assumption 2). Note that md 2 V is necessary condition for capacity as shown in (Nichani al., 2025). The additional requirements imposed on the polynomial activation function appear to be artifacts of our three-step GD analysis, and we conjecture that they could be relaxed when considering a longer training horizon.
3.2 LEARNABILITY STATEMENT
Now we are ready to present our main theorem on the complexity of learning the factual recall task. Specifically, the transformer learns the desired mechanism when the signal term dominates the noise and bias terms as stated below.
Theorem 1. Let Assumptions I and 3 hold for Attention-MLP and and 2 hold for Attention-only. The Attention-Mil MLP model achieves Accuracy 1 oy (1) with probability 1 ov(1) (1) whenever
1
(3.1)
NVLdddAL) NVVd(dAL)
Ndym
Signal
Grudient noise
Mean bias
MLP noise
For the Attention-only model, the same holds with the last MLP noise term removed.
5

---

Published as a conference paper at ICLR 2026
Theorem 1 characterizes learnability a function of (V, d, m) and identifies the following terms that impact the gradient signal-to -oniie ratio:
Signal measures the alignment between the key-query weights W W KQ and the trigger Fuig-
2. Gradient noise is due to the concentration error in the update of w KQ'
3. Mean bias arises from the nonzero mean of token vectors {X
ML.P noise reflects the randomness the MLPP weight matrix W Attention-
We make the following observations.
Multiplicative scaling, Note that the parameters (V, N,d, I,, m) interact in a multiplicative fash- ion. For example, the noise and bias terms in (3.1) all decay with (N xd), suggesting that increas- ing the embedding dimensions d can lower the statistical complexity of learning the correct recall mechanism. While the full 5- 5-parameter trade-off can be opaque, Section 4 we focus on specific regimes that lead to simplification of the scaling relationship and validate the rate empirically.
Optimal storage & sample complexity. Recall that the capacity-optimal construction for the factual recall task requires md Z V parameters (or d² Z V for Atlention-only); and as discussed earlier, sample size N V log V is necessary to observe all distinct tokens. (3.1) implies that in the small- L regime, the optimized transformer achieve optimal capacity and sample complexity simultaneously. For longer sequences. however, these two conditions may not be achieved the same time, i.e., one must increase either the network width or sample size beyond optimality to learn the task this confirms the empirical observation in Figure 1.
3.3 STATISTICAL LOWER BOUND
Theorem I provides upper bound (i.e., sufficient condition) the model and sample size for learning factual recall under a 3-gradient-step optimization procedure. We complement this suffi- cient condition with lower bound indicating that the multiplicative dependence the problem parameters is partly statistical; that is, the scaling behavior will be observed any model satisfying the broader conditions stated below. Our lower bound applies to statistical methods that can query the dataset through the attention outputs at initialization, h; == attn (xt, W In particular, we consider queries of the form {h,, h J.N as the gradient with respect to the key-query matrix WKQ depends on these quantities (see (B.1)). The statement given below:
Theorem 2 (Informal), Any method that relies the noisy version of the queries (h;, h;h, Fiml fails, iee, Accuracy f 1 with finte probability, if N $ V L1,
The complete statement of Theorem is deferred to Theorem 4 in Appendix E. We observe that the lower bound does not exactly match our upper bound in Theorem 1, as Signal S Gradient Noise in (3.1) is stronger than the stated lower bound. This being said, Theorem also confirms the mul- tiplicative scaling, hence suggesting the trade-off between capacity and sample efficiency is present in a boarder class of learning algorithms. A stronger computational lower bound for transformers and gradient-based optimization is an interesting problem we leave for future work.
4 IMPLICATIONS AND EMPIRICAL VERIFICATIONS
In this section, we leverage our main theorem to obtain more concrete scalings between parameters, and present empirical evidence on the derived multiplicative rate.
4.1 ATTENTION-ONILY MODEI.
We start with the Attention-only model which gives a simpler phase diagram.
Corollary 1. For the Attention-only model, the bottleneck term in (3.1) is the Mean bias term. Therefore, Theorem 1 is equivalent the parameter size requirement di Z max 82/3 /N4/33
We make the following observations:

---

Published as a conference paper at ICLR 2026
1.41
Thm::
0.99
N Vlog
punnie ove
225
Vl.
parameted size (d.)
$.2.2.04
2.0
5- 1.40
Thhm:
2,00
popper Coe (e3)
2
2
212
vocab size (V)
(V)
(a) N x V
(b) N x V1.5
(c) Attention only, I. x V
Figure 2: Empiriaal scaling of parameter size via three-step GD for the Attention-only model. In and (b), top-left top-right use x vo.5; bottom-left the parameter counts from (a) and (b) for the I. x V case under two sample-size regimes, N = V V and N x V1.5 Line fitting: We identify in the heatmaps the smallest embedding dimension that achieves accuracies (0.1, 0.12 0.15) and perform a least squares fit. The slopes of the fitted lines and their theoretical heatmaps. Differences in transparency (c) are
bottom-right use x
(c), we compare
counterparts are reported
overlapping points.
The condition in Corollary is the maximum of two terms, where df Z V is due to the capacity requirement in Assumption 2, whereas the second term ensures Signal 2 Mean bias and implies a multiplicative scaling between the sample size N and embedding dimension d (i.c., increasing one of the parameters can compensate for the other).
Note that the Mean bias term arises from nonzero token mean, which potentially be allevi- ated by centering the tokens, for instance through an appropriate normalization layer. Exploring the effect of applying normalization in this model is interesting direction for future work.
Empirical Findings. We run the three-step gradient descent algorithm on an Attention-only model over varying V and d, and report the accuracies in the heatmaps (Figure 2). The plots are in log-log scale; therefore, the slopes give the exponent s in d x V*. As shown in the top row of Figures 2a-2b, the slope for relatively small (where d² x V. By contrast, when the context window is larger (L x V), the requirement becomes d x V, which is also reflected in the experimental results, as observed in the bottom panel Figure 2a.
V0.5) matches the optimal capacity condition
In Figure 2b we run experiments with increasing sample size to observe the multiplicative trade-off. As seen the bottom figure of Figure 2b, increasing the sample size from V log V V1.5 reduces the exponent of the parameter size from 2 to 1.42 (the theoretical value is 1.32). Finally, the learnability thresholds for x V Figures and 2b are plotted together in Figure 2c, to illustrate that increasing the sample size can compensate for the number of parameters in the network.
4.2 ArTENTION-MLE MODEL
For the attention- MLP model, the nonlinear MLP layer introduces additional phases as stated below. Corollary 2. For the Attention-MLP model, Theorem translates to md Z V and
MLP noise,
m = o(d2L, L) and m = o(dV)
Signal
Gradient noise,
V ス dL and m Z d
Mean bias.
V = o(dL) and 2 dV,
where
Signal MLP noise is equivalent to md 2 V mmL
Signal
Gradient noise equivalent to md 2 V m
VN
Mean bias is equivalent to md Z m 5
Signal
N

---

Published as a conference paper at ICLR 2026
Thm:: 1.00
Thm:: 1.80
pasueteriee (ea))
1.05
Fit: 1.12
2
N Vlog
N Vis
parameter size (md)
5= 1.49
2
(V)
+
1 0.90
Thm:: 1.50
Thm:: 1.00
Fit: 1.48
Fit: 0.98
(Maaeeie M0z (nud)
200
213
2
25
2
25
vocab size (V)
(V)
wach (9)
(a) N x V log
(b) x V1.
(c) L x V0.6 and m x d
Figure 3: Empirical scaling parameter size for AftentionMLPP model under two sample size regimes,
= log and N x V2., In (a) and (b), top-row uses L x 170.25; bottom-row uses L x Vu.s, (c), we compare the parameter counts from (a) and (b) for the L x V0. case under both sample-size regimes.
Than:
Fit: 1.05
Fit: 1.14
df
1.55
o
parameter size (md)
227
2
Thm: 1,50
Thm:: 1.60
Fil: 1.55
= 1.49
X
221
213
2
2
2
21 20 woch sie (V)
29
vocab size (V)
VOCED 10X (V)
(a) m d?
(b) m 8
(c) = V0.5 and = V log
Figure 4: Empirical scaling of embedding parameter size for the Astention-ML.P model under width regimes, m df and x d (a) and (b), top-row uses v0.25, bottom-row uses x 10.5, (c), we compare the parameter counts from (a) and (b) for the L x 1o.5 case under both width regimes.
The phase diagram for the Attention-MLP model is visualized in Figure 1b. Compared to the Attention-only case, it exhibits additional regimes because a smaller embedding dimension; this can lead to different dominant terms the gradient. In par- ticular, since large I, and d entail a larger magnitude of the Mean bias (as in the Aflention-only setting), increasing the MLP width m and thereby reducing the required embedding dimension may suppress this bias term
trade off and d and thus use
Empirical Findings. We run the 3-step gradient descent algorithm on an Attention- network over varying V and d and plot the accuracies in Figures 3 and 4. We take the nonlinearity to be the mixture two Hermite polynomials Φ = 0.7hg 0.333,, satisfying the conditions in Assumption 3. We run experiments with width width further, we restrict ourselves to the MLP noise-dominan region.
d2 and x d Due to the prohibitive cost of increasing the
In Figure 1a, plot the scaling of the number of parameters (md) function of vocabulary size V for different sequence-length regimes in L. We observe that L x /0.25 requires md x V, which the optimal capacity, predicted by theory. As T increases, need parameters to
8

---

Published as a conference paper at ICLR 2026
achieve the same capacity, observed in the L v0.5 and L x V0.75 cases in Figure 1, where the slopes agree with our theoretical predictions as well (see also Figures 3a and 3b).
We further test the effect of sample size in Figure 3, where use L V and x d. We plot both heat maps in Figures 3a and 3b, and the fitted lines for L 10.5 together in Figure 3c. Note that we state the plot in terms of parameter count, which scales as md x d8 the slopes from the heat map are scaled accordingly. We observe that increasing N from N x Vlog V to N x V1.5 reduces the network size the optimal level, aligning with our theoretical prediction. The heatmap versions of these experiments are shown in Figures 3a and 3b.
Lastly, we probe the width scaling by keeping the sample size N V Vlog V and L V0.5 fixed in Figure 4. Here, we observe that we can reduce the embedding- dimension requirement by increasing m Theorem although it increases the total parameter count overall, seen in Figures 4b and 4c, since width must grow proportionally more than d to achieve the same accuracy. This is also consistent with our theoretical prediction.
mabbeddin-ddimension
4.3 BEYOND EARLY PHASE OF TRAINING
While our theoretical analysis focuses on particular three-gradient-step training procedure, we empirically observe qualitatively similar multiplicative scalings when the transformer model is optimized beyond the "early phase". Specifically, we train our Attention-onf model using Adam (Kingma & Ba, 2015) with mini-bater gradients. In the experiments, ization in both the attention and output layers and set the learning rate to 0.005. We use a batch size of N/22 (except in the last experiment, where use [N/16]), and run the training for 16 epochs.
layer normal-
We highlight the following observations:
Capacity improvement with multi-pass training. In the top row of Figure 5, we plot the heatmaps for L and N x V log V. early training the slope is suboptimal; notably, by the end of Epoch it closely aligns with our theoretical prediction. Moreover, training the network additional epochs improves the capacity condition to a near-optimal level, as shown in Figures Se and 5d.
Effect of sample size. In the bottom row of Figure 5, we plot the heatmaps for T, = V and N x V1.5. We observe similar trajectory in capacity, while the overall capacities improve compared to the small-sample regime, showing the multiplicative dependence on sample size N.
Effect of sequence length. Figure 6, we plot the heatmaps for = V0.85 and N = V log V. We observe improvements in capacity over multiple epochs, while the capacity is larger than in the L V setting at every stage of training, which shows the effect of the sequence length L.
Fic 2.07
Fit: 1.72
Fie: 1.38
Fit: 1.24
stze
(a) Epoch
(b) Epoch 2
(c) Epoch 8
(d) Epoch 16
(V)
(c) Epoch!
(f) Epoch 2
(g) Epoch 8
(b) Epoch 16
Attention-unly model under two sample size regimes. Top now (a-d): N x V log V. Bottom row (e-h): N x V1. The model uses x V and is trained using Adam over 16 epochs. We observe that the slopes at the end of Epoch I match our theoretical predictions, and the capacity improves the number of epochs increases.
Figure 5: Empirical scaling the parameter size

---

Published as a conference paper at ICLR 2026
Effect of batch size. Finally, in Figure 7, repeat the experiments from this section using the same learning rate and architecture but with a smaller batch size [N/16] J. As before, we consider L V in sample-size regimes, N x V log V and N x V1.5 We observe behavior similar to the larger batch size setting, but with improved slopes in Figure 7. This suggests that smaller batch sizes may improve capacity in practice.
Overall, these experiments suggest that the multiplicative relation between the hyperparameters mains throughout training. However, the exponents depend on the iteration number and batch size. Understanding how capacity evolves during training remains an interesting open question.
PROOF OVERVIEW
In this section, we outline the main ideas behind the proof of Theorem 1. For ease of exposition, consider the Attention-only model under population dynamics with orthogonal embeddings.
The main idea of the proof is the following: We observe that the recall task is achieved with near- perfect accuracy if and only if the attention mechanism can distinguish informative tokens. Once this occurs, the remaining task reduces to learning linearly separable problem, which is well under- stood. Therefore, the proof focuses on the attention scores in (2.1) and characterizes the conditions under which the mechanism selects the informative tokens.
The pre-softmax scores evaluated on fresh sequence Xins with the key-query matrix given by the first gradient-descent iterate W KQ given by
scores == (ztrigez + Xin
W KQ KOO EOOS.
(5.1)
By substituting the exact expression for Wky KQ into (5.1), we analyze scores. For intuition, we present the simplified expression below (see ().1) for the full expression):
scores 7 Zerig 12e
x
Te,eZ Zn (V(1) Zout (p: v. [v))
(5.2)
NL
Informative
7 in
Zi X; Z (V(O) Zone (pi (Iv)).
(5.3)
NL
Non-informative
Here ((1) denotes the first iterate of the value matrix defined in (2.3). The informative term in (5.2) captures the alignment between the trigger vector in the fresh input and the one encoded in the learned weights w and therefore contains the position information of the informative token. By contrast, the non-informative term in (5.3) reflects correlations between tokens and does not contain information about the token position. Thus, the proof reduces to characterizing the conditions under which the informative term in (5.2) dominates the con-informative term. We present the proof sketch for the finite-sample setting with non-orthogonal embeddings, as well as how each noise term in Theorem arises, Appendix B.1.
CONCLUSION
In this paper, we derived precise asymptotic rates for learning with gradient descent on transformers trained on simple recall task with random embeddings and finite samples. Our analysis and exper- iments reveal a rich picture of multiplicative scalings between various problem parameters, showing that parameter count is not the only important factor controlling capacity when learning with finite samples on large noisy sequences. Our results suggest that finer control the data distribution may be necessary for learning efficiently at optimal capacity, for instance by ensuring sequenees are less noisy and more informative, hoping that the discovered mechanisms are robust to harder settings. This is reminiscent of the procedures used for long context extension in LLMs, where most of training happens on shorter sequences, but the final models are extended to work with very long sequences, and empirically do well on retrieval tasks such as "needle-in-a-haystee
(e.g.. Gemini Team, 2024). which resembles our theoretical setup. Analyzing similar scalings in more structured data distributions and architectures is thus an interesting avenue for future work.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT
The work of M. Soltanolkotabi was partially supported by AWS credits through an Amazon Faculty Research Award, a NAIRR Pilot Award, and generous funding by Coefficient Giving, and the USC- Capital One Center for Responsible Al and Decision Making in Finance (CREDIF) Fellowship. M. Soltanolkotabi is also supported by the Packard Fellowship in Science and Engineering, Sloan Research Fellowship in Mathematics, NSF CAREER Award #1 846369, DARPA FasNNICS program, NSF CIF Awards #18 13873877 and #2008443, and NIH Award DP2LM014564-00.
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws. arXiv preprint arXiv:2404.05405, 2024.
Daniel Amit, Hanoch Gutfreund, and Haim Sompolinsky. Storing infinite numbers of patterns in spin-glass model of neural networks. Physical Review Letters, 55(14):1530, 1985.
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 (NeurIPS), 2022.
Gérard Ben Arous, Murat A. Erdogdu, Nuri Ment Vural, and Denny Wu. Learning quadratic neural networks in high dimensions: Sgd dynamics and scaling laws, 2025. URL https: //axxiv. org/abs////// 2508. 03688.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems (NeurIPS), 2023.
Vivien Cabannes, Elvis Dohmatob, and Alberto Bictti. Scaling laws for associative memories. In International Conference Learning Representations (ICLR), 2024a.
Vivien Cabannes, Bertin Simsek, and Alberto Bietti. Learning associative memories with gradient descent. In Proceedings of the International Conference Machine Learning (ICML), 2024b.
Nicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee, Florian Tramer, and Chiyuan Zhang. Quantifying memorization across neural language models. In The Eleventh International Conference on Learning Representations, 2022.
Lei Chen, Joan Bruna, and Alberto Bietti. Distributional associations in-context reasoning: A study of feed-forward and attention layers. In Proceedings of the International Conference on Learning Representations (ICLR), 2025.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Prov- able training dynamics and feature learning in transformers. Advances Neural Information Processing Systems, 37:66479 6656777, 2024.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pp. 5413-5452. PMLR, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovie Stephan. Learning two- layer neural networks, one (giant) step at time. arXiv preprint arXiv:2305.1822 2023.
Mete Demircigil, Judith Heusel, Matthias Löwe, Sven Upgang, and Franck Vermet. On model of associative memory with huge storage capacity. Journal of Statistical Physics, 168:288-201 299, 2017.
O. Duranthon, Pierre Marion, Claire Boyer, Bruno Loureiro, and Lenka Zdeborová. Statisti- cal advantage softmax attention: Insights from single-location regression. arXiv preprint arXiv:2509.21983 /9366 2026.
11

---

Published as a conference paper at ICLR 2026
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Satfield-Doddss Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposi- tion. arXiv preprint wwXXi:2209.106333 2022.
Gemini Team. Gemini 1.5: Unlocking multimodal understanding across millions of tokens con- text. arXiv preprini arXin:2403.055500 2024.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. arXiv preprint 1rXiv:2012.149133 2020.
Gaurav Ghosal, Tatsumori Hashimoto, and Aditi Raghunathan. Understanding finetuning for factual knowledge extraction. arXiv preprint arXiv:2406. 144.45550.,478,, 2024.
John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 7954)2555,-5588 1982.
Yixiao Huang, Hanlin Zhu, Tianyu Guo, Jiantao Jiao, Somaych Sojoudi, Michael I Jordan, Stuart Russell, and Song Mei. Generalization or hallucination? understanding out-of-contex reasoning in transformers. arXiv preprint arXiv:2506. 10887, 2025.
Zhengbao Jiang, Frank F Xu, Jun Araki, and Graham Neubig. How models know? Transactions of the Association for Computational Linguistics, 8:423 2020.
know what language
Harish Kamath, Emmanuel Ameisen, Isaac Kauvar, Rodrigo Luger, Wes Gurnee, Adam Pearce, Sam Zimmerman, Joshua Batson, Thomas Conerly, Chris Olah, and Jack Lindsey. Tracing attention computation through feature interactions. Transformer Circuits Thread, 2025. URL https: / (ttransformrrcirrnn
pubs 2025 (tt22///ttentioonk.... soinnn-kkiidxxx..oo index.nexx. html.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR, 2015.
Dmitry Krotov and John J Hopfield. Dense associative memory for pattern recognition. Advances in neural information processing systems, 29, 2016.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In Proceedings of the International Conference Machine Learning (ICML), 2023.
Pierre Marion, Raphael Berthier, Gerard Biau, and Claire Boyer. Attention layers provably solve single-location regression. In Proceedings of the International Conference on Learning Repre- sentations (ICLR), 2025.
Robert McEliece, Edward C Posner, Eugene R Rodemich, and Santosh S Venkatesh. The capacity of the hopfield associative memory. IEEE transactions on Information Theory, 33(4):461-48, 1988.
John X Morris, Chawin Sitawarin, Chuan Guo, Narine Kokhlikyan, G Edward Suh, Alexander M Rush, Kamalika Chaudhuri, and Saeed Mahloujifar. How much do language models memorize? arXiv preprint arXiv:2505.24832. 2025.
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In Proceedings of the International Conference on Machine Learning (ICML), 2024.
Eshaan Nichani, Jason D Lee, and Alberto Bietti. Understanding factual recall in transformers via associative memories. In Proceedings of the International Conference on Learning Representa- tions (ICLR), 2025.
Samet Oymak. Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention prompt-tuning International Conference on Machine Learning, 2023.
Yang Peng, Yuchen Xin, and Zhihua Zhang. Matrix rosenthal and concentration inequalities for markov chains with applications in statistical learning. arXiv preprint arXivv2508.0437 2025.
12

---

Published as a conference paper at ICLR 2026
Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models as knowledge bases? arXiv preprint 1w..:::999.00/000 2019.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint (xXiv:2008.02217, 2020.
Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of a language model? arXiv preprint arXiv:2002.08810, 2020.
Jonathan Scarlett and Volkan Cevher. An introductory guide fano's inequality with applications in statistical estimation. arXiv preprint (rXiv:1901.00:0 2019.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Proceedings of the International Conference Machine Learning (ICML), 2021.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding train- ing dynamics and token composition in 1-layer transformer. Advances Neural Information Processing Systems (NeurIPS), 2023a.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying mul- tilayer transformers via joint dynamics of mip and attention. arXiv preprint 10X0522310.00535, 2023b.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers lear in-context by gradient descent. International Conference on Machine Learning, pp. 35151-35174. PMIR, 2023.
Nuri Mert Vural and Murat A. Erdogdu. Pruning is optimal for learning sparse features in high- dimensions. In Annual Conference Computational Learning Theory, 2024.
Zixuan Wang, Eshaan Nichani, Alberto Bietti, Alex Damian, Daniel Hsu, Jason D Lee, and Denny Wu. Learning compositional functions with transformers from easy-to-hard data. In Conference on Learning Theory (COLT), 2025.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
Nicolas Zucchet, Jörg Bornschein, Stephanie Chan, Andrew Lampinen, Razvan Pascanu, and Soham De. How do language models Iearn facts? dynamics, curricula and hallucinations. arXiv preprint arXiv:2503.2 1676, 2025.
13

---

Published as a conference paper at ICLR 2026
CONTENTS
1 Introduction
1.1 Our Contributions
1.2 Related Work
2 Problem Setting
3
Main Results
3.1
Technical Assumptions
3.2
Learnability Statement
3.3
Statistical Lower Bound
Implications and Empirical Verifications
4.1
Attention-only Model
4.2
Attention-MLP Model
4.3
Beyond Early Phase of Training
5 Proof Overview
10
Conclusion
10
Additional Experiments
16
B Preliminaries for Appendix
16
B.1 Proof Overview: Empirical Dynamics with Non-Orthon Orthogonal Embeddings
18
B.1.1 Decomposition of output layer weights
18
B.1.2 Characterization noise terms
18
B.2 Preliminary Results: Characterization of Good Events
20
C Proof of Theorem
29
C.I Attention scores and their asymptotic scaling
30
D Proof of Theorem 3
30
D.I Concentration bound for S1
D.1.1 Concentration bound for 0
31
D.1.2 Concentration bound for
43
D.2 Concentration bound for
46
D.2.1
Concentration for
D.3 Concentration bound for
53
D.3.1 Concentration bound for
E Lower Bound
57
14

---

Published as a conference paper at ICLR 2026
F
Auxiliary Statements
60
F.1
Gaussian matrices and related statements
60
F22 Multinomial distribution and related statements
63
Miscellaneous
71
G.1
Rosenthal-Burkholde inequality and corollaries
72
15

---

Published as a conference paper at ICLR 2026
A
ADDITIONAL EXPERIMENTS
Fit 1.50
Fil: 1.37
Fic 1.21
Fit 1.11
PaLLOOO slae
(a) Epoch
(b) Epoch
(c) Epoch 8
(d) Epoch 16
Figure 6: Empirical scaling of the parameter size for the Attention-on model with N x V log V and L x V The model is trained using Adam over 16 epochs. The slopes smaller than in Figure which consistent with the shorter sequence length.
Fit
Fit: 1.49
1.16
Fit: 1.21
(a) Epoch 1
(b) Epoch 2
(c) Epoch 8
(d) Epoch 16
20
250
20
Wode sice (V)
Wobth size
woblh size (V)
woutb www.bize00
(c) Epoch!
(f) Epoch 2
2g) Epach 8
(h) Epoch
Figure 7: Empirical scaling of the parameter size for the Attention-only model using a batch size of [N/16]. Top row (a-d): x V log V. Bottom row (e-h): x VL.s
V), trained with Adam
B PRELIMINARIES FOR APPENDIX
Additional Notation. For vector æ E RVV we use diag(z) VVV denotes the diagonal matrix which has the same diagonal entries with æ, while for a matrix A, diag(A) E I denotes the column vector whose elements coincide with the diagonal entries of A. For a random variable w, Em denotes taking expectation with respect w and keeping the remaining independent terms fixed. Similarly, we use E[ w for conditional expectation, conditioned on w. We use E as an indicator function, which takes values {0, 1} depending the event holds or not. We use C to denote any constant in the upper-bound, which might depend on ø. We use poly p.g (N. d, V, L) denotes a polynomial function of (N, d, V. L) whose degree depends on (p, q) polynomially. For vectors w, i and scaler variable 17
t á + O(π) to denote ||2 w//oc O(n).
Since we do not use positional encoding in the model, without loss of generality we can fix the infor- mative index We will denote the rows of W... with ( {wx) ] m For compact representation the attention with the trigger, define
1. We define the sequence of non-informative tokens N, [0.,2.... _R,2
Zin 1in Zuigg] and X;
T N
] 0
1
RLx((+1)
16

---

Published as a conference paper at ICLR 2026
With this notation, we write the iterates in three-step GD. Let
Pr.i = p(x;; (t), W kQ) and Cooii = [x (x,2,n W KQ E
We have
N
v
-ont
Σ
P: po, ;x Win))
N
i=1
W = Zin
Σ
X, (ddia((on))oonn.an.... X; Zn Win
KQ
N
diag (o'( (Win X; ao.)) (V(") (D: (B.1)
For notational convenience, we define the noise due to finite width (which we defined equivalently (B.11))
FW in: Zin X., Xj) = Will diag (o' W X, (1) Win x,
Wi diag Φ5 ( Wi, Zin x 1n) ( Zin 1
E....
(B.2)
For the terms arising in the expected value term in (B.2), we define
Osyy Eu oz (+wT Zin X, 11. 6' (z Zin XJ x [n)] Byy=EE (0" ( Zin Tin X 1 00 ( w Zin X 12) (
Moreover, we make the following definitions to simplify the notation the following:
A
Zin
aug (x, - [
- -II))")
LN
Σ aryla, (1)) [v) Z
(B.3)
LN
A2,ir = Zin
Uij N
LN
N
(a, - + - -IrI Z
LN
(B.4)
As,5r
1
Σ
N
(2, - (vv) T
1
Σ
N
ary - tlv) Zin 1v 1,
L2V2
N
N
j=1
5=1
and
S
1
[x, - (1,) (dv) (1
[a, [z) (4v)T)
(B.5)
LN
IN
N
S2 =
Σ (NJ - +1v1 - ((v)T)
LN
==1
x
Σ
( - +1v) - -(vI
(B.6)
LN
j=1
S3 L,V2 2
Σ
x
(+v)
1
Σ
(x) 1V
(B.7)
N
j=1
17

---

Published as a conference paper at ICLR 2026
B.1 PROOF OVERVIEW: EMPIRICAL DYNAMICS WITH NON-ORTHOGONAL EMBEDDINGS
In this part, we consider (5.2)-(5.3) with empirical loss and non-orthogonal embeddings. For no- tational convenience, without loss of generality, we assume = and II Iv (accordingly, P. = x. Our goal is to show how the Signal, Gradient noise, Mcan bias, and MLP noise terms arise from the dynamies of the first gradient step.
The analysis proceeds in two steps. First, we show that the first iterate of the output-layer weights v admits a natural decomposition into mean, bias, and noise components. We then show how this decomposition gives rise the terms appearing in Theorem 1.
B.1.1 DECOMPOSITION OF OUTPUT LAYER WEIGHTS
Both the informative and non-informative terms depend on V We show that it can be decom- posed as
N
v = Zout
[1,11 - Tm
(B.8)
NL 1=1
Zou
1
(Iv - -Iv--iviV) b)vlt)+
1 VN
E( a. [I
3
3)Zm
Z
(B.9)
VL
VLV
Mean
Noise
Bias
where the noise term is given by
N Σ (a,11 1=1
1 1 - ((((v=-ivi))) bv11T)).
E =
LN
Here, the bias term arises from aggregating tokens at initialization; specifically, the aggregate- averages + I in (B.8) concentrate around their mean + L grows, and this effect appears as the bias term. The noise term captures finite-sample fluctuations of tokens around this mean. We explicitly factor out the typical operator norm 1/VVLN in (B.9) so that the remaining matrix =\t has constant norm on average. i.c., E||
B.1.2 CHARACTERIZATION OF NOISE TERMS
Signal. Using the mean component in (B.9), the informative term in (5.2) can be written
Informative 7| Ztrig 12e1
a
Z (V Zont(p. - 11pp:-+1pp (v))
NT
1=1
7 71|ztri|
N Σ
1 Zin (Iv - A 1v Zonr (2ui - lv)
1
7 T in
VT2
1=1
Signal
VZ2
negligible terms.
The first term is due to the mean component; the negligible terms are due to the bias and noise in (B.9). Standard concentration arguments for Gaussian matrices can be used to show that the leading term scales as v12. which gives us the Signal term in (3.1). The detailed derivations are provided in Section D.I.1.
Gradient Noise and Mean Bias. For case of presentation, we focus on the large- L regime where we can use the following approximation due to concentration
1
T, Zin X;X,
Z
a la-
(B.
18

---

Published as a conference paper at ICLR 2026
Let Zin denote arbitrary row of Xin- Using (B.10), approximate the non- non-informative with
N lv)
Non-informative 2
in in Zi EZ out
Zot
LNV
N
Gradient
N
Lin 2 Zin
Zout
Σ
ll
negligible terms.
Vd
N
Mean bias
The first term arises from the noise component E and determines the scaling of the Gradient noise term. The second term comes from the bias component and yields the Mean bias term in (3.1). We hide the contributions from the mean component the negligible terms, since they are smaller in magnitude. The fluctuations of each term be bounded stated in Theorem using standard concentration arguments. The detailed derivations are provided in Section D.1.2.
MLP noise. In this part, we consider the Altention-ML.P model. The scores in (5.1) can be defined the same way for this case as well. Let {w> Im denote the rows of Win. ins where Wy ~ [0,, Ia). For illustration, we work in the large-L regime and adopt the approximation in (B.10).
define the MLP-noise term as the deviation of the scores from their expectation with respect to the randomness in Win
MLP-noise scores E..,, (scores.
Under the large-I, assumption in (B.10), the scores admit the approximation (see (C.1) for the full form)
MLP-noise ~ Iin Zi N2d
Σ
FW(W i Zi. X1 X,) - Zout (I),1 - lv).
i,j=1
Here FW(W in Zin, X,, X2) denotes the noise induced by the finite width of Win defined
m
FW(W Zm, X1, X,) = m
wro' W Zin X,L1 @ fww Zin X,11)
E [uuo (oo) Zin X:1 1 ( www Zin ; 11, (B.11)
For large L, standard concentration arguments imply that 1 Zin X, 1 12111122 L-1 - 0. Hence
('ww Zi X, w Zmx
(0)'(0)
bonzero constant
where Assumption ensures (00) (0) * 0. Since E[wa]] = 0, replacing the dedependent factors by this constant yields, have
FW(W.. mn Zm. x, x,) ((0)
m
m
Substituting this into the expression above gives
(0)
N
MLP noise 2 Am
dm
ipptt
Z we
Zone N
(x 11v
dym
Here, the terms can be bounded in the displayed equation using standard concentration arguments, which yield the scaling the MLP noise term in (3.1). The detailed derivations are provided in Section D.3.
19

---

Published as a conference paper at ICLR 2026
B.2 PRELIMINARY RESULTS: CHARACTERIZATION OF GOOD EVENTS
We start with characterizing "good events" which we will use in the proof of Theorem 1.
Lemma 1, We consider E N. and V > N > V > and x Va, and d x V for some c1 E (0. For the following we define, maj (1 R., + +. We define the following events:
(EI) Let 23 == Z Ek and 2 = (2) + (l1g))))) We have
(E1.1) +11 2 12 S and max A |Z ll2 I and max za)1 S
log V vi
(E1.2) 1 Z 1|2 S 2 and 1 vV Z Zm /lo 1
log V
va
(E1.3)
1,4 1.
S log V V and 2 Zin 2 Zm lv
Cx log ((t)* and
2/z ding Z 2m)
- Ck log V /
(E1.4) For all i [N]. L Z X |n e x, 11 log
log V 1x
Vd
(E1.5) For all i E [N], 1 -1 Zi X 11| - L + Ck log V X; 12ll2 n/
(E1.6) For all i E [N]. [ Zn Z x 1 $ (c
1, Ck log V |x, 1ull2
vd
(E1.7) For all i.j E [N]. | in Zin X lr. mayl < 1 x, 1r. misl +
CK Ixla,x 11 V.
b
Vd
(E1.8) For all i [N]. || I N. NZ
Cx N. 1 + log v I,,, 1 122).
vd
(E2) We have
(E2.1) For all i,j E [N] x - misl S CK log² V
VVAL'
(E2.2) For all i E [N], | 16 S log and |X 1 2 +
(E2.3) || E 2ıl|2 - - + s Cx
log² N
NVV and || EN 0; 1 1vllz - :/ S Cx
log? N
NVV
and N Li=1 x; +lvll= < (e+l)
(E2.4) Σij= |1, m0; - 11 S
4N²
V and N
(1 [1., - +) s v and for any h E [V].
2
Σm 1 =, (1 - - 1
CN²
V
(E2.5) ||S| 112 S LAV and [ur(S)) - (1 } 1 (1's + ++) S
Cx log INNVV
(E2.6) ||S21|2 S
Cx Colog and (tr(S2) - (1 - } L=N S Ox log NLV
NVVV
(E2.7) -Cx -Oklog" V
v212 11 = S3 1 1 N 1v 1v 1 Ox log" V
1
1 v=r 1
NVV
NVV
(E2.8)
NT. EN
(NJ - +1
+ 111 = + + Cn log"
I1=
VNV
For any K O, there exists a universal constant Ck > depending only on K such that
1
P( 2 VK and P[(E2) 2 VK
Proof. For (E1):
By Proposition 4, we have 1 Zi. -T dall s 2 V
VVd and by Proposition 6, we have maxegv #x ll2 ||zuiglla $ 2 with probability at least - CVVeep((cllg' V).
By Proposition 6, VV Zim 1y S 2 and + 2 Zin 1 v 1lo s
2 V with probability least vd
1- CC-eexpp-cllo'' V).
20

---

Published as a conference paper at ICLR 2026
By Propositions and 7, we have Jt 4
2m
|
2 2 V vd and
Z Zia "n Zin 1v
Ok log v with probability at least
V). Moreover
1
1 Σ 1zll(zz zx)
1
1
24 diag Zm)
| 1z:lb(zz Zrrg))
V 26
ex(-OxOOk
where we used previous items to bound the last term. For i + k. by using Lemma 3, we have for p S
d(d 2
(d 6p - 2) S d p
B 11 [2,, 2x) < d 1T (2p)) s d p2 p2
d3p
Therefore,
B|||2112 4p 1izi, zx) 200 s 1/2 VT.
By Proposition 15, we have for 2 A p /
1 T Zi 2 Zm.
20
< Cd
P
V 12/2
V VK.
By using p = log V. we have the bound the statement with probability
By Proposition 6 with probability least 1 VK
1 Zm X
11
- l Z Zi X
1z|+ + 1 kiip Zia x 1zl
e X 1. Cx log V | 1r
vd
By the union bound, the item follows.
By Proposition 6 with probability at least VR
V 11ll d V = L - Cx log VI 11 d
1 | 2in x, L1l S 11x 11 + Ck log V
By Proposition 7, with probability
least 1
CN exx((cclog2)) V),
have
|z] Zia Z Zin
X, s 2y (ef Te]x,11 1 + Ck log v ( for all [N].
va
By Proposition with probability at least
Vx
L [1, Zm myj
1 1 1[
1 X 1 my + Ck
|X] 1 A11l1221 X, Inlla log V
vd
For the last item, let nig I_,, Neek. We have
N N,N,N,N... (1 11=12 Zurig 1 ))z
L
+ Σ
V
ny [z}z]]
1 _.)zz..
j=1
j7k
By Proposition 12, we have
V Σ n
2n 1/4 1a)zk
1
197 2
E
2
C(p 1)8
Σ
n (2
-aa z
21

---

Published as a conference paper at ICLR 2026
C (D - 11 16-1112.
Therefore, with probability 1 VR, we have
L + log V 1 1z 122.
Nn N. S Cx (no + d
vd
For (E2):
By Proposition 9, we have the first item with probability
yx.
By Corollary 3, we have | 1 1lllao S log L with probability 1 1-
N Vx
for large enough
L. For the second part, we define fig == 11 [ xee- We observe that
X 11
Σ Plan > > 0] v (1
= -1(1 14) (1-
L 2V
+o(l/v)).
k=1
By McDiarmid inequality, we have
I. I I||x Aello - 1111 2V + +o(L/v))| > VL v| s 2exp((22 -2og²² v).
which gives the result.
Let n = Ei De. where E [[n] = N
v lv. By Proposition 9 with probability
1 - vx+ we have
1
1 || NN" n -1 ) 1 =
1 N
1 >1 (1
1 V
1
log 2 V
(1
S Ck
N
v
N V
Lastly, by Corollary 3, we have || 2 n bl I (
We have
N
NN
1
2 4
(1 - ) + 22
1 v.
V =
12
= 2
-
i,=1
2
N
1
222
1
= (1
1x==,
+
V (1
V
i,f=1
= (1 2
1
| V (1
2N2
1 +
c
By the previous item, the statement follows. Moreover,
N
(1 =ex= - + (11 - +)
1==1
= (1
V
Σ
N
(Le,=er- - =))²
2 2N
+ V
(1
N (La - +)
i=1
1=1
= =N 2
2
(exi
Eki
1 N
2N2
(1
1
Ek
1 N
Σ x
1vv)
N
V
N
i=1
C
< V2
The events for S1, S2 and Sy follows Proposition 10.
(F2.8) follows the second item in Proposition 9.
22

---

Published as a conference paper at ICLR 2026
Proposition 1. We consider the parameter regime in Lemma 1. supk, k221 (k)) (0))())(()) The intersection of (EI) and (E2) implies the following events:
Let Φ
(RI) For all i,j [N] _ 1 r. x 2 Zm 1, -misl Cx
(oog vi
log² ) vVAL)
(R2) sup,, 0' (0)² v |3, - (0)))(()) s /misj Ck
loz va
+ Ck log² V
vVAL
(R3) Let Ae,er = - - (')) S. for * {1,2,3}. We have
supi.rENN] 1 12 Cx NdL + VaL VALPALNA)
supr.rE[N]] 2, S
CKVV
+ VNEWANA
dvNI
.N
Agwe Zin lv 1 Z such that V-T.7
We have Az,ir
(x)(
1 NL VNVABALN
1
sup i,rE[N]
As,irl
NL + ANVADAND
N
(R4) For all i,r E [N}
We have
(0))
1
1 1
/1 S (k)
1
A
A
+
(1
d
N
NdT3 Vdl2 V A LIALVD log? V
logV
+CKO(O0) L2V3/2Vd L2NVVd
We have
0'(0)
1 1 1
1 Id S dvNL NL VVIVARAAN
CRVV 1
A2
(1
d
v
L2N
logV + log
+ CK
NLVVd
NVLVd
ViI We have A3,ir )'(0)" I Zn Iv = AZr
vD Zin Iv 11 Zj such that
Cxd((0)log''
A S
NVV
CK)(
1 NL VNVABALL + (NL + VNVAIAALD
N
Proof. We have the following arguments.
By (E1.7) and (E2.1), we have (R1).
For (R2), we assume (R1) hold. Let w/ = 1, w Z X 1 1. Zin X Lzllz- We write
Ew [o' (1 Zin X,
11.1112100 / ø (1 Zi, X. 1 (23ww 60))
(o)
Ew [ ulv!
Σ
Z 1 1111 Zm Lell
of ((++1) (0)
w,v=l
(6)
+ 1 x,
X 1z (1) (0)
2 Zin
P+
Ew [ (0) ulv!
Σ
2
1 11
2 X|
1
0,0=1
+ is 14--322
23

---

Published as a conference paper at ICLR 2026
3 1 (mij + CK log V
log" V
Ck
vV W O
vd
13
where we used Taylor expansion of and Ew [w is odd for jointly Gaussian (7,, 7z) in (b), and (R1) in (c). Similarly,
Z 12] 0 in (a). 2| = 0 a +
B[0" (11 Zi
E
X; 1 w2 @ L Zi X; 1s [l22;)]] - (()){)2
kg
[w]We] 0 ((x+2) (0)0)00))
Zin x T 1 r 11:l T, Zin X T 1el1
wlu!
1 1
[x Tm in x
(20))
E Wwpwil
1 x, 12l061 1 In ,
_ 1 11l
0 (u (( (0) (0)
w!u!
0,y=1 even レ+ック2
2 L mmy + CK log V vd +
log" V
+0 1
$
CK
VV +
1
For (R3), we define
N
Å Al,ir ==
1
(aaj - o (0)))) - (1)) - )
LN
1
0((0)(zx, ( )
x
LN
j=1
(z)))(z, [ )
LN
N Σ 5=1
(ary ('()))()) - [[v)(2z- - )Lv)
x
LN
+
Σ
(os - ((0)))
- -
LN
1
Σ (ar - o(0))) b,) - )v)
x
LN
We have
(0'(0)ssup,laa sup,lou- (0(22
|| x
|S +
suploy - Φ ('(0)21|8))) ll
LN
if1
(d) (O)
1
1
NV VAL
AL
where we used (R2) and (E2.5) (d). (E1.1), we have
| = |Z 2 S CD)
1
NdI.3
V V VALAAAA
Moreover, we define
2g,ir
Σ (ay ('0))()) - +Lv1 100, - -
NL
24

---

Published as a conference paper at ICLR 2026
N
d'0))aa - + - =111
x
NL
Σ
き (0) bly1 -1
NL
j=1
1
(or) 1 ('0)) ( - -VIvII-))
NL
N
any (ay- (00) - -1
NT.
==1
N
(Urj - -003) )
x
NT.
j=1
We have
N
A2 12 S 210
Σ
(or
1100)/((00--1,1I11 - -b1 1|
NL
j=1
11 +
(ai - ('())) [ - +1v1I_1 2
NL j=1
- -4))) -
(a.
NL
1=1
N
- -00)) - (0) -N - =1v
NL
We observe that
Σ aai - ((())) +1))) =1v1I=1)711
NL
(au ((0))) 11v) (N. - 1 1 1
NL
N
(00))()))(
-
- - -
NL
j=1 J+1
C
SUP kaij Φ o'(0)²
N
Σ (z, (a )1v)(, -
NLVL
VL
N
1=1
siitds aij 2((0)2
N
[N] +L - -1
VI
NT.
f=1
3f1
GM
C
C 1
NLVL ALLAAAA Lvd
where we used (R2) and (E2.1) in (c), and (R2), (E2.5) and (E2.8) in (f). Then, by (E2.6), we have
C
1 1
1
1
|| |22ar112 < VNVL
WIVABANN + C2
NL
NLE "VLVABAA
25

---

Published as a conference paper at ICLR 2026
Therefore, by (EI.1)
CVV
22 = 12,
dUNT. NL; + VVIVABAIN
Lastly, we define
N
N
As,i
1
ass - (x()) - }dv) T
1
Σ
(0f((,)- -
N
N
j=1
N
1
Σ
((0)(xx; - }1v)) T 1
Σ
(rr - o'0)) - tlv)
N
==1
j=1
Σ
- -0
(ay -
1
Σ ( tlv))
N
N
f=1
We have
N
N Σ
As.irl S 2'0))
1
Σ (z) - {1v)
-
1
1
(()))()))) - -
N
N
j=1
+ 200
1
Σ [zj- ( N
Σ
[a))))(0))){; -
N
N
N
+
Σ
(ai - (00)) - +(v)
||
(ar - o'())) - - Σ
N
N
1=1
1=1
S SN VAZSALNE
CO'(0)2
2
NT. VNVAPALW + NT VNVABALW
(2.12)
N
We observe that As,ir
Aair Zin 1 v 1 1 -h and by (B.12), the last result follows. V
For (R4), we assume (E1.1), (E1.2), (E2.3), and (E2.5). We write
6'(0)a (1 - 1)
1
A1,5r
)
d
12
N
-1, + 00(0)²
Sn Zm
-
tr(S1)
Ia
1
7
+
1.))
V
d
d
L2
We have
A 2'(0) (1
1
1 1
+
M Lu
d
L2
N
(1 (1+ (N + +)1
|
[tr(
$ | + +0)( log V
Vd
d
$ 2
1
Vdl V VALAALVd + +Cxo'O0)* Ld
1
log V
log" V L2NVVd
Ndl3
Similarly,
|| A2
2'(0) 4 d (1 V L Iall2
1 L 1
$ $1142 200( log V 122le
(r((2)) - -(1 2 N
Vd
d
26

---

Published as a conference paper at ICLR 2026
CKVV 1
+ NL VVL AWAWA + CK) (
1
log $ V NLVVV NVLVd
log V
dVNL
Lastly,
A3 d'(0)' 1
V2 Z In Lv 1 Z
1 N = 3,ir + *0(0)
E(x; (a. v 1v) 1/
1
Zin Silvil v in
N) 2
N
/=1
By (E2.3), we have
CK2 log 2 V
NVV 2 Zm 17 ZT1 1 (1
Σ az V lv) |v)
N
1) 1
Zin 1v 2
N
N
V2
j=1
CK2 Io V 1
NVV V2 Zm 1llly 1
By (R3), the result follows.
Proposition 2. We recall that 2₁= 2k + 1=1iig Given that (EI) holds, the following statements
hold:
(PI) We have for + j and any E [V}.
C Vd
E (La,=e,
21z Zm X, X 5
(P2) For any & E [V].
Zin X X, sc( [ 12)
d
(P3) We have for + j and any E [V].
C'oog d²
E - - 2m x, x, Z Zm 1 1 2 X X, 2, 24 Zm
<
(P4) For any E [V],
Vlog²
L 12)
E [1 2, Zm x,
X 21,21
Zin] 9 SC d
( +
d
(P5) For notational convenience, let
L
1
s= 24 Zm
x X;
V Iv
Z 7m (X]
For any i, k E (V].
Co V
and E| S C'log v 4 L2
d + 2
Ei<|Znn| S
VVd
Proof. For the first item, we have
E [1 m my;
11zz Zin X
x, 2 in Zi, X X, X, 2x [
(a)
E (1z;= v 12 Zi Z Zz Z Zin
E 2/z Zin x Z Zm x Z zb Zinl
v3 2 Zia
Zin 2 2%.
in
27

---

Published as a conference paper at ICLR 2026
C
Vd
where we used the independence of the rows of X in (a). For the second item, we write
L
E| X x, Z X X 2 Za Znl = V 2 Zin diag
Z, Zm Z Z
LL 1)
+
Zin
7 T Z 2
V2
L
12²
C
d
d2
For the third item, we have
E (1æ;=
v 1zz Zm X; x, in Zm 1m in Zi
x Z T Zin
in
E [1,, - +t)z Zin - Zin 11 v 1 Znn 200) Z zu [in
- E B 2n a. 2 mm 2n Zz Lmm]
1
43 Zin 2 Zi 1v 1/ in
Zin Z,.
S C log V
d²
For the fourth item, we have
E [2] Zin x, x, Zin 1 1
2 Zin X X; K in zR Zim
EE [z--
a,
Z Zin 1v -v Zi Zm ZuzzolZ Lin]
L(L -
1)
+
V2
2
T
Zin
Z
Zin 1v Zin Ziz
LV log' V
V2Vlog" V
d²
d3
For the fifth item, we have
E
2 Zin
T X,
X.
[v)
2 Zin X. 1 [Zin]
1
E
z Zin X _ x, _ Z X X, Lv
V22 Zin Z Zi
L2
=
Lz Zin E x. 1 Zin a0 T Zin
lv
V2
2
T
Zin
2 T
in
Zin (B..133
By (BI), we have
CL V VVd
(B.13) S
For the second part, let nj = e X1 We have
z Zm
[
X, 4 Iv)
in Zin (X.
1 T
v vl 11 11
Σ (ng- - -xT (Σ (nu -t)zil)
v v
(ng t)
11z 2, Zj.
(B.14)
1=1
Let S = (s such that $y == Proposition 8 to bound second moment of (4). We bound each term separately below.
[ 2,2
21 + य.. x 2,). We will use the third item in
28

---

Published as a conference paper at ICLR 2026
We have
tr([Tv -oa--llo (v]))s) = tr(S) V 1TSiv/
= -v Zin Elar a Z Zin
7 Zin Zi Zin lv
C'logVe VVV
Moreover,
tr((Iv - - S(Lv 1 - tr((²) - 1SL 12++
2
1
(LV)
V2
We have tr(s²) 2
CV² lox²
and
de
V
V
1
e, Sll
Σ
2
252
21
+
= ====z z,, Zin 1v + 2 Zin Zin Zj.
2
2
2m1
1=1
Therefore,
1 [r((Iv - 1 1v1 10cliv-vvvi11 [v(b)ss
CV2 log V
s
d²
Moreover, (diaa(())B S CVlog v
Therefore, by Proposition 8, have
E [z] Zin (x,X, X; 1. Iv) Iv) Zm X 1 1v1))11)
2 |zm] [ s CClog" CClo'2V ( 12)
C PROOF OF THEOREM
We consider the following technical assumptions for the subsequent proof. Assumption (Technical conditions). We work under the following conditions:
Permutation. Without loss of generality, assume -v.
Learning rates. Take = oy (1), chosen sufficiently small so that we can write p. = ++v++0(m)
Activation. We consider polynomial activation with degree of p+ satisfying:
- (0(0) o" (0) #0
The smallest non-zero Hermite component of has index yo i.e, q" min(k 0| BloZ)),, (2)) # 0}, for Z (0,
Since the learning algorithm does not assume any structure in the ground-truth permutation, we may, without loss of generality, take it to be the identity. This simplifies the notation in the analysis below. The learning rate 17 is chosen sufficiently small so that the network output remains close to its initialization 1.2 1y, which simplifies the analysis of the three-step gradient descent algorithm. The assumption the activation function is technical and is needed for the analysis of the three-step gradient descent dynamics; however, we believe that such an assumption would not be necessary for general multi-step training.
29

---

Published as a conference paper at ICLR 2026
C.I ATTENTION SCORES AND THEIR ASYMPTOTIC SCALING
Let X be an independent copy of input sequence. By using the technical condition above, we de- compose the attention scores to three terms 81, 82, 83 R4:
XZ W
(1) KQ
N
17 N
XZ
7m
Σ
aijx
oij
(IL-
in
i-=
x X [ - )1v T Zont Zou((z - $ lv)
N
+
17
XZ Zin
Σ
(I, -1,
2
1, - &1 1v) (ai)(Zzo Zout Zour + 1v)
1
XZ Zin Σ
(1 4
N
1,/=1
FW Zin X, X, (v) T Zout (xi
(C.1)
+ Onn'>poly(d,,N)
= =s 52 83) Opoyyyo NV)
The following theorem characterizes the scaling of each term. We recall that (e er. denotes the standard basis vectors in IR_,
Theorem 3. With probability least (1), we have the following:
11=1
1
For all I [L], (ei, 81)
VT2
VNV13/2d
( LY1/2'
||121100 S NVLd(LAd) NLd(L d)
|| z
Ndym
We first make an observation that we will frequently rely on in the following:
Proposition 3. For any p N, we have
E| 1 V A2 E B[| A3, 112] S poly p.p. (d, V, L).
Proof. By Proposition 13, we observe that asy S poly,. (d, V, L). Therefore, we have
| v 12 V ||A S poly,. (d, V, 4)
from which the result follows.
D PROOF OF THEOREM 3
We observe that
Z Zin T Zn zmevezee) X. + Zmev+ieT)
= =x Zin X + e Zin X. + X2 e + [e.ee
(D.1)
In the following, we will consider x; = Cus for fixed l E [V]. We will write
[I + X
xT.x) )
2k + 1 Zerig Zkz
(D.2)
30

---

Published as a conference paper at ICLR 2026
and
(e Zrig Zi X T + | 2e = (Zk, Ztrie) el = 1MM
(D.3)
In the following, we will consider the event.
Event (E)))
D.1 CONCENTRATION BOUND FOR
By (D.1)....)--D.33), we can write that
(1,,8)) = N
Σ
aizzz Zin
X, (IL L 1111)) X
iJ=
x 7m x 10 | 1lv Zone Zout 1 1vv
N
Pki
Σ
Oiy (e1 11,
X 2 Zm 12 - +Iv)² T. Zont Zout $ 1l)
N2
1,==1
0 + 4
We will analyze 0 and y separately. We define
B = B,, + B + Bi3,
where
B =
azy - )L))(2) - tLv)7)
NL
j=1
B
0ig (NJ
1 -1 +1v)T)
NL
B = V NT
1 1l
Σ aijlor - =vv)
D.1.1 CONCENTRATION BOUND FOR 1
We define
C =
T, (a: +1v)z Zin X (Ir - 11111) X.Z Zm
[Fi 1z Zin X X, X Zmn
1 22 1z Zin X. 1 1 X,2, Zm
=C.,
=C.2
By Chebyshev's inequality, with probability 1 - (v))
N
N
tr B out Zone C,) = tr ir(Zour
Σ
C,B.Z.
N
N
out
i=1
tr N
1 N
C.B.
+
1
log V
C.B.
Vd
N
31

---

Published as a conference paper at ICLR 2026
Bounding 02:
We start with bounding 02 term. We have
02 C
+
Vall
Va
1
C
1
Σ
C C,B
Vd
N
We have
D C,B,11
S
2 N
E 1CC C+.))
D C,B.2
E B
Nº
En C,B.3
N N
S
2
We define the scalars
('(0)* (1
L +
0'(0)
1 3
2'(0)*
12
83
NV2
First, we will bound the first two terms. Let * E (1,
Bounding first two ferms. For i + r, by using the definition in A and A in (B.3)-(B.4), we have
1 (1 (1a,=0, }) Zn X, X;2, A., Zi X, X,
tr B B C
1,
(I [m==,, }z Zin x x,
X Z
_
+ 2 (1,=00 1 Zn X x
(A. t. Ia) Din X,
t,
<
2 (La,= レ) 7 X, 2m X 7
+ 2 (10,=0, + ) t.
X Z 122| X Z1
By (PI) in Proposition 2, (EI) implies
tu 1
Ct.
(1 10.=n,
+tzz Zin Zinx X X, X, Zin X
Zin
5
(D.4)
1 Vd'
Moreover, by using (R4) and (P2), we have
1
E
(L..=x. - +)]lA,, [ X2 Z2, Zmn
1
('(0)2 1
Ndls + vata VALSALVa) + +0((0) LAV3/2Vd + LeNvVa) ==1 1
logV
C
Vd(LAd)
.V 1 + VIZENNAANA +0'(0)4 dVNL
NLă
logV + log2V NLVVA
NVLVd)
==2 2
1 NA
LA + Clog V NV3/211//23/ (L A dj3/2'
(D.5)
On the other hand,
tr Br, B., CT.
1 (1 (12,=R, 12, 2i X 1
1,1, A. Znx 1,1, 1, Zi
1a (1æ,=E, - 122 Zin X X.
Z
x.
32

---

Published as a conference paper at ICLR 2026
L Lu.= 1 -t, X, 1 X, 21
- - 2n X
1
1
+ 14
(L2,=0, - +)]lA,, - t Ialle
X,
4
|
X 1 X 2,
By using (E1.4), (E2.2), and (R1), we have
tx E (lz,=e, -oa--t-ie Zin ,,
1
Cl log² V VLVALA Ad
Moreover, by using (E1.4), (E2.2), (R1), and (R4)
[1 (I.-.. -
Lalla 1 Zu Zin X, [x, 2 1
Zin
T
C
(')0)
NaL + VaL VALAANT
+0'(0)1
log v L2N.Vd log V NLVVV NVLVd
Ing V L2V3/2Vd log
VL2(LAd)
VV
ETN VVL VALAAWWA +00(0)1
+= 2
dvNL
Clog V /3/2/V/L+ LVd VAR/NDALAAV PRALD / (L d/3//2
C
(D.6)
On the other hand, for i = r, by (R4),
iCC B C.) + tr (C2, B B.
12 Zm X
A X X Zi
1
(1
})
zu
Zin X; 1
x1 4.
z 2
L
< 4.
(1
Zin
X
X,
2/ 4
1,
L
in
By using (E1.4), (E2.2), (R1), and (P2)
+ tu
7,2
Zin
X
X,
ZT T
in
Zi
122 Zin
T
E
1
2m
(
Ct.
Ct.
(D.7)
1
T2
LAd
Therefore, we have by (D.))..)))))))))))))..). and using N < VL and L < V, we have
N Σ C.B...
1
Zim
+E
Σ
2 F C
Zm
C
C
5 I3SCAA A Clog V
C
NV3/2//d(dA IT) A/d3/22
C
C
$ NNALAAD3JUADD
(D.8)
TVSI dj3/2'
33

---

Published as a conference paper at ICLR 2026
Bounding the third term. We have
N Σ
C
I
2 N
Σ
urCC,,
B
N
1=1
We recall the definition As,rr in (R4):
N
aiy (x, - x1v)
T
Σ
aryla, blv) 2'(1)*
N
N j=1
N
1=1
We have for +
(r(Ci, B B
CT,1) + Br,
r.
(Im +)4 Zin X. x, Zn
(LIEE +) 7m 1
$ (1æ
} Zin
2 2 x, 2,
t3
L4 (Ia
Zm
V2TA [1m.ma,
Zin X:2 X| X,
Ag,ir
V2L6 (Im,=.r
For the first term, by (P3),
la
CO'(0) log V
E
(1,,=,,
1 2 Zin X X, 2,11 im Zin
x, X, S
L
NV2I4 d²
For the second term, by using (E1.4), (E1.5) and (E2.2)
t3
E
(10
+ ) Zin X, 1 1 2 2
2 Zin
T
NV3LA LAd (LV )
o' d'O0)'
1
For the last two terms, by using (E1.1), (E1.4), (E2.2), (RI), (R4), and (P2),
(1 (1z,=0, - +) ir 7 X X -
2
Z
X 2 T
in
2z 11 -m +
(IZ
in
T
26
12
in X. 1 1,, 2, 2z 21112
52
C
1
V
0'(0)* log² 2 V
o'(0)2
+
+
NT VNVAIANNa
V3
d2
R2
Ta
NVV
N
L
11
V
V
2
+
+
V3L4
d
dp
d
Ld
NL
VNVADALI
C
C
&' (0)) log2 V WALLALP
(D.9)
NV3 A
NV
For i = r, by using (R4),
tr,C,,, B 1.1 + tr, B.s
1
(1
Zin X X X As, Z X,Z,,
V
34

---

Published as a conference paper at ICLR 2026
+ L (1 1 1)
1
Z X 1 1 Zn 1 1 X
2t3
2tg
L2 1 Zin 2n X
x
- 2ol 1
L
Then, by (P4), (E1.4), (E2.2), and (E1.5), we have
1 E (1 I X X, x 9Zzn
+
t3
E
(I
in
Zin
T
Co' V
d²
$ NV LAd
+
(D.10)
2 VL
Therefore, by using (D.9)-(D.10) and using L < and < VL, we have
N
+
N
1
Σ
C
2
Imm
< (RaL(dAL)(dAN))
(D.11)
==1
Therefore, by (D.8)-(D.1),, we have
C'oog V
C'log/2v
02 S NTTAU INVAa)AAAAAATA)APP dj3/4
Bounding 01:
We have
1 N
1 N N
01 Lr
1
C,B.,)
+ lr
Σ
C.B.,
+Ir
Σ
C.B.3)
N 1=1
N
1=1
1=1
11
: 92
We have
N
On
N2
Σ
Oig
(1
1b)zz
Zin x, (T
L
1-=1
6(022
Σ
2/ Zin
Σ 11 (I - x,
Zm,, 11v17)11
N2
j=1
(1 +
Σ
a3jzz Zin x_ (Ir. 7 -
N2
x 2 (X,] - [1r
('0))'())) F21
N
x X (Ir
(X] - +1lv1))1
=: 11a + 2 11b + 0 1c-
We start with the last term. By using Hölder's inequality.
1i1el= S 22
1
1 b1
sup |asi - |oc)-('(0)}] [
sup *]j[[]]
124
Z
X
(I
xx - +1vv[)
35

---

Published as a conference paper at ICLR 2026
Colog V VI A
(D.12)
where we used (E1.8), (E2.4), (E2.8), and (R2) in (D.12). Next, we consider &116:
1011l = (1-+)= N
Σ
T Zi
(jj X, 2 Zin
[ay x, 2 Zin x 1r. Zin]
NºT.
1=1
(1 7)
+
Zy Zmn
011 X x Z Zi x 1
Zm
2L2
(1 )
0gj2k Zi x 1n
N2LV
j=1
(1 }
N Σ
O3j24 Zin Zn 1,
2m [x] 11
N2
/=1
For the first summand,
(1 +
E
2/6 Zm
22
2 1zn
11 Zinl )
(55) X
-R x Zi X
Zm
(1
E
2
Zin X
X 2 T
Zin x 11) 1zm
L3
in
CO)(O)
E
Zm
X
X,
|mm
(D.13)
L3
where we used (R1) in (D.13).
By Chebyshev's inequality and (P2), with probability - ov(,,, we have
2 Zm (1- Σ
(,
Zm 11
- [a(aj x, Zid Im X3 1 lmml)
N2L2 j=1
2 log V NN LLd
For the second summand,
(1 1
N12 2 Zi E
XI 1r Zim]
(1
())
2
Zi
E
X
X Z in
Zin X X1 Zm/
NL2
(1 (1- +
Z Zin
E ((0))
x
2,
Zin X, X, Zin] Lv
2
)0)
(1 1 )d (0)
2
T Zm
E | 2
Zin lv +
24
T Zin 2T Zin lv
NL
NV2
(D.14)
(1- $)
24 Zin E.
- ('n))) 2, Zmn]
(D.15)
NLV
(1 }
x, + Zin ( 1 [lv1])n, /m| (D.16)
Z Zi E [a01
NI2
+)
zz Zma E [an - -'(0))
- + Z Zin Zi
(D.17)
NLV
36

---

Published as a conference paper at ICLR 2026
Colg V NVVa(LAd) N13/3 d))
where use (E1.3) to bound (D.14); (E1.6), (R2) for (D.15); (P5), (R2) for (D.16); and (E1.8),
(R2) for (D.17).
For the third summand,
Jjj24 X X, Z Zin Lv
N2LV
= NV2 2g Li Zin Ln + 0'(02 N
o' (0)2
Σ
Zin X X,
ttv
Z Zin
N
j=1
L Σ (a ('0)))( Zn x X Ziu 1v
==1
The first term:
d(()) NV2
Clog V NvVdi
Zi
Zi 2 lv I
The second term: By using
E
2
Zin
{ X,
Iv
Z Zn 1r) i
1=1
N
= Σ
ZB
Zin
X X,
+vv 4 Zzn Zm 1v) )
iZ/u
LVN
d²
j=1
Therefore, by Chebyshev's 's inequality, we have
0'(02 2
N
6((0)2
Σ
2 Zin
-
2 Zm 2n
N2LV j=1
NV/2
Finally,
N Σ
(0 (o)))((00))2) 2n |, x, Zm
N2LV
1=1
(زا - 00 X
NVV
NL
/=1
NVVLd'
where we use (E2.8) and (E1.1) and (R2). Therefore,
NºLV Σ a3zz Zin X 2
Zinlv < <Clog V
(NVVad + NVVLLd
NVVd!
For the last summand,
(1- })
Σ
Zi X X
Zin (x] - +
L3
#=1
(1- })-
Z Zin
1
Ø=1
37

---

Published as a conference paper at ICLR 2026
x L ( - +iv1E)1
Zin - v I [mn] )
E
(1 (0)²
+
2k 7m E 1
- [Zm]
NL
(1 t)
+
NL3 2/ Zin
E [an - -100))
Z v Iv [v (D.18)
We have
[
1 -- {(r1[)1,)
) 2 | 2,n2
2/ Zi E
1 2ha s012 Va
12² 1
S C1 Zm E
Moreover, by using Proposition
E x 1, X Zi
/ 1 11 |mn]
L E
=E [
1
1 1
|mn]
[
Zm Z [Zmo]
= =LE [ - 1 +1 Zun + (
1)
tr
2T Z
2L(L 1)
I - Z Lv
in in Z lv.
V2
1/3
L(L - 2)
+ V2
Lastly,
1
NI.3 2/ Zm E. - ( 2m (X v blvr Zin]
E Ilaan- ( Zin x 1z
2
Zin
(X] [ Zin
NI3
Cloog V
where we used (R1). (E1.4), (E1.5) for the last inequality.
Therefore, by Chebyshev's 's inequality, with probability 1 - ov(1), we have
(D.18) S Colog (
1 NI2VLAd + NLVVa)
1
Therefore, we have
1nvl S &Cloog (
NVV(LAdVa + NLPLLAd)
(D.19)
Finally, we consider 9 1iaa:
O11a = (1 +) +) 20)2²
N
1
N
Σ
1 a - tta,0,) ) 2 Zin - +Iv1[]
Zk in
NT.
NT.
j=1
179
2'(0)2 N
N
2
T Zin
1
(1 (1, -
Zu, +1v1
NL j=1
NL
+ 20)))2 N Σ
N
Z/z
Zin
(1 12 -
(1a-= FNN, N;
2 Zin - -
NL
NT
j=1
i7)
20) N NL j=1
N
Σ zk
Zin
Σ
(L,,=
D. (NJ -
NL
38

---

Published as a conference paper at ICLR 2026
20)2
N Σ 2 N j=1
in
1 NL
Σ
N
11 0 -
I,
2 - 1
Daa + Dab + Due + Oad + Dee
For the first summand, we write
1 1 )'(00²
1 (1 L' NL
Σ 2
Zin
Zin
Σ 1 =
Zmn v1
NL
Daal
+ (1 1 0'(0)2 N NL
N
Σ
2/
Zin
NL
1 1
) Zm
v
1
j=1
(
aa2
We have
|0a11 5
20)3) N
N
10:=0 - +1
sup izj
1m zJ Zm Zin + [vv
N2
Ø=1
O((0)
(11 b)) suplz exe 2 Zmn }lv)
N2
10y=el
;
0(002 N
1 w;=es - ++1) }
suplzs Z exei Zm + [lv)
N2V12
1=1
1fj
Cloog V VL2Vd where we use (E1.1) and (E2.4).
Moreover, let
lea2 =: (1 1 21())2 N
Σ lea2.5-
L NL
3=1
We have E[Dao2,5 Zin] = and E[0 an Don [in] = 0 for j + j', and
E 202. [uu]
122
CL
N NL
N
$
2k
Zin
2 Zin
mmz= R
Z zz Zin
NL
izj
CL
1
N
5 C
dV2
E
2
Zin
NL
Zin
NL
Z
Zin
Ld
where we use (E2.8) and (E1.1).
Therefore, by Chebyshev's inequality with probability 1 1- (v), we have
Maaz1 Clog V
VNVL3322
Therefore,
Clog V
ClogV
1aal <
V12vd
VNVL3/2
39

---

Published as a conference paper at ICLR 2026
Moreover, for the second term, we write
20(1)²
N
N
|uabl= -
12,=0, - +1)
22
j=1
X sup 1e 2 Zmn + 1 2 Zm (X] [1v1))1 inj
(0(0)²
Σ
1 =
11(1 -[, })
N2L3
j=1
x suplz Ze (ex + X 1 Zi b1v11)
8'(0)2 N N (NAVL3 1=1
Σ
1 1z,=el 11
suplzz Zm(exd] + X [ ZX] tivl[)1t
x
Cloog V VL2d
where we used (E1.1), (E1.3), (E2.4) and (R1). For the third term, we write
Jac 2'(0)2 Σ Zm N, N
I-
1
Iv
2
Zm
1
}) $ 1lv)
NL2
N
171
(L) 1)
2/ Zin Zi Zin
1
Σ
c
- =vv)(z )
Σ
N
(z;
#lv)
NL
N
j=1
6' (1 (1 1)
N
N2 z4 Zin 2, Zi
Σ
[(x; (x. +1v
= 0act + 0ar2 + Dac3.
By using (E1.1), (E2.3), (E2.5),
C
C N2
E EiU22 aara Zin
N2
2
Z.,
Zi 2, 2
Moreover,
C
C
E 1 acc Zin]
N3 2
Zin 2,
Zin Z TT
Therefore,
C'log V NVVLd
Clog V NVNLd
acc2 5
Moreover, we have
E(02 1
N
L L-
L-1
Σ
Zz Zm
N
N
Iv
Z Zim N N;
V Iv
ZT 2z Zin
V
40

---

Published as a conference paper at ICLR 2026
N
L
$
E
1n 2 2x mm
N
V
1=1
C L-1
E
T 2/
2m
2
1m
2 in 24
Lm
PALELEN
V2
i=1
S LV2
Therefore, by Chebyshev's inequality, we have
Cllog V
1aa1| S VNVLåd
For the fourth term, we have E[ded Zinl = 0 and
E P ad Zm]
N Lizs
C
E
(le= Zi
L
w==1
N N. N. (N, -Iv
3 ) imn
C
Σ
E
(z)
N,
1 ) [m
VN
Aj=1
C
E
J 2 Zm,
N.
VN
where we used (P2) in the last step.
Therefore, by Chebyshev's inequality with probability 1 1 (1), we have
Clog V
|aal S NVVTAVLAd
For the last term, we have WB Zia] = 0 and
E 10 [Zn.]
C
N
$
1i4s (12;1122=0 - 1zz Zin
N
i,) 11
x 1IINN - 1 - ) im
C
N Σ 11=1
Zm N. L. )- 2 1nmn
1mm)
L
E
2
C
N
E z Zin 1, N,2, 1
N+L+dV
iJ=1
C (1+ d 1t)
L
NPL4AV where we use (E).4) in the last step.
Therefore, by Chebyshev's inequality with probability 1 1 -vo we have
ClogV V
100l S
41

---

Published as a conference paper at ICLR 2026
Overall, we have
|0nia| S Coo (
VLLVd + VNV VNVVLd(Lnd)
(D.20)
Therefore, by (D.11)))...1)))(0..)))) and using N & VL and I. < V, we have
1
101l S C'ogg V
VIAVa VNVVTA(LAd)
1
C'log V 1
+ Cooog V
NVV(LAdVV NIIVLAd VINAVLAdVABAWA 1 NVV(LAdVa VL2(LAd)1/2)
1
CologV
(D.21)
Finally,
012
1
Σ
0iy x - - Zin X (Iz -1. X
N2
(,51
1
Oij (1 mi 1mmmn ト Zin { Xi -v 2 Zinlv
N2T2V
f,=1
N
1
Σ
Oij [1 [m m. - +) Zin 2 ZinZin lv
N2LV2
+-==
Oksy (1y(1n, (In, - t)z Zin X. 1, X Zn lv
N2L3V
For the first term,
N
Ongy
1 (I - 1 Zin [ X; +vv 2 Zi
N212V
f,j=1
1
la my - b1) sup (3(0))) 2x, X, - klv) Z Zinlv
N2L2V
2'
mi my b)) sup (2z
Zin (x, X; [v) Z ZinIv
N2T2V
1
1 m. m.
+1
sup
((()))((0)))2 Zi, (x X; +Iv) Zin Zin 1v
N2L2V
in
2'
N
(z) - 1v4 Σ
sup 2 Zm X x, - 1vv 2 Zn
NL2V
VI2VV(LAA)
where we (E1.2), (E1.8), (E2.3), and (R2).
For the second term
N Ciyy (lz,=
1
v |2 Zin Zin Zin lv
V
(=
1=1
N
N
2==0; - 11
C'log
sup 2 Znn 21 Zu lvl
N2LV2
v3/2 Ld3/2'
i.j
i=1 j=1
where we used (E1.3), (E2.4), and (R2).
42

---

Published as a conference paper at ICLR 2026
For the third term,
N
Σ
ai
( - - Zi X 11 X Ln
N
Σ 6,=1
12,=0, - )
sup/ar
X, 1 x, "
Av
N2L
in
(VL.
Vz
V2L2 ALAd
where we used (E1.4), (E1.5), (E2.4), and (R2).
Therefore, by Chebyshev's inequality with probability 1 - (1), we have
Clogg V
C'log² V
122 S 2
(D.22)
VLVVd3122
By (D.22))(1D.22), overall we have
01 Cloog V
VT2Va + VNVI3/2 + NVVTaVLAd + NTAVV(IAAV)) )
Clogg
+ Clogg V
NVV(LAdNa NVVa NIVLLAd) WIALLAVAABANN
CClog V
NVV(LAdyV + VNVI3/2 VLLPd(Ad) + ViVTa(LAd)
ClogVV
+ V
D.1.2 CONCENTRATION BOUND FOR yp
We recall that
2 Zou Zour(x; 1v) Σ Cny(er (1r) x, N. 1n(2, =lv)
Φ N212 1,=1
In this part, will focus the term
22 aig (el
Σ
+11))² T X X 11(21 t1v) Z Zoue (x;; + lv)
N
aittr Zin 1r(20; - (l) T Zoue 1v)
N2
WJ=1
Σ aistr Z X 11(2; + ZT Zont (2: - -1
N2
P1 1p2.
For the first term, we write
20)22
N
Σ
tr Zou - + 1 2 X LL( - $ +1v) Zome)
P1
N2
1-=1
1
Σ
Zoue(aii o'(0)3((z: - Av Av) 2 Zin Le(2
1
V [v) Zoo1
1
N
41=1
43

---

Published as a conference paper at ICLR 2026
P11 + ff2
We start with the second term. By Proposition 5, we have
N
1 ) Z Z Zm 1
(ay - (10)) mm==
912
N
11=1
log²1 V
Σ A,=11
N
(ai - -'(0))
1v [v)x Zin x,
v 1v)
1
L2L2Vd
Let M = {(ai)- & ('(0) )o Z Zin X, 1 11) J.je[N] We have
N
1 [v) Z X LE(20)
1 |1
(ii
o' (x:
ij=1
<
M
F
Σ
N
(a.
+v)
1
-v)
N s V |
i=1
where we used (E2.5). Moroever,
N
||| =
Σ - ((0)))
2 Zi, X3 1n)²
1 1=
N
auy - 0'(0)9p2 +
Σ
0i (0)22
suplzs Z X, 1
#/1=
i,
N
N
[1+ 7)
V VARRALY
L
where we used (E1.4) and (F2.4). Therefore,
1 (ai (0))
V )o T 2 Zm X, Lx(2,, - 1 [1
Σ
N212
1,11
NVL3/3//LLAA VALPAAN V 1
1
N
VILVLAd
Moreover,
N
1
Σ
(ay - -00))) t Z ZinX1 LL
L
-0.=2 - b1 sup (aa - )'(0)))z, Z X 1 VIRVLAd i.jE[N]
2
1-=1
where we used (E1.4), (E2.4), (R2). Therefore, |412| S VIAVLAd
Next, we consider |p2)|.
By Proposition 5,
22= = 20)22
N
Σ
+)1,,
Zu, d1
JJ=1
d' (0)2
N
Σ
-
Zm le(2, - tdv)
Vd N
1-=1
Σ (aij 2 Zin, 11
N2
44

---

Published as a conference paper at ICLR 2026
N
Σ
(ag -
1 - t1v)
d N
w==1
P21 + P22 P23 + 4-24
For SP24, we define M = (aij - o' (0))) x, 2 Zin X 11. ( Similar to above, we have
1
M P.
NVL3Vd
We have
N
| p =
Σ
(ay- - ((0))) 2n x 12)²
1==1
N
lais - 8' (0)2p)
supl iØj
[ X. - Zin X 1
i/j=1
Σ aa
2'(0)9p)
x 2 X 1
N2
N
S
+
AV A LVd)3
122
where we used (RI), (R2). Therefore.
N
VN
|sp2a|
NVI3VO
(VALPA LVdjo12
VIDVLAd
For 4P23. we have
N
1 +
14pal I
1 = -
supl (ay - 2,0
N
its
NL supl(on Ø' -00000))]]] 1.1 1
VL2VLAd
where we used (E2.4), (R1), (R2). For the first two terms, we define
Vo=
1
xj
11(0;; - 1 +1T)
Vo,1 NL Σ
1 [1)
NL
j=1
1 N
1 TT) Vo.3 =
1
N
1 [Iv)
Vo
[NJ + [I L 12;
(x;
NL
1=1
1=1
We have by (F2.5)
0'(0)2 1
Z Vo 2
d'(0)2
20)
1
1:p22| S Vd L
LALN x V12d
NV Vd
Lastly.
0'(0)2 1:221| I, Ir Zin Vo 2m)
0'(0)2
Z (Zin
20)
1
NVLVVd
5 NL VLVVd
Therefore, 14221 S VLRVLAd
Lastly, we consider PP11 By Proposition 5, we have
lpnl
= (0) V T 0,1 in Zia Vo.1) + o3 (0) (r 2 2m Vo.2) + (00) (Vo,, 2,, Zi Vo.3)
1
1
L
+
Vo,
+
VoT
Zi
Vo
+
2 Zi Vo.3 F
Vd
45

---

Published as a conference paper at ICLR 2026
For the first term, by (E2.5), we have tr V,, 2 Zin =
Vo, 4 1
VI2
For the second term,
N
tr( tr(Vo,1 2 im Vo,2)
NL Σ
1 - [-v)
Vo
Zi Zm (N] - +1I-1)1--11
j=1
We have
E ( - )2 Zm] 2 V2Ld
0,1 Zm
By Chebyshev's inequality,
(V., 2,, Zm Vo.2) - VNVLL/VV
The third summand: We have
(L - 1)tt(( Vo _ Zi Vo,3) S L Z Vo.11211 Zi Vo,312
where we used that Zn Vo.3 is 1-rank. By (E2.5)-
T.
V In 1221 Zn Vo.3112 2
VLLVd VVLN S NVLVA
The fourth summand: We have by (E2.5)-EE..7)
1
C VL2d
Vol 2 Zm Vo.1
Vo Vo.1 - S
F Vd
F
The fifth summand: We have by (E2.5)((2.77)
Vo. Zi Z
S tr Voli Z In Vo Vo.2 2. mm Zn Vo.1)
V
1
NLd Zin o. Vo, 2 <
NId V212
NVI3
Therefore,
Vol1 Z Zin Vo.2 VNV133/2
S
F
The sixth summand:
11)
Vo
T
in Zin Vo.3
S VVN
Lv 2
Zm Vo,1 Vo. 2, 2n
VENL2d
Therefore,
L-
C
Vol Zij mn Vo.3 $ VLVNd
vd
Therefore, we have
1= (1)
ō(1)
8 Pul
VL2
VNVL3
D.2 CONCENTRATION BOUND FOR $2
In this section, we will use 3 = x" (0)) We have
oj 82
Zin Zn X, 2n, 1lv)
2 Zont (3.1 11
N
ij=1
1
ByzI Zin 1, 1 1,
1x(xx -., 2 TZ Zout 1 1v)
L
N
+ N
PM
Σ Sy(e1 - - x, - T ZT Zoue (x, - Lv)
-
negligible terma.
46

---

Published as a conference paper at ICLR 2026
D.2.1 CONCENTRATION FOR
We will write K as follows:
N
Σ
-
Zm
x, x Zin X 1,11 11
K
N 1
x (2; |/ [v) T Zout Zout $ lv)
3
N Σ
+ 2
2k
Zin
T Iv)
zn Zil az + 1v)
+J=1
x (2; - \ 1v)
T Zout Zout (2: - 1 lv)
+ N
Σ
Zk
T
Zin
N N. - 1 Iv) Z Zin (x; + (v)
1J=1
x (2; - 1lv) T Zou Zoue (20, - +lv)
N Σ
+ N2
2k
Zin
T T [tv Z Zi (N₂ - -1v1
i-1=
x (0; +1lv) T Zou Zoun - lv)
N Σ
+ N2T.2
2k
T.
Zin NN., - 11) 2 Zm
(111[-1)
i,j=1
x - 1lv Z Zoue (20, - -/ lv)
-
B
+ N2LV 2 Zia Ii Z 10 - 1v) T Z Zone (z) - 11v
Σ
1,1=1
k1 + K2 + K3 + R4 + K55 + Ki.
By Proposition 5, we have
N
1
K1 N
Σ
(Big - (1 -1
Zi X N,, 1r.
iJ=1
log"! V
Σ (Big - B) 11 Zm X2,,XXXXX,,,.. X - 1occ-b-v))7 |1
N
=: K11 + K12.
We have
7/4
1
Σ
(Bgg B) -mm= (1)===))-)))
Σ>
(2g X Zin x
T1x71111 11)22)
N21.
N
N
(Big - (1x, +) 2)
N
Σ
Σ (z) Xi X.
NELV
j=1
1=1
N
Σ B (B), +) )))
1
Σ
2ZL²
N
22 Zin (X, X, - -. - -bIv)))ae)))
X
Σ
(3y -1)1))))-+)))
N25V
J=1
47

---

Published as a conference paper at ICLR 2026
N
Zm 2 (X; - +1v (1r1))11)
X
2
Y=1
VN
N3/2
1
L VV VVL
LV
L
V
VAIBALA
d AVLA d3/2
VN
N3/2
+
N3/2
L
V VALPALI LVd VaVLAd
1
VTAVa(LAd) NIADAALA VIAPAVAAVV AR where we use (P4)- (P5), and (EI.)) Moreover,
N
K122 2
N3/212/0
Σ
x
i=1
N N
N
Σ
(B:5 ('(0))) 2 X in X, _ 1212)
1=1
N
VI
1
NULVVV VAL LVd
L)
vd d3/2
+
1
21 WIALLAYABBAL VNVI3I2d(LA))
Therefore,
|mil S VI3VA(LAd)s NIARAVLA A
By Proposition 5, (1- ) K2 N
N
Bt tIv) T
j=1
3 LL 1z Zm [ T - +lv Z Zin (2, + 4lly)
x
N
>
N2
2
1=1
Σ (x, $ (z Zm [I:0] (zv) (2. + 1v)
X
Let n4 = 10 S Nar e. We have
N
N
1
Σ
(a, - t(v)|
Σ
(z; - +1v)2 Zi xx:o] - +Iv 2, Z x
N2
J=1
1=1
L-1
(z) - ++4v / - #
N
Zin Zx + 2 Zm
L
N
5=1
n2 (2|ZT Zim (e,e, 2, 1)/
2
N
N
Σ
(x $ 12) Zen (. +Iv) Z.n 1v
NoLV
N
+
Σ -
11 Zin 2n, - |v Zun
N'LV
48

---

Published as a conference paper at ICLR 2026
where we used (EI.1) and (E2.3). Moreover,
N
2 1
E
blv /2 Zm [ Z Z -v Lv
1
1
E
Z
Zi (1,0)] + - 1v / in in (2,00) T + 21,zx Lmm]
N
N
1
(1 (1=== - } Zi (a:00) - 2 in 2m 1 12 21 Zin
(02 (
[v) 2 2s. Zin
N
N
E
2/
Zin (x,0, } Iv) 2 Zm 1v 10 Zm [2 [v) 2 2x Zm
N
NV
N
V
N
NV
d2'
where we (E1.1)-(E1.3) Therefore, by Chebyshev's 's inequality, we have
1 |R2| S VLAVa + NLVa
Moreover, by using Chebyshev's 's inequality
1
blv)
K3-
N2L2
N
T2 Zin (N. N: -V-Tv) Z Zi (z, + (v))))
x
N
1
>
ay [(v)
22L2 d
2
N
ax blv)zz 2m (N. N, - k + (
X
i=1
We have
- 122 Zi N N; - Zin +, Zin
E
(1
N
N
E
112 Zin N N₂ - 4 Zin Zin (2, ( (v)
Zim
j=1
N2
[NT N1 (0) |mm
E
2 Zin
(D.23)
V
We have
N NN,N, N1 - (v 2 Zin (2) + ( 2
CL d
+ L3
43)
E
Zk
Zin
Zin
<
where we used (E1. (EI.1)-(FL)) (FE1.3 and (E1.8). Therefore, we have (D.23) S No (1 + 1.3 ) Also.
N E
+Iv)zT Zin NN. N. - -L/lV) 2 Zin (z: + ل 11) ) || Lmm
x
1=1
CL < 2
(1+ 1
Therefore, we have
NVVIL3d NL + 1
|Kal S
NVV (L2322
49

---

Published as a conference paper at ICLR 2026
Moreover, by Chebyshev's inequality
(Lv) T
Σ 2
14
2
N
Σ
(x;
=1v2z Zi [ -
- -
2m (N. - ++v1[-1) Zm
1=1
+
Σ
x
}1v)11
22
J=1
N
a 1v)z
+1,) Zin [,,0e [1v 2 Zin (N, - -b1v1
i=1
We have
-2),,, Zm 20 x; - +1v Zm (N. -+Av(I-1)4-11) |
E
-0.=2,--
N
E
Σ
[1 -=y=
12 Zin x...o. 1 Ivy Z Zm (N₂ - Izm
i=1
1=1
N2
T [Iv - Zm (N. -))))))))))) Z
E
Zu
Zin
(D.24)
V
We have
E
2/
T Zin
I T - Iv Zmn -bIvII-)) Zin
<
CL
d
24
Zin E [ - ttv) - Zin x2,2a- - +lv) ZZn
Z S
CL
ds
where we used Proposition 8. Therefore, (D.24) S
No1
Also,
Vaz
N
2
Σ
+1v)zz
Zin x,,0 - 4 2 Zm (N. - ((v1 1
1=1
N
= d
CL
Σ
2
Ziu
E [2,2, - - ) Zm
Zi (.,.e - -b 1zm ZT S
CNL
d2
1=1
where we used Proposition 8. Therefore,
1 Inal S (NVVLARd 213/2//12)
Moreover, let
7/ = 24 Zin
N N: -4/11 2 2 (N. -IvI
E 2 Zin
N, N: - 4. Zm (N. - - Zm
By Proposition 5, we have
1
Σ x + (vv)
[ (z. {1v)7)
Kg
2
50

---

Published as a conference paper at ICLR 2026
log-1 V
N
N
Σ (ag - t1v) 2
Σ
(a., - +1vm%%
L
2
j=l
1=1
N
Σ (2) {1v)
2
/=1
D v 1 Zin E [ N; N.,NN-4 11
2 Zm (N," - -1 I (mm ))
log" 2 V
N
Σ j=1
- (a
N
1 Zi R (N. N; 1
Zm Zm 1 -- 1
a;
By Proposition 2
N Σ
az my - t11))²
-))))//
= =EB B Σ 1 my - F)7r)* 2
E
2 ) Σ
2(1 +)
N N
B
E Bhill
V
1=1
j
N2 1 +
1.2
2 V
d
d
Then,
N
Σ x 4
1v)
$
Moreover, by Proposition 8, (E2.3) and (E2.7), we have
Σ x 11 Zn. (N,
L-1 Iv Zu) )1
2
LVN N
VV
Therefore, by Chebyshev's inequality, we have
|Ks|
NL)Vd(LAd)iz NLd(LAd)) + NLVVV
Lastly, by Proposition 5,
N
K6 L
Σ
(1,,
1)2 Zin 2, 2 X 1
1==1
B
N
Σ
(I.
111)21 Zi 2
X,01x1111110 12(2;
N3LV
1,==1
3
4 |z
Zin 2T 2in x 1
NºLV
N
N
Σ
NL Σ
(x: - 11 2n 2T,TZ,Z X
VVd
N
2
j=1
51

---

Published as a conference paper at ICLR 2026
We have
N
2
T
Zin
-in Zin X 12(2; - b1r)T
NL
1=1
$ 2/ Zin
nn Zin
Σ
[lv1 -
NL
+ 2/
Zin
2n mm lvl
Σ ( - 11
N Σ
( + VV
Zi
!
NT.
d3/2/N
1=1
CV
V NTd3/2
Moreover,
N
E
N2LV
Σ
(1
mn
12z Zin I Zin x; 11.) 2 iZin
1/=1
1
Σ
N
(1,=ay man - + Zin Zin Zin x, 1r. )² [Zin]
N1LV2
j=1
2
Zin Zi Zin x, 11. )
N+LVV2
IE
Zkk
iZin
j=1
2
N
N
(1 2: ma - 1 Zin Zin Zin x 11) [Z!.] (
E
N
i#1
2
2m x, X1 2 |u
(2 In Zin
N4L2V2
in
f=1
N
2
+v 1 [
E
T T. 2 2k
Zin
N2V3
2
+ C
E
Zin
2
Zin X 11. [Zn
N
VdV
where we used D.2.1 in the last step. We have
N
) | xx1x) 11)
E
2 2in
2
Zin X
N4LVV2
Cllog² V Ld
N3 2 Zi
Zim
1 balli + & 2 Zi, 2inz S
in
Therefore, by Chebyshev's inequality, we have
Isol S NL
Overall, by using N & V.L
1
|K| < NV Ad) NLd(LAdj d2 + NLVVd VNVIA" VNVLAVd
1
52

---

Published as a conference paper at ICLR 2026
D.3 CONCENTRATION BOUND FOR Sy
We have
N
el
83
N
2/ Zin X x
in
T
1
( ( w Zin x, 11 1 Zin x ,,,1), 1n)
x
k=
E
12 mm x 1 1) //
L ux Zia x 12)) +11) 2our Zou (ar $ [))
N
Σ
2
Zin x 1 1,
2
iJ=1
1
w
( Z X, LL) / Z 1e)
E wiy (
2 Zin X 11) / + w Zin Lin [lv)² 2 (2) + lv)
N
Σ
(el L 11) X
N أ-
1 m
o
w ( Zin X, 1 Zmx (
V + negligible terms.
D.3.1 CONCENTRATION BOUND FOR V
We define
N
1
ū= tr
(z. - 11
X,
x, X, L Z X,, L1)
NL
1=1
1 N
Σ
o Fw Zu 11)x - 1v) Zout Zout)
x
N
= tr
Σ
-11= ) Zin X T X Z in ww ( ( u Z 1L)
NL 1=1
1
Σ
o ( Z X 11) 11v)T)
x
N
log² 2 V
N
+
Σ 1=1
[x,-
1 1
Zin X T X, X o ( T X 12)
Vd
NL
1 N
Σ
T Zin 11
11) X \lvv)l1
11 + D2z where we used Proposition 5 for the second step. We define
s(1) = (0) + tr²(t) and d'(1) =: (0) + (t) and (1) =: (0 + (1).
and write
1 (0) (0) tr NL
Σ z. X X X
T wk
1
Σ (xy + )
53

---

Published as a conference paper at ICLR 2026
N Σ
+ ((0)tt
2.2k
I, Z wy
N t=1
Xi LLV2 ( uw Z X,
1
Σ
x - )
N
+ +0)))t
4i4K mm N1 N N. 2
N
Z
1 ( F X, 11)
Σ
x )
N
N
1r
+ 2 2m X,
X Z in
w ( w Zin X [1)
NT.
1=1
1 N
(+W)
Z 11) w Zin 1L(00 )Iv)T)
N
=: in + 1/2 + 1/1 + 1/44-
In the following, we bound each term separately. Let new = Ei [N] x; = e.
We have
11 =
The
2
Zm
ew + + Z wR
L
N
N
w
V
+ 2
Z
in
N
N
Σ
N,, - Z wk
KEft Anw
We have by using Lemma 5 and Proposition 8,
V
2 Zin
Nww
1
Σ
(N.NN, - ---'iv
2 w/x 2 Z
NL
N
tEffi Sow
=
E
Zin
Nww
Σ
-m 1mm
NT.
N
SE(i1 dno
V
= 2 Zi
Byy
V n2m
N2
N
w=
x IE
[N Ni - Iv)
2 Zin [N] Ni - |zmn]
< N2I2V
C 4
Zin FF.
[NNT - - vv) 2 Zin (NT N₁ - -r 1Zin Zin Zu,8
C
N2Vd(LAn d)
Moreover, by using (E1.1) and Proposition 8,
E
w
Thw N
zk
Zin
e - [v) 2 wr )
2 2iZ.n |
L
N
E
Zin
Thw N
I nww
e
T +
(
nw
1 ) Iv ZTn 1 1mm
1,2
N
N
V2
V
Tba
nww
Rww
_ - L- Iv 11 1m0
<
L.d
E
N
N
e am
+
IN
w=
54

---

Published as a conference paper at ICLR 2026
1
CV2 d
E
sup WEIN]
Th
+
C
V
Ruw
C PN
N
N
d
N
w=
Therefore,
27 1in
S
N2
Moreover,
N
C
ZinX; X, 1 11 (u F X 1I) 2
0
Σ
az T
Zm Zi Zn wx Zin
N
NI
im1
We have for any E [N].
T 2, Z ur up Zin X, 1 ( Zm X, & VI((I + 1 v/
2
Then,
N
KiZy 20 DT wn wx Zin X, 1 f Zin x 11.) /
N
N
Σ
LL,=ee + 11) 2 Vdl3
3
2=1
Then,
B 1 NVdL3
Moreover,
a,4 Zm N. W. M Zn X, L X x,T1x11 12)
N
NT
We have for any E [N],
2 Zin N N W. Zin X, 11 ( Zin X 11)
T
$ VL Z,, N. N. I vi N 1 + 1 + 4
Then,
1
0. Zi N; N. w. Zin X L2Vi I (u Zin X, xT
NL
2
N
+ 11 13 $ VLd(LAd) +
N
[
Σ
x 1
No
N
3
==1
We have
N
E
Σ
x L |
V3L + NVL
N2
Then,
Bi Zil S NVLd(LAd)
55

---

Published as a conference paper at ICLR 2026
Lastly, we have
A 12) Zin X _ x; Z in wel a' (tw} Zin x 11) |/2 2
|iia| S
a3
NT.
N
Σ
Z 11) (
1 T Zin X - +vv)
x
j=1
By using the derivations in the two previous items, we have
+lv) 2 Zin X x, X wa (o (t Zin X, (.,xI11)) In)
a
NI
1
TiZy Zn T Z w w Zin X; 12)
<
NL
1
Z.Z In N Ni wh / Zm x ALUI (e X 10)]|1
2
N
+
+00(0)
Σ NL
x - 1x N N in 20
i=1
N
+ +0(00) 2 VVLd(LAd)
NL Σ (a. ビ1 Zu, N. wk
We have
N
1 7i N N in 20k 1 7mn
E
NL
N
1
E
-Lm--m- - 12 Zin N T N, 2 Zin N, N. N. Zin
N2
v.j=1
I.
1,2
$ NL v 2i Zm diag 2 2mn 2u +
NL V2
2
Zmn 2 T in
Zin 2
Nd(L A d) (D.25)
Moreover,
N Σe 3=1
11) 20 X L - 112
2 W 2 11
NL
N
=
=
NL
1
Σ
w Zin
X T 11(x; - -+vv
y=1
+ NL Σ
22,4 T _ Zin X 11) 11)(w1))
(w Zin X 11)²
We have
( 1. w Xnn 11) (w)) X, 16) S L.
Therefore,
N Σ
u (wwo) X, 11) (wm) I 12)2115 S /
NL j=1
VVL.
Moreover,
1
Σ
w X Lz(2 - +
NL
J=1
56

---

Published as a conference paper at ICLR 2026
N
Σ
Zin [x]
- +1r11) - - 2
NL j=1
N
+
wi Zin 1 v/
Σ
By b1v)"112
j=1
We have
blung Zin 1 v/
N (0)
- -b -
Clogg v VVN
Moreover,
Σ 1 [ - - ~ T
NL
N
=w Zin
{ - (
- +
NL j=1
N
[x] - ( - )1v)T)² T
1
2 w. S NL
NL
j=l
Then, for N < VL
Cloog 16 V A
B Zml S
On the other hand, we have
1 1 N
|221 S
Σ
cc.
X X, Zn wx u (2 Zm x; 11)
Vd
NL
1m, tlvill1
x
y=1
Note that
N
1 N
Σ o
Zn, (tu 11 (zz) - - = |1/00||
(t
Σ (z) - 1 |r)
N
N
==1
j=1
1 N
0 (tw Xn Le) u X, Lo t2v\11 - 1
Σ
Therefore by (D.25), we have
1
E
12 Zm S N2deLLad)
Therefore, we have
Bir|Zin] = 0 and Variance(v/Zza) S 228m
E LOWER BOUND
To prove a lower bound, we construct a Bayesian setting with the same likelihood distribution in our setting. In particular, the ground truth permutation is chosen from the set of permutation matrices:
:: = (P E fo, 13 V II is a permutation matrix}.
We describe our Bayesian setting game between Envi ronment and Lea Tearner as follows:
57

---

Published as a conference paper at ICLR 2026
At the beginning, Environment samples P. Un),(H), probability vectors without revealing them to the learner.
Learner observes + 1 channel that generates words from the {er. e2, T. ,ev
sequentially for / 1,2, N with distributions:
At every round, Environment randomly picks a channel &
Label: Channel 0 generates ~id Unif(V)
Input: Given & and pt, Channel & generates Xt,, P.P.
Noise distribution: Channel [T,] (4r generate X,, Unif(V) independent of Channel
0.
Let ((X,, ISNN be the dataset. We study the Bayes estimator with loss given the representation of the past: S = f(D, (I:N):
p arg P Zm))
(E.1)
PEH
In the following consider the empirical mean and covariance of embedded words the given data, i.e,, S = (M, E., Pi)lEs,, where
1 He -= X L2
Tn Tt and E, 1 _ X,
OR
G..
VIL
where ((g:, Gi)JEeN i.i.d. measurement noise with distributions gr N(0, La) and Gt.ij Gr.ji with GLas No, (Ittm)) i.dd for i <.
Theorem 4. The following lower bound holds:
P.P 1 I(N)
C de o Log (logv) v)
1A
+
H Loog
We information-theoreti argument to prove Theorem 4. For the proof, let (AA and (AA)) denote the entropy and conditional entropy of given C; let [A; B) = H(A) (AA and I [AA; B(C) (A(C) H (A C) denote the mutual information between random variables A and B and the conditional mutual given C, respectively. We let DRL. denote the Kullback-Leibler (KL) divergence. We start with an auxiliary statement for the proof.
Lemma 2. Let A, B, C, D be discrete random variables defined on the same probability space. The following statements hold:
In general, H(ABB,O) H(AB)). The equality satisfied if and only A C.B.
we have [(A, BIC, D) < I(A,B(C)
Set g(A,C) C) be a measurable function of (A,C). If 1 AI(S,C,)) then TIA; D) T(S; BJC, D).
Given, μ. p' E Rdd positive definite E Bdd and supp(A) Rdd, we have
DKL(N(H A, (N + +A,E)) < 5(m -') (4 - n²).
Proof. We have
P(A.C P(A C) = [EO1E log P(A H(A H(AB E log P(A
I(A, BB).
Since the mutual information is non negative, the first item follows. Moreover, since (AA,C =0 and only if A C.B. For the second item, by using the first item,
1 (A BCC, D) HBBC,, D) HBB A, D) < H(BIC) H(BIA, 1 (AA
For the third item, since S is a function of (A, C)), we have
T(A; BJC, D) T((A,S): BlC, D) H(B|C, D) H(B)AA
58

---

Published as a conference paper at ICLR 2026
DB H(B
I(S: ([IBS.C.D))I(SBBBC....)
Let f denotes the Gaussian pdf with and covariance Σ. For any a E Rd, since t - t loo is convex
pla)s
Σ p(a)s( p- a)) --a) log
[aEsupp(A) p(a)s(00 -0 a) p(a)f(x - a) log f(x p a
aEsupp(A)
Σ
1 11
aEsupp(A)
Therefore, we have
DKI + ( + +,E)) S
E DLL)DRL((( + a, Σ)I a, E))
(A)
= DKL(N(p, (H Σ)),
where the last inequality follows the invariance of KL divergence in the second line to constant shifts. The final bound follows the known formula for the KL divergence between Gaussian distributions.
The proof of Theorem given in the following:
Proof of Theorem 4. Since we assume Zi is known by the learner, we will fix in the following without explicitly conditioning thte terms on it. Note that we consider the Bayes decision rule in (F. 1) and use Fano's inequality (Scarlett Cevher, 2019) to lower bound its error probability:
T(P.; 5) log
P[P+ 4 P. Zin]
2 1
(E.2)
log H
We have -P [(P:, P.))ee)) = (p:h))) I(P:: ((m), Eclhswk(phss.)
T(P.; (1m E.))ssN((Pissn
N
Σ T(P.: (m: E.),(( En) (peleen)
Given fixed Zin- we observe that (m, E.) 1 {mm Eubbuce (pu Thereforf, by Lemma
P.. (pelenN and (e)) 1
N I(P2;S) < Σ
I(P:: (e &
Σ I(P:: (Mr. E.)p.))
E
1=1
Xi Pr. where Xi, is function of (P.,p)) Therefore, by Moreover, have P. L. (pi, E.) Lemma 2,
I(P.:S) <
I(X (Mr, E.)lp.)
1m1
We have
((Xasi (mi E.)p,, Zin)
1 V
(TR (14 E.)
(6)
<
Σ
V
DKL (PL (M Ei)l (
V2
1=
where Ih. E.) denotes the distribution of (st, Σ,) Xt,, = Po denotes Po = $ Ek=1 (m. E.)
pt
and (b) follows the convexity of KL divergence in its second argument. For k + j. by the last item of Lemma 2, have
C d
C d
C d d
DKL (Ik ( 11
zz
2112
|
0.
o
L
d L
of
L
59

---

Published as a conference paper at ICLR 2026
Therefore, have
I(P.:S) N C d
C d 03 L
0
Moreover, we can write
((P::
T(P:: (Xi)e (
N
[P.: X (pa ()SNN)
t=1
I(P.; (i.slp.))
where the first incquality follows data processing inequality, third and fourth incqualities follow the first and second items Lemma 2. We have
I(Px; Xe, Pp, ()) = (X,,, Pe. E1)
((e,, L,, P.) = log V.
log V
Therefore, we have [(P2; & log V. Finally, we have
C C d
S) S N logi ^
ox L
0
The result follows from (E.2).
F AUXILIARY STATEMENTS
F.1 GAUSSIAN MATRICES AND RELATED STATEMENTS
Lemma 3. Let ~~ Ia). We have B|||= 1 = d(d + .). (d 2h: 2).
Proof. We observe that ||2||2 x3 By using the moment formula for chi-squared distribution, have the result.
Lemma 4. Let NO, Ia) and S Rdxd be a symmetric matrix. For
P [1 Sz tr(S)) 2 [| 21| < 2e
Proof. We note that x Sz tr(S) has the same distribution with Σd (Z) - -)), where Z, Ma (o, 1). By using the
lemma, we have the result.
Proposition 4. Let E VVV be symmetric positive semidefinite matrix. Let
M = ST.
For poly(d) > V >> d, We have
PP
M
tr(S)
Ta].
max
||S|| Vd
log V, ||S|| log v)] < eepp(cclog V).
d
Proof. Without loss of generality, we can assume that S is diagonal, i.e., S = diag(s),".:, We have
tr(S)
(ziz) 1 Io).
M
Ia
d
i=1
60

---

Published as a conference paper at ICLR 2026
We have
V
d [a)) I d I
1
E
81
2 -
1=1
Moreover, for p Σ 2
E
2,
[ SEIL
1 2P.
By Proposition 15, we have 2 5 psps
E
IM
tr(S)
12
S C
CVVVVoog
11 Vd
Vog
|5 loggo
For 20 log" V, we have the result.
Proposition 5. Let RVV be a square matrix and let M In SnSS4p S For poly(d) > V >> d. We have
||5 Vd
[l]tr(M)- tr(S))
log
S exp(-cloog V).
Proof. Without loss of generality, assume that S is diagonal, i.e., S diag(s1, (8v) We have
tr(M) - (r(()) Σ 81 (11:111111) 1).
We have
422237
E
0| 1))1 s exp
) |>| s 4|s:|
d
Then,
d
[exp()(r(M)) -r(m))))rr(s)))) S exp
|^| 5 4||
We have
tr(S)) 2 log? V || |F
P F
I <exp(-clog V).
Vd
Proposition 6. Let S E XV be square matrix. For E Sd-1 and M = Si. | max
in we have
p P[[(TT Mu ir(S) Tu)
[|sym(S)] rr,)]]r,,ll1 btllt//l//yl httl//symm S) }]
d
S
Proof. Consider d Z), where g ~ N(0,, J.v). We have
Mu = / (wv )) Sg = - sym(uv" ") @ sym(S)p
By using Proposition 11, we have
Elg symmmu @ sym S)gg tr(S)u To.
Moreover,
dv g sym(uv' sym(S)g tr(()u =a A 1)
i=1
where 9i N(O, 1). By using the subexponential concentration, have the result.
61

---

Published as a conference paper at ICLR 2026
Proposition 7. For u, E RVV we have
P T 7, Zin 2i Zi u
T 10(1t d -
V-
C|| log V
VV
V d3/2)
)
S
V).
Proof.
Without loss of generality, we assume that and have a unit norm. Let
1
v1 (u w) (1v - vo T )r.
TT =
We have
v 2 Zm 1n Zi = (u 6) 7 Zin 7 Zin / (u "TT'"""" 2 Zin 2 U1
Without loss of generality, we consider v el and 01 = e2. For the second term, we write
2, = en and let 2 {zilk=s and g Vdveo( 2).
e, 2 in mn 7,T em = (1)=1 + || > 22 + 2 2ō x2
1 T
=[|=1 11 || 22 22 79
sym(z1z2 Iv-
We have
By Lemma and Proposition
5 log V
5 log V vd
log V vd
> 12 1|
S d
and | - 1/s
and
|z| x2 <
< 1 - 6exp(-cloo V).
By Proposition 11, we have
[sym(z)z2) @ Tv-2l12 S ||11||
[sym(z1 Iv-2llF 5 /VI z1 [|2 12122
tr(sym(z1 @ Jv-2) = (V 2) 22.
Therefore, by Lemma 4, we have
(V - 2) 2 22
log VV + log² Vi 10521]] V]
7 sym((zz2) @ Tv-29
S 2|
1a/l|//2122 22|12
d
d
2 eexp eexp(_clog* V).
By union bound of the precious two items, we have
P e in
Zin
2 ez S 2og V
d/2 +
V
vB]
d )
2 18exp((ccloo V). (E1)
Next, we redefine the notation: Z = { y
We write
2jn 2T,
V = ||211 + 2
222
- I1 21
V-1
(1| 112 - 1)
d
d
d
By Proposition bla, we have
-z= v la 11)=1 < log V |
VV
P 인 7z
< 1 2 exxp(cclog? V)
d
d
By using the first item above, we have
|||2| Zmn 2
V V-1
6log
V
+
S - 22exp(-clog V). (F2)
d
/3/2
The result follows (F.1) and (F.2).
62

---

Published as a conference paper at ICLR 2026
F.2 MULTINOMIAL DISTRIBUTION AND RELATED STATEMENTS
Lemma S. Let (" ,yy E Mult (N; "
For t E BY,
V
(Σ
N miel)
exp
Then, if = v. i
[V].
We have
-R
II mi(m= 1)... (n: k; + 1)
N(N 1)...N (...NN-KK- -+++))
where K == Ei= ki.
VR
By the previous item, write
E
[n]]
N
N(N-1)
V
V/
(
1 )
(V- -V2-
-
NVS
E [n]]
N + 3N(N-1) + NON ((N-
V
[nf]
N 7N(N-1) 6N 11 NON -LI--IIN--- 2NN-11N--1
V
V
- For # /. E [n? NIN-1) + 2N(N-D(N-2) MNN-1NN--)NN
V
19
V
- [2 n?)² N2 INN-D
Proof. have
Let 2j sampled from -e ,ev with "....,pp))) We have n. -Ei e, xx We
E cxp Σ t:me)
exp Σ (txz;))
(E exp (t, (1) (t)
piet)) "
(Eneet) N
im
The later statements can be derived by using z; e and taking derivatives of both sides with respect (21, ... ",,)).
Proposition 8. Let n == (mI.... ,n) ~ Mult(L, lv) and S E VVV be symmetric matrix. The following statements hold:
We have
L E[w] Sæjæιæ]] Mb-2
E(diag(n & 1 1v Sdiaggn 1v) T. Ela: Sm S.S.
We have
Ejnn T Sn] = 2L(L-1) V2
S1v + LE [ Sxn] + (((L-1) tr())+ 44-11(4-22 (044-211700 44-1)1-2111$S1 1T[SSv1 lv.
We have
E [[(n= +1v) S(n (ov)
V diag(())
L
2 Slv + 1
v2 (LTSSV) Lv 2
L(L 1) (r
(Iv [lv 2
V2
L L(L 1)
(r
(Iv-
+lv S(Iv (1v1b)s)
V2
Proof. For the first item, observe that
e, (ddiag(n) Sdiag())e:: E[nyn:]S.. (6g + L(L
L
1(L-1) 1) 1/2
Sijn
63

---

Published as a conference paper at ICLR 2026
from which the first equation follows. For the second equation,
L ej (n + Sdiaga - [Lv)le, = [B((n + (n.
- -11ss, = -
L
و.g
Sij.
V
For the second item, we have
[B|nn' Snl); Σ
Syk Ein
jk
LLL -1)(1 - 2)
L(L ((L-1)((4 -
L(L - 1)
Σ Szk) +
+
(
Sik +
SEX)
V3
V3
V2
iF
if
L:
31((1 - 1)
T(I. 1)(I, -2) 1)(1.-2)
+
+
Sii
V
V2
V3
L
Sn + L(L - 1) tr(S)
2L(L 1)
Σ
Skk=
L(L 1)(L - 2)
Syk)-
V
V3
jk
For the third item, we have (n - 1-t11 1v) = Σ (x; convenience, let
1 v (v) in distribution. For notational
8ij : (x. - [v)] S(a; V lv).
1
Then,
L
EE [(n - -vv) s(n - ( = Σ
B[ayss]
i.A.
By independence, only (i, 1) where each index occur even times contribute. The possible cases are:
All four indices equal (i j k 1): There are L many terms here with contribution
E[ 2
V diag (Iv v Iv 1v1l) V |
1
1
1 diag(S)=
2 S1 + 12 (1T.51))r1l(
Two distinct indices, both pairs diagonal and k. and i + k): There are L(L 1) many terms here with contribution
lv V2 tr (Iv (Iv- 1 - }s)
Su = E[su] 2
Two distinct indices, paired off-diagonal: (i = k and and + j): There are L(L 1) many terms here with contribution
4 (v) - [ S(x2 V [v) - 1 - 1
2 B
E
Lr
[(
1 tr (Iv v 1v (oduv-11,11 1v1))s)
v2
v
Proposition 9. Let V > L. There exists a universal C > 0 such that the following holds:
. .Let -Letmij mij = (1- (1--)) 11)1 i= + &. For K > 0 and 2 log V,
-E +I| 1r - maj < pa + +
64

---

Published as a conference paper at ICLR 2026
P| 1y - mush 2 CK2 log V VVAL $ V Vx
For K > and 2 Vo
E (X. [ - sc V
1 {ntL
NL -i (x
+1 - (1v1)) - - (1=+vv1
CK
lox-V + VNV
log" V. N
(1+ log21
S
vVAL
Proof. Let va i.idd copies of æ₁ We note that X,1 = Et. va in distribution. For = j. we have
11,,,,11111
1
2
Σ
10.=00 =1+ (L 1) 2
a. V)
L 1stersh
Define
Y =
mm mm mu +) and Fi = oVi,,..
Given that
To 1m ~ Binomial(1- - 1, ) = B| s SCWDy v P), p log V. (F.3)
where we used Corollary 3. As for the quadratic variation
L QL = Σ
1 =
3 -
2 For p 2 log V, by using triangle inequality,
1
$1
V
2
(a) 1
V Σ
(1 - 1
1-1 Σ
Øir
1-1 Σ
Bir
pp lv
1=1
1=V+1
r=1
p-
1 L
Cpp V Σ
I.?
= Cp V
where we used Hölder's 's inequality in (a) and Corollary in (b). By using (F.3) and Proposition 15, for 2 log V, we have
L
E
Yu
1
VV +r²).
L
1=1
By using = log V, we have
> CeK² log² V VVAL
Yv
<
L
VK
65

---

Published as a conference paper at ICLR 2026
Hence, we have the i case. For i + j. we have
1 1 4 v
1
,r
V
1=1
We redefine the martingale difference sequence
V 2
Conditioned on X,, we have (Y,, 3r are i..d and
[| = 0 and -
- lv
1
-
By Proposition 15, for p 2 log V, we have
1
vP + pz
+ 0²
E T 1=1
VLV
By using p log V, we have
P +|
-
4
CK2 Co" V
1
1
VVVL
VK
For the second item, we define
[x 1 -
1
1 (vrl
Y
V (Iv v(Iv
and QN = N E Y
We have
VL [X]- - (lv 11 21 [XT 11))11 =
QN NE
21
= -E (1 1 7
1 (XT t1 -x
[11 (x] - (1v11) (1 -1-
(1 +
+ NNE
(+ ( (lv12)*
1 - - [
(0) CN Iv
2 QN.
where we use Proposition 14 in (c). Therefore, we have ||Q. 2|2 S CN
V Moreover, we observe that
(x. 4 Lv 17] 11 1 2
Y la <
By using the first item,
200 + 1 (x. 1 lv ( 1,1 s1+0( p² + ).
VV
Therefore, by using Proposition 15, we have
l
N
|
S c (VoVlosV
I V NV + Vo
ps + 2T))
E
N
vv
1=1
By using p log V. we have
p
1
n
Y
>
CKKlogg v((
NV
1
+
N
1
11+
VV
log²
)) s VK
1
66

---

Published as a conference paper at ICLR 2026
Proposition 10. We consider S1. S2 and Sy defined in (B.5), (B.6) and (B.7) in the regime V
N > V and L x V E E (0, 1). For any K > and V 2 (1),(1), the following holds:
1. We have
1 1/V [1r(())
2 3 +
CK log" V L2NvV
E²
2
IP
||Sil|2
I
1,2V2
VK
We have
P (r((2) - (1 }
1
1 -1
C
K log V
or || C
K log" V
$
4
L2N
LNV
NLV
VK
3. We have
CK* log V 1 Iv 3 S3 N NVV
CK CK V NVV
A. 1vll s
1
7K
Proof. We define n, n,= K( S N 2j = eill We observe that
tr(S)) = (1 /
2
1
Σ
m
v2r and || 2| S sup
nf L
L2N2
i=1
By using Proposition 9 and Corollary 3, we have the first item. For the second item, we write
S. = (1 - +)
Σ
N
[
1
I
L
v
j=1
2
Σ
- (NN] - +1rl (NN - -+111--))
+ T2N2
jek
=: S2 S22
We will analyze S21 and S22 separately.
Bounds for S21: We have
1 L
tr(S21) - (1 _ )2
L
1 L N
= V T2N2
V DI
v I
1
1-1 1 1-
112
=Yı
We have E[ 3 and by Proposition 9,
Cp²
B[J 17
VV WAL
Therefore, by Proposition 15,
E (r() 1 1
C
PN
+ +NN Dd
2
LN
(
V
VVAL
By using p = log V. we have
([(((Sa)) -
1 1 1
1--1
K log 3 V LNVNV
1 VK
(R4)
V T2N
Moreover, by Proposition 9, we have
P S21 -(1 } L 1-13
12N V (Iv - 1r171/1
> C
Klog" V LN
1/44//0 N
1
(1+
log" V
VNV
1
S
(F.5)
VK'
67

---

Published as a conference paper at ICLR 2026
Bounds for S22:
We write
N
(r()
2
1ay=m.
2)1Z-1(N
-bIvIZ-1 -
(NN -b
N
k= j=1
N
22.
2=
Let Fi = o (Ni:k) and 22 0. We have
E E [x 4(I. 4(1-1) 1)
k-
2
E
1 -ay= -) - 11
Fk-1
LNN4 V
2
j=1
= =1 -1- )) 4(T. 1)
Σ [ 1 11
k-1
L V2
j=l
Then,
N
4(I- 1) N 4-1
ΣΣ INN - 1
QN
=
E F = V L k=
k=
1, -(1--) 1) N-1
= (1 v)
L-N4V2 Σ (N- - - I
Then, for p 2 log V.
E
[[ $J$$ <
5
Σ
N
E
|
1
11
2 (a)
5L--3 N
ΣB E
I 1
L3NaV2
LoN822
k=1
5p
(F66
T2N2V2
where we used Hölder's inequality in (a) and Corollary 3 in (b). By using Proposition 15, we show the following:
To bound EE 12, for p 2 log V. by using the conditional independence of (ι,) k-1
we write
E
Y2 2,
1
N 1:k1 2x
(NT -I -
LNZ
VV
Cpk
- -1.II~))2-))1
I (NN
LN2
Therefore,
C
VPVk
(NT - (NT 1 -I
E 1 5
+ pkž) E
LN2
VV
I
5
C
VPVk
+ pkž
E
N
1T,,,,
1 17s
LNZ
VV
Cp²
1
pkb
(F7)
LN
VV
A
where used Proposition in (c).
68

---

Published as a conference paper at ICLR 2026
Then by using (F.6) and (F.7), we have for p log V
(1)tr(S22)11))
p
+ N'
Cp
s NLV
S
L V VV
NLV
Then, we have
P [tr(S22)1>
CK log² V
$
1
(F.8)
NLV
VK
To bound ||22ll2, we define
k-1
V )sym (INJ I - 1
Y =
===
We have
E
2 Σ
- (N - 1 -
B
x INT. 1-- Fr
+
2
Σ
E
INT
1 -
x (NJ 1 1 - -
21
[1 13, Iv
[N]
V2
k-1
2L
- - 1 I
Σ [NT
V
Therefore, we have
N
QN
Fx-1
N-1
2L
1 -11 Tv
(N -
V2
k=1
N
N-1
(N ( -I - --
2NL N-1
(NT 1 -1
V2
k=1
N2
1
N-1
Σ (NT
- -
V
LN
R=1
Then,
2NT.
N-1
( 1)
$7
E QI
<
E
(NT
1v
11
1/2
k=
69

---

Published as a conference paper at ICLR 2026
N2
E
1
Σ (N₂
N-1
1
-
1
.)
V
[[L
- 1)
2
k=1
2N
E
1
(NI
--r
1
[2]
V2
N2
E
1
N-1
Σ N v
(N1 1 -
8/
V
[ ([--) 1)
k=1
CN
p²
CL2N2
1
P
P
(1+
p² )
1
+
(V
V2
VV VVL
V
NV
N
CN2L2
p +
p²
1
V2
N/V
VVVL
To bound
BI]lY
[22 we observe that
We have
E
[(11) )) (NJ (
11 - ) (px.) No.
IINT - + - - -4 - -TT
Moreover,
(N,7 - (N
E
[1((1
-
27k: N16
$2
(12;=0x -
23 " - [|NN] -11
s CVE - -1v ill (VE+ +D) (1
)
p²
VV VVAL)
By Proposition 15, we have for p= polog V,
E|| BI CI - 1 - (VV
Lk VPV
+ +DDVKA))
which implies
k ||| s CLO (vVV EV +
Therefore, for p = log V. we have
E 221 S C Vp
D5/2
pd
NLV
LNVVV
LN
Therefore, we have
K3/2 10g V
P 1|2122212 C
(F99
NLV By (F:4), (F.5), (F88), and (F.9), we have the second item. For the last item, we have
VK
S3
1 N
22 Mral
||
N
(z,
- lv) |r)
22
N
J=1
By Proposition 9,
Σ C
N
1v)
1v)
1
CK2 log 2
5
N 1=1
N
N NVV
VK
The displayed equation implies the third item.
70

---

Published as a conference paper at ICLR 2026
G MISCELLANEOUS
Proposition 11. Let A gaxd and E RVV. Let M = @ B. We have
M 2 = || and |M =
and tr(M) = tr(A)tr(B).
Proof. The Frobenius norm and trace straightforward. For the & norm, let A Di daui'
and B =: プド= of " We have
d v
d v (4,, = ΣΣ
M Σ E
o, (u, @ aūj)(, ,,)
i=1==1
,=
For any (i.5) # (V, we have
(n4 @ (uy @ uUy)
= (v, @ ( @ .yy) = 0.
Therefore,
M 2 = maxo;o = max oi max oz.
y
(
Proposition 12. Let (0, Ia) and Ph RRd [0, 2o) denotes degree i polynomial which takes nonnegative values. For 2 1, we have
B[)P<()))1 < (8(p 1j) [[Px(=)].
Proof. By hypercontractivity. it is sufficient to prove that (L)
EP(=z) S 85. We have
E| (=)))(=))) < (x)) Px)) < 2 Px()) [ (x)
which proves the result.
Proposition 13. Let and [O, Ia). For L > and u. u E S we have
[k/2]
E Hos
(VI""
1 w u))
Hes
1
w
T
v)
k!
(2i
1)!!
(L 1) (us, 0)k-22
VI
2
2i
Proof. For a E R, we have
[k/2]
Hey (ax)
Σ
k!
(a² - -1)aa
Hex 2
k
2iil(k
24)!
1=0
Therefore, for a = 1/VL, we have
E
Has
w
Tu)
Hee
(1
[
VL
VL
|k/2|
k!
E
Σ
(a² 1) k
2i H en- (w Tw)
2iil(k: - 2i}!
1=0
k/2
Σ
2' - 21}! (a² -1)'a Hen- (w Tw)
k!
x
1=0
[k/2|
Σ
k!
2
(a² 2) n - 28) 2k: 21)! (u, x)*
2kil(k
- 21)!
1=0
1k/2
k!
(2i
1)!!
(k)
(L 1)26
(u, (u, x)
2i!!
2i
1=0
71

---

Published as a conference paper at ICLR 2026
Proposition 14. Let Z be random variable and X be d x d symmetric matrix valued random matrix. We have
B[ 3 B|ZX]
= E|Z
B
Proof. We observe that
Zia X
2 La
ZX X
E La XX
EiZ
ZIa
X
00
00. 0.
(G.1)
ZX
X
By (Ben Arous al., 2025, Proposition 24), we know that (G.1) is equivalent
E[
E|
Since X - vX is monotone in matrix order, have the result.
G.1 ROSENTHAL BURKHOLDER INEQUALITY AND COROLLARIES
We will rely on the following inequality:
Proposition 15 (Peng et al., 2025, Theorem 2.1)). Let {M& LN be a d-dimensional symmetric
matrix valued martingale adapted to the filtration {Fu) Let YR == M M be its corre- sponding difference sequence and the quadratic variation is defined as
QN
F
For any p 2 2, suppose
E
1Q
< and
sup E kE[N]
|Y=11]]
< 00.
Then it holds that
| 12
S c (VpVloge
|
#7
+ (p logd)g sup
Y 111 +).
kE(N)
We have the following corollaries:
Corollary 3. The following statements holds for general L, V 0:
1. For X Binomial(L, }), we have
L V + pl ( E E| - kglpl' SC
2. Let c (cp,... ,cv) ~ Multinomial(L,, F1v). For p 2 1, we have
E||c|| C (+))
pL Y/2
L
p
3. By following the notation in the second item,
V >> L, we have for 2 2r + 1,
2e
log f--1
L
log L-2
P|||c|0 2 log I.] <
log
If > V, we have
eL V
|eello0
$
2 L/V
72

---

Published as a conference paper at ICLR 2026
Proof. The first two items are direct consequence of Proposition 15. For the third item, using 2k S (w 1)-- ++1) and linearity of expectation
E Cau 1).... (Cer - k 1)]
P|llello 2 kj S
Plee 2 k] S
k!
10
w
L( 1) 1)... (L k + 1)
k!Vk-1
For V > L, by choosing & llog [], the result follows. For > V, by choosing i: 1 J. the result follows.
0
73


Reviews:
[{"review_id": "9FwodtvU4Z", "reviewer": "Reviewer_iRpb", "summary": "This paper presents a theoretical framework for understanding how transformers acquire factual recall when trained on finite data with non-orthogonal embeddings. Focusing on single-layer attention and attention-plus-MLP architectures, the authors derive precise conditions under which gradient descent can identify and retrieve informative tokens within a sequence. They show that learning efficiency follows a multiplicative relationship among vocabulary size, sequence length, embedding dimension, MLP width, and sample size. Unlike prior analyses, this work establishes learnability even when the vocabulary size is larger than the embedding dimension. Empirical results corroborate the theoretical scaling laws and demonstrate relevance to practical settings.", "questions": "- The clarity of some assumptions needs to be improved. For example, what do $\\eta \\approx 0$ and $\\gamma \\gg 0$ specifically mean? Why are these conditions not expressed using big-O notation?\n- Similarly, it may be non-standard to write $\\Omega(V \\log V) \\le N$ and $V \\ge \\Omega(1)$. When using big-O or related asymptotic notations, should readers assume that these are taken with respect to $V$?\n- In addition, how did the authors name or interpret each term in Eq. (6)? It is recommended to provide readers with more intuition or explanation for these terms to enhance understanding.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This work provides a clear problem formulation and rigorous proofs supporting its theoretical claims. The technical results are overall a good contribution . Beyond theory, it offers numerical experiments that closely match the analytical predictions, showing the precision and tightness of the derived results. The findings also yield practical insights into transformer training, key architectural parameter choice, the importance of the MLP component.", "weaknesses": "- The clarity of some parts needs improvement (see Questions). It is recommended to include a brief intuition or proof sketch to better convey the key insights underlying the proof. The current presentation may be less accessible to a broader ML audience.\n-  The gap between theory and practice remains substantial. The definition of the factual recall task requires further clarification. In the current setup, the model identifies a marked token within a sequence and maps it to another token through a ground-truth permutation $\\Pi^*$. However, this formulation appears disconnected from practical scenarios, particularly regarding how the key information is represented or “marked” in the token features. It also seems misaligned with the setting described in the cited work [1].\n- As for the theoretical aspects, it is recommended to include a quantitative summary of prior works. For example, specifying what learnability results/bounds have been established in those studies that analyze orthogonal embeddings.\n\n[1] Nichani et al., Understanding Factual Recall in Transformers via Associative Memories", "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": "r8bztQ3WPJ", "reviewer": "Reviewer_BKC3", "summary": "This work presents a theoretical analysis of the training dynamics of transformers on a factual recall task. Specifically, it considers two model architectures: (1) a single attention-only layer and (2) one attention layer followed by an MLP layer. The study examines how these transformers learn to map an output token to an informative position within a sequence. The authors quantify model accuracy after three steps of gradient descent and establish learnability guarantees as functions of vocabulary size, sequence length, model size, and the number of training samples. From an empirical perspective, the results provide practical guidance on selecting the hidden dimension and sample size, while also elucidating the role of the MLP component in recall tasks.", "questions": "It appears that the considered factual recall task can be learned exactly after just three steps of gradient descent, provided that the embedding dimension or the sample size is sufficiently large. Could the authors offer some intuition for why this occurs? For example, even in the orthogonal regime, $d > V$, why can the model achieve perfect recall accuracy with only three-step GD?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "+ The paper is clearly written, with precise notations and formal language, and includes detailed technical assumptions.\n\n+ The theoretical results appear sound and solid. To the best of my knowledge, this work provides the first theoretical analysis of training transformers to perform recall tasks, particularly in the regime where the embedding dimension can be smaller than the vocabulary size.\n\n+ The theoretical findings are further validated through numerical experiments conducted in controlled settings, which support the correctness of the theory and demonstrate its potential generalization to more practical scenarios.", "weaknesses": "- Some assumptions may require further justification. For instance, in Assumption 1, it is not entirely clear why it is necessary to set $L = V^c, c \\in (0, 1)$? Under this assumption, the sequence length is necessarily smaller than the vocabulary size, and the motivation for this specific scaling is not well explained.\n\n- It is also suggested that the authors present their results in a more comparative manner and provide readers with additional background context. My understanding is that the main contribution of this work lies in analyzing the learning dynamics of transformers on factual recall tasks in the regime where $d < V$. However, it is somewhat difficult to determine whether the theoretical settings in this paper are consistent with prior works. For example, regarding the choice of model architectures, learning tasks, training algorithms, assumptions, and the final results. Without such contextualization, it becomes challenging for the general audience to assess the true contribution of the paper and identify whether the new result stems from new assumptions or from novel proof techniques.", "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": "htmAV68wr3", "reviewer": "Reviewer_nvK4", "summary": "This paper investigates the gradient-based learning of the Transformer architecture for identifying informative tokens within the context. The theoretical result provides conditions on hyperparameters for successful learning, and the empirical result validates this.", "questions": "1. What are the useful implications of the theoretical results? Specifically, what is the takeaway for practitioners who train Transformer language models? What is the main contribution to the deep learning theory or language model theory that is not available from prior works?\n\n2. Does adding an MLP layer on top of the attention layer hurt? It seems that this makes the theoretical condition for successful learning more difficult to satisfy.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper addresses an important problem of the learning dynamics for Transformer models.\n\n2. The paper complements its theoretical predictions through experiments.", "weaknesses": "1. The results in the paper are limited to a simple toy model (easy task, single layer, 3 simplified gradient steps on a fixed batch). Although additional experiments exist beyond this toy setting, it remains unclear how the results will generalize to more realistic settings.\n\n2. The setting assumes that the input contains a trigger embedding that marks the informative token, which makes this toy setting less practical. It is difficult to imagine a realistic task where this assumption holds.", "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": "A6gCbBooht", "reviewer": "Reviewer_yVha", "summary": "This work establishes a learnability statement on a specific factual recall task. Within a sequence, there is a token position which is marked, and the goal is to retrieve this token with the attention mechanism and then use the MLP to recall the correct label of the input, which is assigned randomly. This task is analysed in the first few steps of learning dynamics in the large limit of tokens, and under various (scaling) parameters such as embedding dimension, sample complexity and hidden layer size. The authors provide both a learnability criterion under these hyperparameters and a statistical lower bound based on the accessibility of attention queries that holds for more general algorithms (but is not very strong or specific). \nThe main takeaway is that denoising the correct item and memorisation are more difficult for longer sequences, where a larger embedding size and/or more samples are needed to reach a similar performance as for small sequences.", "questions": "- Do you have any insights in your model on how strictly non-orthogonal vectors would do, and how close to non-orthogonality your embeddings actually are? Since they are random and not learned, I would expect that they are still \"quite\" orthogonal. What would change in your experiments if you could actually learn them too, and therefore impose better-than-random, learned, possibly non-orthogonal (or more orthogonal...) structure? Do you have an intuition for that? I think an empirical ablation like this would help contextualise the result a bit better.\n\n\nSome suggestions/clarifications on intuition:\n- In the description of the problem setting, you say \"can then recall the correct label via an associative memory mechanism.\" (L.126). Later on, L.144, you say the goal is to \"learn the target function (permutation)\". I understand that the VxV matrix is essentially the random input class->output class mapping, but it took me a moment to see that this is where the memorisation comes in; you might want to make this more explicit.\n- Do you have an intuition about how increasing the sequence length L and the vocabulary length makes the task harder? Longer lengths make detection harder, and larger V memorisation - but if intuitively the two functionalities are located in different models (L.125), why do they both scale linearly in all of the other parameters, and not independently?\n- Do you observe that $W_{KQ}$ learns $z_{trig}$?\n- L.219 \"during the second gradient step,\" -> \"after the first step\" (?) (the object reference is a bit ambiguous)\n- What do the different alpha (transparency) values of the regressed information mean? It is there in almost all figures.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The presentation and clarity of the writing make it easy to follow the work.\n- The dynamical analysis of the first three steps already gives a good insight into the general behaviour and the type of solution that is learned to fulfil the recall task.\n- The illustration of Thm1 in terms of comparing the signal with the noise at various network components is helpful and intuitive.\n- The empirical findings validate the theoretical results closely, even when the theory does not hold anymore for longer time dynamics.\n- This work offers plenty of interesting directions to follow up on for the community.", "weaknesses": "- The model is not very close to large-scale models, but it has some resemblance to factual recall tasks and comes with the advantage of being analyzable.\n- see questions", "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)