Pretrain–Test Task Alignment Governs Generalization in In-Context Learning
Authors: Mary Letey
Jacob A Zavatone-Veth
Yue M. Lu
Cengiz Pehlevan
Venue: ICLR 2026 Poster (2026)
Keywords: In-Context Learning
Task Alignment
Spectral Bias
Pretraining
Linear Attention
Primary area: learning theory

Abstract: In-context learning (ICL) is a central capability of Transformer models, but the structures in data that enable its emergence and govern its robustness remain poorly understood. In this work, we study how the structure of pretraining tasks governs generalization in ICL. Using a solvable model for ICL of linear regression by linear attention, we derive an exact expression for ICL generalization error in high dimensions under arbitrary pretraining–testing task covariance mismatch. This leads to a new alignment measure that quantifies how much information about the pretraining task distribution is useful for inference at test time. We show that this measure directly predicts ICL performance not only in the solvable model but also in nonlinear Transformers. Our analysis further reveals a tradeoff between specialization and generalization in ICL: depending on task distribution alignment, increasing pretraining task diversity can either improve or harm test performance. Together, these results identify train-test task alignment as a key determinant of generalization in ICL.

Published as a conference paper at ICLR 2026
PRETRAIN-TEST TASK ALIGNMENT GOVERNS GENERALIZATION IN IN-CONTEXT LEARNING
Mary Letey, Jacob A. Zavatone-Vetts Veth, Yue M. Lu, and Cengiz Pehlevan
Harvard University
(maryletey, jwwvatonevethheeeeeee...... p www.ahaaaaaaeeeee (yuelu,
aaaaaass.aaaaaa harvard.edu
ABSTRACT
In-context learning (ICL) is central capability of Transformer models, but the structures in data that enable its emergence and govern its robustness remain poorly understood. In this work, we study how the structure of pretraining tasks governs generalization in ICL. Using solvable model for ICL of linear regres- sion by linear attention, we derive an exact expression for ICL generalization error in high dimensions under arbitrary pretraining-test -esting task covariance mismatch. new alignment measure that quantifies how much information
This leads about the pretraining task distribution is useful for inference at test time. We show that this measure directly predicts ICI. performance not only the solvable model but also in nonlinear Transformers. Our analysis further reveals tradcoff between specialization and generalization in ICL: depending on task distribution alignment, increasing pretraining task diversity either improve harm test performance. Together, these results identify train-test task alignment as a key determinant of generalization in ICL.
INTRODUCTION
Pre-training on simple next-token prediction enables Transformer models to acquire a remarkably broad array of capabilities, from language translation code generation and mathematical reasoning (Achiam et al., 2023; Anthropic, 2024; DeepSeek-Al AI et al., 2025; Vaswani et al., 2017). Among the emergent abilities that enable Transformers to flexibly execute a myriad of tasks, their capacity for in-context learning (ICL) is particularly striking, as it allows for test-time task execution without task-specific pretraining (Von Oswald et al., 2023; Wei et al., 2022). In other words, ICL reflects the ability to emergently meta-learn learning algorithm during pretraining, which is then applied to learn from data within context at test time (Akyürek et al., 2023; Raventós et al., 2023; Zhang al., 2024a).
For ICL to effective, the tasks encountered at test time must not be totally unrelated to those encountered during pretraining,
there is free lunch. Though substantial theoretical attention has been devoted to the question of why and how ICL emerges and how well the resulting algorithms perform, this key issue of how pretraining tasks should be selected to enable ICL in the real, test- time world remains undercxplored (Lu et al., 2025; Zhang et al., 2024a;b). This motivates the central question of our work:
Central Question
How does mismatch between the structure of tasks seen in pretraining and the structure of tasks faced at test time affect the ability in-context learners to generalize?
Here, investigate how pretrain-tess test task alignment affects generalization in simple model set- ting: ICL of linear regression. Our main contributions are as follows:
We give a precise mathematical analysis of the performance a simplified linear Attention mod- ule learning to do linear regression in-context. We model pretraining and test task distributions

---

Published as a conference paper at ICLR 2026
with arbitrary covariance structure, generalizing previous works this model that assumed iden- tical task distributions (Lu et al., 2025; Zhang et al., 2024a). In this solvable model, ICL general- ization is governed by particular alignment metric between pretraining and test task distributions.
Though derived for a simplified linear model, we show that this alignment measure predicts the generalization performance of nonlinear Transformers trained to do linear regression in-context.
Finally we show in the solvable model that it is not always optimal from a generalization stand- point to pretrain on the same distribution of tasks that the model will encounter at test time, as echoed in kernel regression (Canatar et al., 2022).
In all, our work sheds light on how pretrain-test task alignment shapes the performance of in-context algorithms. It reveals how the inductive biases of Transformers can result in optimal task misalign- ment: rather than teaching to the test, different curriculum of pretraining tasks may better enable a Transformer to learn the algorithm that enables it to generalize well at test time.
1.1 RELATED WORK
Empirical studies of ICL. Empirical work has shown that LLMs can learn diverse tasks from ex- amples alone, with performance improving predictably with scale (Achiam al., 2023; Anthropic, 2024; DeepSeek-AN AI et al.. 2025; Vaswani et al., 2017: Von Oswald et al., 2023; Wei et al., 2022). Several studies document how architectural components, such as attention heads or MLP layers, are recruited during training to implement ICL. (Kratsios & Furuya, 2025; Reddy, 2024; Tong & Pehle- van, 2025; Zhang et al., 2024b). Various works have also focused on understanding what algorithms transformers can learn to perform, including gradient descent, Bayesian inference, and compression. (Ahn et al., 2023; Cole et al., 2025; Elmoznino et al., 2025; Lee et al., 2025; Liu et al., 2024; Ma- hankali et al., 2023; McCracken et al., 2025; Shen et al., 2025; Singh et al., 2023; Wurgaft et al., 2025: Zhang et al., 2023). Others have explored the role of task diversity in shaping generalization, showing that diverse pretraining induces transitions from memorization generalization (Raventós et al., 2023). The specific effect of data structure and its role in ICL emergence has also been studied empirically, notably by Chan et al. (2022) highlighting the importance of anisotropic data for the emergence of ICL abilities. Our work provides theoretical counterpart to these empirical investigations, as we model the generalization effects of structured task distributions.
Theoretical studies of ICL. Theoretical work has flourished in simplified Transformer models, particularly with linear or kernelized attention. A number of studies show that these architectures can implement classical learning algorithms, including kernel regression, ridge regression, or gra- dient descent, purely from in-context tokens (Akyürek et al., 2023; Bai et al., 2023; Li et al., 2023; Von Oswald al., 2023; Zhang al., 2024b). These insights have advanced our mechanistic un- derstanding of ICL as algorithm emulation. However, most of these works make restrictive assump- tions. Commonly, data is drawn from isotropic Gaussians, train and test distributions are matched, and generalization is studied only in the infinite-sample or population limit. Even recent studies of sample ICL retain these simplifying assumptions i.e. without full generality of training and testing task distribution (Fu et al., 2023; Li et al., 2024; Lu et al., 2025; Zhang et al., 2024b). A notable exception is the work of Goddard et al. (2025) that studies tasks sampled from separate por- tions of spherical task manifold. Our work advances this "task generalization" frontier by deriving an exact expression for the ICL generalization error in the presence of arbitrary task covariances and finite-samen sample regimes. This allows us to explore how strk-utructure mismatch affects generalization, a setting largely absent from prior theoretical models.
Train- Test task alignment in other settings. Outside of ICL, the idea of train-tess task alignment in linear regression has been studied in the context of out- oito---distribution generalization under tar- get vector and covariate shifts. particular, past works have studied which measures of alignment between train and test feature covariance matrices determine out- outso--distribution generalization in high-dimensional ridge regression (Atanasov et al., 2025; Canatar et al., 202 Patil et al., 2024; Tripuraneni et al., 2021). These works build on substantial body of results showing how the gen- eralization performance of ridge regression is determined by the structure of the training feature covariance -rriiiaceeprriccipall directions with larger population variance are learned first-and and the align- ment of the task vector with those high-variance variance directions (Advani et al., 2020; Atanasov et al., 2024; Canatar et al., 2021a; Dobriban & Wager, 2015; Hastie et al., 2022). Our work brings this
2

---

Published as a conference paper at ICLR 2026
perspective into the ICL setting by analyzing the effects of the pretraining covariance, and pretrain- test misalignment on generalization. We show that even in simple linear regression settings, struc- ture mismatch induces rich, nontrivial behavior, reinforcing the broader principle that distributional alignment between pretraining and test-time data is a key driver of generalization.
2 MODEL SETUP
We begin by setting the solvable model used to derive the results of this work. This setup builds on previous works that apply a reduced- reduced-parameter version of linear attention linear regression ICL (Lu et al., 2025; Wang et al., 2020; Wu et al., 2024; Zhang et al., 2024a).
ICL of linear regression. We consider an in-context regression task: the input to the model is sequence of the form {=r, y1, 22, 12. By, ye, 12+1), and the required output is the matching VE+1 corresponding to ø4+1- This input called a context, and l the context length. We consider the case that the relationship between and y approximately linear:
y/ ( w)
for noise and task vector w. Thus, the model needs to form an estimate of u using {a V/, 22 3 ) and then apply to 1 to estimate y/+1-
Pretraining data. The pretraining data batch will contain sample sequences of the above form, i... for I,..,, n, the 1th sample sequence {I1, 22, 3/, ........ Le, ye, 2t+1} related by the approximate linear mapping from (1), s/ (ae) wh0 + e.", where now wh the task vector corresponding to the f-th sample context.
We will sample the pretraining batch as follows: For i 1,...,, and μ = I.....n.
x! miid (O,, Ia/d), whth munit (t1 ... ... &
of Miid N/O,,) p);
where 1j miid N(0, Cirain) for j = 1,...,k. k.
The parameter here is called task diversity. Note that < n, the pretraining batch contains some tasks repeated across the contexts. In this way, control both the amount i of actually unique tasks seen during pretraining, as well as the structure the task distribution using Ctrain- This distinguishes tasks (Fu et al., 2023; Li et al., 2024; Lu et al., 2025; Raventós et al., 2023).
setup from previous studies which typically focus isotropic or matched
Linear attention. We will study the performance of the linear self-attention block (Wang et al., 2020) on this in-context regression task. The input to the linear self-attention model is an embedding matrix Z made up our context sequence. Here, following the convention of Wang et al. (2020); Wu et al. (2024); Zhang et al. (2024a), chose to embed {x1, v1, 22, y2. ........ Lt, ye, 2++1}
7
I1
22
Dg
Ø(+1 0
((((+))((+11)
12
Vf
where 0 in the lower-right corner is placeholder token for the wr-1 wish to predict. The model's output (Katharopoulos et al., 2020; Shen et al., 2021; Wang et al., 2020) is given by
= = Z + VZ(KZ)
(4)
for value matrix V (d++))) and key, query matrices K, Q such that K STQ E R(d+1)x((+1)
Following the positional encoding in (3), the linear attention model's prediction for 3e41 is
i A441.41-
(5)
Previous works from Zhang et al. (2024a) and Lu et al. (2025) have shown that the output 7 = A441,841 of the model can be reduced to give simpler, analytically-tracablee model. Writing the attention and value matrices as
Vi U2
U12
M
M
m22
KTQ.
(6)
m
Mi22
3

---

Published as a conference paper at ICLR 2026
the predictor expands
8 ÿ = oF mi+i U22 Mi
++1 Σ
m2l)
Σ
yi + 222 Σ 05 + Mi1
m. U21 + m2 Σ y. 221 (7)
The last two terms in (7) depend on 121. While these terms in principle could contribute to the pre- dictor, prior work argues they may be relatively weak signals for estimating w in high dimensions: 2 U21 carries task information, and y U21 provides only one-dimensional projection of a and w (Lu et al., 2025). For analytical tractability, we therefore make the Ansatz ty = O, which yields reduced model that we analyze throughout. We justify this choice in Appendix H comparing the parameter-reduu reduuced model analyzed below with full-pat parammarretter trained linear attention.
With this simplification,
rewrite the simplified model's output
Ge+1 = tr(THz)
(8)
for parameter matrix
= U22
MTT/d
m21] E (d+1))
(9)
and data matrix
Hz 20+1 Eig
Eige vil Rx((11))
(10)
Reduced model optimization. Given batch {af .f..... o ++1 V 17-1 of pretraining data (as explained above), we can find finite-sample optimal parameters by minimizing MSE loss on next-output prediction with ridge regularization, giving
I" arg min T
Σ - WW41-trr(Tie (((((")))))) + + ")t((TT"),
(1I)
where À > 0 is a regularization parameter, and H is defined by (10) for the pth context. We will focus on the minimum- norm predictor, i.e., on the limit where x - 0 (Hastic al., 2022).
Test error. We finally wish test the pretrained model on general task to see if the model can genuinely perform in-context regression. The test distribution Ptest is then
at Lid. (O, Ta/d), & Mid. N(O,p))
w lest lidd NO, Creat)
(12)
We will measure ICL performance by the average MSE error of our optimal estimator de+1 tr(I* T) over the test distribution,
(TC = Pr [(ye+1 tr(
(13)
We highlight here the generality of the test task distribution through the new matrix Ctest, allowing study the interaction of training and testing task structure.
Data parameters in the high-dimensional limit. dimension d, context length f, pretraining batch size and task diversity k. As is standard in the theory of high-dimensional regression (Advani et al., 2020; Atanasov et al., 2024; Hastie et al., 2022: Lu et al., 2025), we will consider scaling limit where all four of these parameters are taken to infinity in a way such that the following ratios remain constant:
We have introduced four data parameters: token
n
k
a.
=T.
and
=
(14)
d2
d
Considering this limit makes the model analytically tractable, but preserves interesting phenomena that are present at finite size. Going forward, we will refer to a, t, and k as the context length, batch size, and task diversity parameters, respectively.
4

---

Published as a conference paper at ICLR 2026
Pretraining task quantities. Before presenting our formula for this ICL test error, it will be help- ful to first define some lask-distribution quantities, which depend on the pretraining task covariance Curain and task diversity parameter These quantities effectively tell us how well we recon- struct Orrain from the k-sample task covariance that the model sees during pretraining,
R
(15)
JE[K]
Because the pretraining tasks (tI..... b&} are random, Rg is random matrix. However, in high dimensions, the following deterministic quantities capture the relevant information about R&: let the deterministic matrix F (2) and deterministic scalar Mn(==) be defined through the implicit equations
Fx(2) = ( -1 =Ma(s)) K
Cerain + zIa
(16)
Mx(=) (r(F)())).
(17)
Then, we have the high-dimensional equivalence (RR + =zaa)= Px(2).
As define formally in Definition of Appendix B, this equivalence holds in the sense that the (normalized) traces of (Rx + += and (x(2) multiplied against test matrices coincide in the high- dimensional limit. Here, z R4 is noise threshold parameter suppressing smaller eigenvalues of Rs.
Intuitively, (z) and (2 give us information about how much signal in Curain can be recovered after a finite number & of pretraining samples, filtered by noise level z. As K + co, Rg approaches Curain and we therefore fully recover the original distribution of tasks. We will use (2) to refer to the derivative of F. (2) with respect to 2; this gives measure of the sensitivity of this covariance recovery matrix to the noise level 2. These quantities will play key role in our discussion of task alignment.
Given this setup, we compute a formula for the ICL generalization error EICL, in terms of the data parameters a, K, the pretraining task covariance Ctrain. and the testing task covariance Crest- We present this formula and discuss its implications
following sections.
Notation. We write tr[A] tr( (A)/dd /d for the dimension-scaled trace of matrix A. We use nor- malized Frobenius inner product between two matrices, i.e. (A, B) = tr(AB (r / We abbreviate the normalized traces of the task covariances Ctrain tr [Crriin] and Cteat tr(Crest]. We denote high-dimensional equivalence (as in Appendix Definition 1) by =
ICL TEST ERROR DEPENDS ON TASK MISALIGNMENT
3.1 TASK ALIGNMENT DETERMINES GENERALIZATION IN THE SOLVABLE MODEL
We state the main result of our analysis of the simplified linear attention model, which an analytical formula for the ICL error (13) in high dimensions. Here give only informal statement of this result and focus on its implications; the formal statement is given by (C.7) in Appendix C.
Result: High dimensional formula for the ICL error E(T*
The ICL test error (13) of the simplified linear attention model set up above given by EIC1(1") escalar AArain Otest) + Emiselign (Cirain: Creat) eICL EteL(Crrsing Creat) (18)
(
in the high-dimensional limit, where
misuligu (Otrainy Crest) = (Ctestt IC)
(19)
measures the alignment between Ctest and
K k= q= (g) -o))
(20)
5

---

Published as a conference paper at ICLR 2026
Appearing in this formula are effective ridge variable À and effective noise variable given by the solution of the joint equations
MM (a) =1- for 1, == 0 for T
(21)
a= (p + Carain)/ + X.
(22)
where M,(()) is determined self-consistently as in (17). Finally q is pretraining-errov term given by (C.8) and escatar Atrain Crest.) is given by (C. note that both y and Escalar only depend on Crest through its trace Crest, and Escalar only depends Ctrain through its spectrum Atrain-
We first support this result through Figure 1, which shows agreement between our theory curve formula EICL eICI(Cusaie (Curann: Crest) and numerical simulations of MSE error (I') given by (1I) and the test distribution (12). By comparing the left (EICL) and right (Cmissalign) panels of Figure we that the behavior of eICl is highly dependent on the emisalign term: for well-aligned training and testing distributions, low and decreasing emisalign in k immediately leads to a monotonic decrease in etcr. in K: for worse-aligned training and testing distributions, eTCL. can be nonmonotonic or even monotonically increasing in k. This is intriguing. as one would naively expect additional task samples to improve in-context performance, but this is not true in general: whether additional task samples are helpful for ICL depends on the alignment of the pretraining and testing distributions.
Crest Crrain
Idx 1/120
Idx 79/120
idx 99/120
109/120
Idx 120/120
CoitallennCoo Coraln Chust)
0.8
30.6
0.4
0.2
0.0
0.5
1.0
1.5
2.0 0.0 0.5
1.0
155
2.0
K k/d
k/d
Figure 1: Theoretical etcr (left panel) and emisalign (right panel) curves plotted against numerical simulations of E(T*) computed directly from sampled data. We choose Cursin with uniform cigenvalue distribution:
Ctrain x diag((d , - 1, ... 1]) such that tr(Ciesin) Crest Ctrain (red curves) with testing on single task directions, i.e., the "idx i/d" labels
d. We compare
correspond to rank-1 test covariances Ci = e spiked at index i. In this way, Clest captures the strongest task direction of Cerain and Clest captures the weakest. Purameters: d = 120, a = 2, 4, p 0.01. Shading represents =std of numerical simulations. We calculate the simulation values of emisaligo in the right panel by subtracting Escalar from the MSE simulation values E E(I*).
To gain intuition for this misalignment error (19), useful analogy is to consider the simplest matrix "misalignment" measure (Crest Or Firstly, this measure captures misalignment that can occur from misalignment of the eigenvectors of Crest and Ctrain: see Corollary 4.1 and its proof for further discussion. Going beyond eigenvectors, because the eigenvalues of C are ordered opposite to the eigenvalues Cerain- this measure effectively captures the relative strength of signal directions between Ctest and Otrain: if Ctrain and Crest share the same eigenvalues, alignment will be max- imized when these eigenvalues appear in the same order, and minimized when these eigenvalues appear in opposing orders. This is precisely why (Crest Crr1 can be interpreted as misalignment measure, as even this very simple Ansatz shows how misalignment arises from mismatches in the relative weighting of signal directions between train and test tasks. This relative strength argument holds for our alignment measure, where instead of (Crest Crrli)) we use (Ctest Here K obviously depends Cusain they share the same eigenvectors and importantly, has the same property as I that its eigenvalues are oppositely ordered to the eigenvalues of Ctrain (we prove this claim in Appendix D for > 1 and study it numerically for < 1; conjecture that it holds for all 7).

---

Published as a conference paper at ICLR 2026
However, this ceneenallueorderring argument is incomplete for the ICL setting. It assumes that Curana can be fully learned, but with finite context length, finite task diversity, and label noise, this cannot be the case. Thus, the model cannot access the full covariance structure, only a partial version, whose resolution depends on both sampling, and noise. This is precisely what the resolvent terms (-c()) and (.).) in K can capture, as explained in the setup. This is analogous to the alignment measures that emerge in ordinary ridge regression under covariate shift, which capture how much of the training feature covariance can be resolved from finitely many samples (Atanasov et al., 2025; Canatar et al., 2021b; Patil et al., 2024; Tripuraneni et al., 2021). Furthermore, the effective noise o must play a key role in alignment, as the model does not know a priori that this is an "in-context" learning problem, and must learn to decouple the tokens from the task for each context in order to extract the task information implicit in that context. The ability of the model to do this depends on both the label noise (p), the context length (a), and sufficient number of contexts (A). In fact, the form of a = (p + tr(Curamul//a a + À is also familiar from ordinary ridge regression as an effective noise-to-signal term: the optimal ridge regularization parameter balances the variance due to label noise (p) with the estimation error from having finite data (Atanasov et al., 2024; Canatar et al., 2021a; Hastie et al., 2022; Patil al., 2024). At infinite sample size, the optimal ridge is simply p. the label noise. However, at finite sample size, the effective regularization is increased due to the sample estimation error, and becomes + o where of characterizes variability or complexity of the regression task. In the same way, our model has to resolve the statistics of the tokens æ over samples (e, measured by a), and so we expect terms familiar from linear regression to appear in our formula.
Test on powerfaw
TTes on low ranis
Test pretrain
0.2
0.5
10
1.0
-gA Gor fos
des
Ortnakge (E(ss:))
(Greef (00))
(Gug Frabe)
CKALGp, Guer)"
ercr,Crrain, (Ctrain: Crest) ) against alignment measures: emisalign (Cirainy Crest): ). Creste F.. tr(Cteat Crrin) and 1/CKA(Curann Crest) from left right. Ctrain is fixed be a diagonal matrix with powerlaw spectrum Curain x diag([1 ..., d' with power p = 0.9 and tr(Ctraia] = 1. Grest is varied over a range of different covariance matrices that are simultaneously diagonalizable with Curinn specifically power spectrum with different powers (circles connected by solid line), and low-rank covariance matrices C, diag((d/r)1, Og-r] (triangular markers connected by dashed line). Changing the power of the powerlaw tests or the rank of the low-rank tests will make them either more or less aligned with Creain- Parameters: d = 120, a 2, T = 4, p = 0.01.
Figure 2: etcr
In summary, our aligmment measure is motivated both by arguments regarding spectral ordering and relative strength, as well as having all components necessary to capture finite size effects, In Figure 2 we compare our alignment measure to three others: the population matrix measure (Cerain Crest ) a simpler version of our alignment measure involving just the resolvent (Crrain FirrainFx(Or Fx)())), and finally the Centered Kernel Alignment (CKA) measure (Kornblith et al., 2019). This comparison illustrates how different measures emphasize different aspects: (Cirain (a)) accounts for finite-sample res- olution effects but not as specifically as (Ctrain does, while (CtrainCC1 and CKA both miss saitee effects altogether, CKA instcad being designed to detect nonlinear representational similarity. We perform this comparison by first noting that, trivially, Emisalign is monotonically lated to cICL: larger Emisalign implies larger eICL. Figure shows CICL for fixed Ctrain and varying Crest plotted against the above matrix alignment measures between Ctrain and Ctest- We see that (C) and (Ctrain Fx(a), (a)) are the strongest drivers of CTCL. (i.e. most monotonically related, with X)) obviously being perfectly correlated), while the performance of (Crest 2 and CKA as predictors of excı, is lacking.
7

---

Published as a conference paper at ICLR 2026
powerlow on pretrain
Thanstommer ICL error 12
(Cux))))
[())
CKMC, Cua!
Figure 3: ICL test loss nonlinear transformer against different alignment measures. The setup of the covariances is identical to Figure 2, the only difference is that here ICL error is computed as the MSE on the test task as performed by a trained two-layer architecture with softmax attention and MLP connections. Our measure Cmisallgn achieves the best correlation with ICL error: the Spearman coefficients (measuring monotonicity, over all test covariances and averaged over the different values) are 0.99 (ours), 0.98, 0.96, and 0.39 from left to right. Parameters: d 20,
2, = 4, = 0.01.
3.2 TASK ALIGNMENT PREDICTS GENERALIZATION IN NONLINEAR TRANSFORMERS
We further support the predictive power of our alignment measure defined by émisalign for ICL error in two-layer transformer architecture with softmax attention by Figure 3. See Appendix H for details on the architecture setup. We see that even for an architecture far from the linear attention considered by the theory, our theoretically-derived alignment measure is still the best predictor of how changing the test distribution affects ICL error, compared to other alignment measures.
MISMATCHED TASK DISTRIBUTIONS ARE OFTEN OPTIMAL
We have identified the relevant measure of alignment between the pretraining and testing task distributions for linear transformers learning perform ICL. It is then natural to ask whether for a fixed test distribution Crest. is it always optimal to pretrain on the test distribution, i.e., is ercu Cect,Crrstt Ctest) ) S eICL. (Cerainy Crest In this section, show that the answer this question in general no, it is not always optimal to pretrain on the test distribution.
4.1 COVARIANCE ALIGNMENT AND EIGENSPACE STRUCTURE
We first show that it suffices consider the case in which Ctrain and Ctest co co-diagonalizable, as the ICL error is extremal this case. Note that we can focus on Emisaligu- as other terms the ICL error (18) contained in escalar (C.10) are independent of the eigenvectors of these two matrices.
Corollary 4.1: Eigenspace alignment extremizes misalignment error
The misalignment errоr Emisalign (Cirain: Cuest) (Crest, K.) is extremized when Crest and Cirain are co-diagonalizable. Concretely, letting \ (Ctest) 2 2 >e(Crestt and A (K) 2 Aa(KC) be the ordered eigenvalues of these two real symmetrie matrices, we have
Σ A((Crea))e-stt1(3 & emisalign I Σ (((eee)))))()))
(23)
5=1
with equality in either the upper lower bound if and only if Ctrain and Ctest are co- diagonalizable.
To show that this result holds, we observe that, as Crest and IC are real symmetric matrices, the desired bound (23) is simply restatement of Ruhe's trace inequality normalized by 1/d, where equality holds if and only if Ctest and K. oo-diagonalzzale. We give self-contained proof of this inequality in Appendix E (see also Marshall et al. (2010) or Li (2020)'s expository blog post). By its definition (20), K. and Cirain have the same eigenvectors, the claim follows.
8

---

Published as a conference paper at ICLR 2026
Therefore, the misalignment eπror be minimized maximized if we assume that Curain and Crest are simultaneously diagonalizable. We therefore restrict our attention to the co-diagonalizable setting, and write the ordered eigenvalues of Cerain and Crest as X (Ctrain) 2 2 da (errii)) and X1 (Crest) 2 ... (Ctass)) respectively.
4.2 OPTIMAL TEST COVARIANCE FOR FIXED PRETRAINING DISTRIBUTION
Before presenting results optimal pretraining structure for fixed task structure, which is a more practical question answer, begin first with simpler question: For fixed training distribution Crainn is it always optimal to test on the pretraining distribution, i.c. is EICL (Ctrain Cerain) S CICL Ctrain: ) The following result answers this question in the negative:
Corollary 4.2: Trace-constrained optimal test covariance
For fixed Crran and fixed Ctrain Ctest and 1, have that eICL (Ctrains Ctest) is minimized by the single-index spike covariance with eigenvalues
(An(Ctest)
Ar(Crest)
Ad(Crest)) = [dCtrain
0
i.e., all test signal is aligned with the largest eigenvector of Ctrain-
To see that this holds, note that EICL (ocL(Cuest (Curain Ctest) is a linear function of the cigenvalues of Crest as eissalgnn (Ctrsin Cvest) is linear Crest and Cscalar (Atriin Great) does not depend on Ctest when Cuest = Ctrain Furthermore. the constraint Cuest Ctrain is a simplex in the space of Ctest eigenval- ues. The minimum will thus be attained along the simplex, the vertex corresponding to the lowest eigenvalue of K. As argued in Appendix D. this corresponds to the largest eigenvalue of Cerain when T 1. As further discussed in Appendix D, we conjecture this to be true for all
This result illustrates bias towards exploiting low-dimensional structure seen in pretraining to opti- mize performance test time. these cases, where we are optimizing test structure Crest for fixed pretraining structure Cirain, observe that the best generalization is achieved by concentrating all signal into a single shared direction, collapsing the task manifold into single dimension. other words, it is easier learn generalize the pretraining structure than to generalize over the entire pretraining structure learned from lim- ited samples. This is familiar from spectral bias results of ordinary ridge regression (Canatar al., 202 la;b) Note that, had our training distribution been isotropic (Ctrain Ia) then all trace-fixed test covariances will perform equally; utilizing anisotropy is crucial.
degenerate low-rank test structure highly aligned to
4.3 NON-OPTIMALITY OF PRETRAINING ON TEST STRUCTURE
We now investigate the original question of the optimal pretraining task covariance Ctrain given particular test task covariance Crest. We will provide an example showing it is in fact not optimal general have Ctrain = Crest- We particularly emphasize that whether or not a particular task structure is better for pretraining depends strongly on the task diversity k.
Consider the case of power-law task distributions at both training and test time, i.e.
Otrain & diag
Perain
d
Perain
Crest OK diag 11 Pract
d d
(24)
with the normalization constant chosen such that tr(Curaiul = tr(Crest]. We can study the effect of changing Cursin relative Ctest on eICL by changing the exponent Pirain relative fixed Prest.
Figure 4 shows that, for low pretraining on Ctrann with higher spectral power relative to the test power can improve ICI. error. We see that focusing pretraining on a low-dimensional subspace, effectively creating strong inductive bias, can improve in-context learning performance when train- ing data scarce. The model generalizes better by overfitting to fewer directions, rather than learn- ing more directions weakly. However, increasing the pretraining power too much relative to the test power will worsen ICL performance, as the pretraining task set is now too nww-dimensiona to cover enough variation at test time. We highlight that the potential advantage coming from increasing training spectral power weakens as soon as x is large enough for the model to able to resolve enough task directions during pretraining.

---

esclarr and Published as a conference paper at ICLR 2026
Figure 4: Heatmap of theoretical ICL eπror given by (18) for simultaneousyy diagonalizable powerlaw task covariances Cirain and Crest with spectral power Prrain (variable) and Prest (fixed). The æ-axis shows task diversity K and the y-axis shows the difference Ptrain Ptest between task spectral powers. The colourbar shows the % improvement in error by training on Ctrain instead of Ctest- This shows that increasing spectral power in the pretraining tasks can markedly improve ICL error on the same test distribution. Parameters: d 100, Pvest 0.9, a 1,7 = 4, 0.01.
Train power Test power
% Improvement in Error
12
08
04
0.
-0.4
0.3
k/dd
CONCLUSIONS
We have developed a framework of in-context task alignment, showing that our derived measure of error from task misalignment is a robust predictor of ICL performance, even in nonlinear architec- tures. This derivation builds on a model in-context learning of linear regression, which we extend to fully generul task covariates. Previous works have shown that the performance of ICL models can suffer under task shift, particularly in cases of low task diversity (Garg et al., 2022: Goddard et al., 2025; Zhang al., 2024a). We highlight that, while task misalignment is provably driver of ICL error, our fully-general covariance framework can be used to show that task misalignment can indeed be utilized to improve ICL performance in cases of limited data.
Looking forward, our analysis provides a rich model of ICL with many further avenues of inves- tigation. A more detailed analysis of the interaction between pretraining-specffic misalignment error émisalign could be used to derive a general heuristic for optimal pretraining, with potential implications for practical settings. Beyond task alignment, additional highly relevant phenomena can be investigated: we highlight (1) a generalized learning transition in task diver- sity (empirically studied by Raventós et al. (2023); see Appendix F) and (2) test-time scaling (see Gozeten et al. (2025) and Appendix G). We leave exploration of these phenomena to future work.
AUTHOR CONTRIBUTIONS
All authors conceived the project. M.I.L. performed all computations with assistance from J.A.Z.- V., Y.M.L., and C.P.. M.I.L. performed all numerical analysis and experiments, and wrote initial draft of the paper. J.A.Z.-V., Y.M.L., and C.P. supervised the project and contributed to review and editing of the manuscript.
ACKNOWLEDGMENTS
M.I.L. is supported by Graduate Fellowship from the Kempner Institute. J.A.Z.-V. supported by Junior Fellowship from the Harvard Society of Fellows. C.P. is supported by an NSF CA- REER Award (IIS-2239700, DARPA grants DIAL-FP-038 and AIQ-HRO01 12520041, the Simons Collaboration on the Physics of Learning and Neural Computation, and the William F. Milton Fund from Harvard University. This work has been made possible in part by gift from the Chan Zucker- berg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXivv2303.0000 087744 2023.
Madhu S Advani, Andrew M Saxe, and Haim Sompolinsky. High-dimensional dynamics of gener- alization error neural networks. Neural Networks,
2020. doi: https///oii.og 10.1016/j.eumme12220......222 URL sutps://weet sei scwencedi com/ swtsscieeencceeeeeeeeeee atticlle/p11//3888300000000
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, 2023. URL htttt:://aaaiiiieieee abs/2306 00297.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in in-context learning? investigations with linear models. In The Eleventh International Conference Learning Representations, 2023. URL wwwpreeeeeeeeeeeeeeeeeeeeeeeeeee net/ forum?
idd00g044ByN41.
Anthropic. The claude model family: Opus, sonnet, haiku, 2024. URL https://api.
semanticscholes org/ wtagg//orpssii 268232499.
Alexander Atanasov, Jacob Zavatone- Veth,, and Cengiz Pehlevan. Risk and cross validation in ridge regression with correlated samples. In Forty-second International Conference Machine Learning. 2025. URL https: opwnoeeeiiiieeeee .nt/ Ltte://ooommmmmmm/22220
Alexander B. Atanasov, Jacob A. Zavatone-Veth Veth, and Cengiz Pehlevan. Scaling and renormalization in high-dimensional regression. arXiv, 2024.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisti- cians: Provable in-context learning with in-context algorithm selection. In A. Oh, T. Nau- mann, Globerson, K. Saenko, M. Hardt, and Levine (eds.), Advances in Neural Information Processing Systems, volume 36. pp. 57125-57211. Curran Associates, Inc., neurips.. ec wwwppp/pppppeeeeeeeeeeeeeeeeeeee www.raapperroooreeeeeeeece .d
2023. URL
fi
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature Com- munications, 12(1), May 2021a. ISSN 2041-1723. doi: 10. 0723///4167--221--1
URL
http://dx doi. or0/10...
11100.0338 1038/ 44444447-222 21211-21111 3103-1
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Out-of-distribution gen- cralization in kernel regression.
M. Ranzato, A. Beygelzimer, Dauphin, P.S.
Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Process- ing Systems, volume 34, pp. 12600-12612. Curran Associates, Inc., 2021b. URL https://proceeding neurips. wwtp:pcppppppppeeeeeeeee httpe//eeeeeaapee//////Ee.. 691 691cco.d665011877777/11133333333
Paper. pdf.
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Out-of-distributia generalization in kernel regression, 2022. URL httph::tp///.c...t. https///aawwxw../aaaaaee abs 2106. 02261.
Stephanie C. Y. Chan, Adam Santoro, Andrew K. Lampinen, Jane X. Wang, Aaditya Singh, Pierre H. Richemond, Jay McClelland, and Felix Hill. Data distributional properties drive emergent in- context learning in transformers, 2022.
Frank Cole, Yulong Lu, Tianhao Zhang, and Yuxuan Zhao. In-context learning of linear dynamical systems with transformers: Error bounds and depth-separation, 2025. URL https: //arxvv.. org/abs/ aaaaab//550... 2502. 08136.
DeepSeek-AL, Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang. Xingkai Yu, Yu Wu, F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang, Bochao

---

Published as a conference paper at ICLR 2026
Wu, Bei Feng, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang. Han Bao, Hanwei Xu, Haocheng Wang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Qu, Hui Li, Jianzhong Guo, Jiashi Li, Jiawei Wang, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, J. L. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Liang Zhao, Litong Wang. Liyue Zhang, Lei Xu, Leyi Xia, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Meng Li, Miaojun Wang, Mingming Li, Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang. Qinyu Chen, Qiushi Du, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, R. J. Chen, R. L. Jin, Ruyi Chen, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang. Shuiping Yu, Shunfeng Zhou, Shuting Pan, S. Li, Shuang Zhou, Shaoqing Wu, Shengfeng Ye, Tao Yun, Tian Pei, Tianyu Sun, T. Wang, Wangding Zeng. Wanjia Zhao, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang, W. L. Xiao, Wei An, Xiaodong Liu, Xiaohan Wang, Xiaokang Chen, Xiaotao Nie, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xinyu Yang. Xinyuan Li, Xuecheng Su, Xuheng Lin, X. Q. Li, Xiangyue Jin, Xiaojin Shen, Xi- aosha Chen, Xiaowen Sun, Xiaoxiang Wang, Xinnan Song, Xinyi Zhou, Xianzu Wang. Xinxia Shan, Y. K. Li, Y. Q. Wang, Y. X. Wei, Yang Zhang. Yanhong Xu, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Wang, Yi Yu, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yongqiang Guo. Yuan Ou, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yunfan Xiong, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Y. X. Zhu, Yanhong Xu, Yanping Huang, Yaohui Li, Yi Zheng, Yuchen Zhu, Yunxian Ma, Ying Tang, Yukun Zha, Yuting Yan, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhenda Xie, Zhengyan Zhang. Zhewen Hao, Zhicheng Ma, Zhigang Yan, Zhiyu Wu, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Zizheng Pan, Zhen Huang, Zhipeng Xu, Zhongyu Zhang, and Zhen Zhang. Deepseck-rl: Incentivizing reasoning capability in Ilms via reinforce- ment learning. arXiv, 2025. URL https:
12948.
Edgar Dobriban and Stefan Wager. High-dimensional asymptotics of prediction: Ridge regression and classification, 2015. URL /etpss///ariv..
abs abs/1507 03003.
Eric Elmoznino, Tom Marty, Tejas Kasetty, Leo Gagnon, Sarthak Mittal, Mahan Fathi, Dhanya Sridhar, and Guillaume Lajoie. In-context learning and occam's razor, 2025. URL https:
//arxivv wwwoor//abbaa /2410 14086.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimiza- tion methods for in-context learning: A study with linear models, 2023.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: /ttp////eeeeeeeeeeeeeeeeeeeeeeee net mete/oounm forum? id= f12
Chase Goddard, Lindsay M. Smith, Vudtiwat Ngampruetikom, and David J. Schwab. When can in-context learning generalize out of task distribution?, 2025. URL https: http://axxvvvooo/e abs/2506. 05574.
Halil Alperen Gozeten, M. Emrullah Ildiz, Xuechen Zhang, Mahdi Soltanolkotabi, Marco Mondelli, and Samet Oymak. Test-time training provably improves transformers as in-context learners, 2025. URL https://arxixs wrg.org/gaao... abb 22003253.02 11842.
Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan J. Tibshirani. Surprises in high- dimensional ridgeless least squares interpolation. The Annals of Statistics, 50(2):949 -009949998666 986, 2022. org/10 1214/221AA02213
doi: 10.1214/21-A022133 URL
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5166-5165. PMLR, 2020.
Simon Komblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited, 2019. URL hwwptttts///eettt wwtp::/wwwwww.w////p//o/ooooss/t /oo.b////99..0011111 00414.
12

---

Published as a conference paper at ICLR 2026
Anastasis Kratsios and Takashi Furuya. Is in-context universality enough? mips also universal in-context, 2025. URL pattps://arxxe org/ httpg//bb/2502...02 23322.3777
Jin Hwa Lee, Andrew Kyle Lampinen, Aaditya K Singh, and Andrew M Saxe. Distinct computations emerge from compositional curricula in in-context learning. In Workshop on Spurious Correlation and Shortcut Learning: Foundations and Solutions, 2025. URL https: /ppereviee....
net/ net/forum forum? dd=oo5TNikeee kek..l.
Yingcong Li, M. Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning, 2023.
Yingcong Li, Ankit Singh Rawat, and Samet Oymak. Fine-grained analysis of in-context lin- ear estimation: Data, architecture, and beyond. In A. Globerson, L. Mackey, D. Bel- grave, A. Fan, U. Paquet, Tomezak, and C. Zhang (eds.), Advances in Neural In- formation Processing Systems, volume 37, pp- 138324-138364, Curran Associates, Inc., 2024. URL. https://prroeet ngs.
www.innn.uneiiieeneneeee neur cc/
file/
.dd..
Yingzhou Li. Von Neumann's trace inequalities. online blog post, 2020. URL: https: //yingzhouli. html.
Toni J.b. Liu, Nicolas Boulle, Raphaël Sarfati, and Christopher Earls. Lims learn governing prin- ciples of dynamical systems, revealing an in-context neural scaling law. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 11997-111777. As- sociation for Computational Linguistics, 2024. doi: 10.18665//1/2224emmlpp www.eemmlllmminnee.. main.842. URL
http: /stp://ddet dx doi org/ 10. 111.1..3//1...2211.111 emnlp mainn.ain...m 842.
Yue M. Lu, Mary Letey, Jacob Zavatone- Veth, Anindita Maiti, and Cengiz Pehlevan. Asymp- totic theory of in-context learning by linear attention. Proceedings of the National Academy of Sciences, 122(29):e200559912,, 2025. doi: 10.1077/pmas.
URL https:
// /ww. pnas hwrg//doi/aba/
1073/ 10000//paaas 2502599111
Arvind Mahankali, Tatsunori B. Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention, 2023. URL
attps://arxiv.ort oaaiii.rrr///oo abs////bb/337.. 0357773777
Albert W. Marshall, Ingram Olkin, and Barry C. Arnold. Inequalities: Theory of Majorization and Its Applications. Springer New York, NY, edition, 2010. doi: https://doi.orr/10.. 1001/
978-0-387-68276-1.
Gavin McCracken, Gabriela Moisescu-Parte Pareja. Vincent Letourneau, Doina Precup, and Jonathan Love. Uncovering a universal abstract algorithm for modular addition neural networks, 2025.
URL htapps///axxvv.. org/ abs /205. 18266.
Pratik Patil, Jin-Hong Du, and Ryan Tibshirani. Optimal ridge regularization for dutti--dis prediction. Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 39908-39954 PMLR, 21-27 Jul 2024. URL https: httpp:://prcceednngge .ww. press// v235 215.//ati12a..........
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diver- sity and the emergence of non-bayesian in-context learning for regression. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neu- ral Information Processing Systems, volume 36, pp. 14228-14246. Curran Associates, Inc., wtw.:eeddddinn...enueeeeeeeeeeee neurips ec cc/paper files/ hailes/paper/203// Cer--Coofeeencce file://e0022222aa444400032772662 oor. .df..
2023. URL https::/prrceeddnng
Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in in-context classification task. In The Tweifth International Conference on Learning Representations, 2024. URL /ttttss://ppenertttt net forum? hf
13

---

Published as a conference paper at ICLR 2026
Zhaiming Shen, Alexander Hsu, Rongjie Lai, and Wenjing Liao. Understanding in-context learning on structured manifolds: Bridging attention to kernel methods, 2025. URL https: //axxis.//rrvv... htt/abb//5005 10959.
Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities. In Proceedings of the IEEE/CVF winter conference on appli- cations of computer vision, pp. 3531-3539, 2021.
Aaditya K. Singh, Stephanie C. Y. Chan, Ted Moskovitz, Erin Grant, Andrew M. Saxe, and Felix Hill. The transient nature of emergent in-context learning in transformers, 2023, URL https: //arxiv. www. wtaaa///1111 08360.
William L. Tong and Cengiz Pehlevan. MLPs learn in-context on regression and classification tasks. In The Thirteenth International Conference Learning Representations, 2025. URL https:/ http:////pereeeee.emmm net forum? i witmmmddddbitttttt
Nilesh Tripuraneni, Ben Adlam, and Jeffrey Pennington. Covariate shift in high-dimensional random feature regression. arXiv, 2021. URL https: / arrivv org/ abs/2111.1.1. 08234.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvint- sev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient de- scent. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learn- ing. volume 202 of Proceedings of Machine Learning Research, pp. 25151-35111 PMLR, 23-29 html.
Jul 2023. URL https: http//pccceennggge/e/ee
mlr press / v22///22////////2222/// www-osswllle
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large lan- guage models. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https:/ cpenreview.eeeeee/eeeeeee net wttp://wwwwee.tttiieeeeeeee forumiid-y2kSSs22:
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In The Twelfth International Conference on Learning Representations, 2024. URL https://
Daniel Wurgaft, Ekdeep Singh Lubana, Core Francisco Park, Hidenori Tanaka, Gautam Reddy, and Noah D. Goodman. In-context learning strategies emerge rationally, 2025. URL https:
/ //arxiv.ccmm wrw. abs /20.../50......
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in- context. Journal of Machine Learning Research, 25(49):1-55, 2024a. URL http: jm.r.
brg//papers/v252233 1042, html
Ruiqi Zhang, Jingfeng Wu, and Peter L. Bartlett. In-context learning of a linear transformer block: Benefits of the mlp component and one-step gd initialization, 2024b.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization, 2023. URL org/abs/////eeoooooo.ggg/eee aa////////300 19420.
https: //appsi//rrxv..../ org/
14

---

Published as a conference paper at ICLR 2026
SUPPLEMENTARY INFORMATION
In the following sections, we provide proofs of the results presented in the main document. We begin in Section with some preliminary simplifications of the ICL and IDG population risk for general T. Then in Section B we will set up and complete the random matrix theory calculation necessary to justify our main result, equation 18, which will be done in Section C. Section will cover the proof of the thenomena-basee Corollary 4.1. Then, we provide additional analysis of transitions in the behavior of the ICL. error as function of the task diversity &. in Section F, and of shifts in the test-time context length in Section G. We conclude with details of our numerical experiments in Section H.
For convenience, we include the following table detailing quantities appearing throughout the main paper as well as this supplementary section.
Symbol
Mcaning
Def'n
d n lestt k Z Zy.
Task and token dimension
Sec. 2
Number of pretraining contexts
Sec. 2
E,
Pretraining and testing context length
Sec. 2
Pretraining task diversity
Sec. 2
Embedding choice of a context (context μ)) sequence
(3)
T, Hz
Reduced model parameters and data
99, (10)
A
Ridge parameter
(II
Optimised I' given data sampling Hz....., Hz.
(I))
EIC((T)
ICL test errors for parameter
(13)
Otrain
Training task covariance (this is finitely sampled by k:)
(2)
Crest
Testing task covariance (no sampling restriction)
(12)
P
Label noise in both pretraining and testing Asymptotic samples parameter between n and d2 Asymptotic context length parameter between & and d Asymptotic task diversity parameter between l and d Pretraining task sample covariance
(2)
T
(14)
o
(14)
k
(14)
Rk
(15)
X,
Rescaled ridges and effective noise
(21), (22)
Fz(2)
Deterministic equivalent of resolvent of R& at ridge
(16)
Mn(12)
Stieltjes transform of R₁ at ridge Deterministic equivalent for EicI.
(17)
eICI (Ctrann Crest)
(18)
emisalign ( (Crrains Ctest)
Pretrain-Test misalignment term in EICL (Cerains Ctest)
(19)
K
Pretraining-dependent comparision matrix in Emisalign
(20)
Table 1: Quantities appearing throughout the main document
Symbol
Meaning
Def'n
Token covariance matrix (taken to be Ia) IDG test error for parameter r' Linear response matrix in simplified SDG. ÊiC Quadratic response matrix in simplified EIDG, Eicl Pretraining task sample mean and covariance Resolvent of feature Hz covariance used in I Parameterisation of resolvent for convenience Extended parameterised resolvent that contains I"
Result
(T)
(A.2)
Amg. AIeL
(A.8), (A.9)
BIDG, BICL
(A.8), (A.9)
by. Rs
(A.10)
G
(B.10)
G(x)
(B.11)
Gox((T) (x)
(B.14)
Gli)]
Leave-one-out extended resolvent Token- and noise- averaged version of Gexi (T)
(B.19)
(T)
(B.39)
Deterministic equivalent of I
(B.52)
eIDG (Atrain)
Deterministic equivalent of εDD (I*)
(C.6)
eICL(Crssia Cuet)
Full explicit formula for (18)
(C.7)
escalar( CaaslarrCrrainn (Ctrain; Ctost)
Ctest -isttrisoc term in CICL (Cirainy Crest)
(C.10)
Table 2: Quantities appearing only in the supplementary information, particularly Section B
15

---

Published as a conference paper at ICLR 2026
Notation Here vec(-) will mean the vectorization operation under the row-major convention: for a d1 d2 matrix A, vec(A) ) is a vector in R41d2, formed by stacking the rows of A together. We will use this together with the matrix Kronecker product where by standard results have
vec )) = u @ v
(25)
(u & (w @ 8) (uw @ (us T)
(26)
We will also use the notation [M] o to mean the principal minor of a matrix M (i.e., first row and column removed). We use the following normalized trace
tr = tr(A).
d
We use = when defining new quantities. We use = for deterministic equivalence, which describes the large-dimensional limit of random quantity (either a scalar or matrix) that concentrates to deterministic value as dimension approaches infinity.
A GENERAL ERROR FORMULAS
Here we will set up the definition of various errors more generally than in the main document. As this is an in-context learning setup, there are different types of "generalization" that can be studied: generalization over the tokens and generalization over the tasks. We will define test errors to capture these separately.
In general, we wish understand the performance and behavior of this estimator I which is pretrained on data from the pretraining distribution Ptrain. when tested on new data. Namely, we will analyze the average performance of an estimator ý = Ir[T, Hz) as function of given fixed under the MSE loss, which is the natural loss for regression test. The most general this expression can be
LMSE (T) = Ep,.s [Wee4 - (t(TTTTH))") H2))²
(A.1)
for a general test distribution Pust detailing how to sample tokens tasks w, and noise at test time.
We will consider two different testing regimes: the in-context learning (ICL) test, where the model sees new tasks w, and the in-distribution generalization (IDG) (or in-weights) test, where the model sees the exact sume tasks used in pretraining . )x Explicitly, we define these test distributions and corresponding error functions as follows:
where each t, "iid (0, Cuain)
(TG = Eppo [We++1 tr(T, H2))
(A.2)
PING == ot mid. V(O,Ia/d) n" unitt (tr.... ,t), o ~us. N(0,p),
where E [6] and E [m]. and
and
(T') = Epac (we+1 LtrT Hz))
(A.3)
PICL == a. mid (O, Ta/d), r miid N(0, Crest), & Mia N(O
where [Rues]] and μ E [n].
Notice here that we've introduced two different context lengths: é for the IDG distribution, which the same context length as the pretraining setup Purain given by (2), and lest which the con- text length testing time. This allows to later explore the effect of pretraining and testing sequences of different context lengths.
We assume that both task covariance matrices Ctrain and Ctest well-behaved in high dimensions, specifically that tr(Ctrain], tx [Creat] = (1). This ensures the task signals are not over or under amplified d
Lemma 1. Simplified test losses. Consider IDG and ICL test distribution and corresponding error functions as given by (A.2) and (A.3). For fixed parameters I e [dx((+1) and data sampled ac- cording to test distribution PIDG PicL, the corresponding errors ennG and exc1. can be expressed
16

---

Published as a conference paper at ICLR 2026
EIDG( p + tr[ERR] - 2tr[TA t
(A.4)
(T) ~ ++ te|t - 2 2tr[TAAcul + tx| tr
(A.5)
where
ADG
ER&E (tr[ERe] + p) Eba]
(A.6)
ArcI.
ECrest E
0]
(A.7)
DRA + & (ur[ERe) p) (tr[ERA) + p) Ebk
BIDG
(A.8)
(tr[ERR] + p) (Eb&)
(tr[ERk) + +p)
ECrest +
Eraae (tx (tr[2Crest) +
0 (tr(2CCrest] p)
BICL.
(A.9)
0"
for by and Rg defined by the pretraining task sample {wi ,.., wkk as
by =
tj JE[R]
Σ
RS Σ ttt
(A.10)
JE(A]
The "=" in these formulas comes from ignoring terms of order 1/d and weaker, i.e. high- dimensional treatment.
Note that Result 1 is a more general result than the setting considered in the main paper, as Result allows for general token covariance E. The rest of this work will take = Ia computational tractability.
Remark 1.
Lemma 2. Fix the task vector w that defines y/s = w 0; + for context Z. Denote the conditional expectation over only ais c, holding fixed, by B.g, Then we have the following
Ex, y/++1 1] =
(A.11)
Ez,,, H =
(A.12)
Ex, [b2+1] tr trlLww ] I p
(A.13)
Exc [1 Hz] = d [oww Ty
(tr(Zww *1 p) Ew]
(A.14)
Ewow E + $ [Ir(Eww 7 1+ +
(urw ] + p) Ew
Ex,e(veo(HHz) voo( (zz) 1
(tr[Eww 1 + (Ew)
(Ir[Ewww 1 (+)))
(A.15)
where recall
Hz xz4 [ Eise V:0 E ui E Rdx((++1)
(A.16)
Proof. Equation (A.11) follows immediately from the linearity of ye+1 in e and a;, which are both mean-O random variables. Similarly for (A.12), Hz is linear in ø0+1-
For the conditional expectation of y241 in (A.13), simply expand
Ex, [W²+1]] = Ez,c [(w" Øe+1 + (e4+1)(( + EE+1)]
(A.17)
Eg,e /ww T Teai ed41
(A.18)
1 jw
EEw +
(A.19)
tr trww T] + p.
(A.20)
17

---

Published as a conference paper at ICLR 2026
For (A.14), have
d
Ex,t
11-13102412
ye-
T
Ex,e (0+11 c<+1)(w x, =,
(A.21)
ise
ise
Ex, 184+10/41 ax
(A.22)
ist
(E/d) ww (E/d)
(A.23)
ist
Eww"" Σ
(A.24)
and
En,e
Σ yet 1
1 Σ
En,c [2041 + (4 x, +G)((a w + G)]
7 ise
i56
(A.25)
1
Σ
Ex,e [ w(w x + EεE,)]
(A.26)
8
1 Σ (E/d)r((w" (E/dl/w (F//d)) + p)
(A.27)
ise
d (tr] Euow
[ p)
(A.28)
as required for (A.14).
Finally, for (A.15.) first will be helpful to note, by Isserlis's theorem / Wick's theorem, that
2
Bmcc a wws aıæ]]]
Eww EE d
tr tr|_ww JJ..
(A.29)
d²
will also useful to rewrite
(1
1 Σ Eise yeas iier u?
vec(Hz)² (041 2
[ige
Eigevi)
by converting between vec() and matrix Kronecker. The components of this Kronecker product are independent and can therefore be averaged separately. Clearly
(0 E
so we focus on the second matrix in the product. This matrix will have blocks that sum over two l sums, leading both ((1) and (1//E) terms in the final expression. We will ignore terms of order 1/E as we eventually will only use this formula in a proportional limit of E, + oo such that t/d (11), and thus 1/6 is negligible in this limit.
18

---

Published as a conference paper at ICLR 2026
We proceed boek-by-blockk Firstly,
d²
de
Σ
ya
Σ
E
(2, w <e)(w" + c))o]]]
12
WSE
d²
Bmg [ w + e:) My ej)ar]]
(A.30)
1
d2
Bu,
[ ww xa + 111
d
(62
1)
E Ex [x ww x222
(A.31)
d
1
2 TE+ + de
tr/[wo
Nnww T + g02)
d
+
(6e 1)
Swow
(A.32)
p2
12
Inow E + d
+ i
& (ur[EwwrT] } 1) E
(A.33)
Ewo""
d E+
(ur(Eare } p) E.
(A.34)
&
Secondly,
Ex,e
Σ y
= l
e 82
Li [11]]] 31 + - E) pz Wix, [01111] Exe. [32]
-f2-))
(A.35)
12
iust
1
Ed (d JE + 2Eww TE) w + PIw
(1-2) (Ir(Enowo J] p) Ew = (Ir[Ewow 1 + p) Ew.
(A.36)
(A.37)
Finally,
8 [uil + -- 2''Ene
E
Σ v
Eu,c Bu Bue [v2]
(A.38)
1
2 (Earrlu T Taxx "ww" Toraa wax T w) + 6ppEn (wo "xx" T w/] 3p²)
+ (1- 1) (tr(Eww 1 p)²
(A.39)
(tr[Eww 11 2+p)
(A.40)
Combining these pieces gives (A.15).
Proof. [of Lemma 1] We begin by noting that for both IDG and ICL test errors, can expand
e(T) E,,,w, (1e41-tr(TTH)) tr(THH))) = Ew Exz,e (getz ( (r((H)))) "
Enp Ean (yC+1 - vec veo vec(Hz)))²) "]]
En EE,, | - 2 vec(T) T Em [B.,, [We+1 veo (x))]]
+ vec Ew Ex,e [vec(H2) vee (H2) []] vec(T)
(A.41)
where the difference between IDG and ICL error comes from the different task distribution over which we take E,,. Note this E,, is shorthand for "expectation over task distribution," " which will be different distributions depending if we using PIDG PICL.
19

---

Published as a conference paper at ICLR 2026
Now the above Lemma simplify these terms drastically. For the IDG distribution, have
E [0]
bb
En www
Rg
Therefore,
Eu Fe,, 11+1 = tr [ERR]+ +
(A.42)
1
ER (tr(ERR) P) Ebs] 1
Eu [Lize, [Be++ [z]]
d -ADG
(A.43)
ER + & (tr[ERa] + +p)
(tr[ERk] + p) Σby (tr[ERe) + p)²
Ex, [veeveee)) (ooc()z)
(tr[ERR + P) (2bx)
BIDG.
(A.44)
For the ICL distribution, we have
Eu (m)) = O,
En (wow 1 = Ctest
and so
Eup Ee,, 112 = Ene [tr[Eww 1 [ p] = -Ir[ECeest]] +
(A.45)
Ew [E [ye+] Hzll
[EOte E 0] =AAICL
(A.46)
ECtest + (tr[ECCest]] +p)
Eu [Ben, (gec(HZ)) veo(Hz)
0
(tr (tr[SCteal] + p)
E 60 BICI.
(A.47)
Upon noting that
vec(L)) B) vec(I") tr (ΣTВT")
substituting these components into (A.41) gives the required results.
B RANDOM MATRIX THEORY CALCULATION
Note that while Result is offered for general token covariance Σ, all steps and results in this section and the following section assume Σ I.-
The setup and structure of this formalism and corresponding results will follow Section SI.3 and SI.4 of Lu al. (2025). There will be key differences caused by the different task statistics (O, (. Creat) that we use in this work compared to the isotropic tasks used in Lu et al. (2025). Since the token structure is the same between the work of Lu et (2025) and our work here, we will not prove various етrоr bounding claims or approximations rigorously here, as such bounds would be identical to what can found Lu et al. (2025). However, care will be taken to highlight the distinction and differences required for handling the additional complexity introduced by our consideration of non-isotropic tasks.
We will use to denote two scalar quantities which converge in probability in the high nigh-dimensiooal limit, or for two matrices to represent deterministic equivalence.
Definition For two d (possibly random) matrices Ag and Ba. we write Aa Ba if tr(Ma(Aa Ba) + O in probability d + oa for any sequence of bounded spectral norm test matrices Ma (Atanasov et al., 2024; Lu al., 2025).
As noted above, will attempt to control rates of convergence.
We begin by considering the closed-form optimized parameters given by (11). This optimization problem can be solved explicitly
vec(T") = 1 AI + Σ vec Hz vec (Hzz) T
Vly vec (ec)H
(B.1)
20

---

Published as a conference paper at ICLR 2026
Ultimately we wish to characterize the ICL and IDG enror of these parameters, which given by Result as
(TT p + tr(Ctest 2r Accl + (rccc (T*) GTDG(T*) p + tr| Rr) ADbal + tr[r" BIDG (T*)
BICL
(B.2)
(B.3)
for
AIC1 [Crest
0]
(B.4)
Crest + (tr(Crest) + +D)
BiCL =
(B.5)
0
(Ir[Crnst] + p)²
ADDG
[R (tr[RR] + +D) byl
(B.6)
Bs +
(tr [tr[Re]] + pp)
(tr[Ra] + bk (tr[Ae] + p)²
BIDG
Oterain
(B.7)
+ p)bb p)) bj
Again, are distinguishing between pretraining and testing context lengths, use
Ottrain /drain/d,
atest = /deat/..
These expressions for m and its corresponding errors depend explicitly the randomness found in the particular sample of the data 2, w., C; we wish to able to characterize the typical perfor- mance, and thus need average over this disorder. We will do by following random-matrix style computation that computes a deterministic equivalent for the parameter matrix I* and its cor- responding ICL and IDG errors, valid in high dimensions in the proportional limit
Lrain d
Otrain = ((), leest
Oltest @(1,, ā = O(1),
n de
=T= (B.8)
To do this, we will express necessary quantities terms of resolvents.
RESOLVENT AND EXTENDED RESOLVENT SETUP
I' above be rewritten
vec(T*) G (Eueta Uu vec(H,)) /d,
(B.9)
where G is the resolvent matrix
G G= (Euetn) vec(H,, vec(H,) /dd(HH)
(B.10)
We will find it helpful explicitly include additional matrix Bplaceholder (d) (that
G(T)) (Epeln) vec(Hp) vec (H,) //d nT Bllacehooldrr + rAI)
(B.11)
for non-negative scalar т. Notice that G(0) = G. Eventually Bplaecholder will be explicitly related to BICL or BING in way that will allow us to more easily compute the PT or [ term in the error formulas.
We can write G in a cleaner and more useful way by concatenating yı and vee Hμ) into an extended vector
Mu/d vec vrec(H,)/vd
E +d++
(B.
We also extend Bplaceholder
0
Bext
(B.13)
Bplaceholder
We then define an extended resolvent (with very similar structure as G) to be
1
Gexi (T))
(B.14)
Epe(n) Z,zi Bext T\-
21

---

Published as a conference paper at ICLR 2026
We that 2, and Bex have block structure, and expanding the block inverse see
d Gex((T) -c(T))((T) +(m))((()(())))) (T)
[
-(T))q (x)
(B.15)
for
(
1
G(x) (Enein) Un voc(Hp)) Rd(d+1)
(B.16)
d3/2
and c(z) defined self-consistently by
+
1
Σ Yullo Hu)) G(T) vec(Hl)).
(B.17)
d3
KE
M.PE(M]
At we can now see that Gext explicitly contains information about the optimal parameters r* we re considering, as
(()) Vd
vec([").
(B.18)
We will thus proceed by computing deterministic equivalent for this extended resolvent.
DETERMINISTIC EQUIVALENT FOR THE EXTENDED RESOLVENT
We will closely follow Lu et al. (2025), omitting details that remain unchanged between the formal- ism in their work and ours here. Intuition-based headings will be given in red.
The computation proceeds by defining "leave-one-our" version of Gexts
ext
Even zu + π Bext + TAI
(B.19)
By construction,
Gext (Ewe(n) Zn12/4 + TBxtt + +Nl) =1. I.
(B.20)
extt 212 Goxe( Boxt + TXI) I.
(B.21)
ME[M] 1 2j. G.G.
TT Averuge over disorder in x. e, which concentrates. The term II Gext= z is well-behaved, specifically as argued by al. (2025), it concentrates around conditional expectation when the task w and Glu) is fixed. We thus have
TT G ~ (ww
(B.22)
where
ir (IC F((w,)))
x" (w)))
(B.23)
and
Tww 1 (p tr/ww 1) Ia (p+ tr(ww TI) w
(p + tr 1]) w +ottrwww T])
E(w)
(B.24)
Glu) Replacing 2 ext zy in (B.21) with x" (ww,) gives
Σ ME
Glu) Zuz + Gext Bext + TXI) ~ I.
(B.25)
1 (ww)
Next, we will also average the term _ on on the left-hand side of (B.25) over the remaining disorder in a, holding w fixed, and replace z,²ح with this conditional expectation. Doing so introduces some small error which is bounded in Lu al. (2025), From this we have
ch4 Gext B.n.c (= + Gexi(T Boxt + TAT) T.
(B.26)
(ww,
22

---

Published as a conference paper at ICLR 2026
We already have enough to compute Bx,e (2 I using Lemma 2 as
Ez,c 22
E
d? T.
(B.27)
where because it will appear in subsequent equations we define the matrix
p + tr(ww"]] r 1 va vec ([rrw] T
va vee (www T (p+trlwwed Duol)
(B.28)
(p (((
I E(w)
Now substitute this expression into (B.26) for the conditional expectation Enn [2 giving
1
Glul Y
GexI Bext =AI) ~
(B.29)
1+ x²
where we recall n/d?.
"Leave-one-out" terms behave like their full-sum equivalents. In high dimensions and for large n, there is negligible difference between Evzp and En. Thus, we replace G by Gexet and x (ww) by
x(wp) tr (Ieexlyo - 11 E(wy)))
(B.30)
d²
So finally we have the expression for Gext
Gext
1 1+ x(w)) ME(N)
Σ
T BBex Thl ~I
(B.31)
TI
Exploit finiteness of training task set. So far we are summing over n task vectors, but really only n/k of these are unique. Thus, we can simplify (B.31)
T Σ
1
Gext
T + nt Bext + TAl
~ I.
(B.32)
Te x JEIN
Indeed x(t;) is also self-averaging in lú. and as argued by Lu et al. (2025), concentrates to its mean
Rave
1 1/2
Σ x(ty)-
(B.33)
JELR
We can thus simplify our expressions by substituting (B.30) into (B.33) and performing the sum over WT, m. Wg. Here will the that
(o (o+t((t,t)))= = [r(Rs) Σ JE(K]
(B.34)
Σ
(p+ tr
1 - tr(Htx]) by
(B.35)
jE[K]
(++ (p+ tr(
(B.36)
T JC(
where for (B.35) and (B.36) have used that tr(t,t,] RR( tr tr(Rg] tr(t,t]] has mean and variance ((1/d). We thus have that
Σ E(t,) BIDG
(B.37)
K
JELN
and so
Rave db tr (IGeulyo [I @ BIDGl)
(B.38)
23

---

Published as a conference paper at ICLR 2026
Final formula for extended resolvent in terms of training task sample. Gext (T) is asymptotically equivalent to
The extended resolvent
(e(x)= = 1
Purain
Va vec ([Rx Pursinb))
]
+ r Bext + TXI
Xave
Lia vec ([R
Pirain bkl)
La 8 BYDG
(B.39)
with Xave defined self-connistently by
Xave Σ
x((;)
1 E d
ir (GGoxlyy Ul @ BUDGL)]
(B.40)
jE
with
by
Σ tj. jE[k]
R4 K
Σ t,t
Ptrain p + tr(Rx).
(B.41)
5E[K]
RR + Perain La Perain bv Pirain
BIG
(B.42)
Ptrain bg
Simplifying self-consistency variable X. X is defined terms of the dominant component [Gext]yo of Gext- To simplify this, choose ansatz for Bext
0
0
Bext
0
Ta 69 Brest
where Biest will later be BiCI. or BING depending on which error we compute. Then (Gox)/y can be expanded as
la & BING 7la Bust + TAIa Ia+1
1 X=
BING + Brest + TN4+11 T
= I Ø
(B.43)
1 1t
Ia 2 FE(T))
(B.44)
where
FE(T) = 1 +1 BIDG - n Busst + +T TAId+1
We can use this to replace Rave with X= defined self-consistently by
Xx 2tr (1 BIDG BBeast + AT(A+1)
BIDG
(B.45)
Relating (T) to a deterministic equivalent for I* Lu et al. (2025) shows that two key quantities tr[[" A and (r( B[ in E(T*) can also be computed from the extended resolvent G As this is purely matrix algebra claim, and independent of the particular task structure hidden case. We thus include this result here without proof and
within r** generalizes immediately refer to reader to Lu et al. (2025) for a derivation.
Lemma 3. [From Lu et al. (2025)] For any matrix Rdx(d+1),
[r A = c(0)vd [0
vec(A)"] Co)e(0))
(B.46)
where ej denotes the first natural basis vector in of +d+1 For any symmetric and positive semidef- inite matrix B [g(++))((+1)) if we set
Bplaceholder Ia 8
(B.47)
in (B.13), then
B(
(B.48)
dr
((TT)
==U
24

---

Published as a conference paper at ICLR 2026
Using Lemma 3, find the deterministic equivalent for I" from
tr(r> A <(0)vd [o
(B.49)
(0)
c(0) tr (lRa Purain (BwG + X(1 + -A
(B.50)
d (RRg Pirain bul (BING A(1 xo)
(B.51)
and thus
1 = RR Perain bu] (BIDG + A(1 + xo)1)
(B.52)
SELF-CONSISTENCT EQUATIONS AND FINITE TASK SAMPLE AVERAGES
So far, everything is left in terms of the quantities bEs R&. BIDG that depend on typical sample of k task vectors W1 ... Wg. What remains be analyzed is the behavior of this sample and its sample statistics, i.e., how does the finiteness of affect 1** The following steps will deal with this.
Characterization of Xo We see that we almost have characterization for the deterministic equivalent of I' but we need understand what X(1 + Xo) term doing. From (B.45) we know Xo is defined self-consistently from
~ tr
T
BING + ATId-1)
BIDG
(B.53)
d
1+ Xo
TXo
- tr
(BIG + A(1 + xo)lo+1)
BIDG
(B.54)
1 Xo
d
d+1 1
A(1 + xo) tr (BIDO+
x(1+xo)
+ A(1
+ )
(B.55)
Remember BDGG is a (d+ 1) x (d+1) 1) block matrix with upper d x d block given by R4 + Purain/ We wish express the above resolvent just terms of Rg.
Working heuristically and block-inverting inverting the BIDG + A(1 + x) resolvent we have
tr (BINO + A(1 + Xo)
-11 -
~ tr(Fn(o))
(B.56)
where
o = Ptrain
+ X(1 + xo)
(B.57)
a
FR(a) = (Ru + ala)
(B.58)
Sample covariance Rg and the Stieltjes transform. We thus have
TXO d+1 X(1 + xo) ur(FR(oo))
A(1 + xo)MM (Ptrain X(1 xo)
1 + Xo
d
where (x) is the asymptotic equivalent of the Stieltjies transform of the Wishart resolvent FR(o) = (Ry + ala)
Following Atanasov et al. (2024) for correlated feature wisharts, we have
1 Mx(a) / s trr((Crrain + soT)-1)
(B.59)
where is defined by the self consistency equation
11
tr( tr(Cerain (Cirain + soTl
Mx(a)
(B.60)
K d
K
25

---

Published as a conference paper at ICLR 2026
and also only depends o, ni and Ctrain In other words, we have self-consisteney equation for (x) as
Mx (())) Crrain + old) |
- =
(B.61)
Similarly, we can express Fn in deterministic form
(Rx + -oaa) Fe(o) = ((11= σ Malot) Ctrain + old)
(B.62)
Note that
Fo(o) = -Falo) (Ra +ola)=2 do
d
o ((=M.l))+ =
_ do
Mu(o)) Cirain + +la) (x(
(B.63)
Given particular problem with particular Ctrain, each of these can be computed.
Self- Cofsistesstt equation for Xo Recall to characterize Feg need to understand X(1 Xo) We now know this variable is defined by
X.
Ptrain
Ar = =1 -1-7 + 2)
XM
(B.64)
In the ridgeless x - 0 limit this equation becomes
Ptrain
(B.65)
While this equation has nontrivial solutions for < 1, its only solution when T >> 1.
Final deterministic equivalence for I in terms of Wishart resolvent. Recall from above that
Tog (BIG (
Ptrain
We can use the trick that
R2
Ptrain by] S (BIG + Actet1 - (Ptriin
+ +AC) +Ac)let) let1)
(B.66)
a
where S is almost-identity matrix that selects top d x (d+ of (d 1) (d+ 1) matrix. We can then simplify
m S ss = SI-oFRR(o 0)
+ +i) (P Ptrain
(BiN + +X)-
(B.67)
a
(B.68)
We stress that this is an approximation that comes from the full block inversion of the resolvent of BING, but is robustly handled in Lu et al. (2025).
CHARACTERIZATION OF QUADRATIC TERM IN ERROR.
The only remaining error term left to approximate is those with the form IBBTT This is what we introduced æ and Bext for.
ve d 1 (T 0)
5. tr (lal (res'B
(B.69)
d c(
for Bplaceholder = La 8 Brest. The definition c(T) comes from the scalar term in the block inversion of (T))
cT(TT) (e(T) = -c (T) Is (()))
(B.70)
with
T TPuain
+ Ar
tr R
Perain FE(xx) R Perainbs] 7) (B.71)
1 Xx
(1 X 2
26

---

Published as a conference paper at ICLR 2026
with
1
(11 1
T
Xx ~ jth
BIDG + + AT
BIDG
(B.72)
-11 BinG + t Brest + TAla+1
-1
Fr(x=) = 1 +
(B.73)
Now following through the calculus, get
d X= dr
Fb(xx) = =-F((x>) Brees (1 1xxx)
BIDG
FÆ(xx)
(B.74)
dm
d
= jtr tr ( ( ( BING
- d Xx
(B.75)
dn
Plugging (B.74) into (B.75) and simplifying for xo gives
TX6
1 tr (Brest BoD F Fr BIDG)
1/ tr( Brsst Fo) 2 Lr(BreaF
(B.76)
(1+ xo)
21 (r(Po) X22 (r( (
for
Fo BIG + (la+1)
x= x(1 x
can finally simplify the quadratic term as
1 di c(TT) ( 0) d
tr( Boeat
Xo (1 (1+ x
PPrrinn tr(TTa ATG) - d (T(TTn(Tan))
(B.77)
DICTIONARY OF VARIOUS DETERMINISTIC EQUIVALENTS
Lemma 4. The following is summary formulas. Their derivations follow simply from the above discussions and characterizations. Let
Fo (BING + )tari)
== + xo)
and 1og as given by (B.68). Then
(r(] Ma(o)
(B.78)
tr|F) (A)] ~ MM
(B.79)
tr| tr|Tēq A ~ tr(Curaiul - - o2
(B,80)
À (r([ (Txx [ Ā (1 20 Mx -2M((())
(B.81)
tr| 11 4Tu ~ tr(Crest] Cresst))
(B.82)
tr[Bnn
-iM((o
(B.83)
tr|Bing Fol = Mx(a) (M)(o)
(B.84)
tr tr(Bict,Fb ~ tr (Crese + Prest Ia F2(())
(B.85)
atestt
Crest + Prest Io F2((0)
tr(BIct ~ tr
(B.86)
arestt
urfley BiDG (T )gg) ) ~ ir(T eq ATDCl - Ā (t([To((Tog))) (Tog)7]
(B.87)
tr(T" (Ic T] tr
(Crest Prest Id
(Ta - 20 Fx (2F
(B.88)
atest
27

---

Published as a conference paper at ICLR 2026
C PROOFS OF ASYMPOTIC IDG AND ICL ERROR CHARACTERIZATION
We finally have all the necessary components to write down the deterministic equivalents of the IDG and ICL errors given by Lemma (or (B.3) and (B.2)). This will comprise our main theoretical result.
Result:
High-d High-dimensional formula for the IDG error
Let
- Cuaann + 2la) 2la)" Fx(2) x
(C.1)
Fx(2) -Fx(() ((=Ma(e)++ 1 x -,((es) Cirain + +la Fx(z) de
d (C.2)
Fx(2),
Mx(z) 1 (r((a(=)),
Mx(()) = ir(F(()))
(C.3)
and
iM, (a) - A2o
(C.4)
a= (p + Cerain) (amiin/(rrraai) + X.
(C.5)
Then we have the deterministic equivalent of (A.2) as
GIDG(T") ~ , p - o2 - İ(1 - MoM
(C.6)
T- -1-12220 - 27Mx(() - X2M)(o))
GEIDG(Arrain)
Result: High-dimensional formula for the ICL error
Take F", a, and À as defined as above, and = eIDG eIDG (Atrai) The deterministic equivalent of (A.3) is then
(Ict(1*) ~ p
p + Grest
(1+ (4) + (q )))M((())
Otest
+ q tr Px()) + (q) 0²) tr [Ctest F.())
(C.7)
EICL (Curain Crest).
Proof. Using Lemma 4 and (B.77) in (B.3) gives deterministic equivalent for EIDG(T*)
(TD) p Ctrain 2(Ctrain + + o² Mx( Ctrain
Mm,(( - 1(1 -o)
1 - 27Mx(0) - (2)
T +1 27Mx(O) - (2)
Ptrain - (Ctrain +o MM - x(1 220Ma(0)-0MMM.0
p 0 )2M((o) X(1 2o.w(o) (004(o))o2M0 o2M((o)))
1- --) (2
+ 1 - 22Mx(o) - {P)
5 (o) - 1(1 - -2
-
5 (1 2)Mx(() (zM)(o)
required.
28

---

Published as a conference paper at ICLR 2026
Using Lemma 4 and (B.77) in (B.2) gives
(IC ~
+ Crest
2 Crest - o tr (CoesFF((o))
+
Crest p Crest Chst La
(4a-20F((0)- -2
Cbest
Crest + p +tet Clest Id]
(Fr(o) + +XE((TO))
+
Chest
M²) - X(1 20.M((o) - M2
T (1 22M,(0) (2)
Replacing
q p pto d² (x) X(1 - 20.M((o) ( -2M
(C.8)
T (1 - 22M - (2M)
and gathering Crest terms we get
CICL (Ciraiu Crest) = p+
P + Crest
(1+ + (a) -
Csest
+ (Crest; (o) + (g\-
(C.9)
as required. The remaining terms only depending on the test distribution through Crost, and thus not containing any structural information about the testing tasks, we call
Escalar (Atrains Ctest) =
Gress
(1+ 20) + (o) -2)
(C.10)
Otyest
and the remaining deeeddeen terms form the Emisalign (Ctrains Ctest) contribution.
Note that the above results define self-consistently in (C.4) for general ridge A. For the main document we have explicitly taken the ridgeless limit, although in general this necessary and can work from (C.4) and use this À in all formulas.
Remark The ridgeless + 0 limit of (C.4)
SMx
Ctrain + P + 2) = -
(C.11)
Ottrain
exhibits different behavior for and < 1. For 1 in the ridgeless limit, À 0 explicitly.
We can solve some of the self-consistency equations here explicitly in the isotropic case and recover previous results.
Remark 3. In the ridgeless limit X + O with isotropic Ctrain Ctrain Ia. have explicitly
Mx(v) V Ctrain
Ctrain
V(
Ctrain
Ctrain K
4Ctrain
(C.12)
K
K
Ā M
@ Oerain
T<I
(C.13)
0
T>1
For Orrin (2025).
1, Crest
Td. and Otrain Otest this recovers the theoretical results of Lu al.
29

---

Published as a conference paper at ICLR 2026
D ORDERING OF EIGENVALUES
Here we analyze the spectrum of the Creain-dependen component of the misalignment error term (19),
K= <F= (y)-o2))(((0))) Given the self-consistent definitions of (o) and (7) in (16) and (B.63), write the eigen- values of K. as
An Au to (+ -²² B, (Crrain) +
(xC)
Andi (Crrai))ain) +
for
A, = -+
1
-MMoo()
a
(D.1)
o BR = Ma(o)) M,(o).
(D.2)
k
Conjecture: Eigenvalue Ordering
We wish to show that
(1) S A2(1)) S... Aa(K))
for
X1 (Cerain)) A2 (Ctrri)n))
)a(Crrain)
i.e. the ordering of eigenvalues K. is opposite that of Ctrain-
We will separate this into two cases. For r > we provide rigorous proof. For T < due to the unclear sign y/2 - q we cannot currently provide rigorous proof with the same methodology. Instead, we provide numerical plots supporting our claim, with the goal of proving this formally for the t < 1 case in subsequent iterations this work.
Lemma. We have that - S (.M < 1 and and M2, 1-
Proof. Rs is a positive semi-definite matrix and so each eigenvalue of aRRe + aTa) is non- negative increasing function of a 2 0. Therefore a tr(( + =1a) increasing in a and so
lim a tr((Rs, + TTa) "1 axx d
is an increasing non-negative function of a. Thus we have
< lim aM,(o) (o)
o
Also as Rg is positive definite it is obvious that
aM)((a)= = lim a trr( + -1o)0 1.
d
By equivalent argument,
8² M, < lim (²M < o
It's helpful to first write the self-consistency equation
((1-1+ o Malo)) Cerain + o
M.(aa = tr
(D.3)
for (o) equivalently as self-consistency equation for the "renormalized ridge" & of this Wishart resolvent, defined by ā @ur((Ctrain + + (M)(o))
30

---

Published as a conference paper at ICLR 2026
equivalently
δ =
(D.4)
1 tr (Crrain + o)
In the - 0 limit, 7 limits to 0 when K > 1, and limits to the unique solution of
σ
(Crrain + +
(D.5)
K K
for K < 1. For our purposes, this means that
lim aMa(d)=
Ri
(D.6)
0, R>1
Differentiating with respect to a gives
lim 0² (,) a-
[K 1,
<<1
(D.7)
0, k>1
S K-1
and so 1 - x (Ma(o) and (x
With this lemma in hand, can prove the main claim. For 1, have Ā 0 and (p + Ctrain/ and
Xe Axde(Cerainl) +o a o² BBA: +
A, (Ctrain)rain) +
Without loss of generality, let's compare A (IC) and (x())). First, we can see from the Lemma above that A, 2 O, and thus we have that
Ag A1 (Ctraii) to
(Ctrran))))) +
Furthermore, have
A, oBB 1 ((x
We reason about this quantity follows. We have that A.. BB as result of this lemma, and so (Ax a Bx & (Ax - a Bx) Rearranging gives
Badi +1
Badg(Crrain)
Agdi(Crraina +
+ a
and are done as we've we've shown (KC) (xC) for (Cerain)) > A2
For F 1. Ā * 0 and so the negative-definite contribution of (a) complicates the above argument. We provide preliminary numerical evidence for this case as follows. We compute numerically from its self-consistency equation (21) to compute AC, of which we plot its eigenvalues against the cigenvalues of Cursin- As demonstrated schematically Figure 5, for values less than 1, it is consistently the case that the cigenvalues of K. negatively correlated with the eigenvalues of Ctrain- We emphasise that this is heuristic check, and we plan to prove this rigorously future iterations of this work.
E TRACE INEQUALITIES FOR MISALIGNMENT
In this appendix, give self-contained expository proof of Ruhe's trace inequality in the form we require for Corollary 4.1. We actually present two lines of analysis: first a direct proof of the inequality, and then an alternative analysis based Riemannian optimization.
We begin by recalling the general setup: let A and be x real symmetric matrices. By the spec- tral theorem, they orthogonally diagonalizable, with non-increasingly ordered real eigenvalues
31

---

Published as a conference paper at ICLR 2026
0.1
0.25
0.5
0.75
0.9
0.5, 0.5
0.5,
100
2.50
50-
Eigenvalues of R.
2.25
25
2.00
1.75
7.0
1.50
155
125
1.0-
0.5
3.00
0.25
0.50
0.35
1.00
1.25
1.50
L.P5
2.50
0.50
0.29
0.30
0.15
1.0
3.25
1.50
1755
2.00
100,
0.5
100,
100
1.5
10-
Eigenvalues of R.
1.0
0.0
0.6-
0.4
200
0.5
0:00
1000
2.25
150
177
1.00
0.50
0.23
0:30
0.79
1.00
150
1.75
200
Elgenvalues of Crain
Eigenvalues of Grain
Figure 5: Demonstration of opposing eigenvalues for r < values, at a range of k and a values. Curann here the same as in Figure 2, i.e. powerlaw spectrum.
(A) 2 A2(A) 2 ... 2 Aa(A) and A1(B) 2 AB( 2 .... Ad(B), respectively, We want to show that
d Σ
d
(A tt((A S Σ (A))
(E.1)
j=1
j=1
with equality attained in either bound and only if A and B are co-diagonalizable.
Writing the eigendecompositioon of A and B as
A = O and B == OBABOT
(F.2)
respectively, where OAOT la. OnO Ta [Ma]e Ar(A)osj. and [An]y ABB) we see immediately that it suffices to consider the case in which is diagonal, with non-decreasing elements, as
Lr(AB) (TAAOOAOT) where O=O]On.
(E.3)
and the matrix OAgOT has the same eigenvalues as B. We can see that A and B are co- diagonalizable if and only if OA is diagonal. Morcover, can be attained only if OA TT is in fact equal to AB up to pertmuation of its diagonal elements. What remains is show that they are in fact bounds.
see that the claimed bounds
E. DIRECT PROOF OF RUHE'S INEQUALITY
We now give a direct proof of the claimed inequality. We remark that the version of Ruhe's 's inequality proved in Marshall et al. (2010) is not sufficient for our purposes, as it assumes that both A and B
32

---

Published as a conference paper at ICLR 2026
i.e., that (A( 2 0 and XB( 2 0. We instead follow the proof for general Hermitian matrices outlined in Li (2020)'s 's blog post.
positive
The strategy of this proof is to proceed by induction the matrix dimension d. The claim clearly holds for = 1. as then A and B are scalars equal to their only eigenvalue, and both bounds collapse. Now consider d 1. Assuming- ordered eigenvalues, we have
noted before-ther that A is diagonal with elements given by its
tr(AB) = Σ (A)
(E.4)
1=1
> (A((A) +
+ Aa(A))Bss
(E.5)
(A)(A) *c(AMBB + Σ {o(A)BBJ
(E.6)
We now observe that
(A,(A) -A) {d(A)BB;, (rB
(E.7)
where A is a (d- 1) (d- 1) diagonal matrix with رزj Xy(A) (A and B is the (d- 1) x ()) principal submatrix of B given by discarding its last row and column. The ordering of the eigenvalues of A implies the ordering
(A) Aa(A) 22.... ... (e-1 (()))))(-AA2)))) 2
(E.8)
of the eigenvalues of A.
Thus, tr( tr(Aßr (B is the trace of product of (d - 1) x (d - 1) real symmetric matrices with having its eigenvalues in non non-decreasing order along the diagonal, so on the induction hypothesis we have the bound
-MA))- S Su(A S
(E.9)
By the Poincaré separation theorem (also known as the Cauchy interlacing theorem), as B is principal submatrix of the real symmetric matrix B, we have the incquality
Xy+1 (B) S (B) & (B))
(E.
for all j [d - 1).
Combining these results and using the fact that Az(A) - >a(A) 2 0 by our ordering assumption, we thus have the upper bound
d-1
d
tr(AB)= S
Σ
(A)(4) - (B(A))A,( Σ \a
(E.11)
1=1
j=1
d (4) - (A(4)))( Σ Xa(A))
(E.12)
5=1
d-1
(B)A),((B - -2(AA)
Σ
1,(B) xaa(A) Σ Bij
(E.13)
Σ (AAA),)(B) - Aa(A)
Σ Az(B) + Aa(AA)
Σ
A3(
(E.14)
1=1
Σ (B))
(E.15)
1=1
33

---

Published as a conference paper at ICLR 2026
AB( Here, the first line uses the induction hypothesis in the form re-stated above, the second line uses the upper bound from the Poincaré separation theorem, and the remaining three lines are just algebraic simplification. Similarly, we have the lower bound
as Bi
4-1 EA((A)) -A(A))a--B
Lr(AB) M
Σ
(a(A)
(B.16)
4-1 (2((A) XcAA)AA-s+((BB)
d Σ
Σ
Xa(A)
(B.17)
j=1
d-1
(AA)A -d(AA) Σ
Aa-j+1(B). + Xa(A)
Σ
(B)
(E.18)
==1
j=1
E ((A)) - Aa(A) Σ 1₂(B) Aa(AA) Σ Nr(B)
(B.19)
1=2
(AA))
(E.20)
By induction, we therefore conclude that
(B) S Σ (A)A)),(B),
(E.21)
as desired.
E.2 RIEMANNIAN OPTIMIZATION
We now give an alternative perspective on this result based on the properties of orthogonal matrices. Consider
f(O) = (AAA
(E.22)
function the space of orthogonal matrices, where ÅB is diagonal matrix with non-zero elements given by any permutation of the eigenvalues of B. Near the identity, we can write any orthogonal matrix
1 252 O(1)) T= 2
(E.23)
where S is a skew-symmetric matrix (ST = S) and is a small parameter. Substituting this expansion into f(O), we have
I f tr(A - A + O(t)),
f(O) = tr(AAAn) ttr((AA, +
(E.24)
where [X,Y] XY YX is the matrix commutator. Irrespective of the permutation we have chosen to define ÂB, [AA, = 0, so f(O) is stationary at the identity.
We therefore consider the quadratic term
H tr(A 2AASĂBS A
(E.25)
as it will determine whether this stationary point is a local minimum, local maximum, saddle. As [A An] O, we can immediately simplify this to
H = tr(AAA A^SAES).
(E.26)
34

---

Published as a conference paper at ICLR 2026
For brevity, define a, (Aa)a and b, (Ag)s. Then, expanding in components,
Σ (a,b;
(B.27)
s==1
Σ a;(b; b
(E.28)
as by skem-symmetry Sys -Sij. Using the fact that the terms clearly vanish, we have
=
11-Eaate Ea(b. - 6,) Σ abbb: - b,)
(E.29)
By re-labling indices i ++ j and using the fact that S3 Sty
Σ aib by) Σ ag -b;)s
(B.30)
so upon re-combining terms
H H=- aa)((;))bb (0)][b--b)) -0--bb)) bb) icj
(E.31)
This now depends on the ordering of the eigenvalues, and thus on the permutation of indices. If the eigenvalues of both A and B are in non-increasing order such that a; 2 aj and by 2 by whenever i j. we have < O, and the identity a local maximum of f(O). In contrast, if they are oppositely ordered, then H 0 and the identity is a local minimum. Otherwise, the stationary point is a saddle. These results are consistent with Ruhe's 's inequality as proved above.
F PHASE TRANSITION WITH INCREASING PRETRAINING TASK DIVERSITY
As have seen throughout, the task diversity parameter is crucial because it quantifies the quality of the pretraining data: higher i implies a richer span of task variations, enabling the in-context learner infer the structure shared between pretrain and test sets more accurately. Understanding thus sheds light on the implicit algorithm performed by the model. Previous work from Lu et al. (2025); Raventós al. (2023) in particular have investigated the performance of ICL as s: increases. In particular, Lu et al. (2025) showed, in the isotropic task case, that there is a phase transition at = 1: for K < the model has insufficient task diversity to be able to generalize within tasks, and is forced to memorize the training tasks, while for &. 1, the model has seen sufficient tasks to generalize in-context efficiently.
The settings considered in both Lu et al. (2025); Raventós et al. (2023) focus isotropic tasks where the pretraining and testing task distributions are the same. Here, we verify that not only does this phase transition in task-diversity remain in the presence of structured pretraining task distributions, but further is independent of the test distribution (see Figure 6).
Corollary F.1: Phase transition in task diversity
In the proportional limit of a, + such that a/7 for 7 (1) fixed, there is phase transition in eICL at k = Vank(Crrain)// Specifically, for rank (urain) // then
(p= (1
p=>) tr Ctest (2) Otrain + Ta)
p+oesum
K < r
lim eICL (CrraCrs Crest) 7,00900
+
1 P p+eeram
tr
Crest (2 + Ia)
K > r
(F.I)
where and xy are defined self self-consistently as solutions of
tr
(x) Ourain + +1a) -1 = - K
(F2)
tr
(z2 Curain + Ia) T.
(F:3)
35

---

Published as a conference paper at ICLR 2026
Max-aligned
Max-aligned spiie
Test on pretrain
Test pretrain
0.2
Test aligned
Test on aligned powerlow
Test on isotropic
Isotropic
Test unaligned powerfaw
Test on unaligned powerlaw
0.
0.50 0.75
1,03
1.25
1.50
3.75
2.00
225
0.25
0.50
0.75
1.50
3.75
2.00
225
(a) Full-rank phase transition
(b) Half-rank phase transition
Figure 6: Theory curves with corresponding numerical simulations showing phase transition of ICL error in for range of test structures. For (a) Ctrain is the same as in Figure and is full-rank; for (b) Cresin = diag diag(2la/2 O4/2l, thus half-rank. Parameters:
d = 80, 80, = 80, 0.01. Tests for each are done on Ctrain La, powerlaw (spectral power 0.5, aligned and unaligned with Ctrain), as well as the rank-1 Ctest optimal from Result 4.2.
To prove this result, we first suppose that Ctrain is invertible. We start by writing the self-consistency equation
Ma(a) Ma(a)= - 2 Cerain + toaa
(F.4)
for Mx equivalently as self-consisenncy equation for the renormalized ridge o, defined by
ō ttr[(Ctrain
equivalently
(F.5)
(Cirain + 5) -
In the + limit, limits to when x > and limits to the unique solution of
ō (Cerain + +0)
(R6)
k K
for K < 1. For our purposes, this means that
[1 11-K, K3
lim aMM(())
(R.7
o+0
[0,
K >
as well as
1 +
aMa(o)
lim FF((o) lim o-0
1
Cuain la
(F.8)
(xCrrain + +aa
-1
K <<1
(F19)
0,
K>1
and finally
lim ('F
((00))((((rttr -(rCtrain +
+1)
K < K.
(F.10)
(,
>
for x defined by tr (m) Craan + Ia) -1] = -1-k.
Using these characterizations of the 0 limit, and noticing for problem that if 7 0o (so ă 0), then must take the -> limit of eICL, being careful to use
T + Curain 1
(F.11)
7
36

---

Published as a conference paper at ICLR 2026
to enforce the proportional a/t limit. Doing gives the required answer.
The above proof requires explicitly that Ctrain is invertible. is not, we have to more careful with handling the - limit of (D.4). One potential branch of this solution is
+ # tr((Ctrain + + +
(F.12)
R
or equivalently
ΣΣ á
1
7(d
(F.13)
d
where r is the rank of Crrain- This is what causes the split in behavior of the solution K r/d (which was previously 1). For K < r/d, this is solvable at nonzero ō and we end up with the familiar solution branch of
aM, (a) tr +
(F.14)
For n. >> r/d, there is no longer sensible solution at nonzero ā (since cannot be negative) and so we're forced to take + 0. This gives
1 ā Nō d
oM,(o) न (Crrann + +
(d)
(F.15)
0+0
This concludes the proof of the claimed result, and recovers the previous result for invertible Ctrain-
G CONTEXT LENGTH SHIFTS
Training a
Guan isotropic
Gras powerlaw, power -
Coun powerlaw, power 1.5
ICL etroor
2215
10
16"
10a
10'
Test-time l/d
Figure 7: Theory curves with corresponding numerical simulations showing monotonicity of error test-time context length for variety of task structures. The dashed line corresponds to training a. Parameters: d = 150, a 2, = 4, k = 1, 0.01. Testing is done on the same distribution as pretraining.
We have thus far assumed that the pretraining- and test-time context lengths are the same. Our gen- eral result in (C.7) allows for an atest that differs from the training context length a, but shows that this shift only affects a single scaling factor in escalars where a factor of 1/a is replaced by 1/ 1/Qtest- Despite EICI. not being monotonic training context length, as demonstrated by Lu et al. (2025), it is monotonically decreasing in test context length: testing on larger-context prompts can only de- crease error, and testing on shorter-context prompts can only increase error. This is demonstrated in Figure 7 for variety of different task structures. Recall that aee only appears in eICl through escalars as part of an effective noise term: having longer test-time contexts will give the model a bet- ter estimate of the token distribution, leading better predictor for the task corresponding that context.
37

---

Published as a conference paper at ICLR 2026
H
EXPERIMENTAL DETAILS
Code for all simulations can be found on Github.
H.1 LINEAR SIMULATIONS
Linear simulations are done by sampling pretraining and testing distributions as described by (2) and (12), and computing numerical I* by (B.1). We simulate the ridgeless À 0 limit of this theory by taking À = 0.00001.
To justify the Ansatz of taking the value component 12 trained linear attention behaves similarly in ICL error to the Ir(T"H") model, with the discrepancy decreasing as d increases. For the full-parameter linear attention experiments here, we initialise V, K, Q with LeCun initialisation (i.c. so 221 0), and train for 10000 epochs with BS d2 using Adam W with LR 0.0001.
include Figure 8 demonstrating that
de irter altoisn
6-32 Inear sttenson
incer attenzin
repuced gamma
KLL wner
9.55
Figure 8: Comparing trained linear attention with full non-zero V, K, Q parameters to reduced T* model at various d. Here p = 0.01, 5, a = 1. Errorbars and shading show +std over 20 runs of each experiment.
We also include Figure showing that explicitly initalising 22 (i.e. the chosen Ansatz) im- proves the match between trained linear attention and our reduced I model, as expected, but this improvement not extreme.
ihi cation
0-
5.20
0.73
1.50
1.25
1.50
200
Figure 9: Comparing reduced gamma model at "low dimension" d = 16 trained linear attention with both LeCun initialisation (U21 * 0) and V = la+1 initialisation (221 = 0). Indeed the latter achieves a closer match to the reduced model, expected, but the difference is not stark. Same parameters and training proceedure as Figure
38

---

Published as a conference paper at ICLR 2026
H.2 NONLINEAR ARCHITECTURE
The architecture used for Figure 3 is a transformer formed of two transformer-blocks, constructed follows. The input to the architecture is the Z matrix (3). We apply causal mask to ensure that each token attends to itself and prior tokens in the sequence. Each transformer-block made of first applying single-head softmax attention this masked input, with residual connection, and then normalised; this is then passed to a single hidden layer MLP with GELU activation, followed by final residual connection and layer norm application. The final logit computed by a dense layer projecting the output of the two transformer- blocks into scalar.
The model is pretrained with data sampled by (2). We form Z-matrices from these samples, which are again the inputs to the architecture, and the model is pretrained to predict the corresponding 1 value by minimising MSE loss between the final logit and ye+11- Training done using Adam W for 1000 epochs with batch size 16 and learning rate 0.0001.
Testing is done by sampling a batch of n contexts from (12), and tracking MSE between the true M1+1 and the model output. This test sampling is repeated 500 times to average out noise.
39


Reviews:
[{"review_id": "B1FTtBQQJn", "reviewer": "Reviewer_ZJRu", "summary": "This paper studies how the structure of pretraining data affects in-context learning (ICL) using a simplified linear attention model for linear regression, where training and testing tasks come from Gaussian distributions with different covariances. The authors derive an exact analytical expression for ICL error and propose a new alignment measure between pretraining and test task distributions that accounts for finite-sample effects, task diversity, and noise. This measure correlates extremely well with actual ICL performance, even in trained nonlinear Transformers. Finally, they show that pretraining on the same task distribution is not always optimal—when task diversity is limited, models generalize better if pretrained on a lower-dimensional (power-law structured) set of tasks.", "questions": "See weaknesses above and the following:\n\nI’d appreciate the authors’ view on how Section 4 relates to the “task-scaling vs. context-scaling” perspective (e.g., Abedsoltan et al., 2024: https://arxiv.org/pdf/2410.12783). That work finds that Transformers (unlike MLPs) benefit from both (i) increasing the number of pretraining tasks and (ii) longer context at test time.\n\nIn your Fig. 4, my understanding is that context length is held fixed, while task diversity (κ) and spectral shape (p_train) are varied. Could you clarify:\n\n1) If we allow context length to grow (while task diversity remains small), do your results predict that matching the pretraining distribution to the test distribution becomes optimal in the large-context limit? Or does some degree of structured misalignment remain beneficial?\n\n2) Since Fig. 4 already shows that increasing task diversity removes the non-optimality, how would you expect context scaling to interact with this effect? Would larger contexts further reduce alignment sensitivity, or play an orthogonal role?\n\n3) Finally, It would be interesting to understand whether your framework can jointly explain task and context scaling effects.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well written and clearly states its problem; overall, it is an engaging and enjoyable read.  \n2. The analysis in Section 4 is particularly insightful and well structured. The authors examine two complementary cases—fixing `C_train` to find the optimal `C_test`, and fixing `C_test` to find the optimal `C_train`—to understand when pretraining and test task structures align or conflict. The finding that pretraining on the true test distribution `C_test` can be suboptimal in low-data regimes, where a lower-dimensional or more anisotropic pretraining distribution yields better ICL performance, is especially thought-provoking.  \n3. The proposed alignment measure is well-motivated and technically novel. It quantifies how well the pretraining and test task structures align while explicitly accounting for finite-sample effects, task diversity, and noise—factors often ignored in simpler population-level alignment metrics.  \n4. The paper provides strong empirical validation: the proposed measure correlates almost perfectly with actual ICL generalization error in both the analytically solvable linear model, outperforming other common similarity metrics such as `C_test C_train^{-1}` and CKA.", "weaknesses": "I think the paper is beautiful from a theoretical perspective. However, from a more practical ICL standpoint:\n\n1. The paper’s main results are construction-based. It would strengthen the work to show whether a linear attention model trained with stochastic gradient methods actually converges to the analytically derived (constructed) solution. This matters because, while constructions can represent any algorithm in principle, demonstrating convergence under realistic training dynamics would make the results more compelling for practical settings.\n\n2. It remains unclear how additional architectural components—such as multiple layers or MLP blocks—would affect the proposed alignment and the observed “non-optimality of pretraining on the test distribution.” In more general Transformer settings, would the same phenomenon hold, or could deeper architectures mitigate or amplify this effect? Some discussion or intuition on this would be valuable.", "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": "b2ZkPFa8ky", "reviewer": "Reviewer_vZ3u", "summary": "This paper studies a linear self-attention model trained on an in-context linear regression task, where the task vectors in training and at test time follow different distributions. By deriving the test error in terms of the data statistics in a high-dimensional limit, the authors quantify how the information in training is useful at test time. The analysis also reveals that the optimal pretraining distribution in the sense of minimizing the test error is not the test distribution in general.", "questions": "- Section 4 makes it clear that, given a test distribution, it is not always optimal to pretrain on the test distribution. But it's not yet clear to me how one should design the optimal pretraining distribution that minimizes the test error. Could the authors clarify whether this question is addressed in the paper, or if it remains an open problem?\n- In line 369, what are MLP connections? Is the architecture softmax attention and MLP with skip connections?\n- A very minor suggestion: it might be helpful to avoid using the word \"attention\" in its psychological sense (lines 38, 407) in a paper that studies the self-attention model.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written and mostly clear. Frequently providing intuitive interpretations for the equations was very helpful for me to follow the theoretical results and their implications.\n- The paper studies an interesting setting in which the in-context task distribution is different in training and at test time, providing analytical expressions for the test error in a high-dimensional limit.\n- The finding that it is not always optimal to pretrain on the test distribution is interesting and somewhat surprising.\n- The alignment measure introduced in this paper predicts the test error of softmax transformers more accurately than the conventional alignment measures, justifying the applicability of the theory developed in the simplified setup.", "weaknesses": "- I have several questions regarding the motivation to drop two terms in Equation (7). \n\n  First, [1] showed that $v_{21}$ remains zero if initialized at zero with  $w_\\mu\\sim \\mathcal N(0,I)$ and $n\\to \\infty$. Their proof relies on the mean of $w_\\mu$ over $n$ sample sequences being zero. However, in this paper's setup with a finite (and possibly small) $k$, the mean of $w_\\mu$ isn't close to zero. Hence, using [1] to motivate setting $v_{21}=0$ in line 165 may not be applicable.\n\n  Second, I'm not sure why $x_ix_i^\\top v_{21}$ is considered to contain no task information (line 161). One reason linear attention cannot achieve zero training/test loss, even when the noise $\\epsilon=0$ and the sequence length is sufficient to recover $w$, is that the model relies on the in-context covariance estimated from the training set to approximate $\\sum_{i=1}^l x_i x_i^\\top$ for new test sequences, rather than using the true $x_i$ covariance of the current context. So $\\sum_{i=1}^l x_i x_i^\\top$ should contain essential information for robustly inferring the task vector. The limitation instead seems to come from expressivity -- single-layer linear attention lacks the expressivity to properly utilize the information in $\\sum_{i=1}^l x_i x_i^\\top$.\n\n- In Equation (11), does ridge regularization in the reduced model match ridge regularization in the original model? Training the reduced model regularizing $\\Gamma$ may yield different models from training the original linear attention regularizing the value, key and query weights. If they are not equivalent, it would be important to clarify this distinction.\n\n[1] Zhang, R., Frei, S., & Bartlett, P. L. Trained transformers learn linear models in-context. 2024 JMLR.", "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": "lGqo56wTjn", "reviewer": "Reviewer_4Nn9", "summary": "This paper  investigates how aligning or misaligning train and test over task diversity contributes to or minimizes ICL error.  The authors use a simplified linear attention model for ICL of linear regression tasks that enables them to formally characterize ICL generalization error and that gives them a new way of measuring alignment between training and testing data sets with respect to tasks.", "questions": "Are you predicting only for $y_n$ given a fixed input sequence $x_1, y_1, ... x_n$ ?\n\ngoverning the input?\n\nIs generalization or testing as in equation 12 is just limited to shifting the slope, as the distribution for test x input and scalar epsilon are identical to train?\n\n\nHere are a few typos and grammatical suggestions:\n\nl, 391: increases/ aggravates/exacerbates but not extremizes\n\nl 431  all the signal", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Understanding generalization error in ICL is clearly linked to pretraining and this paper provides a novel approach to thinking about generalization in terms of task diversity, whereas prior work has looked at other generalization problems.  This work can in principle give us a better understanding of how pretraining can affect ICL.  Intuitively task diversity is a dimension on which real ICL as opposed to memorization can emerge.", "weaknesses": "However if I understand correctly, task diversity here has apparently quite a limited definition.  as in equations 2, 3, 18 for example task diversity in the linear equation set up discussed here is primarily the variance in the matrix W in the linear relationship y_i = Wx_i +b for the inputs in the task.  You do not vary $d$ the dimensionality of the tokens in C train nor do you vary the distributions for x_i or b.\n\nThus, the authors propose an alignment metric for train and test that predicts ICL error on a limited testing suite when the main parameter of variation is the slope of the linear relationship between x and y in the input context.  As the linear relationship here is in d dimensions, in fact each dimension m can convey a \"different\" affine equation, $y^m = a^mx^m$.   Thus task diversity comes down to how many linear equations are encoded in the matrix W.\n\nThis is potentially very interesting.  If we knew that language generation tasks are in general expressible in terms of m linear relationships using d dimensions, then this gives us a general approach to ICL.  However, even for simple tasks and simple attention layer only models, the outputs of the attention layer formulated as a mathematical function become highly nonlinear \n\nThis is a partial step but only a partial step in learning how ICL really works and how it relies on pretraining.  The reason for this partiality is that that the divergence of test and train is very limited.  Only in the case in which we are varying C train with respect to task, in the very specific sense you have defined, do, for instance, the elements  F(z) and M(z) measure the degree to which we can capture the information in C train.  \n\nBut generalization should also cover generalizing from the test data inputs for $x$ and  $b$, for the linear function learning task as in Garg et al 2022.   Naim and Asher (2024), for instance, take a detailed look at what happens when we shift distributions for inputs, slopes and intercepts (x, a, b) in a one dimensional linear function learning task.  The results show that transformer models with softmax and normalization don't really compute functions like regression and develop systematic errors as training and test distributions diverge.   Another paper, Naim and Asher 2025, shows formally that transformer architectures have inherent limitations in computing linear functions.  However, they also show that the same transformer behavior applies to polynomial functions and even some continuous functions.  There should be some discussion of how generalization on task diversity interacts with these other forms of generalization, espeically since the observation that adding task samples does not always improve performance matches Naim and Asher's observation of degradation of predictions when extending in context examples after a certain point. \n\nMore generally, omitting a discussion about how other dimensions of the problem affect distribution vitiates at least to some degree the authors' claim that \"these results identify train-test task alignment as a key\ndeterminant of generalization in ICL,\" \n\nThere are also some specific worries.\n\n\n\nIn formula 4, you do add but not Norm, as is standard in transformer architectures.  With linear attention in addition there is no exponential. Thus it looks like formula 9 becomes the output of a 1 layer 1 attention head model, which Naim and Asher 2024 show doesn't really have much predictive power. \n\nIn equation 8 you drop certain terms that do contribute to the model output values for given tokens.  So I have worries about how these analytical results transfer actually to working transformer models with multiple layers, multiple attention heads and non linear scoring functions.\n\nYou are working with linear attention and with restricted assumptions on how the calculation of values in the attention layers is made (see remarks above).  So while it's really interesting that the measure works on a more standard transformer as in H2, we can't transfer the mathematical analysis very readily to the non linear transformer of H2. \n\n\n\n\n\nThe last section is interesting  but given that only task diversity seems to be changing, then it's not clear why this is so surprising\n\nIn the conclusion you speak of \"fully general task covariates\".  But the entire exercise is built around a linear relationship between inputs x_i and labels y_i.  So how does this provide a fully general task variation?  In many generation ICL tasks it is not at all obvious that there is any linear relationship between the input context and the generated string.  The linear relationship is a very specific ICL restriction.\n\nOmar Naim and Nicholas Asher. Re-examining learning linear functions in context. arXiv:2411.11465\n[cs.LG], 2024\n\nOmar Naim and Nicholas Asher.  analyzing limits of icl \nhttps://arxiv.org/pdf/2502.03503, 2025", "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": "uP7ai5uGJP", "reviewer": "Reviewer_cB3E", "summary": "The authors present a theory of in-context learning of linear regression in a simplified one-layer transformer, and calculate the generalization error for in & out of distribution tasks.", "questions": "In addition to the questions above, I have a few more:\n- From Figure 1, it seems that $e_\\text{misalign}$ captures much of the behavior of $e_\\text{ICL}$. In what regime is this true?\n- In many figures, specific values for the thermodynamic parameters $\\alpha, \\tau, \\kappa$ are chosen. Are the values chosen part of a typical regime? I might imagine that some of the conclusions drawn will change as these parameters are varied (as an example, the authors note that previous work identifies a phase transition at $\\tau=1$)\n- What are the main obstacles to generalizing the theory beyond linear problems & Gaussian data? (I.e. what goes wrong in the calculation)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The authors provide key steps towards a more complete understanding of ICL in transformers, through their analysis of a simplified transformer & problem setting. \n - The subject matter considered is relevant to the state of the field and in particular to AI safety (Understanding distribution shifts in transformers is important for safety concerns)\n - The work is, to my knowledge, original\n - The paper is presented clearly, with all mathematical objects given appropriate intuition\n - While I have not carefully checked the calculations, the results presented follow earlier work I am familiar with and the methods used seem plausible.", "weaknesses": "In my opinion, the main weakness of the paper as it stands is that not enough is done to explain the limits of the toy model considered. While the authors do show empirically that their model explains behaviors of non-toy, nonlinear, transformers trained on linear problems, it is not yet clear how far one can stretch this model. \n\nIn particular:\n- The data is still quite simple: Does the model hold up for non-gaussian data (with a linear rule)?\n- I imagine the model will break if the in-context regression problem is nonlinear. Is this correct?\n- What can we learn about real data from the alignment measures proposed in the paper?\n- The authors claim that mismatched task distributions are often optimal. Is there a real-world example of this phenomenon? (I.e. does the theory at least make a qualitative prediction about LLMs?)\n- What further steps must be taken to bring theory & practice closer together?", "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)