Understanding the Role of Training Data in Test-Time Scaling
Authors: Adel Javanmard
Baharan Mirzasoleiman
Vahab Mirrokni
Venue: ICLR 2026 Poster (2026)
Keywords: Language models
Learning theory
Chains-of-Thought
Inference compute
Test error
Primary area: learning theory
TL;DR: Theoretical framework to study test-time scaling, Chains-of-thought reasoning, and task selection for training in the setting of in-context learning with linear self attention.

Abstract: Test-time scaling improves the reasoning capabilities of large language models (LLMs) by allocating extra compute to generate longer Chains-of-Thoughts (CoTs). This enables models to  tackle more complex problem by breaking them down into additional steps, backtracking, and correcting mistakes. Despite its strong performance--demonstrated by OpenAI's o1 and DeepSeek R1, the conditions in the training data under which long CoTs emerge, and when such long CoTs improve the performance, remain unclear. In this paper, we study the performance of test-time scaling for transformers trained on an in-context weight prediction task for linear regression. Our analysis provides a theoretical explanation for several intriguing observations: First,  at any fixed test error, increasing test-time compute allows us to reduce the number of in-context examples (context length) in training prompts. Second, if the skills required to solve a downstream task are not sufficiently present in the training data, increasing test-time compute can harm performance. Finally, we characterize task hardness via the smallest eigenvalue of its feature covariance matrix and show that training on a diverse, relevant, and hard set of tasks results in best performance for test-time scaling. We confirm our findings with experiments on large, nonlinear transformer architectures.

Published as a conference paper at ICLR 2026
UNDERSTANDING THE ROLE OF TRAINING DATA IN TEST-TIME SCALING
Adel Javanmard Baharan Mirzasoleiman") Vahab Mirrokni Uniiersiity of Southern California, UUniversity of California Los Angeles, Google Research a jjavanmauuse. edu, baharan(cs ucla.edu, mirrokni@google com
ABSTRACT
Test-time scaling improves the reasoning capabilities of large language models (LLMs) by allocating extra compute to generate longer Chains-of- Thoughtss (CoTs). This enables models to tackle more complex problem by breaking them down into additional steps, backtracking, and correcting mistakes. Despite its strong cesormnnce-dmmossttee
by OpenAl's ol and DeepSeek RI, the conditions in the training data under which long CoT's emerge, and when such long CoTs improve the performance, remain unclear. In this paper, we study the performance of test-time scaling for transformers trained an in-context weight prediction task for linear regression. Our analysis provides theoretical explanation for several intriguing observations: First, at any fixed test error, increasing test-time compute allows us to reduce the number of in-context examples (context length) in training prompts. Second, if the skills required to solve downstream task are not sufficiently present in the training data, increasing test-time compute can harm performance. Finally, we characterize task hardness via the smallest eigenvalue of its feature covariance matrix and show that training hard set of tasks results in best performance for test-time scaling. We confirm our findings with experiments large, nonlinear transformer architectures.
diverse, relevant, and
INTRODUCTION
Scaling test-time compute enhances inference in large language models (LLMs), by enabling reason- ing with long chains-of- of-thughtt (CoTs). This allows models generate more intermediate reasoning steps for complex problems, evaluate multiple options, and backtrack to find more accurate answers, all without changing the model's parameters. There been recent body of work on this idea (Snell et al. 2024, Welleck al. 2024| Muennighoff et al. 2025 Yeo et al. 2025). with OpenAl's ol (OpenAl 2024 and DeepSeek RI (Guo et al. (205)) demonstrating strong reasoning performance with consistent gains from scaling test-time compute. However, our understanding of the training data properties that support test-time scaling remains limited.
Training diverse and difficult data has shown to be beneficial to enable test-time scaling complex reasoning tasks, such as mathematical competitions (Muennighoff et al., 2025), medical reasoning (Huang et al., 2025b),, and code (Yu et al. 2025). Difficult examples cannot be answered by the model being trained or other more powerful proxy models. However, the precise notion of difficulty and the relation between the amount of compute training and test time remains unclear. In particular,
often identified as those that
Does increasing the test-time compute always improve the downstream reasoning perfor- mance?
(ii)
Can increasing the test-time compute lower the requirement training-time compute?
(iii) What are difficult training examples and why are they beneficial for test-time scaling? Addressing this question requires a rigorous understanding of the effect of training data and its properties on the performance of test-time scaling.
In this paper, we theoretically study the performance of test-time scaling for transformers trained on an in-context weight prediction task for linear regression, where the goal is to predict the linear weight

---

Published as a conference paper at ICLR 2026
vector from the sequence of imput prompts. This framework has been used previously for analyzing the mechanism underlying training CoT (Huang et al.] 2025a). During training, the model performs direct coneeeext-learinng and outputs its prediction of the weight vector. At test time, the transformer performs CoTT' and generates multiple intermediate steps before arriving at its final prediction of the weight vector. Our analysis yields several intriguing findings: First, fixing the test error, by increasing the test-time compute we can decrease the number of in-context examples (context length) in training prompts. Second, if the skills needed to solve the downstream task (corresponding to directions in the data covariance matrix) are not sufficiently represented in the training data, increasing test-time compute can harm performance, effectively causing the model to overthink. Finally, we characterize hardness of a task based on the smallest eigenvalue of its feature covariance matrix and show that training on a diverse, relevant and hard set of tasks during training yields the best performance for test-time scaling.
Our main contributions and the organization of the paper are discussed below:
(a)
In Section 3 We study in-context learning in transformers with single linear self-attention (LSA) layer trained via gradient descent. Despite the problem's 's non-convexity, we show that gradient descent, when initialized randomly but suitably, converges to a global minimum, which we explicitly characterize. Our analysis allows for general feature covariance. During training, the model engages in direct in-context learning, but at test time we employ chain-of-thoughh (CoT) prompting to let the model generate intermediate reasoning steps before producing its final output. We demonstrate that, with CoT prompting at test time, the transformer effectively implements a multi-step (pseudo-) Newton's 's method for loss optimization. Notably, this part of our contribution extends the results of Zhang et al. (2024) by incorporating CoT dynamics at test time, and of Huang ct al. 2025a) by accommodating general feature covariance.
(6) By analyzing the expected estimation error of in-context weights for test prompts, in Section 3.4 we introduce a measure of task hardness defined by the ratio the smallest eigenvalue of the feature covariance matrix to its trace. We interpret the eigenvectors representing different skills relevant to the task, with the corresponding eigenvalues indicating the strength of those skills. Under this interpretation, hard tasks are characterized by a long-tailes spectrum of skills, while easy tasks correspond to having only a few well-balanced skills.
This analysis leads to two key consequences: (1) For a fixed test error, increasing test-time compute allows us to reduce the required number of in-context examples (i.e., the context length) training prompts. (2) We derive test-time scaling laws for our ICL setting, capturing how test error depends on test-time compute and highlighting the role of factors such as context length, feature dimension, and task covariance structure in shaping the overall trend.
(c) In Section 5.1 we study setting with T tasks, where each task is specified by its feature covariance matrix (interpreted, as discussed in part (b), as the set of skills required for the task together with their relative strengths). We extend the analysis of Section to this multi-task setting and characterize the estimation eтror of the final CoT output. Based on this characterization, we formulate quadratic optimization problem to determine the optimal task selection probabilities, demonstrating that training on a diverse, relevant, and sufficiently hard set of tasks yields the best performance under test-time scaling. We validate theoretical results with experiments both Linear Self-Attentio (LSA) models and the more complex nonlinear transformer architecture GPT-2.
RELATED WORK
Recent work has highlighted several phenomena relevant our study. First, it has been observed that simply increasing test-time compute and reasoning depth can, counterintuitively, harm performance, a phenomenon termed overthinking. The empirical study of et al.] (2025) suggests that LLMs tend to overthink simple problems by generating unnecessarily long outputs, and underthink harder ones, by providing shallow or incomplete reasoning that overlooks critical steps. In Wang et al.) (2025), it is argued that exploring more reasoning branches may degrade system efficiency as many branches may be trapped in overthinking. Second, Recent work has explored the test-time scaling paradigm (Snell et al. 2024: Welleck et al. 20244,, with OpenAl's ol (OpenAl 2024) and DeepSeek RI (Guo et al.] (025) demonstrating strong performance through reinforcement learning on millions of samples and multiple training stages. MMennighoff al. 2025) proposes simple framework,
2

---

Published as a conference paper at ICLR 2026
which involves training only 1,000 samples with next-token prediction and controlling thinking duration via a simple test-time technique, and show that achieves test-time scaling and strong reasoning performance. Finally, prior studies data mixtures emphasize the importance of balancing training corpora with sufficient coverage of topics matched to downstream tasks, as imbalanced data composition impair generalization (Xie al. 2023 Nguyen al.) 2024). However, prior work has been largely empirical, whereas we develop a theoretical framework that rigorously analyzes test-time scaling and Chain-of-Thougee effectiveness, overthinking, and principled strategies for task selection during training.
3 IN-CONTEXT LEARNING
In an in-context learning (ICL) scenario, a model is presented with instances of prompts of the form
P. = (I1, for (z1)....., In, h (In)), with z; drawn i.i.d from distribution Dr. and ht sampled inde- pendently from a distribution over functions in a given function class. The goal in-context learning is to train model so that when given test prompt P, (I1, h. ................................ Tm. hr (1m), Fm++1) with an independently sampled hyr,, it is able to make a prediction on 2m+1 that is close to hy (xm+1) Therefore, key distinction from traditional supervised learning that in ICI., each prompt has its own distribution. For example, in linear regression, hr (x) = (wr, 2), where each prompt has its own ground truth Wy. Thus, ICL the model must generalize not just across data points but across
distributions, and be able to infer the correct predictive rule on the fly for each new prompt without modifying its parameters.
3.1 IN-CONTEXT WEIGHT PREDICTION AND LINEAR SELF-ATTENTION
We focus on ICL for linear regression task, where each prompt P. = (I+,1, Vr.1.... Ir.n Vr.,) with Vr = (W=, (w,)) where Z+,; N(0, A). Wt N(0, Ia). Most previous works on this setting focus on prediction without directly estimating the weight vector of the test prompt (Ahn et al. 2023a Zhang et al., 2024 Mahankali et al. (003). Here, we take similar approach to Huang et al. (225aa) and consider in-context weight prediction where we require the model to directly estimate the wight vector of test prompts. To this end, we adopt the embedding used by Bai et al. (2023); ; Huang et al. (2025a) which includes the weight-estimation:
Fr.
X, Vr Odxn 01xn
0
F,
Vr.
(3.1)
0
0
Wp
vitp
where wo E 12d is an initialization for the weight estimate.
We next proceed by describing the transformer architecture. We consider one layer self-attention with residual connection. Let E' be embedding formed from the prompt. A self-attention module takes as input an embedding matrix and outputs matrix of the same size,
fauca (E; WK. WQ, Wv, WP) = E WpWyE. vs (WKE) WoE
where v is an activation (e.g. softmax) that is applied column-wise wise. Following Gatmiry al. (024); Huang et al. (2025a); Zhang et al. (2024); Ahnn et al. (2023b), we consider Linear-Self Attention (LSA) where the activation v is the identity mapping. By defining W := W Wo, V WyWv and
Sneaa--Sel-AAttnntion
0 (W, V) arrive
ETWE P
FLSA ()
(3.2)
The estimation of the transformer for w, is given by the last token of the output sequence, namely wi fL.SA (E 0)) 1422:2d-11111 which is obtained by restricting the last column of fL.SA (Ez00) )) to
entries [d + 2 2d + 1]. We assume rhp = for simplicity.
We learn the parameters of the transformer by minimizing the following empirical loss over B independent prompts:
B
2(0)=
1
Σ
JISA (Es; (Oa, 0, wr, 1)112=
(3.3)
2B
3

---

of the ground Published as a conference paper at ICLR 2026
We consider the behavior of gradient descent-trained networks the population loss induced by the limit infinite training prompts:
(())
lim L(0) - En Ew -t. 1m.. (I (): - (0,, 0, wr. 1) B-300
(3.4)
Our first result shows that with suitable initialization and step size, gradient descent converges to global minimum of I(0),, which explicitly characterize.
Theorem 3.1 Consider the linear self attention network over the population loss ization
3.4
with initial-
0
0
0
0
cl
0
V(0)=
(
0
C
(0)
Va1
0
0
0
0
0
0
0
for some real-valued c. Also define
r (1+
1 b Rdxd
A br(A) n
(3.5)
We run gradient descent the population loss with constant step size n S 11((c² lop)) W21(t) cc The gradient descent converges to a global minimum of the loss given by
We also fix
0
0
cl
0
0
0 0
V.
W.
(3.6)
0
Note that chain-of-thoure reasoning is not employed during training; however, discuss in the following section, the model engages in chain-of-thought reasoning at test time. In contrast, (Huang et al. 2025a Theorem 3.1) consider the setting of isotropic Gaussian features (A incorporate chain-of-though reasoning during training by generating intermediate steps through gradient updates on the linear regression objective. Also, the result of (Zhang et al., 2024, Theorem 4.1) does not apply to our setting, since it works with a different embedding and trains the model by minimizing the expected prediction loss function.
and
3.2 TEST TIME CHAINHOF-THOUGHT
During test time, we observe a test prompt of the form P = (I1, (Wrest, 11), Ims (Wrest Tm)) of possibly different length than the training prompts, and Wtest may be never seen before. We let the transformer to generate k steps before it outputs the final prediction truth Whest- Specifically, we let E, be the embedding at the i-th step of generation, and have fI.AA [E+)(d + 2d + -1] as the prediction of the next link in the chain. We then append it the current embedding, as follows:
T
Tm
0
0
X
test
0
0
0
E,
y1
Vm
0
0
3test
0
0
(3.7)
0
Wi
W1
Odxn
WC
U/y
0
01xn
with wy == FLSAA (iji)))d2:22411-1- The final prediction is given by Wx+1- In our next proposition,
give explicit characterization of the recursive updates of w;.
Proposition 3.2 Consider the LSA model with parameters V. and W. given by (3.6) and assume test prompt of the form P = (x01, (Wrest, L ...,, Ims (Wrest, Im))). Initializing the test time CoT with wh = 0, we have
W1-1 = W/
Xtest X. test (w: Wrest))
(3.8)
YT1
where Xtest = [2al... pool.../am.. [m]]]]] E RRdxm Therefore, the final output (after step of generation) is given by
Wh++1 1 -11 ---r-1 Xrest X test )")
Wrest
(3.9)
4

---

Published as a conference paper at ICLR 2026
Remark 3.3 Consider the quadratic loss [(w)
2m Prest X best W 1
X test w) 622 with Yerst X
teat Whast-
The gradient of the loss is given by Vl(w)= = 1 test (Brest
Xtest X tast (w Wrest ) and the expected Hessian is given by E [[W)((w)) E Xteat = A. Treating r, given by (3.5).
as a regularized form of A, the update (38) is (pseudo-) Newton's method for optimizing the loss.
3.4 HARDNESS OF A TASK
We define task by the covariance matrix of its features (A), so different tasks have different features covariances and for each task, we have many prompts with features generated from N(O, A), but cach with its own wy. Now suppose we perform direct in-context learning on task and then use it to predict labels on queries from the same task (without CoT). Our next result will bound the expected estimation error and we use that to define a measure of task hardness.
Theorem 3.3 Consider the LSA model with parameters V. and W. and assume test prompt is of the form P = (2), (Wrests I Ti)...... Imm (Weest: Tmm)): ). Initializing the in-context learning with wo 0, - Xtest X test w with Xtest [x1] (ell.../mmm [mm]] IRdxm, We
the estimate of will be given by i x2 have
1 ne ( (T+ -T tr(A ()(A-)))) Whest
T Exsoss (11ii- Weestt (12) < test
where the expectation is with respect to Xtest. we obtain
Taking expectation with respect to Weest NOO,
d
(1 1+ tr(A)
) d
(1+ tr(A) tr(A))
[[I 2 Weest (13)
(3.10)
Amin (A)
Amiu (A)
m
Based on the above result, we define the hardness of a task, with features covariance A, via the following measure:
tr(A)
(A)
(3.11)
Amin (A)
Note that it is invariant to scaling of and would be higher A has some small eigenvalue as more data is needed to learn these directions. Our next results bound the expected estimation error under CoT during test time.
Theorem 3.4 Consider the setting of Theorem 3.3 and let wk+1 be the model estimate for the target 3.3 task after generating steps during test time. Also suppose that m Q(k2 d) and that eigenvalues of A are upper and lower bounded by positive constants. We have
B(I Weest 12) t O O(k/ d/m))
where the expectation with respect to Xust [8il... [mm] and Weesst N(O, I)-
Corollary 3.5 Under the setling of Theorem 3.4 and letting Amiu (A) > 0 be the minimum eigenvalue of A we have
-2k
B(I Wrest s d 1+ Hard
(1+)
The above corollary is also consistent with our measure of hardness: the estimation error of wk+1 increases with Hard(A), J. addition, if we want to get the estimation error below some target level E, harder tasks require longer CoT at test time (larger k).
Note that in Corollary 3.5 it was assumed that is full rank. features are coming from subspace of lower dimension), then one cannot estimate Whest along those directions, as we do not see any information about them during the process. This course is not an issue if the prompts during test time coming from the same task, those directions do not contribute to the predictions. In these cases, by restricting to the relevant subspace, hardness of the task can be defined similarly where Amin (A) is the minimum "non-zero" eigenvalue of A.
rank deficient (that is the
5

---

Published as a conference paper at ICLR 2026
An interpretation of the hardness measure is that each eigenvector of A corresponds to specific skill needed for solving examples from that task, with the corresponding eigenvalues indicating the strength of those skills. An easy task is that relies nonzero eigenvalues similar magnitude), while a hard task draws on many skills, reflected in a long-tailed spectrum. The proposed measure captures this intuition quantitatively.
few dominant skills (a small number of
Test-time scaling. Our result in Corollary 3.5 provides test time scaling for our ICL setting. Note that the computational complexity during test time increases as we allow for more steps of thinking: Specifically, it is O(kd² ) as the matrix 7 - X can be computed once, and each step thinking involves multiplying it with the current estimate. Our result also captures the role of Amin, tr(A) and the prompts length during training and the features dimension d in shaping the test time scaling law. Another observation is that at any fixed test error, by increasing we can decrease the length of prompts during training. In Figure I illustrate test-time scaling for several choices of prompt lengths (n) and task hardness.
Remark 3.5
200
--n-00HH40
n-10, H-20
0-n-10.H-03
150
nm20
n=20 H=20
n=20 n=20_1130
Extiration
100
50
50
100
150
200
250
Computational complexitty (kdd)
Figure 1: Test-time scaling for the in-context learning. Here, n is the number of in-context examples (context length) in training prompts, and H is the task hardness.
4 TASK SELECTION FOR TRAINING
We consider set of T tasks with corresponding covariances . Ar. Similar to previous sections we draw infinite prompts (B + oo) but here each prompt is selected from task with probability R, 0, where Σ; n, = The goal of this section to derive optimal choice of {m)) and build insights about this choice.
Theorem 4.1 Consider the linear self-attention network and the population loss 3.4 under the multi-task setting, with the same initialization given in Theorem 3.1 Redefine I as follows:
=
n
Σ Aete + n
(2 Σ
AZTE + Σ tr((c))eTT))
Σ ACTe))
(4.1)
n
KETT
EETT
EETT
EE
Then a similar statement to Theorem holds true.
Consider target task with covariance Σ (which can be different from any of the tasks during training).
For a prompt from it, P = (æ1, (Weest, #1)...., ........ Im. (Wrest, Zm²)) we Xtest = [x1]... ... ........... | e [dxmm and Σ = 1 Xtest X X Initializing with wp = 0 and allowing chain-of-thought of length k,
A.9 implies that the L.SA estimate of Weest is w (T 1 r-1 Whest Therefore,
B(Ilwa+1 Wteet
E[I(( - r-1 E) wrestl)) 11?) E[tr((1- - (4.2)
where the second step is by taking expectation with respect Wrest. In the next proposition, derive a prompt instance independent upper bound for the estimation error in terms of the population covariance E.
Proposition 4.2 Suppose that m = ()kPd). Then,
E[tr((T-STre) - tr(T)tr(T - (1 (()) (4.3)

---

Published as a conference paper at ICLR 2026
The optimal choice of tasks selection probabilities is the that minimizes the expected estimation error during test time given by (4.2) We instead use the upper bound given by Proposition [4.2] and focus on the term tr((7 r 1 ET 1/2)2k) in (4.3). which captures the effect of thinking and the
dominant term with an exponential rate. This results in the following optimization for choosing task selection probabilities:
min [E
Bitr((11 - r 1/2 1/
(4.4)
subject to
Te 0. [TE Te =
FE[T]
When is the test time thinking useful? We observe that the effect of thinking at inference time is captured by the term tr((1 - -1-1//11--// Depending on the cigenvalues of r-//2//-1/2 this term may shrink or grow as k increases. Intuitively, if each eigenvector of Σ (representing the skills required test time) is sufficiently represented in the training data so that strong along that direction and r-1/25r-1/2 remains mhel--then additional thinking improves performance. In contrast, if some task-relevant directions are underrepresented in the training data, and thus not well learned by the model, increasing the amount of test-time thinking can degrade performance, effectively leading to overthinking.
Remark 4.1
4.2 OPTIMAL CHOICE OF TASK SELECTION PROBABILITIES
We next analyze the optimization 4.4) to argue that choosing a diverse, relevant and hard set of tasks during training results in best performance for test-time scaling.
Diversity. A key observation that we must select a diverse set of tasks so that the spectrum of T adequately covers all directions in the target covariance E. Failing do so causes r-1/2 to be large along uncovered directions, resulting in higher test error that may further amplify with additional reasoning steps.
Another important notion is the relevance of the selected tasks to the target task. Recall the expression given by (4.1). When d < n, and noting that the eigenvalues of A are 0(1), I can be replaced by i := ExE(ry AgTe, which is a convex combination of (Ac) [EE)) Hence, minimizing ur((T r- 1/2 in effect corresponds to approximating Σ with a convex combination of the task covariance matrices and so tasks which place high weight on directions well represented in (i.e. relevant ones) are desirable.
Relevance.
Hardness. The other factor in task selection is the hardness of tasks. We argue that when the target task hard (as often the case where models are compared on difficult benchmarks), our proposal favors selecting hard tasks during training. Without loss of generality, by scaling features we can assume that tr(Ae) = 1, Vl and tr(E)= = With this normalization the hardness of task captured by the minimum eigenvalue of the corresponding covariance matrix. Now, invoking the test error given by tr((!- r r 1a--//rr-1// 1/2)), the absolute eτror along minimum eigenvectors of Σ contribute more towards the error. Given that the target task is hard one, mmin (E.) small and in the next proposition we show that to estimate Σ well on this direction by a convex combination of available tasks, we need to select some hard tasks (those with small minimum eigenvalue).
Proposition 4.3 Suppose that (o Omin (Σ) & and define D E Juiin (Ae) S A(e Tmin (.)) Note that D corresponds to tasks with small minimum eigenvalues (hard tasks), since both E and mmin(2) are small. Then, EteD te 1/2. In words, at least 1/2 of task selection probubilities are on hard tasks.
Further simplification of task selection procedure. Note that the optimization problem (4.4) is inherently nonconvex, which motivates us to turn to simplifications that transform it into a convex and tractable form for large-scale problems. We make two modifications: 1) As discussed before when < n and since the the eigenvalues of are ( I' can be well approximated by I ΣeE(r ArTe, which is convex combination of (Ae)ee(m) [m]- 2) The objective in (4.4) seeks make r 1112251-122 close to the identity matrix. Instead, we minimize || EI-E-'rje R same goal but through different formulation. With these consideration, we propose the following
| Σ I-'fi which pursues the
7

---

Published as a conference paper at ICLR 2026
alternative optimization for choosing task selection probabilities (
min (et)))
--
Ae
(4.5)
RE[T]
subject to
Σ (E[T]
Ft =
te z 0, [
This is quadratic optimization problem and can be efficiently solved at scale.
5 EXPERIMENTS
In this section, we conduct experiments to validate theoretical results.
Setting. We conduct experiments in two settings. First, we consider a transformer with a single linear self-attention (LSA) to confirm our theory. Then, we consider large, nonlinear transformer architecture namely GPT-2 to validate our conclusions. In both cases, the data distribution follows in-context weight prediction task in Sec. 3.1 where T ~ N(0, A), w, N(O, La). We choose the token dimensions d = 10. During inference, we let the model to output multiple steps before returning the final predicted weight vector. At each step we concatenate the embedding with [0a, 0i, 1] as in Eq. (3.7) and input the concatenated embedding matrix to the model. The predicted wg will be returned after & steps of COT. We report the average results and error bars over 10 runs.
Transformers with a single linear self-attention (LSA). We train the transformer architecture in Eq. (3.2p on the synthetic data generated as described above. We generate 5000 examples, use batch size B = 1000 and run Adam with learning rate 17 = 0.001 for = 1000 epochs. For the results reported in the main paper we follow our theoretical setting Sec. 3.1 That is, we initialize transformer weights (V(0), W (0)) according to Theorem 3.1 where 1/31 (0) is set randomly with entries independently and uniformly drawn from [0, 1). and we set c = 1. We do not perform CoT during training.
We report additional results with random initialization and training with CoT in Appendix B
Large, nonlinear transformer architectures. We use decoder-only Transformer architecture (Vaswani et al. (017) from the GPT-2 family (Radford et al. 2019), consisting of 12 layers, 8 attention heads, and 256-dimensional embedding space. In total the model contains 9.5M parameters.
This architecture takes as input sequence of vectors in its embedding space and predicts the weight vector within the same space. We apply this architecture to prompts of form (Zr,1 1r.1." Zr.ms Dr,m, Wo, 1) in the following manner. In line with (Garg et al. 2022). we map each yr,, to the same dimension as Fr,s by appending zeros, and map Er,i, Vr.; into the latent embedding space of the Transformer through (learnblee) linear transformation. We get the predicted Wr as the model output. Similarly, we map the model output, i.e., w, from the latent embedding space of the Transformer to d-dimensional vector through another (learnable) linear transformation. Training is performed with batch size of 64 over 25k total steps. The model weights are randomly initialized, and CoT is applied during both training and inference. We use curriculum learning (Garg et al. 2022) to speed up training.
Larger test-time compute reduces the requirement on training-tine compute. Fig 2a show the test error vs length of CoT (k). For the LSA model, we use n = 10, 20, 30 and for GPT-2 we use 20, 30, 40. We see that by increasing the test-time compute (k)), we can decrease the length of prompts during training to get similar test error.
When more thinking hurts. For training, we sample prompt inputs from N(O, A) where a skewed covariance matrix with eigenbasis chosen uniformly at random and i-th eigenvalue proportional to 1/5. For test, we sample prompt inputs from N(O, Ia). We normalize the inputs so that their expected squared norm is equal that of inputs encountered during training. Fig 2b 2d show that when some of the directions the target task are not sufficiently present in the training data, allowing for more thinking during test time would hurt the performance. An interesting observation is that when the model is in overthinking regime (Fig 2igg2]] [d] larger prompt length n yields a higher test loss, while
when the model is not overthinking, larger n reduces the test loss (Fig 2a 2ch.
8

---

Published as a conference paper at ICLR 2026
n-
10
n-
10
Test Loss
Text LOSS
10-2
207
n-:10
D-20
10
(a) LSA
(b) LSA
20
30
40
Tect 1066
Test.css
0.08
0.04
0.06
0.04
0.00
0.02
(c) GPT-2
(d) GPT-2
Figure 2: More test-time compute reduces training-tine time requirements for (a) layer transformer and (c) GPT-2. However, insufficient task coverage in training data makes longer CoT's harmful for (b) one-layer transformer and (d) GPT-2. For GPT-2, the errorbars are std of 10 runs. For LSA, negligible start from the fixed initialization in Eq. (3.6). Same value for is used in training and test.
5.1 TASK SELECTION
We design an experiment illustrate that our method prioritizes diverse and hard tasks. We consider multi-task setup with four task types, where each type is defined by two parameters and B. These parameters respectively control the decay rate and the support size of the eigenvalues. Specifically, eigenvalues are proportional to i for i E [B] and zero elsewhere. The positions of nonzero eigenvalues are uniformly shuffled within [d], and the eigenvalues are scaled to have unit sum. Here, B captures task diversity, while captures the task hardness (with larger a producing smaller nonzero eigenvalues, corresponding to harder tasks according to measure (3.11).
The four training task types Easy-Short (a 0.2, B 20), Hard- (a 0.8, B 20), Easy-Long (a = 0.2, = 100), and Hard-Long (a = 0.8, B 100). The target task is set with 0.8 and and the support of cigenvalues, We then solve the quadratic optimization problem (4.5) to obtain task selection probabilities Fę, for 1,..,, 200. Fig Ba displays these probabilities, colored by task type, with solid lines indicating their average per type. As shown, harder and more diverse tasks receive higher selection probabilities, while easier, more concentrated tasks are weighted lower. Fig 3b further plots selection probability versus task hardness, confirming that harder tasks are indeed favored, consistent with our theoretical analysis in Section 5.T.
1000. We generate 50 tasks of each type by randomizing the eigenbases
To demonstrate the improvement achieved by our task selection procedure, Appendix B.1 reports simulation results showing test error as a function of the test-time thinking length (k) for different task selection strategies. The findings indicate that our optimal task selection procedure effectively prevents overthinking.
5.2 EVALUATION ON REAL REASONING BENCHMARKS
Finally, to demonstrate the validity our conclusions on real reasoning tasks, we train Qwen 2.5-7B- Instruct the OMEGA dataset (Sun et al. 2025). We chose two tasks from the OMEGA dataset, namely GCD and polynomial root reasoning. These tasks are designed such that training on one does not benefit the performance the other. We fine-tune the base model (Qwen-Base) with RI.

---

Published as a conference paper at ICLR 2026
0.03
0.03
Stort Easy Long Hard-Long
0.025
aheiio probability (m/
0.025
prohshility
0.02
0.02
0.015
0.015
0.01
0.01
0.005
0.005
0
100
150
250
Index
Harduess of task
(a) Selection probabilities for different task types
(b) Selection prohabilities vs. task hardness
Task selection in a multi-task setup (a) Each color corresponds to a task type with solid lines indicating the average selection probability per type. As we observe harder and more diverse tasks receive higher selection probabilities, while easier, more concentrated tasks are weighted lower (b) Task selection probabilities versus task hardness. As see harder task are favored in the selection.
Figure 3:
separately on the training data of GCD and Poly. We call these models Owen-GCI and Qwen-Poil We evaluate both models on the test data of GCD. Table shows that, as expected, for the harder tasks that require longer reasoning, all models have lower performance. However, we see that while shorter test-time thinking (CoT length less than 1k characters) yields much better performance (+44.69%) on GCD for Qwen-GCD compared to Qwen-Base it yiclds a slightly lower performance on GCD for Qwen-Poly, compared to owen-Base (-1.39%). Interestingly, when models reason for longer at test-time (between Ik and 2k characters), Qwen-Poly has a much lower performance Owen-Base, while Qwen-GCD outperforms Qwen Qwen-Base by 11.2%. This confirms our theoretical results that when training and test data are aligned, more thinking helps. But, insufficient task coverage in training data makes longer test-time compute harmful.
(-6.37%) compared
Table I: Average accuracy GCD for Qwen2.5-7B Instruct (Base), Base model fine-tuned on CGD (Qwen-GCD) and Base model fine-tuned on Poly (Qwen-Poly). For all the models, the accuracy on examples that require longer CoT is lower (compare the second column to the first column). This confirms that examples that require longer CoT are generally more difficult. The % in () shows the fraction of test data with the corresponding test-time CoT length. Notably, shorter CoT's (0- (0-1k) considerably improves the performance of Qwen-GCD (75% versus 30.39%) and slightly harms the performance of Qwen-Poly (29% versus 30.39%). Longer Co''s improve the performance of Qwen-GCD (38.4% versus 27. 2%) and significantly harm the performance of Qwen-Poly (20.83% versus 27. (.).
CoT length
[0, 1k)
[1k, 2k]
Qwen-Base
30.39% (30% data)
27.2% (70% data)
Qwen-GCD
75% (15% data)
38.4% (85% data)
Qwen-Poly
29% (32% data)
20.83% (68% data)
6 CONCLUSION
In this work, we provided a theoretical and empirical framework for understanding in-context learning in transformers, showing that cfain-of-thouuht prompting at test time enables models to emulate multi step (pseudo)-Newton) 's method. By introducing principled notion of task hardness based on features covariance spectrum, we derived scaling laws that clarify how test-time compute, context length, and task diversity interact. We proposed an optimal strategy for task selection in a multi-task training that shows training on a diverse, relevant and hard set of tasks during training results in best performance for test-time scaling. We also validated our findings on both linear self-attention models and GPT-2. We will conclude by discussing some limitations of our work which pave the way for future directions of work. Our theoretical analysis in this work is limited to linear regression tasks and single-layer linear self-attention, and an important direction for future work is extending these results to nonlinear data generation settings and transformers with nonlinear activations.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
AJ was supported in part by the Sloan fellowship in mathematics, the NSF Award DMS-2311024, an Amazon Faculty Research Award, Adobe Faculty Research Award and iORB grant form USC Marshall School of Business, BM was supported in part by the NSF CAREER Award 2146492, NSF-Simons Al Institute for Cosmic Origins (CosmicAI) and NSF AI Institute for Foundations of Machine Learning (IFML).
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances Neural Information Processing Systems, 36:45614 45650, 2023a.
Kwangjun Ahn, Xiang Cheng, Minhak Song. Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Lin- ear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082 2023b.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36:57125-5722 2023.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? study of simple function classes. Advances in neural information processing systems, 35:30583-3000 305933, 2022.
Khashayar Gatmiry, Nikunj Saunshi, Sashank Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn implement multi-step gradient descent for in-context learning? arXiv preprint arXiv:2410.08299, 2024.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et al. Deepseck-rl: Incentivizing reasoning capability Ilms via reinforcement learning. arXiv preprint arXiv:2501.12944 2025.
Roger A Horn and Charles R Johnson. Topics in matrix analysis. Cambridge university press, 1994.
Jianhao Huang, Zixuan Wang, and Jason D Lee. Transformers learn to implement multi-step gradient descent with chain of thought. arXiv preprint arXiv:2502.21212, 2025a.
Xiaoke Huang, Juncheng Wu, Hui Liu, Xianfeng Tang, and Yuyin Zhou. ml: Unleash the tential of test-time scaling for medical reasoning with large language models. arXiv preprint 000:2504400889 20255.
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. arXiv preprint 2.Xi:23077.0356.. 2023.
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang. Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling. arXiv preprint arXiv:2501.19399 2025.
Dang Nguyen, Wenhan Yang, Rathul Anand, Yu Yang, and Baharan Mirzasoleiman. Mini- batch coresets for memory efficieefficient language model training data mixtures. arXiv preprint arXiv:2407.19580 2024.
OpenAl. Learning to reason with lls.. http:///ppnaa..ooommma////mmmmmo
2024.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Charlie Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can be more effective than scaling model parameters. arXiv preprint arXiv:2408.00 03314, 2024.
11

---

Published as a conference paper at ICLR 2026
Jinyan Su, Jennifer Healey, Preslav Nakov, and Claire Cardie. Between underthinking and over- thinking: An empirical study of reasoning length and correctness in Ilms. arXiv preprint arXiv:2505.00117. 2025.
Yiyou Sun, Shawn Hu, Georgia Zhou, Ken Zheng, Hannaneh Hajishirzi, Nouha Dziri, and Dawn Song. Omega: Can Ilms reason outside the box in math? evaluating exploratory, compositional, and transformative generalization. arXiv preprint S.IV:22506188800 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. Compressed Sensing, 2010.
Yuhang Wang, Youhe Jiang, Bin Cui, and Fangcheng Fu. Thinking short and right over thinking long: Serving Ilm reasoning efficiently and accurately. arXiv preprint arXiv:2505.13326 2025.
Sean Welleck, Amanda Bertsch, Matthew Finlayson, Hailey Schoelkopf, Alex Xie, Graham Neubig, Ilia Kulikov, and Zaid Harchaoui. From decoding to meta-generation: Inference-time algorithms for large language models. arXiv preprint arXiv:2406.168335 2024.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoe V Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36:69798-6993 2023.
Edward Yeo, Yuxuan Tong, Morry Niu, Graham Neubig, and Xiang Yue. Demystifying long of-tnnof-thougt reasoning in llms. arXiv preprint arXiv:2502.03337, 2025.
Zhaojian Yu, Yinghao Wu, Yilun Zhao, Arman Cohan, and Xiao-Ping Zhang. Zl: Efficient test-time scaling with code. arXiv preprint arXiv:2504.00810 2025.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024.
12

---

Published as a conference paper at ICLR 2026
A PROOFS OF THEOREMS AND TECHNICAL LEMMAS
A.1 PROOF OF THEOREM 3.1
Lemma A.1 Assume an initialization of the form
0
07
0
0
cl
0
V(0)=
0
0
(
0
0
C
V31
0
0
0
0
0
0
When the linear transformer is trained under gradient descent. Then V(t) and ((t) have the following form:
0
cl
0
0
0
W21(t)
(t)
V(t)= =
W(O)
(t)
0
0
0
0
with Vo1 (4) E TR ddd and (1) E IR.
Lemma A. very similar to (Huang al. 2025a Lemma C.2). The difference is that here the features covariance is non-identity, while we do not do CoTT during training.
Given that several blocks of (t) and (tt) remain zero across the gradient updates, reduce
the loss function in simpler form. Define the shorthand (I)) Va1 (1) Rdxd and w(l))
(1) E R. Invoking (A.3) we can rewrite the loss as follows:
L L(((t)) 2
(IfIA 0(1)).-11 (0a, O, wr Dillee
0a
0
1
0
2
upp
7
VaI(1)XX (crio + (2)(()
WT
1 (I(( - 1) 1
E 2
(A.1)
with Ā = XXT /n. As we see w(t) does not provide additional degree of freedom in minimizing the loss since it appears as the term (l))((1). This clarifies fixing w(l) - along the gradient updates.
We then have
L(O) B (ICVVPTT I+ 2cVA)]
Lemma A.2 Let E Rddx with columns drawn i..d from N(O, A). For any deterministic matrix A have
XXT
XXT
n
1 (A(A+ +
E
AAA+
n
n
n
71
Using Lemma
A.2 we have
B(A3)
A :r(A)A:=IT IA.
Therefore,
(()) 21r (VTAKT)
d + tr(VA)
(A.2)
2
This is convex in V' and so gradient descent with a fixed step size $ 1/ converges to its minimizer, V LI. We have
-
13

---

Published as a conference paper at ICLR 2026
can take L c² |pP To find the minimizer of L(0), we set its gradient zero,
2 (TAVT + VTA) + cA 0,
which is continuous-ting Lyapunov equation (Note that and A commute and both are symmetric). Hence, it has unique solution given by V = -
As the final step, show that V, and W. also global optimum for the population loss, without making the specific structure imposed by gradient descent as described in Lemma A.1
We continue by computing the output of LSA, and recall that up 0.
E. W E.k:
VE,
X
DOU
V
Y
E
W
Odxn
01xn
XXXT
Xy
O
W14
1
UX
yy
0
W21
n
Oaxn
0ax1
0dx1
Oax1
W34
O
01xn
0
W44
XXX W14 + Xy W24
1
yX
W11 + yv W21
n
Waa(t)
Therefore,
fLSA ET, 0)( - (0,, O, Wr 1)T
[XX W14 Xy W247
0
0
3XX W14 + yy W24 0
V
1
WT
W44
XX XX w. W24
0
XX
W14 + w. XX
w. W24
V
n
WT
WA
The loss function is the expected squared norm of this object. To minimize it, make its first, second and last entries zero by setting the corresponding rows in to zero. This gives us
L(0)
2
E
(V/1 + Vs2w )xx (W14 + W2aw;) - wy + Vsa W s4
2
n
1 -V W4 +(ai - 1)w, + Va2w, AW14 + Va2w AW 11/AW1 +
E
2
82
2 (VVSIÁWSs 1)w, + Vasw AW (2)
E
2
(V5. + V22WTT 1) 1)w-
E
where the penultimate step holds since the cross term is an odd function of w, and so its expectation is zero. The other eliminated term is squared and non-negative. In the last step used that wi AW14 is scalar and can be replaced by its transpose. In addition W24 is also scalar and can
14

---

Published as a conference paper at ICLR 2026
commute with Å. Note that V32 and Wi4 do not offer any flexibility in minimizing the loss function as their effect can absorbed Va1 W24. Therefore, we can set them to zero. Hence,
I (I(Va T),, (
min L( 2 min
VaL. W24
Observe that the right-hand side is of the form (A.1) and hence its global optimum (reached by gradient descent) serves the global minimum of the loss.
A.1.1 PROOF OF LEMMA A.1
To prove this lemma, we prove that when the irrelevant blocks are 0, the gradients of the loss remain zero on those blocks and they never update the corresponding parameter block.
We do induction on f. Suppose that the claim holds for t. We start by computing the output of LSA.
E W E
VE,-
0
0
0
07
X
0
0
0
y/
0
F.J
W
V31 (t)
(t)
0
Odxm
e0t
O
0
0_
[01xm
Oaxn
0g7
X
0
cuoo
O1xn
0
W21 (t)
V Xt)
0g
Odxn
tbo
0
m
01xn
0
01xn
0
Oaxn
0g7
01xn
0
TX up + W22(t)
(t
0g
01xn
0.
0.
1
(A.3)
V (1 (cun W (1)w.)
n
where we used that y = X w... We proceed with calculating the derivatives of the loss:
V(((t)))) Vv (IILs)(E): (0a, O, we, 1)1l(a)
E (Er;A(Er0(()))))) - (0ds, O, Wr, 1)) Erli, -1] WT E.E.
n
E. W W. -1] (0,, 0, w. 1000 01 Fy -1 W (VF,
W E, E
E
(A.4)
We note that
E. S.W,ER 1: -1]
VE,
(0a.0,u0 w. 10,
0.
0
(A.5)
(t)xX (cuo + W2x + n1 -wr)
In addition,
XXT
Xy
0
0
E,G, 1-T WT E, ET
[cup W2a(t) 0j
yXT
yy
0
0
7
0
1
[ewoXXX + Wxx cúoo
+ W2a(O)y)
0
0).
(A.6)
15

---

Published as a conference paper at ICLR 2026
Plugging in from A.5 and (A.6) into A.4
get the following structure for the gradient of the
loss:
0
:
0
0
VL(0(())
(A.7)
vvs.)((()))
Vv ((0tt)
0
0
0
0
where
Vv,, L(O(t))= E [V)i (cuo (2a(t) + n(wb w+)) (ci0 XX T W21(t)yX }
n.2
V12 I(((t)) E (Va (cupp W2a())ur,) n(ūo w-)) (cur +W2((t)v
n²
Recall that we set un = 0 by which we obtain
Va2 L(((t)) = E (VVa()) -nI)w, -TTww, x w.) = 0
n
(1) E n
since it an odd function wo, N(O, I). (Note that the population loss is non-random, due the expectation in its definition. Since the initial V(0), W (0) are non-random, the trajectory (tt) and ((1) are non-random.)
We next proceed with calculating the gradient with respect to W. We have
V WwI(((t)
Vw (IVIsA(FF;:00))))-1 (Ods O, w/r 1) 11(.))
2
[ETETVV
(0,, 0, Wr, 1)) Er (: -1]
7
E (VE,
E. S.W, E.E (0a, 0, 20 - uv. t.., 0) Erli, -1]
E
(XXX (XX] (Vxx (cup + W24(1)wz) + n(wio wir) )) yx Va) (t) (Vai(t)XX (cwo + W24 + +n w ))
1
E
20 1]
n²
0
where the last step follows from the following equation and simple algebraic calculation:
XXT T Y
Xy
0
0
E,E,
yy
0
ErG; -T]T-11T- = 10
zio
1]
0
wh
0
1
0a
VE,- __ W.
0
(0a, O, Wy - 0)
Val (1)xXXXxxxx + W24 (t w, ) n(wo - wT)
0
Recalling that ui 0 we simplify Vw follows:
- nT)wr yx (t) (t) - nI) wi
XXT Val (t)
1 n
V F L(0())
E
00 0000
(A.8)
We have E[XXX Vis T (t) + nz nn]=w, = 0, since Vaa (t), W2a(t) are non-random and w, is zero and independent of X. Hence, V Vw is the only non-zero block.
16

---

Published as a conference paper at ICLR 2026
A.1.2 PROOF OF LEMMA A.2
XX XXXXXXXXX
Σ D
Azyz]
n
There are ntn 1) terms where + j and n. terms with j. Since æ; and x are i..d., let denote either of them. Thus,
E [ AXX n²
(n(n - 1)E(xr [AE[ '1 nE 22:2 Axr '])
n²
The first term is the second moments, For the second term we use Isserlis's theorem, by which we have
(E xx
Axx 1)
Au (B(x, B[zr,]] + (x B[z,r,] B[x,,, B [x)
Assuming E(rx" ] = A, we get
E
ArT.
A(A + A TJA+A TH(AA).
Therefore obtain
XX AXX
AAA (A(A+ A A
+ 71
n
A.2 PROOF OF PROPOSITION 3.2
Recall V. and W. given by (3.6), as the estimated blocks transformer after training, We next rewrite the updates for w; is more explicit form:
00a
0
E.
W WFS..-1)
fLSA(E,, 0*) (----
VE.
W;
CW
1
VE,E
C 0 0
m
TXX
Xy
Cw
YXT
yy T
V
Wi
m
Oaxd
0mxm m
rO
O
Tw ;
0
1
0
0
VXXXX
VA1 Xyy T
W;
0
0
Recalling that Vs1 P-1 /e obtain
Wi+1
W/s
r X eess XT
trst wi
r-1 X test Verst
m
m
W/
r
X X
test (wee Wress))
m
Rearranging the terms, Wi+1 Whest (T PT X X.T (Tss)) Wsest) which results in
mn
Wk+1 = Wsest + (I - p-1 Xlest X ) (wwo Wteest)
= (I r-1 Xrest X Wrrest
(A.9)
which completes the proof.
17

---

Published as a conference paper at ICLR 2026
A.3 PROOF OF THEOREM 3.3
Define the shorthand Å - X m.. We have
[[ Waest (1e = BEIKL - --'-1 (=1A) Wrest = whest 1B( )r- r-1 Wesst
=Weest -1-1 - Al + LE(A -2 2A))
Using Lemma A.2
we have
B(Âr -A) = m 1
AT -r A + (2Ar 2 2A tr(Ar
m A A - tr(AAI
m
m
Using that r'1 and commute and both are symmetric we obtain
B(I1 Wrest (12z) w Weest 21 + m+ 1 r 2
1 -r(Ar-? 2)A A)Weest
m
m7
= (1 - 1--1 A) 11'-1A Wrest
(I' 212 3)
West
Using the definition r = (1 1 ) + 1 ttr(A) it is easy to see that
03 7ー ^ 1 1 (T+
Also since r-1 A-1, we have
r~2 + tr(1l 3 1 + tr(A-1)A
Combining the last two equations, have
tr(A ")a) Wreat ne (1+ 1)2 + _ (I
Exa. (1 - Wreat (1( S Weest
m
Taking another expectation with respect to Werest
N(O, 1) we get
B([1 Wrest (13 -/2 ne (d + Lr(A))(r((A)-))
1
21r(A)rr(AA)))) -1)) +Ur(A))) 1 (d + ir(A (((AA))(r(A))
The claim follows by noting
tr(A)tr(A-)) - d tr(()) Amin
tr(A)ttr(A22) < d (r(()) Aman
where Amin is the minimum eigenvalue of A.
A.4 PROOF OF THFOREM 3.4
We define Å
m E T.oe. After ls steps generation, we have Wk+1 (I- A) Weest
and so
E(I Whest (1;,) = [=)==(((= [E[[(T-T- AA* Wrest [2) ) - (r- A)t). (A.10)
Note that 1 and A commute and both are symmetric. Therefore 1 is also symmetric. We denote by o; the eigenvalues of I ' The matrix I r-1, ' however is not symmetric. We denote by ã; the eigenvalues of 7 r-1 Å. By Weyl's inequality, have la; ā,l T- A) 8.
18

---

Published as a conference paper at ICLR 2026
We then have
6; - (o; + =5jek
2h
2k
1 +
2k
2k
8
1+
j=1
2k 1
S 0
22))
(A.11)
2 2kg y. We next proceed by bounding B( B(A)). Observe that Å ai mi Eie(m] T with I₁ ~ N(O, A). Using concentration bounds on random matrices with in- dependent sub-gaussian rows (See e.g. /Vershynin 2010 Eq. 5.26)), we get that with probability at least 1 2e-
Define Å; 1
A - | S mmax c
CV
m
Vm
Since and have bounded eigenvalues, by adjusting the constants e, C (absorbing |M|lon and [T 'I' into these constants), we also have that with probability least 1 -at
d t = r 1).lpp S maxx(, E= CV
(A.12)
m
m
We define the probabilistic event = ( S CK2 Vd/m). Obviously, [((,)e) S CK2 sil//dy vd/m. We also have ((A.)))) Scke Vd/m P(A 2 s)ds. Note that by definition of A. we have
2k
2k0
Á
< 2k max
s C' max(k6, (
/=1
Note that since cigenvalues of A are upper and lower bounded by constants, so are o:'s. Therefore, we can work with one constant C" that works for all E [d].
By virtue of the above bound, if Å 2 s we have 2 min( E ( ) 2 _ min(s, stt). We next choose t such that for CV d apply the tail bound (A.12
ym we have max(e, =2) I 12 min(s, )े that
In addition, for 2 Ck2 dd/m we have E min(s, stt 2 C CVd/mm and so it suffices to have y/m ym
/mm ydz). max( + min(s, s Therefore, can min( Y 8, x sik, V
E
B
P(A, 2 s)ds
Cka dd/m
100
P 82 bz
mni
CK2 v dd/m
100
2 exp yds
Ck2 /d//
By considering each of the four terms the minimum operator defining t, and following algebraic manipulation, it can be seen that the right-hand side above is O /d/m) and hence,
E E(A)) E ((,))) B(Å.e)) C(k
(A.13)
19

---

Published as a conference paper at ICLR 2026
for constant C' 0 and for all E [d]. Combining (A.10) and A.14 and the bound Å, obtain
E[ Wrost Wrosttleg)
(1: s
E1
Ar
E
2k
2k (1 B(Å,))
0
0, (1 + Ck #
i=1
-r<((T A)) (1 C m
where (b) follows from (A.11 and (c) follows from (A.13) In addition, step (a) follows from the following lemma from Hom & Johnson 1994).
(forr & Johnson 1994 (3.3333))) Let be a given d by d matrix and let be given positive integer. For all 0 we have
Lemma A.3
4
(A
forq l,
where for matrix B, (B denotes the singular values of B.
Step (a) follows by using the above lemma for A = - (パ-))) m k. 2, = d. This completes the proof.
A.5 PROOF OF COROLLARY 3.5
Recalling the definition of I' given by (35 we have
I- A
(1 (1+ )A tr((a) A
[(n 1)A tr(()) + (r(A)T)
Amin + tr(A)
(n 1)Amin + tr(A)
+ Hard(A)
1 Hard
n
= 1 1 Hard
I.
Therefore,
2k
tr((1 - ())))))))) <d 1 1 (a)d(A)
which completes the proof by invoking the result of Theorem 3.4
A.6 PROOF OF THEOREM 4.1
The proof follows long the same lines as in Theorem 3.1 Under the multi-task setting, each feature is now coming from mixture of normal distributions. The main modification needed in the proof is on statement of Lemma A.2 which is extended as follows.
Lemma A.4 Let X Bdxn with columns drawn i.i.d from a Gaussian mixture distribution, with probability Te from N(O, At). Then, for any deterministic matrix A we have
XXT
XXT
A(Te) Σ A
E(Ae(A A)Ae + Actr(AAe))
=
71
n
n
20

---

Published as a conference paper at ICLR 2026
Using Lemma A.4.
we have
S = B
XX
XXT
ApTe)² 2 Σ (2 + Astr(Ar))Te.
Continuing from A.2 have
c²
tr (vSvT) +2 d
(( = 2
+etr(V( Acte)))
Setting the derivative to zero, obtain
c² VL(0) (V(S+ + DAeTe.
Neme.
Solving this equation, using that S is symmetric, we have
V -=- A -
where the last step follows from the definition of I.
A.7 PROOF OF PROPOSITION 4.2
The proof is similar to the proof of Theorem 3.4 Recall 3 == 1 Xtest 1 Et= T,,;" the empirical covariance of the features in the test prompt. A major difference with the proof of Theorem 3.4 here have to work with r-1 and r- neither of which symmetric. To relate the trace of their powers to their singular values, we do symmetrization step. We write
(1 ---
= -11---)j-----...011--111
= ---1811---12$01/// -1--
Hence,
tr((1 tr((I-
- - h1100--phyyy4r77111mmmm/00011113
< tr((1 [(([--1"030022311111------111111
tr((1 [|([0=-~0280--22100--0---"111110
bb(1--11/229---/9991--1111
<
(A.14)
In the equalities above we used the identity tr(AB) tr(BA). The inequalities follows from the fact that for positive semidefinite matrices A, B we have tr (r(B)) < tr())tt(BB..
We next denote by o; the cigenvalues of / 1/2/ 1/2 and by o, the eigenvalues of I 11//2 Similar to the proof of Theorem 3.4 we have o²# o7* (1 A,) with B(Å,) < Ck /d/m for all i [d].
21

---

Published as a conference paper at ICLR 2026
By continuing from (A.14), get
-
s <tr(T")(r(()) E
str((T")tr(TT)
2 (1 + (())
<tr((~+)(r(T)
oi 2 + Ck /
= tt (
o(1)),
where in the last step we used that //dd/m ((1) by our assumption. This concludes the proof.
A.8 PROOF OF PROPOSITION 4.3
The proof follows from the Markov inequality. Define discrete random variable X which takes values miin ( (Az) with probability we, for l E (7]. We then have
P(X S 2(E + mmin (E)))
Σ (E(T)
Te (omin (A))
Σ
EED
In addition,
B[X] Σ
Tea (Az) S minn Σ TeAe) min (T),
by using the convexity of minimum eigenvalue and Jensen's 's inequality. Recalling from 4.1 we have
n2
n nー ス 2
-
ApT
BE(T]
for n 2 Combining the above two equations, we obtain
E 20 min (T) S 2(0 ()) + e).
Therefore,
Σ Te P(X S 4(8 mmin (E)))
LED
F(X 4(E + Omin (E)))
E[X]
2 4(E + Tmin (Σ))
2 2 2
where we used Markov's inequality in the third step.
B ADDITIONAL EXPERIMENTS
We report additional experiments on a transformer with single linear self-attention, when starting training from random initialization and performing CoT with length & during training. Similar to the main paper, the data distribution follows our in-context weight prediction task in Sec. 3.1 where Ir.i ~ NO A), w, ~ N(O, Iz). We choose the token dimensions d = 10. During inference, we let model to output steps before outputting the final predicted weight vector. At each step i concatenate the embedding with [O,, w;, 1] as in Eq. 3.1 and input the concatenated embedding matrix to the model. The predicted tr will be outputted after lt steps of COT.
22

---

Published as a conference paper at ICLR 2026
Fig. 4a 4b show the test loss during training for 2, 4. For each k train the model with n = 20, 40, 80. The training and test data are generated from Ir,5 ~ N(O, Ia), w, ~ N(O, Ia). We that for fixed value of k, larger yields lower test emrrr, which confirms our theoretical results.
Fig. shows the test loss during training when training distribution is skewed and some directions of the downstream task not represented enough in the training data. from N(0, A) where A a skewed covariance matrix with eigenbasis chosen uniformly at random and ith eigenvalue proportional 1/2. For test, sample prompt inputs from N(O, Ia). We use n 20. We that larger yields higher test error Thus, larger test-time compute hurts the performance when some directions of the downstream data are not enough presented during training.
0.8
0.8
25
0.6
0.6
Te4t
0.2
0.2
1.0
0.0
200
400
600
602
1000
230
400
6:0
1002
(a)
k
(b) 4
(c) Skewed
Figure 4: Transformer with single linear self-attention. (a), (b) Fixing the test error, by increasing k we can decrease the length of prompts during training. (c) When some directions of test are not enough represented in training data, more test-time compute hurst the performance.
B. EFFECT OF TASK SELECTION TEST TIME SCALING
To demonstrate the improvement we get from our task selection procedure, consider the set up of Section 5.1 where we generate Whest with i.i.d entries from (0, 1). During the test time we initialize with wh = 0 and let the model generate the final estimate of Whest after k step generation. We set the prompt length during training to n = and prompt length during the test to m = 500. In the plot below show how the error | wx behave for the following task selection procedures:
1) Optimal task selection: We set the probabilities ne by solving the optimization problem (4.5). 2) Uniform selection: We select tasks during training with equal probability. Easy task selection: or Easy-Long as described in Section 5.1) with equal probabilities. In Figure the estimation error for different task selection procedure plotted versus k. the length of generations during test time before outputting the final estimate.
We select only the easy tasks (Easy-Shop
As we see under the optimal choice the error goes down with k, while under under the other procedures the егтоr goes up with k, indicating that proper choice of tasks during training can avoid overthinking, which can occur under other choices of tasks during training.
100
opoinal selection
80
un
easy tasks selection
errr
40
20
8
k
Figure 5: Effect of task selection during training on test-time scaling
23


Reviews:
[{"review_id": "pfUNIqpnix", "reviewer": "Reviewer_cDdC", "summary": "The paper studies test‑time scaling in a stylized in‑context learning setting. The core technical model is a one‑layer linear self‑attention transformer trained on an in‑context weight prediction task for linear regression with Gaussian features. They find that CoT induces an iterative update that they interpret as a multi‑step pseudo‑Newton method. They define a task hardness measure, derive error bounds that decay with the CoT depth, and argue that more test‑time compute can compensate for shorter training prompts, and that insufficient skill coverage in training can make longer CoTs harmful.", "questions": "What is the test‑prompt length in Figures 2?\n\nWhat is the sample complexity to estimate a usable covariance from a small validation set, and how robust is the solution to estimation noise?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper addresses important questions such as whether more inference compute always helps, whether it can trade off against training context length, and what counts as difficult training data.\n\nProposition 3.2 derives the iterative update which makes the connection between CoT and a preconditioned iterative method precise and inspectable.\n\nThe proposed hardness measure is scale‑invariant, emphasizes tail eigen‑mass, and comes with a narrative mapping eigenvectors to skills and eigenvalues to skill strength, which leads to a clear understanding that harder tasks need longer CoT to reach the same error.", "weaknesses": "Experiments initialize with the closed‑form optimum using population statistics, so they basically hard‑code the solution instead of learning it, which invalidates empirical support for the learning‑dynamics claims.\n\nAll experiments are on synthetic linear tasks, there is no comparison to closed‑form ridge/OLS, no ablation on preconditioner estimation error, no evaluation on real reasoning benchmarks, and no demonstration that the proposed task‑selection improves anything beyond plotting selection weights.", "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": "vZtWAEWxQh", "reviewer": "Reviewer_8CQr", "summary": "This paper discusses the topic of test-time scaling and chain-of-thought prompting for transformers, examining when and why allocating extra compute at inference improves performance.\nMotivated by recent systems (e.g., OpenAI’s o1 and DeepSeek R1), the authors analyze transformers trained for in-context weight prediction on linear regression to clarify the training-data conditions under which long chains of thought emerge and help.\nThey show theoretically that, for a fixed test error, increasing test-time compute can substitute for longer training prompt context length.\nThey also find that if the relevant skills are insufficiently represented in the training data, additional test-time compute may harm performance.\nThe work formalizes task hardness via the smallest eigenvalue of the feature covariance matrix and argues that training on diverse, relevant, and hard tasks yields the best gains from test-time scaling.\nThe analysis connects chain-of-thought prompting to multi-step (pseudo)-Newton’s method and derives scaling laws governing the interaction among test-time compute, context length, and task diversity.\nAn optimal task-selection strategy for multi-task training is proposed and validated on linear self-attention models and GPT-2, with experiments extending to large, nonlinear transformer architectures.\nThe authors note limitations: the theory focuses on linear regression and single-layer linear self-attention, and future work should extend to nonlinear data generation and transformers with nonlinear activations.", "questions": "* It remains difficult to separate gains due to the proposed (pseudo-Newton) mechanism from those due to decoding strategies or exploration effects. Is there anything that the authors can discuss on this point?\n\n* Prompt sensitivity / implementation sensitivity\nIf small implementation choices can flip conclusions, reproducibility is weakened.\nIs there anything that the authors can add information on this point?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* A coherent framework linking theory and experiments  \nThe paper examines the effectiveness of test-time scaling (longer Chains of Thought) with a back-and-forth between theory (analytical results and scaling laws) and empirical validation (linear self-attention and GPT-2 / larger nonlinear transformers), which strengthens the credibility of its claims.\n\n* Mechanistic identification: CoT as multi-step optimization (pseudo-Newton’s method)  \nBy mapping the inference process to iterative optimization, the paper makes why CoT works more transparent; this unifies prior empirical observations and informs design choices (e.g., preconditioning and stopping criteria).\n\n* A theoretical trade-off between test-time compute and training-time context length  \nUnder a fixed test error, the paper shows that increasing test-time compute can substitute for a longer training context length, providing a principled basis for allocating budget between training and inference.\n\n* A principled metric of task hardness via the smallest eigenvalue of the feature covariance  \nThe work formalizes “hardness” spectrally, enabling computable diagnostics that connect to scaling laws, dataset construction, and benchmark selection.\n\n* Clear conditions for when overthinking can be harmful  \nThe paper specifies that when the necessary skill directions are underrepresented in the training data, increasing test-time compute can degrade performance, offering actionable guidance for safe deployment and monitoring.\n\n* Scaling laws that provide predictability  \nBy deriving scaling relationships among test-time compute, context length, and task diversity (via the features’ covariance spectrum), the paper moves beyond observation to testable predictions that facilitate replication and extension.", "weaknesses": "* Narrow applicability of the theory  \nI understand that the authors explicitly state the limitations (linear regression; single-layer linear self-attention) and future extensions (nonlinear data generation; transformers with nonlinear activations) in the Conclusion and Limitations sections.\nHowever, the main claims (CoT ≈ pseudo-Newton’s method, scaling laws, etc.) are not guaranteed to carry over to recent advances in LLMs, such as different model architectures, like LLaMA-base, GPT-OSS-base, and MoE-based models.\nThis is extremely important for the contribution of this paper. It would be better to discuss , as part of the main content, how the proposed method is promising for extension to different model configurations, even if confidence is currently limited.\n\n* Dependence on the task “hardness” metric  \nIn Sec. 3, the paper defines hardness as $\\mbox{tr} (\\Lambda) / \\lambda_{min}(\\Lambda)$ for the task covariance $\\Lambda$, and claims that hardness is based on the smallest eigenvalue of the feature covariance.\nHowever, it seems that relying on a single spectral indicator may be sensitive to preprocessing or representation choices and may not fully capture linguistic complexity or reasoning modes.\n\n* Detectability of the conditions under which overthinking is harmful  \nIn Sec.5, the paper reports observations in which longer CoT hurts.\nWhile the condition is described, a practical method to detect it in advance (a deployable diagnostic or proxy) is not specified.\nIf my understanding is correct, this indicates that there is still no way to predict this in advance.\n\n\n* Generality, scale, and diversity of experiments  \nThe experiments are conducted only on LSA and GPT-2 (12 layers, 8 heads, ~9.5M parameters).\nIt is insightful, but generalization to current large-scale systems and broader task suites (code, long-form reading, knowledge-intensive tasks) remains limited.", "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": "mqkbJ1PjnX", "reviewer": "Reviewer_NuKS", "summary": "This paper proposes a theoretical lens for studying how training data influences test-time scaling. They propose a simple linear regression prediction task in which the model repeatedly predicts the weight vector given X,y pairs via chain of thought. They demonstrate that increasing test-time compute reduces the in-context examples requirement during training. They also characterize factors that influence the test-time scaling curve, such as training data hardness and diversity of training data.", "questions": "The theory predicts that the generalization error decays (roughly) as $\\frac{1}{n^{2k}}$ - does this scaling hold even for extremely small $n$? In practice, if the model encounters too few examples during training, there could be an error floor that test-time scaling cannot overcome. Is there a theoretical or empirical lower bound on $n$ below which test-time scaling becomes ineffective or unstable?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This is a well-written paper. I find the intuitions on task hardness, task selection, and diversity to be particularly helpful in understanding the provided theorems.\n* The definitions of task hardness (as a ratio of the sum of variances to minimum eigenvalue) is particularly interesting and to my knowledge quite novel. As the paper states, \"an easy task is one that relies on a few dominant skills... while a hard task draws on many skills, reflected in a long-tailed spectrum\". In the context of LLMs, task difficulty is usually defined as length/complexity of a problem. The proposed definition provides a new axis of task complexity that relates feature conditioning and data geometry.\n* Although the primary focus of the paper is to understand how training data affects test-time scaling curves, the framework naturally explains when overthinking occurs. This interpretability demonstrates the elegance and generality of the theoretical analysis.\n\nOverall, the paper advances our mechanistic understanding of why test-time reasoning improves model performance, and has important consequences for designing the appropriate training data for language models.", "weaknesses": "The analysis is mostly confined to the synthetic linear regression task setup. It is not immediately clear how well these definition of task hardness translate to natural language reasoning tasks.", "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": "a2TN9sdoOw", "reviewer": "Reviewer_suZT", "summary": "This paper studies the relationship between test-time scaling and training data in the context of in-context weight prediction for linear regression. The authors develop a theoretical analysis based on linear self-attention (LSA), derive convergence properties, propose a definition of task hardness, and formulate an optimization framework for task selection. Experiments are conducted with both LSA and GPT-2 architectures to validate the theoretical trends.", "questions": "See weaknesses.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Provides a systematic theoretical analysis with clear derivations of convergence and error bounds.\n2. Introduces a task hardness measure based on the covariance matrix, which is simple and interpretable.\n3. Offers a formal perspective on task selection, framed as an optimization problem.\n4. Addresses the timely topic of test-time scaling, which has attracted significant attention in the community.", "weaknesses": "1. Overly idealized setting: The core analysis relies on linear regression and LSA, which are far from realistic large-scale model training.\n2. Limited GPT-2 validation: Although GPT-2 experiments are included, they remain confined to the same synthetic weight prediction task, limiting external validity.\n3. Simplistic hardness definition: The hardness measure depends only on the smallest eigenvalue, reflecting the limitations of the simplified setting and failing to capture the complexity of real tasks.\n4. Insufficient experimental support: The experiments are small-scale and restricted to synthetic tasks, without evidence from realistic reasoning benchmarks.", "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)