Test-Time Training Provably Improves Transformers as In-context Learners
Authors: Halil Alperen Gozeten
Muhammed Emrullah Ildiz
Xuechen Zhang
Mahdi Soltanolkotabi
Marco Mondelli
Samet Oymak
Venue: ICML 2025 poster (2025)
Keywords: in-context learning
transformers
test-time training
gradient descent
Primary area: theory->deep_learning
TL;DR: We characterize the theoretical benefits of test-time training for in-context learning and corroborate its benefits for tabular foundation models.

Abstract: Test-time training (TTT) methods explicitly update the weights of a model to adapt to the specific test instance, and they have found success in a variety of settings, including most recently language modeling and reasoning. To demystify this success, we investigate a gradient-based TTT algorithm for in-context learning, where we train a transformer model on the in-context demonstrations provided in the test prompt. Specifically, we provide a comprehensive theoretical characterization of linear transformers when the update rule is a single gradient step. Our theory (i) delineates the role of alignment between pretraining distribution and target task, (ii) demystifies how TTT can alleviate distribution shift, and (iii) quantifies the sample complexity of TTT including how it can significantly reduce the eventual sample size required for in-context learning. As our empirical contribution, we study the benefits of TTT for TabPFN, a tabular foundation model. In line with our theory, we demonstrate that TTT significantly reduces the required sample size for tabular classification (3 to 5 times fewer) unlocking substantial inference efficiency with a negligible training cost.

Test-Time Training Provably Improves Transformers as In-context Learners
Halil Alperen Gozeten
Emrullah Ildiz"
Xuechen Zhang' Samet Oymak'
Mahdi Soltanolkotahi Marco Mondelli
Abstract
may falter. This has motivated two popular approaches: in-context learning and test-time computation. In-context learning (ICL) incorporates demonstrations related to the query as part of the prompt facilitating better inference by the model (Brown et al., 2020; Min al., 2022). Test-time computation methods explicitly increase the inference-time compute to elicit higher quality responses (Snell et al., 2024; Jaech et al., 2024). An important instance of test-time com- putation is test-time training (TTT) where the weights of a model are explicitly updated to adapt to specific test in- stances (Sun et al., 2020; Liu et al., 2021). A gradicnt based TTT approach can be described as follows: Given a test prompt and a pretrained sequence model, we update the pretrained weights by performing few gradient iterations suitable self-supervieed objective (e.g. next-token pre- diction objective). We can then use the resulting model to perform the inference on the test prompt. This procedure is applicable beyond language models, and it is conceptually similar to meta-learning approaches such as model-agnostic meta-leaming (Finn al., 2017).
Test-time training (TTT) methods explicitly up- date the weights of a model to adapt to the spe- cific test instance, and they have found success in variety of settings, including most recently language modeling and reasoning. To demystify this success, we investigate a gradient- based TTT algorithm for in-context learning. where we train a transformer model on the in-context demonstra- tions provided in the test prompt. Specifically, we provide a comprehensive theoretical characteriza- tion of linear transformers when the update rule is single gradient step. Our theory (i) delineates the role of alignment between pretraining distri- bution and target task, (ii) demystifies how TTT can alleviate distribution shift, and (iii) quantifies the sample complexity of TTT including how it can significantly reduce the eventual sample size required for in-context learning. As our empirical contribution, we study the benefits of TTT for TabPFN. tabular foundation model. In line with our theory, we demonstrate that TTT significantly reduces the required sample size for tabular clas- sification (3 to 5 times fewer) unlocking substan- tial inference efficiency with a negligible training cost.
Recently, TTT has found significant success in the context of language modeling by boosting the accuracy and rea- soning capability of state-of-the-ar models (Akyürek et al., 2024; Sun et al., 2024). This success is in part due to the fact that TTT can be naturally integrated within in-context learning: We can fine-tune the transformer model to fit to the labels or chain-of oh-ouuh rationales provided as part of the in-context demonstrations, and then re-apply ICL on the adapted model. In fact, recent work (Akyürek al., 2024) utilizes variation of this procedure and additional data augmentation to obtain remarkable improvements in the ARC reasoning benchmark. This motivates the central question of our work:
1. Introduction
Modern language models can be viewed as general-purpose, addressing diverse range of user queries in zero-shot fash- ion (Kojima et al., 2022; Wei et al., 2022). However, as the complexity novelty of the query increases, such as in com- plex multi-step reasoning scenarios, the pretrained model
What are the provable benefits of test-time train- ing of transformers, specifically, for enhancing in-context learning?
'Equal contribution 'University of Michigan, Ann Arbor UUniversity of Southemn California 'Institute of Science and Tech- nology Austria. Correspondence to: Halil Alperen Gozeten <alperen@ uni.h.
Contributions. As our central contribution, we address this question by providing a comprehensive theoretical study of test-time training for one-layer linear transformers. Focus- ing on prompts following a linear dataset model, we provide
Proceedings of the 42nd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Test-Time Training Provably Improves Transformers In-context Learners
precise risk characterization of TTT with one gradient step update rule. This characterization is established in terms of three ingredients: (i) context length (number of in-context examples during inference), (ii) target sample size available for TTT in Section 4, and (iii) alignment between the pre- trained model and target task in Section 5. We show that, as sample size increases, TTT can alleviate the distribution shift bottlenecks that arise in standard ICL. This also re- veals regimes where TTT with zero small initialization is preferable to TTT with the pre-trained model (i.e. cold vs. warm start). Interestingly, our experiments with the GPT2 architecture (Radford et al., 2019) show that multi- layer transformers exhibit behavior in line with our theory. Our theory and experiments demonstrate that one step of TTT yields significant performance gains with less computa- tion, consistent with recent empirical observations (Akyürek et al., 2024) that only few gradient steps offer substantial test-time improvements. Additionally, while standard ICL requires Q(d) context length under an isotropic task prior, with d denoting feature dimension, we prove that TTT can succeed with o(d) context length by effectively memorizing the target task. Our technical novelty stems from accurately capturing the statistical benefit of TTT during in-context learning. Specifically, while the transformer is trained a single prompt, we characterize how the sample complex- ity benefit of TTT is proportional to the number of target examples within the prompt.
landseape of a one-layer linear attention model and show that the optimized model implements single step of pro- jected gradient descent over the in-context demonstrations, More recent works (Lu et al., 2024; Wu et al., 2023) extend these to characterize the pretraining task sample complexity of ICL. While these works focus on pretraining capabili- ties, we focus on the adaptation of the pretrained model to the target task. Additionally, rather than empirical risk minimization, we use gradient descent for adaptation and characterize its risk when the sample size is determined by the target prompt length. In our theoretical model, each to- ken represents data point containing input features and the target label. Remarkably, TabPFN (Hollmann et al., 2023; 2025) shows that, using this simple encoding and pretrain- ing the model with sufficiently rich data priors, transform- ers can accomplish state-of-the-art classification on tabular datasets via in-context learning. There have been also efforts (Thomas et al., 2024) to fine-tune in-context tabular models like TabPFN at the dataset level using retrieval-based local context selection.
Test-time training. Test-time training (Sun et al., 2020; Liu al., 2021; Gandelsman et al., 2022) and related test- time adaptation (Wang et al., 2021; Niu et al., 2022; Yuan et al., 2023) methods aim to overcome distribution shift bottlenecks. This is typically accomplished by adapting the model with the test example using self-supervised or unsupervised objectives. These methods can work with just a single text example or admit a minibatch of examples. For sequence/language modeling tasks, one can utilize TTT the test sequence via the next-token prediction objec- tive (Sun et al., 2024; Hardt & Sun, 2024; Hübotter et al., 2025). Specifically, during in-context learning, the query is unlabeled (e.g. math problem to solve), but we are provided with related examples and associated labels (e.g. through retrieval). Thus, we can fit to these labels as the source of supervision (essentially fine-tuning the model on the dataset of in-context examples). For instance, the training objective in Akyürek et al. (2024) utilizes this approach boosted by additional data augmentation. Finally, the computational efficiency of TTT is an important consideration which moti- vates our investigation of TTT with a single gradient update.
As empirical corroboration of our theory, we explore tab- ular learning and TabPFN (Hollmann al., 2023; 2025)- - state-of-the-art tabular foundation model pretrained with structural causal model priors. TabPFN is well-aligned with our theoretical setting with similar token encodings but different prior distributions. An important drawback of TabPFN lies in its inference cost, as it uses a full tab- ular dataset as context during inference. In line with our theory, we demonstrate that TTT can convert TabPFN into a task-specific tabular model that works equally well with significantly less data (up to 5 times fewer). This in turn implies substantial inference gains given that the complexity of celtmax-attention is quadratic in the sequence length.
2. Related Work
3. Problem Setup
We organize our discussion of related work into two main areas: in-context learning and test-time training.
Notation. Let [n] denote the set (1.... .m) for an integer n 2 1. We denote vectors and matrices using bold lower-case and upper-case letters, respectively, such as x for vectors and X for matrices. The element x; refers to the i-th entry of the vector x. We represent the zero vector of size n by O, and the zero matrix of size m X n by 0mxn- The operator tr(X) represents the trace of X, Xt is its Moore Penrose pseudoinverse, and ||X]|2 denotes the spectral norm of X. Given x. Rd, the notation [x y] E R represents the
In-context learning. In-context learning has received sig- nificant interest during the past few years (Brown et al., 2020; Liu et al., 2023; Agarwal et al., 2024). This has also motivated works toward a stronger theoretical under- standing of ICL (Zhang et al., 2024; Mahankali al., 2024: Ahn et al., 2023b; Xie et al., 2022; Garg et al., 2022; Li et al., 2023). Closer to us, Mahankali et al. (2024); Ahn et al. (2023b); Zhang et al. (2024) study the optimization
2

---

Test-Time Training Provably Improves Transformers In-context Learners
row-wise concatenation, while [x:y] IRd represents the column-wise concatenation.
k, the distributions (D D. and the pretrained starting point W".
In-context learning and fest-time training, In-context learning is the ability of the model to learn from demon- strations in the prompt, i.e. in-context. Specifically, given a sequence of demonstrations of desired input/output pairs (1). (x2, (Xn.),, yv)) E Rd X R followed by a query input x in the prompt, the model can guess the correspond- ing output query y. Concretely,
Architecture. We study the one-layer linear attention model as the sequence model:
SM(Z. W)
ZWv
In-1d+11 =x W (5)
where [l denotes the entry (n + 1, + 1) of the corre- sponding matrix. Here, the query, key, and value matrices We. Wx. Wy E R(d+lxid+l) are defined as
can define the context tokens 2 = [x, Rd+1 for i E [n] and the query token z = [x; 0] e Rd+ Then the input prompt can be written in the form
W
0axi
0 Oaxe
0dxl
W K
Wv =
0
O
0 Ixl
and we set the data matrix X Inndd and the corresponding
Z = [z1
In =I"
= yi y2
X1 X2
xn x Yi 0
E Dotdd))
labels y R to be:
To estimate the output query y. we focus on a se- quence model SM(Z, W) with Z the input prompt and W the model parameters. We assume that the sequence model is pre-trained with token distribution DrT where (z)) [x:y] iid D.T with the corresponding optimal pa- rameters given by
X
Xn
y y= [M yu]
Here, we collapse the query and key matrices by identifying the top-left d'd block of Wo E R(d-1xxd+1) with a single matrix W E Raxd, and choose the value matrix to retrieve the output of the query token with one-layer linear attention. We note that similar models have been used in prior work (Zhang et al., 2023; Mahankali et al., 2023; Ahn al., 2023a; Li et al., 2024; 2025) for the theoretical analysis of variety of phenomena, albeit not for characterizing the benefits of TTT.
W" arg min Bua,,, [(- SM(.M, w)
During inference, we test the sequence model on another dis- tribution D.TT and observe & samples of this test distribution STT= {(,),))) where y, the label of the query token inside Zj. The main idea behind Test-Time Training (TTT) is to refine the model's parameters using the test data SrT before performing inference. Concretely, the empirical loss of the sequence model on the test set STT with an arbitrary model parameter Win is given by
Data model. We consider a linear model with Gaussian data. Specifically, during pre-training, we assume the context tokens z; = [x and the query input/output vector yx: to be sampled i.i.d. from the distribution DtT (Eg)) Concretely, the outputs are generated according to the linear model y: = x BPT + E., with task parameter ßer ~ N(0, Eg), input features x; ~ NOO. and noise terms & ~ N(O, o.) for E [n].
(s,, E0,--M(S,, SM), J=l
(2)
During inference, we test the sequence model on a new task parameter ßTT with the input prompts generated following the test-time distribution (Brr). This test-time distri- bution is governed by a similar linear setting. Concretely, the outputs are generated according to the linear model yi xx BrT + E.. where the input features x; are sampled i..d. from the same feature distribution N(0, Σ,) and the noise terms are sampled i.i.d. from the same distribution N(0,02) o2).
One can thus refine W" by optimizing the above test-time empirical loss. In this paper, we focus on TTT strat- egy involving a single gradient descent step over Lsm. i.e. WIT := W v L,,,",, (W"). The error incurred by the model can then be calculated using the population loss via
L(W)=
(y SM(Z. W
(3)
We construct the test-time set STr [(Z},))) s=l by fol- lowing the test-time distribution D),(BTT) We first sample (n + k) query input/output pairs ((E, 5:) sid DIT(BTT)) Then, we designate the first n samples as fixed context to- kens across the set, whereas we use the latter k samples as queries in the set. In formulas,
Now, we will define the expected loss of the weights WIT we achieve after test-time training over all test-time training
sets STT:
Lrr := Bsn, [[(WT))]
(4)
Our goal is to characterize the loss Lrr obtained after test-time training as a function of the number of context samples n, the number of data points in test-time training
F 82 3, 8j4n Z,
=
51 32
5. 0
yj = jtn
Vje E [k].
(6)
3

---

Test-Time Training Provably Improves Transformers In-context Learners
Our procedure can be implemented using a single forward-backward pass by putting all examples within the prompt and using suitable attention mask to ensure that only the first n examples (context examples with labels) attend the last k examples (query examples without labels) and vice versa. This allows for efficient parallel training. This also means that we use a fixed context-query split where the same n examples are used as context during TTT. This can be generalized to K-fold procedure where all examples are used as both context and query during TTT (as in Akyürek et al. (2024)); however, we opted for the 1- fold option, which is more amenable to precise statistical analysis.
Remark:
We start with the characterization of W and its loss L(W").
Proposition 4.1. Let = = 0 and let Prr be an arbitrary new task parameter. Then, the optimal pre-trained model's ' parameter W and its population loss are given by
1 n1 n+ 1+++++
d+1
W
L(W") = =1BmTl
x+d+11
The characterization of W is obtained from Li et al. (2024, Corollary 1) and its loss with respect the new task param- eter ßrr is provided in Appendix B for the more general setting with noise term σ². Due to the technical difficulty of analyzing the single-step gradient in the noisy setting and to have manageable expressions, we will focus on the noiseless setting.
Single-step GD for TTT. We now turn our attention to deriving the single-step gradient update over the test set STT- To this aim, we denote by Xcontex and Ycomex the context inputs and their outputs, whereas we denote by Xusin and Fimin the training query inputs and their outputs:
Given k, n,n, and d, the loss of the pre-trained parameters, Ls,, (W*), is a random variable with respect to the test-time set STr- Consequently, the loss L(Wrr) is also a random variable dependent update to the pre-trained model's parameters, we define the expectation of the loss L(WTr) with respect to STT in (4), treating it as a function of the step size 17. The optimal step size is then selected to minimize the expected test- time training loss. In the following theorem (proved in Appendix B), we present the optimal learning rate and the corresponding improvement in the loss induced by test-time training.
cootext = [1
Fcontest = =[5 3n|".
SIT. When applying single-step
Xurin
Ente]
Juain
Fa-l ... Fatk]
The following proposition (proved in Appendix A) charac- terizes the single-step GD TTT update.
Proposition 3.1. Consider the linear attention model with parameters W E Raddd Suppose the test-time training loss function is defined as in (2) and define contess == Xconses Yeontexs E Rdd Then, for any step size n > the new parameter WIT after one pradient-dsseen step from W is given by the rank-1 update
Theorem 4.2. Let n/d = 0(1. Recall the definition of the expected loss LTr(Wrr) with respect to Syr in (4). In the isotropic covariance and noiseless setting (o² = 0). the optimal step-size that minimizes Lrr(Wrr) is
Wr W + 2n Xuain (Yrrain Xrrrii W (contes)) Meontext
In the coming sections, we will start our analysis when the covariance matrices (Ep. E,) are isotropic (Section 4), and then we provide an extension to more general covariance matrices (Section 5). Finally, we corroborate our theoretical results with empirical evidence (Section 6).
d n ~ 2(k d) (n d) 1/Trll?
With this optimal step-size n1", the improvement the loss due the test-time training
k
d²
- k (n dy3 1BTT1
4. A Analysis for Isotropic Covariances
In this section, we study the loss Lrr(Wrr) induced by a single-step gradient update in the isotropic scenario (Ep. Σ,) = II. for an arbitrary test-time parameter Brr- The assumption of isotropic covariance matrices enables us to explicitly characterize the behavior of the loss as a func- tion of n, d, and k. We first analyze how Lr7 (Wrr) varies with the number of in-context examples (n) and the embed- ding dimension (d). We then compare two distinct choices of W as function of the test-time training set size k: (i) W is obtained from (I), and (ii) W = 0xxd- The former represents a test-time training approach, whereas the latter corresponds to training from scratch with the single-step gradient update.
Remark 4.3. We note that in the proportional d regime, the Gaussian approximation in Lemma B. which is uti- lized in the proof of Theorem 4.2, introduces an error of O ). Additionally, during the derivation, we omit lower- order terms (such as 2n n² + 2n), which introduces errors of one degree lower relative to the leading terms. Conse- quently, after carrying these approximations through the entire derivation, the overall approximation error remains at most Onn since the final loss expression is zeroth- order in n, d. Hence, our result is robust to these errors as n, + oo while maintaining the ratio a = n/d. The same reasoning also applies to the non-isotropic covariance case, which will be discussed Section 5.

---

Test-Time Training Provably Improves Transformers In-context Learners
1.0
Empirical Loss (y=
1.0
Theoretical Lass (W" =
Population Loss (Normalized)
Theoretical Loss (y=
Population Loss (Normalized)
Emplrical Loss (W" opt)
0.8
Empirical Loss (y=
Theoretical Loss (W"
Theoretical Loss (y
0.8
Empirical Loss (W" 0)
Empirical Loss (y- 100)
Initial Loss opt)
Theoretical Loss [y= 100)
Initial Loss
0.4
0.6
0.5
0.2
0.4
1.0
4.0
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
a= n/d
y= kid
(a) oon-monooonic behavior the loss L(WT)) after the test-time- training update with respect to a = n/d for various y = k/d values.
(b) Losses L(Wr)) after the test-time tome-training update as function of y = k/d when using optimal pre-trained weights vs. 0axd (null).
Figure 1. Plot of the normalized population losses after test-time training when Σ, = Lp I, o' O, and the task is BTr 1. Solid lines denote theoretical predictions which match the empirical (markers) results. (a): The figure illustrates non-monotonic trend as the ratio ş shifts, see Corollary 4.4. Setting: 300; the ratio å changes between 0.1 and the three lines depict y = 1, = y = where k/d; the step size is selected optimally as per Theorem 4.2. (b): The figure reveals threshold in & at which the preferable initialization switches from the optimal pre-trained W to null model Oad as increases. transition threshold aligns well with Corollary 4.6. Setting: 200; d 400; & changing between 0 and 4n with equal increments.
Theorem 4.2 establishes the improvement produced by the test-time training a function of k, and d. When and d are fixed, such improvement is proportional to ktd Addi- tionally, if k/d = y is fixed, the loss LIr(Wrr) exhibits an intriguing behavior as a function n/d = a, as described by the next corollary (proved in Appendix B).
We plot the behavior of the loss Lrr(WrT) function of n/d for various values of y k/d in Figure la, making the following three observations: (i) As stated in Theorem 4.2, the improvement achieved through test-time training is approximately proportional to 1/(a + 1)² for large n and d. This behavior is evident Figure 1a, which shows that for every value of y. the improvement diminishes n/d increases. (ii) The non-monotonic behavior described Corollary 4.4 is clearly observed in Figure la. The increasing/decreasing regions in the loss as a function of for different values of y are consistent with the theoretical results. (iii) As increases, the improvement from the gray curve by test-time training is proportional to the ratio y/(y 1) for a fixed n/d. This observation aligns with Theorem 4.2.
Corollary 4.4. Recall the definitions of = k/d, a = n/d, and consider the loss isotropic covariance and noiseless setting. Ify > }, then the loss Lrr(WTr) non-nonoen non-monoto a. Specifically, the loss and decreasing
as function of a in the
LTT (Tr) is increasing for
By
1+
7+ for a - Conversely, y < 1.2 then Lrr(Wrr) is monotonic in
In addition to the scenario where we begin from optimally pre-trained model W natural question is how much im- provement can be achieved when the initial weight matrix W = 0,xa (zero initialization). This initialization corre- sponds to training the model from scratch using single- step gradient descent. In the following theorem, we quantify the improvement gained by applying single-step gradient descent from this initialization.
The phase transition points described Corollary 4.4 are approximate but become exact as and d approach in- finity while maintaining the ratios y k/d and n/d. The non-monotonic behavior identified in Corollary 4.4 is observed as a result of two opposing effects, which are the initial loss and the improvement by TTT. As grows against (i.e. a increases), the pre-trained model does better ini- tially and already has a lower loss before TTT, which makes harder to be further reduced by TTT the rank-1 I update unable to correct all directions. On the other hand, when d grows against n (a decreases), the initial loss is high, pro- viding more room (error) to be corrected by rank-1 update, and thus, the improvement by TTT is larger. This intuition aligns with Theorem 4.2, which establishes that the TTT improvement scales as و. Together, these two trends result in the non-monn oon-oontooiic behavior.
Theorem 4.5. Consider the isotropic covariance and noise- less setting (o² = 0). Suppose the initial weight matrix is W 0axa- Then, the optimal step-size that minimizes LTr(W is
n¹ 2(k + () + 4n d)
With this optimal step-size n",, the improvement the loss
5

---

Test-Time Training Provably Improves Transformers In-context Learners
due to test-time training
feature covariance matrix Σn isotropic. This assumption enables us to characterize the loss function Lrr (Tr(TTr) with respect to the alignment between the pre-trained initializa- tion matrix W and the new task parameter Brr-
k
n²
- -(r))--rrrWr
k+d+1 1 n² + 4n 3 Brrll
We start with the definition of two quantities, which will be used to express the optimal step size and the loss improve- ment via the test-time training.
where L(W") = 1187111
We provide the proof of Theorem 4.5 in Appendix B, where we also solve the more general noisy setting with initial weight W = 0axd- If = 0 (1[TTT11) and k grows suf- ficiently faster than d (e.g.. k/d - co), then the test-time training update can reduce the loss to near 0 as k, n, d - oo with a = n/d= = @(1.
Definition 5.1. For an arbitrary W E Rdxd, define A BTT (1 Brr and n 1 [WW]..
Here, the term A represents the misalignuent between the test-time task parameter Brr and the initial parameter W, whereas the term represents the total signal power of the one-layer linear attention system with parameter W. Note that the terms A and B are function of the new task parameter ßTT and the initial weight parameter W.
Armed with Theorem 4.2 and 4.5, we can now establish when it is better initialize with the pre-trained W opposed to the zero (null) initialization, as the size k of the test-time training set varies.
Using the definitions of A and B, we are ready to general- ize Proposition 4. to the non-isotropic and rank-deficient covariance matrices. Note that for the rank-deficien co- variance matrix, we take the optimal W as the minimum Frobenius-norm matrix that minimizes the loss in (I).
Corollary 4.6. Recall the definitions of a = =n/d and y= = Consider the setting in Theorem 4.2 and test-time training described in Proposition 3.1 with both pre-trained and zero initializations. Then, under the isotropic covari- ance and noiseless setting, there exists a threshold y" given (a+ (a+1) 1)²
Proposition 5.2. Let Ix = I, = O and suppose Σp is an arbitrary diagonal covariance matrix with the first r diagonal entries being non-zero. Then, the minimizer W of the pre- pre-training loss (1) that has minimal Frobenius norm and its population loss are given by
such that y < y if and only if it is better to utilize the pre-trained initialization over the zero initialization Oaxd-
by y" x
(ar 22
Corollary 4.6 identifies a phase transition point as a function of y and that distinguishes the region where test-time training outperforms training from scratch. In Figure Ib, we illustrate the loss Lrr(Wrr) as function of y = k/d for pre- trained and zero initializations, with a fixed a = n/d = 1/2. Based on Corollary 4.6, the phase transition point for y is (3/2) 10, which aligns with the empirical results. 5/2
W P== -(n 1)1' tx(E))Z))' L(W") = A B,
where I' = ddigg(,,, Oz-r).
We generalize the characterization of W obtained from Li et al. (2024, Corollary 1) to rank-deficient covariance matrices, and its loss with respect to the new task parameter Brr provided in Appendix C.
The pre-trained initialization provides a significant advan- tage when the test-time training set is small, as it introduces a strong prior that facilitates better performance. How- ever, as the test-time training set grows, this initialization becomes a limitation because the rank-one update to the pre-trained matrix is insufficient achieve near-zero loss. In contrast, with zero initialization, the rank-one update becomes highly effective when the test-time training set is sufficiently large, enabling it to achieve near-zero error. This demonstrates that while pre-trained initialization is bene- ficial in data-scarce scenarios, zero initialization better suited for scenarios with sufficient test-time training data.
Note that the eigenvalues of W" are smaller than 1/(n + 1) as Σg is positive semi-definite matrix. This condition ensures the stability of the linear attention model described in (5) as the magnitude of the SM output scales with the context length n when W is fixed. Assuming a similar criterion for the set of W, we provide the optimal step size parameter and the corresponding test-time training loss.
Theorem 5.3. Let n/d = @(1), Σ, = I, and 0² = 0. Sup- pose that the initial weight matrix W is a diagonal matrix whose eigenvalues are in the interval [0, 111 Then, the
'This assumption is equivalent to more general case where Σ, and Σe jointly diagonalizable, as proved in Appendix C. Joint diagonalizability refers the case where there exists a unitary matrix @ c add such that the covariance matrices Σp = Q and Σ, = QA,0 can be expressed with diagonal Ap and A, matrices.
5. A nalysss for General Covariance
In this section, we extend the previous analysis to the gen- eral case where Σ, and Σp are any jointly diagonalizable covariances. Specifically, we assume the task covariance Eg to be a non-isotropic diagonal matrix while keeping the
"In the joint diagonalizahility case, the restriction W covers the space W that is jointly diagonal with Eg and Σ,
6

---

Test-Time Training Provably Improves Transformers In-context Learners
Thooodiall Loss Ry 100
1:2
Tresrciall Los Ceg wart
ropulation (oommlltedd
Erepiiial les Mae (88" -aag
Pepustatoo Lase (Kormaltra
Emprical Ray ("" see
0.6
Theerstial Lees Razi
Treennieal (aa
0.3
0.8-
Subsoquent steps
Bas
0.6
leital Hest
ir (4p
0.3
Populaation
f-0.2
0.2
0.2
Inttee
0.00 0.25
0.50
0.75
1,00
1.25
1.50
1.75
2,00
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
0.0
y
y-
Gradient Descent Step
(a) Losses after -est-iim--taannnng as a func- tion of = k/d for optimal weights W based on covariances and also 0xzd-
(b) Losses after test- traaainnn vs. y for optimal weights W with worst/best aligned covariances and 0asd-
(c) Comparison of losses after single and multiple steps of gradient descent at test time for different step size decays.
Figure 2. Population losses in single and multi-step update scenarios. I = L, o² = 0 and step sizes are optimal based on the formula Theorem 5.3. W' = opt and W = represent the pre-trained matrix in (1) and zero (null) initialization, respectively. (a): Setting: и 300; d' = 600; k changing between 0 and 4 with equal increments; Σn.1 = diag(/ 0.5 120):: Σpp diag(0.5 1 120)); Pir 250 ; 0.5 220 Solid lines denote theoretical predictions which match the empirical (markers) results. (b): Setting: " 250;
= 500; k changing between and 4n with equal increments; Σạ. = diag(1, 0uy): Σμ² = diag(O, Lun): Br = 0 (c): The figure illustrates the empirical results, which imply that with the initially optimal step size, significant improvement can be gained by single gradient step. Each line depicts a different reduce factor on the step size n7 after each gradient step. Setting: n = 50; d = 100; = 50 x d; oo W is optimal.
Ep 1: BrT=
optimal step size that minimizes the population loss given in (3) after the test-time training update
parameter satisfying ni Wil2 < 1- 5(n, d) for some fixed 8(n, d) > 0 as k/d approaches infinity by Theorem 5.3. This means that test-time training achieves the optimal solution when we have zero or small initialization of W as k/d - co.
A
n¹ 2 2(k d) BBTTI (A
For the remainder of this section, we focus on the op- timal pre-trained W constructed in Proposition 5.2, i.e. the minimum Frrnuum---oonnuu---a norm solution of the pre-training loss (I). When we initialize the pre-trained model weight as W the loss Lrr (Wrr) can be computed by combining Proposition 5.2 and Theorem 5.3. When initialize the weights W Oasa. then the loss (rr(Wrr) is obtained by using the fact that L(0 ||BTTI?. By combining these results, we identify the phase transition point in the non- isotropic task covariance setting, which determines when test-time training outperforms training from scratch.
With this optimal step-size Ti, the improvement in the loss due to test-time training and the initial loss are approxi- mately
k A L(W) A
L(W) (T)
k+dA+
A more general and equivalent version of this result, which applies to any pair of jointly diagonalizable (not necessarily diagonal) covariance matrices Mn and ΣB and to any W that is also jointly diagonalizable with them, can be found in Appendix C. The above approximations are robust and be- come exact as n.d - co while having fixed ratio n/d; see Remark 4.3. Theorem 5.3 characterizes the improvement due to test-time training for a rather general class of W. It shows that. for fixed ||WI the loss improvement increases monotonically with A. In other words, when the misalign- ment measure A is larger, test-time training achieves greater performance gains. Consequently, considering the optimal pre-trained W*, if the eigenvalue spectrum of Σg aligns with the entries of the task Brr (i.e., larger eigenvalues match larger entries), then A is larger and thus yields a bigger improvement.
Corollary 5.4. Consider the setting in Theorem 5.3. Recall the definitions of a = n/d and y = kjd and define A
B
Then, there exists phuse transition
|Brrll""
IBrrlP (c1 -2) - (c1 c2) such that <<y and
point y*
C2 (2c1 + c2)
only it is better to utilize the pre-trained W over the null initialization 0axa-
In a similar way Corollary 4.4, the phase transition point mentioned above is approximate for finite k, and d. and it becomes exact when k, n, and d approach infinity while keeping the ratios and fixed.
Furthermore, we note that the optimal L(W) over all W e d is O(n ) by Lemma 1. This optimal loss behavior can be achieved through test-time training with the initial
We illustrate the findings of Corollary 5.4 in Figures 2a and 2b and make the following three observations: (i) The
7

---

Test-Time Training Provably Improves Transformers In-context Learners
0.7
0.8
Scratch Pre-train
0.6
0.7
Loss (Normalized)
Accuracy
0.5
0.6
0.5
0.4
0.4
0.3
TabPFN TabPFN TTT
0.3
0.2
0.2
200
400
600
800
1000
6 6
10a
2
102 102
(Context length)
(test-time train size)
(a) Accuracy of TabPFN v2 model with and without test-time- training a function of number of in-context samples n.
(b) Losses after sett-imm-ttriiing for GPT2 comparing pre- vre-trained and scratch models function of test-time training set size
Figure 3. (a): Accuracy of TabPFN v2 with and without test-time training. Setting: k = 1000; a changing between I and 1000. (b): Normalized loss after test-time-trrining with pre-training zero initialization. Setting: d 60; 40; k changing between 64 and 512 in increments of 64; o² = 0.01; Ip- diag(0.1 125, 0.5 1o. 1s): L, I. We sample Brr from the distribution N(O, Ip,r) where EhT diag(12,. Joo 0.1 ..2).
phase transition point is monotonic as a function of C1, while keeping C2 fixed, which implies that it is monotonic as a function of the misalignment term A. This means that, as we increase the alignment between the task covariance Lp and the new task parameter Brr. the crossing point of the green and blue curves occurs later than that of the green and orange curves, observed in Figure 2a. (ii) The worst- aligned case over the set of all possible (Eg. Brr) pairs is the one that satisfies BTT ΣaBrT = O. In this case, the corresponding e₁ is equal to I, which implies that the phase transition point in Corollary 5.4 is less than 0 as C2 > O. This means that training from scratch is always better than utilizing the test-time training for the worst-aligned case, which is depicted in Figure 2b as the orange curve. (iii) The best-aligned case over the set of all possible BTT) pairs is the one that satisfies BTT EgBrr/tr(Ep) = 1|Brll?. In this case, the corresponding C1 scales on the order n and C2 on the order of - which in turn makes the phase transition point from Corollary 5.4 approximately n. When k, n. and all grow while maintaining the ratios n/d and y = k/d fixed, this phase transition point approaches infinity. As result, test-time training is always better than training from scratch in the best- best-aligned case, which is depicted in Figure 2b as the blue curve.
on tabular learning by demonstrating how test-time training reduces context length for in-context learning on TabPFN. Finally, we demonstrate how the pre-trained model behaves against the scratch model under the -es-time-training update for the GPT-2 model in the presence of distribution shift,
6.1. Multi-Step Gradient Updates
In Figure 2c, we investigate the benefits of multiple gradient steps by choosing the optimal single-step size 1) based on Theorem 4.2, and then applying different decay factors step size after each subsequent step. We note that reducing the step size each time is necessary to ensure continued im- provement of the loss. The results in Figure 2c confirm that a single-step test-time- training update can capture signifi- cant improvement, whereas additional steps yield diminish- ing returns compared to the initial step. Consequently, the single-step gradient update offers favorable trade-off with less computation yet with the performance near multi-step finetuning.
6.2. Experiments TabPFN
We demonstrate that test-time training (TTT) can signifi- cantly reduce the context length required for a given per- formance on TabPFN v2 (Hollmann et al., 2025). Specifi- cally, we evaluate the TabPFN v2 model on the The Tremen- dous TabLib Trawl (T4) dataset (Gardner et al., 2024) for more comprehensive evaluation, which large-scale high-quality collection of tabular benchmarks. Following the official TabPFN v2 implementation and our theoretical setup, select the datasets containing at least 1,250 sam- ples (with 1 1,000 for training. using an 80-20 split), limit
6. Experiments
In this section, we provide empirical results in light of our theoretical insights from previous sections. The first discus- sion of this section focuses on whether further computation via multi-step gradient descent at test time yields significant improvements beyond a single-step update. Then, we focus
8

---

Test-Time Training Provably Improves Transformers In-context Learners
the number of classes 10 by choosing the most frequent ones, and restrict datasets to 100 randomly selected features. We also convert regression tasks into 10-class classifica- tion tasks based on quantiles to maintain consistency across training and evaluation. For each selected dataset from T4, we evaluate TabPFN v2 with context length n varying from I to 1000. For the TabPFN (blue curve Figure 3a), we directly load the pre-trained model and vary the context win- dow length during evaluation. In contrast, for TTbPFN+TTT (orange curve in Figure 3a), we finetune the model using different context lengths with k= = 1000 samples. As the con- text length decreases, the samples are divided into 1000/n groups, where each group undergoes 50 training iterations. We then report the average performance across all datasets with enough training samples.
three times using different random seeds and report the averaged results.
The results in Figure 3b show that using pre-trained model with TTT is more advantageous when the test-time training set size is small. However, as the training set size increases, training from scratch outperforms the pre-trained model, aligning with Figure 1b.
7. Discussion
We have developed a theoretical framework to character- ize how single-step gradient update at test time enhances in-context learning. In the case of an isotropic covariance matrix, we analyzed the improvement in loss under test-time training as a function of the number of in-context samples (n), embedding dimension (d), and test-time training set size (k). We then extended our analysis to the non-isotropic covariance case, examining how the alignment between the new task parameter (BrT) and the task covariance ma- trix (Ex) affects performance. Our results reveal a phase- transition threshold on k in both isotropic and non non-isotropic settings, beyond which zero initialization (training from scratch) can outperform pre-trained initialization. Empirical results align well with our theoretical findings, and together they underscore that single-step test-time training offers sig- nificant performance improvements at low computational costs. Overall, our findings highlight test-time training as a lightweight, supervised enhancement to standard in-context learning.
In the previous sections, we have theoretically shown that TTT reduces the required in-context sample size for queries from a new task by performing only single adaptation step, thus enhancing the efficiency of inference. As empiri- cal corroboration of this, Figure 3a illustrates that TabPFN with test-time- training allows the model with 200 in-context samples to almost match the performance of the TabPFN without cest-timettrainnng with 1000 in-context samples. This means that test-time training reduces the number of required samples for tabular classification by about 5 times. It is worth emphasizing that the sample complexity benefit is more evident near the zero-shot regime as TTT helps the model memorize new task dynamics by improving the ac- curacy from 0.2 to 0.45. A log-scales version of Figure 3a is provided in Appendix D (Figure 4) for a clearer view of performance gains across different scales of n.
rest-immetraining
As future direction, one can investigate the analysis of TTT under more general architecture and data model. Possible extensions to the architecture model might be to analyze the TTT with softmax attention or multilayer atten- tion, whereas possible extensions to the data model might be to analyze K-Fold vross-validation instead of the cur- rent 1-Fold validation approach. Additionally, investigating TTT in unsupervised or semi-supervised in-context learn- ing (ICL) settings presents another promising direction for extending this work. We hope that our results will serve as foundation for future research on the interaction of TTT and ICL methods.
6.3. Experiments on GPT-2
We demonstrate how the pre-trained model compares against the scratch model under test-time training with distri- bution shift using GPT-2 architecture (Radford et al., 2019), as shown in Figure 3b. Our results are in agreement with Figure lb.
Throughout the experiments, we consider the linear model with Gaussian data following Section 3. In order to obtain pre-trained model, we sample multiple tasks from the distribution with task covariance Σg and train the model from scratch on these tasks until convergence. In contrast, in the scratch setting, the model is initialized randomly using the GPT-2 architecture without any pre-training. During test- time training, we evaluate the model on new task parameters sampled from the distribution Σg,,, where the covariance structure is provided in the caption of Figure 3b. We then apply test-time training to both models by varying the test- time training set size and updating all parameters of the GPT- 2 model. The results are averaged over all the new tasks sampled from Ep.. Finally, we repeat the entire procedure
Impact Statement
This paper presents work on test-time training to provide stronger understanding of the principles of modern Al mod- els. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here
9

---

Test-Time Training Provably Improves Transformers In-context Learners
Acknowledgements
tional conference on machine learning, pp. 1126-1135. PMLR, 2017.
H.A.G., M.E.I., X.Z., and S.O. were supported in part by the NSF grants CCF2046816, CCF-2403075, CCF-2008020, and the Office of Naval Research grant 1I0014241222 M. is funded by the European Union (ERC, INF, project number 101 101161364). Views and opinions expressed are, however, those of the author(s) only and do not necessarily reflect those of the European Union or the European Re- search Council Executive Agency. Neither the European Union nor the granting authority can be held responsible for them. M.S. is supported by the Packard Fellowship in Sci- ence and Engineering, Sloan Research Fellowship in Math- ematics, an NSF-CAREER under award #1 F184339,, DARPA FastNICS program, and NSF-CIF awards #1813877 and #2008443, and NIH AP2LM014564400 The authors also acknowledge further support from Open Philamthropy, Ope- nAI, Amazon Research, Google Research, and Microsoft Research.
Gandelsman, Y., Sun, Y., Chen, X., and Efros, A. Test-time training with masked autoencoders. Advances Neural Information Processing Systems, 35:29374 -23375, 2022.
Gardner, J., Perdomo, J. C., and Schmidt, L. Large scale transfer learning for tabular data via language modeling. arXiv preprint arXiv:2406. 1001,, 2024.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? a case study of sim- ple function classes Advances in Neural Information Processing Systems, 35:30583-30598 2022.
Hardt, M. and Sun, Y. Test-time training on nearest neigh- bors for large language models, 2024. URL https: //arxiv.orl wttp 18466.
Hollmann, N., Müller, S., Eggensperger, K., and Hut- ter, F. Tabpfn: A transformer that solves small tab- ular classification problems in a second, 2023, URL ht http://axxxvv..g//bbbb// 01122000888....
References
Agarwal, R., Singh, A., Zhang. M., Bohnet, B., Rosias, L., Chan, S., Zhang, B., Anand, A., Abbas, Z., Nova, A., et al. Many-shot in-context learning. arXiv preprint arXiv:2404. 1018, 2024.
Hollmann, N., Müller, S., Purucker, I.., Krishnakumar, A., Körfer, M., Hoo, S. B.. Schirmeister, R. T., and Hutter, F. Accurate predictions on small data with tabular foundation model. Nature, 637(8045) 2025. doi: 10. (088///1586-004-000000 //doi org/10. 10. 1038/ 024-
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn implement preconditioned gradient de- scent for in-context leaming. 2023a. URL https: //arxiv. www. htabb//233...00 00297.
URL https:
Hübotter, J., Bongni, S., Hakimi, I., and Krause, A. Effi- ciently learning at test-time: Active fine-tuning of Ilms, 2025. http:: http://atw..ooooaa/////////o httg//bb/24100 08020.
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-context learning. Advances in Neural Infor- mation Processing Systems, 36, 2023b.
Jaech, A., Kalai, A., Lerer, A., Richardson, A., El- Ei-KKihkk, A., Low, A., Helyar, A., Madry, A., Beutel, A., Car- ney, A., et al. Openai ol system card. arXiv preprint a4Xi::241..16720,, 2024.
Akyürek, E., Damani, M., Qiu, L.. Guo, H., Kim, Y., and An- dreas, J. The surprising effectiveness of test-time training for abstract reasoning. arXiv preprint 2xXxv:2411.07279 2024.
Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. Large language models are zero-shot reasoners. Ad- vances in neural information processing systems, 35: 22199-22213, 2022.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Li, Y., Ildiz, M. E., Papailiopoulos, D., and Oymak, S. Transformers as algorithms: Generalization and stabil- ity in in-context learning- In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 19565-19594. PMLRR 23-29 Jul 2023. URI. https: / //proceedings .mr. wwt.ssmww...ll..eeeeeeeeeee www.pppwwsssseeeeeeeee .htl...llll
Chandrasekaran, V., Recht, B., Parrilo, P. A., and Will- sky. A. S. The convex geometry of linear inverse prob- lems. Foundations of Computational Mathematics, 12 (6):805-849, October 2012. ISSN 1615-3383, doi: URL http: /ttp://dx. /xx 012 012-9135-7.
10. oi. wrg/rog/11100
Li, Y., Rawat, A. S., and Oymak, S. Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond. In Globerson, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J., and Zhang, C. (eds.),
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Interna-
10

---

Test-Time Training Provably Improves Transformers In-context Learners
Advances in Neural Information Processing Systems, vol- 37, pp. 138324-138364. Curran Associates, Inc., 2024. URL https: https////ppoeeedngg...rrioggo cc/p aper d58279e75dd223883 ottp:///eeeeeeeeeeeeee------eee
Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling IIm test- time compute optimally can be more effective than scal- ing model parameters. arXiv preprint arXiv:2408.03314, 2024.
pdf.
Sun, Y., Wang, X., Liu, Z., Miller, J., Efros, A., and Hardt, M. Test-time training with self-supervision for generaliza- tion under distribution shifts. In International conference on machine learning, Pp. 9229-9248 PMLR, 2020.
Li, Y., Tarzanagh. D. A., Rawat, A. S., Fazel, M., and Oymak, S. Gating is weighting: Understanding gated lin- ear attention through in-context learning, arXiv preprint arXiv:2504.044 04300, 2025.
Sun, Y., Li, X., Dalal, K., Xu, J., Vikram, A., Zhang, G., Dubois, Y., Chen, X., Wang, X., Koyejo, S., al. Learn- ing to (learn at test time): Rnns with expressive hidden states, arXiv preprint arXiv:2407.0 04620,, 2024.
Liu, P., Yuan, W. Fu, J., Jiang, Z., Hayashi, H., and Neubig, G. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys, 55(9):1-35, 2023.
Thomas, V., Ma, J., Hosseinzadeh, R., Golestan, K., Yu, G., Volkovs, M., and Caterini, A. Retrieval & fine-tuning for in-context tabular models, 2024. URL https://arxxi org/abs
Liu, Y., Kothari, P., Van Delft, B., Bellot- B., Mordan, T., and Alahi, A. Tu+++: When does self-supervised test-time training fail or thrive? Advances Neural Information Processing Systems, 20-22208-2,80, 2021.
05207.
Wang. D., Shelhamer, E., Liu, S., Olshausen, B., and Darrell, T. tent: fully test-time adaptation by entropy minimiza- tion. 2021.
Lu, Y., Letey, M., Zavatone-Vedt J. A., Maiti, A., and Pehle- van, C. In-context learning by linear attention: Exact asymptotics and experiments. In NeurIPS 2024 Workshop on Mathematics of Modern Machine Learning. 2024.
Wei, J., Bosma, M., Zhao, V., Guu, K., Yu, A. W., Lester, B., Du, N., Dai, A. M., and Le, Q. V. Finetuned lan- guage models are zero-shot learners. In International Conference on Learning Representations, 2022.
Mahankali, A., Hashimoto, T. B., and Ma, T. One step of gradient descent provably the optimal in-context learner with one layer linear self attentonn, 2023. URL
https://arxil org/ / 005766.
Wu, J., Zou, D., Chen, Z., Braverman, V., Gu, Q., and Bartlett, P. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv: 231.0839,,, 2023.
Mahankali, A. V., Hashimoto, T., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations, 2024. URL attps://openrevieet nete wwt.terunt/iiiii 8p3fu56.Kc.
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Rep- resentations, 2022. URL https::/ net /forum?id==ddVFCH
Min, S., Lyu, X., Holtzman, A., Artetxe, M., Lewis, M., Hajishirzi, H., and Zettlemoyer, L. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.1 12837, 2022.
Yuan, I.., Xie, B., and Li, S. Robust test-time adaptation in dynamic scenarios. In Proceedings of the IEEE|CVV Conference on Computer Vision and Pattern Recognition. pp. 15922-15932, 2023.
Niu, S., Wu, J., Zhang, Y., Chen, Y., Zheng, S., Zhao, P., and Tan, M. Efficient test-time model adaptation with- out forgetting. In International conference on machine learning, pp. 16888-16905. PMLR, 2022.
Zhang, R., Frei, S., and Bartlett, P. L. Trained transformers learn linear models in-context, 2023. URL https:/ //ts: xiv. xiv.org///ooooooooo........oooo. org/ //sg//////// / 230.2
Petersen, K. B. and Pedersen, M. S. The matrix cookbook, nov 2012. URL http:://ww/t compute dtu. dk/pu /pu bdb/pubs/ 3/////37444fll html, Version 20121115,
Zhang, R., Frei, S., and Bartlett, P. L. Trained transform- ers learn linear models in-context. Journal of Machine Learning Research. 25(49):1-55 2024.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners. OpenAl Technical Report, 2019. URL https: //cdn.openai.......m..omoo.aa..a cww/come/aaooeeeee heae/beeerr-ee lang-age modelss/a nguage_ nguage_models_ wwwelsddoolomoooee are unsupervised mul tittkk_t carners.pdl.

---

Test-Time Training Provably Improves Transformers In-context Learners
A. Proofs for Section 3
Proposition 3.1. Consider the linear attention model with parameters W E R Suppose the test-time training loss function is defined as (2) and define Wemteu comeee Yeomext
R.. Then, for any step size 1j the new parameter WT after one gradient-descen step from W is given by the rank- 1 update
W W 2n Xrrin"
(Yrroin
Xuain
W Mcontess)
Mcontext
Proof. We derive the update on the weight matrix W via a single gradient step using the loss function defined over the set Zuain mann Yurin R Note that SM(Z,, W) X W.ecmeext Yeontexi
X Wilcomtest- Then, for the training set
Zuiin the vector of predicted values under W is:
8 uaiin (W) X Wii context R
Recall the objective given by (2):
L(W)- = Σ
SM,,, w))
Σ
G x Wu context) )
Vtraio Xrain W context l
i=++1
1=++1
The gradient of L w.r.t. W is then obtained by differentiating the summation:
Vw L 2
I W Meontext) Mcontext
rriin (Ftrain Xtrain Wccooeeet Mcontext
This is rank-1 update as gran Vrrain Xiran Wuccomee)) 1) Mcontext pq" rank-1 matrix. Evaluated at W, the update after one gradient step with step size 77 becomes:
W W 77 Vw L(W)
= W 217 Xmii "(rminn X W Wigonteea) Mcontext
This completes the proof.
B. Proofs for Section 4
Lemma B.1 (Validity of Gaussian Approximation Isotropic Case). Let X E Rndd have i..d N(O, 1) entries and let [|w|| = 1. Define q XT Xw. can be written as q = w + g + such that g ~ N(O, Ig/n) and is a residual random variable that obeys
B[[le||]] 59 n²a
Note that represents a lower order term as we have ||w]] = and Blll d/n.
Proof. Set = Xw N(0. I,). Let P = -- Decompose X = www X + PXT and observe that PXT is independent of wwT XT wh Additionally, introduce independent v N(0, I.) and set X PXT + wVT. Finally, set TIXX"Th//// h/||||| and
vix'
g= PXT (h- /nh//Ihll) Vnww" h/1//]/ h-
To proceed, observe that X' has i.i.d N(0, 1) entries and hence g nla) We can write
X Xw X w||h|| P = w||h|| g +r
(1/11 ntww
Let us now set r 2 n)w and investigate B[lle'|| 11 Recalling r = ri - r2, can apply the standard upper bounds
|[|le'1]]] < 3 E[( n) 331+ 3
12

---

Test-Time Training Provably Improves Transformers In-context Learners
Next, from standard facts about chi-square random variable with n degrees of freedom, have B[[ n)² 2n. Similarly, set h h Vnh//|Al)- Note that |[[[/|P]] B[[]A]] vm)²] 2n [[[[a]] 2. This step the fact that
|[[]]]]] ) - which can be obtained by induction argued in Section 3.1 of (Chandrasekaran et al., 2012).
(75) vn++1
Furthermore, considering that X' is the sum of two orthogonal vectors, we obtain
B|l|nil = BX| (h
Vnh/||ADI)]] s BI|XT (h Vnh///bDI]]== d d [[]] 2d.
Finally, set 2 ///////|| N(0, 1). We bound
B[]|r1]]]] s B[I Vnwell"] = n.
Aggregating these, we obtain
B[lle'I']s 3(2n + 2d n) + 6d.
To conclude, set e'/n After normalizing X Xw by and the change of variable - -g/n with g V(0,fa/m) we obtain
n x Xw = -w=
such that B[[le[]]] B(lle" = 9/nn 6d/n?.
Lemma B.2 (Population Loss Expression). Consider the linear attention model as the sequence model described in Section 3 with weight matrix W. Suppose that the new task during the test-tine test-time training is Brr. In that case, the population loss of the model in (3) with respect to this task is given by:
L(W) B D. (B) (0: - x0-x- w'TTT))]
= =B [xx nI, WI, WT Σg n(n WEE,E,,W,,,W + WE Ontr(W"E, E,,,,, o
Proof. Since the target function f is linear, the labels are given by:
y; f(x)) & Brr x; + 8.
We collect the samples into X aod and Y e R" as before, that:
= XBrr +
where & (f1....,En, represents noise the context labels. Recall the prediction by the model:
SM(Z, W) = =
Then, we have the following population loss for the new task BTT:
L(W) BLLyMar. BTT BTT- WXX y)² +
Since = XBTT E, we have:
SM(Z, W) WX X +x =
The error is:
x BTB - xT XB -
WX XBrT) - - W
= =x WX X) TWX
Therefore, the population loss is equal to:
L(W) = Ex (xT (1- W X) - x [TWTT]] 5) +02
Ex Exg (I (xT[I-WX' WX X) x Wx w () x) +02
13

---

Test-Time Training Provably Improves Transformers In-context Learners
We will first take the expectation with respect to query and noise vector 5. Then, using independence, will take the expectation of the resulting expression with respect the context matrix X. Expanding the square gives
(x (1 W WX B) 1 x w
WXTT)) 5) (x" (I W ----XXXXBB x) 2(x" (I W X) WX 5) (x wx'
Now, let's consider each of these three terms individually, First of all, the expectation cross-term
Exgl-2((x' WX B) ( WXT
=) =.
Besides, since the first term does not depend on 5. we have:
Exgl(x (1- W B) B (I WX X) 2, - BT
Consider now the last term T W Recall that & zero-mean noise with B[EE"] L and B[xx"] Ex. First,
condition on
(x (x"xxTg) g) 5 (XW" xx WX kg.
Taking expectation over 5:
Ee((x WX 5) x] t' xx )
Use the cyclic property of trace and the fact that xxx is rank
tr(XW xx WXT) = tr(xx" WX XW x (WX XW xx
Thus, we obtain:
Eel(x' felxxWXTTT) x) [x] o XW")x.
Taking expectation over x yields:
Ex (WX XW )x] trwwtrrwxxx..mmmmm E.).
since [[xx Ex. Combining both expectations:
Bxgl(x WX === - XW X,),
Thus, combining the three terms yields the following expression:
Bri (1 w x) -0_(1--XXXXX B) + tr XWT
Now, we aim compute the loss function averaged over all possible in-context example sets X:
L(W) Ex BTi T (I W X Σ, (I W x) Br tr XW Σx) +02
Rewrite the first expression inside the expectation as
BTT A Σg A BrT
where
A A=
Then
Bx BTT A TA, ABTT BBTTT (Ex(A BIT.
We can write
AEEAA- (1-X XI-X" Ex (1- xx
Hence
ExIA E, A] = Σ, Σ, W X| B|X NMIx Σ, + BI(X"X WT W. (X X)).
We know that 1 X| n Σx; thus, our expression becomes:
Ex E. A] Σ, WE, nE.WT E, BIX X X)).
Using Lemma B.3, we know the following fact:
Ex [x X) M (X x)] = n(n+ ME
14

---

Test-Time Training Provably Improves Transformers In-context Learners
Finally, the overall expression becomes:
Bx|A A]
Σ. WE, n W Ex + n(n + W,,, E,W,, ,,
Adding the constant terms to the expectation gives:
BTT Σ nLr WT. W LS + n(n + 1)2, 1),, #tt(( n E, BTT.
Now that have calculated the expectation of the first term, let's focus the noise term involving trace inside the expectation. We have:
Bx trr(x'' XW XWT),)] ntrrw =,,
As result, we arrive at the final expression:
I(W) EX ((x BTT +5 WXTx)
=BT [2, nI, W., nE,WT I. n(n 1),, E.W,,, +ntr(W'E, E.W,),,, + + Pentr(W'E,WE,, w,,,,) +02.
Lemma B.3. Let Raxd be a random matrix whose rows x; Rd, for i = I....,n, n, are i..dd drawn from N (0, E). Let M& Rald be any fixed symmetric matrix. Then, under these assumptions,
Bx (x
n(n 1)EME + E
Proof. We can write the matrix X as:
x = Σ
ial
Therefore,
11
Σ
x
M
Σ x
ΣΣ
x
(x) Mx,) xJ.
Taking expectation over X, we split the sum into the diagonal part (i = j) and the off-diagonal part (i i):
Xx MX x)}
Σ (xx Mx)
Σ B (x]
diagonal
(ii) off off-dingnnal terms
(i) Diagonal Terms (i= D). Each x, ~ N(O, Σ), so we set
Mx))
Mx,) MBx) x
Recall the standard result, which follows from Wick's 's Theorem (Petersen & Pedersen, 2012)[Section 8.2.4] for zero-mean Gaussian x ~ NN(0 2) and any symmetric matrix A:
x(x² T Ax)x7 = 2ΣAΣ + tr(AE)E.
Substituting A M. it follows that
E [x] Mx)x 2EME t)(
Since we have diagonal terms, summing over = 1,.... yields
B x (x) Mx 12MME+ tr Σ)E].
15

---

Test-Time Training Provably Improves Transformers In-context Learners
(ii) Off-Diagonal Terms (i+ D). For i j. the vectors x; and x, are independent. Then,
My (x] Mx,) x3] = E[x; x [M ΣMM.
Hence, considering that we have n diagonal terms and n(i 1) off off-dia terms, combining (i) (ii) gives:
Ex (X (XT n (2EME + tr( ()( n(n 1)Σ = =1(n + 1) ntr(M
Ex
This completes the proof.
Proposition 4.1. Let Lg = = = 0 and let Brr be an arbitrary new task parameter. Then, the optimal pre-trained model's 's parameter W" and its population loss are given by
|Brlll d
W
L(W"
n +++++
n+ ++d+1
Proof. Recall the loss expression from Lemma B.2:
L(W) = Brr _ nI, WI, - WT 22 n(n + 1)2, W W',,W, WI, + wtr(W = WE, Br wo'nttn'',,,,... WE,) + o
We will simplify the expression under the assumptions Σ. Ep and W ntd=1+00 1, where the optimal pre-trained W follows from Theorem of Li et al. (2024). The first expression of interest is:
Σx nIx WZ,
n Eg n(n + 1)E, W.W ntr(W , WI,)
Substituting Σ, and W yields:
nn + 1) n n+d+l+00 n+d+l+0 (n+d+1+ooy) 1 (n+d+1+o7)
nd
( (d41=oo-m0--G--+++)))mm+o---
(++d+1+02)
d ++d++++07
In addition, the noise term is:
2ntr EE,r, Φ² nd (n d ++
Combining these, the final closed-form loss
L(W) ||BTT|1 (d + 1 o + o2
nd
0² (n+ 1 +0²
d 1++++02p
Finally, setting 2² = in this expression yields the desired result.
Theorem 4.2. Let n/d @(1). Recall the definition of the expected loss Lrr(Wrr) with respect SIT in (4). In the isotropic covariance and noiseless setting 2o2 0), the optimal step-size that minimizes (rr(Wr) is
d
1"" 2(k d)
With this optimal step-size 1ī", the improvement in the loss due to the test-lime training is
k d
L(W" (TT(WT)
k (n + d)3 BBT1
16

---

Test-Time Training Provably Improves Transformers In-context Learners
Proof. We consider the noiseless (o² and isotropic case where Σ, L = Throughout the proof, omit lower order terms like +1 in (d or 2nz in (n²² + 2n) as they're negligible compared to the higher order of the same variable (see Remark 4.3 for error discussion). This convention allows us to simplify the subsequent expressions.
Now, we aim to calculate the expected gain in the loss with respect to (n + k) samples used during the test-time training process. That is, we will take the expectation of the loss function given by Lemma B.2 and compute L(W") -rr(Wrrr Ex,an Xonen [A[] where Af given by:
AL L(W"- - L(WTr)
= BrT" _ [xx nIx W 2x WT Ex n(n W W www.wn E BIT
BTT" [2x nin WTT Eiz WIT" En + n(n 1)2, WT I. W WTT + TE WTT BTT
=Brr" _ [I'x (WIT- W Ex n Es (WIT" w Σ. + n(n ا W Σ. W W Σ. WTT I
+n (tr(W"T mg W Σ tr(W I, WTT Σx) Ex BTT.
(10)
In the noiseless and Σ, h- = case, the optimal pre-trained W becomes = . / as shown in Theorem 1 of Li et al. (2024). Recall that when there's 's no noise and the linear task is Brr, the gradient update given by Proposition 3.1 becomes:
W TT M 217 train "X taai (BTT W X context X context BTT) BIT coonexxt "X
context-
A key technical challenge while using the above update is that the expectation of AL in (10) involves 8th moments of X Coontext To overcome this problem, we will use the Gaussian approximation following Lemma B. by approximating X ooontegs Xxontext n Brr+ vn In where out 1// and reducing the problem to the isotropic case from Lemma B. Also define AW := Wir W" Then, using the approximation and the fact that Σ2 = I, we can write:
Brrll D). This non-isotropic form can be obtained by factoring
n AW= X T "Mrin BIT
n n+d+1 ExB
Vn
111 8 (MBIT' Ex + vng vng x1!22
n+d+1 n
n² n n+
ExBITBIT E n vn
Xuin Xurai
nBIT BIT '2 + Vn 2012
n+ d Σ, Brig 2!/2
1112 V
-
BBT Exs n n+ +++++1
1122 -x gg 1/2
n+
n(d n(d+1)
Bri BrrBraT+ T Vn(d 1)
v n d +
X
"irain
n d Brig"
g
n d + gg
11 d
AWT = =2n n BTT BTT
níd 1)
n 8 8
vn(d+1) 1)
n ++d+1 Brig
V
n
gg
Xuain
Xunin-
n ++d++
Notice that plugging Σ. I in Equation (10), we observe the following difference term in two of the expressions:
W W WIT WTT T Σ, W (W*T AWT) E, (W* + AW)
W AW WW W AWWAAW AW.
Hence, we need to calculate both first-order expectations Ex,.. [AW], Bx.. [AWT] and second-order expectation Excing [AWTAW] By direct calculation, we obtain:
n n BIT Bre n + 1BTTI 12 Exmingg /AWT].
n(d
E Exmin.g [AW] = -2nk
(11)
For the second-order expectation Exm.. [Xrain X X Xrain XX Xxmann. we will benefit from the fact that Ex (XX x)] k(k + 1)Σ² + Atr(E)E where X E Rlxd and its rows are drawn i.i.d from N(O, E). This result follows by plugging M = into Lemma B.3. our setting, Σ, which gives
Ex,.n [ X X Xarai X klk d 1)I.
17

---

Test-Time Training Provably Improves Transformers In-context Learners
Using the above result, we obtain:
Ex
AW AW|
41f B
n(d + 1)
BTT Brr" Vn(d+1
n ++d++ 8
n nvn + d n n
Brig"
n n d + 88' kik + d +
+ d+
n(d + 1)
Vn(d+1) +
n+ d BTT BTTB
n+d+1 BITg"
n + 8BT Vn(dd + 1)
n 88
Ind++1
Vn
41fk((k d 1) 1)
+n+d+1 BrT'
8BT
+ BTT
n d 88
n d
n(d + 1) n1 BIT Brr+
Vildd +
vn
n n+d+1 Brrg"
gBr
n n+ gg
11 n
4nf k(k + d + 1)m
Bg (Vild 1)BrrBri'
+ (d 1)ggrr nBBrgg Vngg' ) (Vn(d+ 1)BrrBrr
d 1)²
+(d 1) BTT8"
- ng BrT" T vnggg)
k(k + d 1)m
41/
Bg [n(d 1ITIBBTII BTT - n(d + B) BIT J gg" + (d 1)'11r111gg"
(n d 1)2
+n Br
n(d= 1)gg' BrI n(d 1) Brr Bri gg" - n(d 1)gg" Bri ngg" gg
41f klk + d 1)n
[n(d 1B BIT" nld 1) BIT + (d 11T1
(n d 1)2
+n +n - -
3n(d 1)IBBTII BTT BTT T nid 1)
= 4/ k(k + d 1)m ((d+ I)(d 3) nd) 1BrTIB Bri" ((d 1) n(d + 2)) |Brll
(n d 1)?
Applying the first-order expectation results from (11), obtain that the first terms of the loss difference are:
Ex BIT nx, (WTT W Σ, BIT] = 21pkn
n(d + 1)
[Brrll
+ d |BTTl12]
n
n d +
nd
2nkn n + nd
+ |BBr 112
(12)
E X BIT n Ez (WIT" W Σx BIT]
2nkn
1B l1.
(13)
n +
The other two terms in Equation (10) are:
Ex [ BIT Tn(n 1) Ex
WtT Eg W
WIT E. WIT) I. BrT]
n(n 1)BITT Ex
W AW AWT AW AW BIT
=-4nkn(n + 1) (n d 1)2
nd
IBBTTI12 477
k(k + d 1 1)n² 1) (nr + 1)2
(n(d 2d 3
nd) |BBTII (Id + 1)2 n(d 2))
(14)
and
Ex BET In (tr(W*T E, W Σ2) tr(Wrr Σ, WTT Σx)) BIT
= n n1BBTll? Exming [tr(- WT AW Ww AW AW)
--4n² 2kkk +d++ D)n²
n²
n (d 2d nd) |B d d(dd + 1)2 n(d + 2)) [|BTTl|]]
4nk
B
(15)
(n d 1)2
(n + + 1)2
Combining (12), (13), (14), (15) in Equation (10), the expected improvement in the loss is the following:
r²d n + ||BrTl13 4nk (n d + 1) |BBTlI 4nk n² + 1)d
n²
L(W) - Lrr(Wrr) = 4nj
(n+d+ 1) |BTTll
407 k(k + 1)n²
[nnn + 2) 2d - nd) (n + d 1) ()( + 1)2 + n(d + 2))
(n + d + 1)²
18

---

Test-Time Training Provably Improves Transformers In-context Learners
Rearranging and approximating by omitting lower-order terms, obtain:
4nk (n + d + 1)2 BBTll 41f" k(k + d + ())nnd(n d)(n² + d)
df
|BTTl11.
(16)
(n + 1)2
Taking the derivative of this expression and setting to yields that the optimal 17" solution is approximately:
d
1j"
(17)
2(k + d 1)(n² + d) + d) 1BBTII
Plugging the optimal 'f'' value in (17) into Equation (16), we obtain that the loss improvement is approximately:
2kd
kd
AL (n+ d 1)2(k d 1)(n d)(1 $1)
1BITh
|
(n d + 1)2(k d 1) d)(1 21
k
d
1BTT11
k d + 1 (n + d 1)P(n d)
Considering that n/d = @(1,, we conclude that
d AL ktd d (n + dys BBill..
Corollary 4.4. Recall the definitions ofy k/d, n/d, and consider the loss (Tr a function of a in the isotropic covariance and noiseless setting. If 2. then the loss Lrr (Wrz) non-monotonic in Specifically, the loss Lrr Orr(Wrr is increasing for a y+1 1 and decreasing for a
y+ 1. Conversely, if < 2. then Lrt (Wrr) is monotonic in a.
3y
Proof. Using Theorem 4.2, the new loss (WT(WTT) is approximately:
Wrr(Wir) = 1Brrll2 n+d k n+ 1)
k
3)
18TTIB
k d
k
where a = + E (0, 1).
Since l/Brrll is fixed, we focus on the second multiplier and define f(a) = a gives:
& ked a². Computing the derivative of f(a)
k k
f(a) = 1
d²
Solving f'(a) = 0 yields amn turning point in (0, 1), so f increases on (0, a and decreases on (aaits 1), making non monoooooe. On the other hand, in the & regime, one finds
ktd
This lies in (0, 1) when
Atd
1, i.e. k <> 4. In that case, f(a) has exactly one
k+d
1
derit & (0, 1),
3k
hence (a) cannot change sign over (0, 1). Therefore, f(a) is monotonic ntd
d which implics that Lrr is also
monotonic in the ratio n/d. Finally, since --T. the threshold until which Lrr increases translates to in terms of a. Consequently, for y > 1/2, the new loss is non-monotonic with respect to a n/d, whereas for /2, remains monotonie. This completes the argument.
3y
y+
Theorem 4.5. Consider the isotropic covariance and noiseless setting (o2 0). Suppose the initial weight matrix W Oaxa- Then, the optimal step-size that minimizes Lrr(W mT)
n¹" = 2(k + d (n 4n 3 d) 1Bm
19

---

Test-Time Training Provably Improves Transformers In-context Learners
With this optimal step-size Hf, the improvement in the loss due to lest-fime training is
n²
L(W" - (rr k+ ++d+1 n² 4n + 3 + 1Brllh
where L(W") = 111r1112
Proof. In the proof, will handle the general noisy scenario, and setting σ² in the final expression will recover the result for the noiseless setting. Similar to the proof of Theorem 4.2, we aim to calculate the expected gain in the loss with respect to (n k) samples used during the test-time training process. During test time, we draw (n + k) total samples ((E,, F1) drawn at test time, each with an associated noise variable (=)) Let's denote Sirain = ( ... &x and
Scontext = Em-1 EnoR We will then evaluate the expectation of the loss function given by Lemma B.2 and compute L(W") LIT(WIT) = Ex M Ronten [A_] where AL is given by:
AL L(W") L(WTT)
= BBr [Ze - nI W. nE,W Ds n(n 1) W",, ntr I, E. Brr
T Ex)
-BIT [Z, nL,,, WTT Ei WT, 2n + n(n WI, + Wrr Di BTT
-o'ntr(WTi' E,WTTE
=Brr" =BIT'InE, (WTI W Σg n Ex (Wri" T W Ex n(n 1)L, (''T Σg W WIT Σ2 WIT) Ex +1 (tr(W*T Σ, W Σ2) tr Σg Wr Σx)) xx Bir onn ttr(W" Σ, W Σ,) tr LrrWrr (.
(18)
Recall that when the linear task is BTT and W = Oaxa. the gradient update given by Proposition 3.1 becomes:
WTT W +2
Xumin TXrrain (BIT W coontext Jo "ycontex) Xtrain Strain (X Ycontext)
W +2
Xuain "Kuain (Brr W Xcontext T(amtt"iaeee "(XcomoxBIT Econtexs) Xuaii 'Suain (Xeontex "Xeoonext BrT Econtex))
W +1 Xunin X BUT Xurain "Sminn (X (Xcontext (X BTT Écontext))
Ximin Xunin Brr + Xirain "iraii ||rr - Xcontexi X goontext + Econtexi Xcontext
2n
Then, obtain its transpose as
WrT 2n/ Xoontexi Xxontext Brr + Xeontest Scontext [BTT" T X Xurain Frrin T Xumain]
Thus, the expression WTT WTT becomes:
WIT WTT = 2 Xcontext Xoontexs Brr + Xeontexi "ooomee BIr Xerain X Arain + Etaii "
Xrain Xurin BIT + X Euraia [Bri" Xcontexa "eoonext + Econtex Xeontoxe
Therefore, we can write:
Exuas Nontu Scotess
WTT W 4n Ex Nesaee Xcontexi Econtest Emin X Xurain Stminn "X Xcontext
Xeontext T Scontext BrT Xurin X Xuran" T X Xurain BTT Econtext TX "xeetextt
Xeontext X BrT Evain Xurain "vaain BuT Xcontexs "xentext
+ X contexi X BrT BiT" T Xmmin " Xurai X Brr Brr X Xeuntexi X
comtext
By plugging M / Lemma B.3 and considering that Σx I, we have the following fact:
Ex
X
X Xurai X = klk d 1)7.
20

---

Test-Time Training Provably Improves Transformers In-context Learners
Applying this identity twice and using the fact that Xxmi [Xain Xuain "] =r(2,)) dl, obtain:
Bxum Eum Xoane Reoeess
WWr''ni
= 4nf (o kndll + o²k(k + d DellBTTll + o² (nddn(n 1)BrrBBi + +1BBTTl(!!
+kkk + d+ 1)1lrrll (n(n 1)ßrr nllBrill?)))
= (o2kn (o'd (k 111 (n(n B) BTT' =nIBmll ( + D(
= kn (o'd (kc + d + (1())TI)) ( (n + B)Brr + 11Brrll2).
(19)
Besides that, we calculate the first order expectations of Wrr as:
Ex Nonenn (BIT" T Mx (WIT W") E. BrT 2nnn Brr Nouuc Xtran Xurin BrrBBr cooteex X
BrT
2nnk2'1B1TT15?
(20)
and similarly,
Xoarn Ex Ee Xccuna Eecerat
BI (WIT" WhT) I. ErBrT 21 BIT nIx
(21)
The next term in the loss difference is:
Xusis finis Xcunta Sonnros
BIT "Tnnnn + 1) (WrT T Σ W - WIT I. Wrr) I, BrT]
-1n -nn=knn kn (n 1)((n 201lBTll + oه) o3 (o'd (k + d D1BBIl1b)
(22)
Using (19), the last term is:
Exran-Sminn. Si Xomer fornra Brr 'n (tr(W"T Σ, w' tr(Wrr' E. Wrr Σx,) Ex BiT)
= ---1fBBr tr(Bx Smin Namesskknns Eamess
Wrr W
kn (o²d (k d 1)1BITI1))
(o²d (1 (n
1)1TTT1)
(23)
Finally, the noise term is equal to:
Ex Emin X Xannes Ecorrss
0 - -a00-tr(WW222 (TW
=n kn o (o' (k 1) (o'd (n (n d D)
(24)
Combining the Equations (20), (21), (22), (23), (24) the expectation the loss difference (18), we obtain the following quadratic expression of ry.
L(W") (rr(W
-4nkn 11 4 kn² 1) (1( 2)1(BTT + o') (o + (k + 1)(
- kn²p (o2d + (k + 1)) ( 2d (n d 1) 11 [BTTT
-
+ (2
1)1
-4kn² - [ ([(r) (
(o'd
+(n + 1)( 2T + 0²) |
4kn" 4 - 7/ (o'd + (k d D) (oAd [Brlll (n22 + 4n + +011Brrll + 20 + d 1))]
Setting the derivative to and solving for gives:
1BTTl
n"' 2(02d (k +k (k+++)))B+1))) (o-d d+ + |BTTll (n² 4n + 3 +d) 20 d 1)(())))
21

---

Test-Time Training Provably Improves Transformers In-context Learners
At this optimal value nf", the improvement the loss is:
k &11rT11
11² 1/Brl11
[BTTII
o + (k+ d 1)1BrTlB (odd [/Brlll (n² + 4n + d) + 20 + d 1)1BI1l?
In the noiseless setting, the optimal 17" corresponds to:
n¹
2(k d+ 1)(n² + 4n + d)1BBTTI
At that optimal value 17²,, the improvement is:
k
k d 1 n² 4n + + d 1BTTI
This completes the proof. As final note, recall that the initial loss is |Brll Thus, the new loss of the system after the rest-time-training update is:
(1 + d n2 4n + 3 d
BBrll
In the proportional n, d regime and when k/d - this gives an error of order O(n -mmathima that of the optimal Wop from Lemma C.I:
4n d
+ 4n + d |Brll
Corollary 4.6. Recall the definitions of n/d and == k/d. Consider the setting in Theorem 4.2 and test-time training described in Proposition 3.1 with both pre-trained and zero initializations. Then, under the isotropic covariance and noiseless setting, there exists threshold y" given by y* pre-trained initialization over the zero initialization Oaxd-
(a + 1)²
such that < y if and only if it is better to utilize the
(a 2)
Proof. Considering that the initial loss of the weight matrix W Odsa is [|Brrll and the improvement given by Theorem 4.5, the new loss after the test-time-tte -imi--tiiiinggggggg update is approximately:
(1 (i-x
k + [/BrilB-
On the other hand, recall that by Proposition 4.1, the loss of W" before the update is d+1 improvement given by Theorem 4.2, the new loss after the test-time-trainnng update is approximately:
μ*æ+ [BrTl] Combining this with the
n (14d
d k
d²
BBTT
k (n + d)3
Let us define 8 : Etd over zero initialization with the below inequality:
A and 0 == m+d Then, we check when it's better (yields smaller loss) to use the pre-trained matrix W*
4
1 1 1-B= > 0 0-80"
=
1
0 (-00881-00
> 0 +0+1
B
d
1 +0
k
k
d
d
ntd
(n d)²
(ar + 1)2
k
> y k/d where a
+ 2d
a+2
Thus, we conclude our argument.
22

---

Test-Time Training Provably Improves Transformers In-context Learners
C. Proofs for Section 5
Lemma C.1 (Optimal W). When the system is noiseless (o² = 0), the optimal which minimizes the population risk with respect the new task Brr given by Lemma B.2 and its corresponding loss are:
ByT (n 2)118TT112
2
Wap
L(W
n 1Br1122
Proof. Recall that when Σ. = and 0² = 0, the loss formula given by Lemma B.2 becomes:
L(W) BTT [ nW nWT n(n 1)W W #nt(WW W) BIT.
(25)
The gradient of this expression with respect to W is:
Vw-[(W) 2n Brr BrT 2n(n + 1)W Brr BrT" 2nBrr BrrW 0.
Simplifying, this gives:
(n 1)WBrrBri BrrW -BrrBri"
(26)
Multiply (26) on the right by Brr. This yields:
WB + BTT.
So ßTT is eigenvector W corresponding the eigenvalue 1/(n 2). Next, consider any vector that orthonormal (or simply orthogonal) to that is v BIT O. Multiplying (26) the right by V:
Brrll Wr =
Wv=0 whenever BrBrT- O.
Hence, all other cigenvalues of W are O. Using the cigendecomposition of W, this uniquely specifies W Brr Bri Solving for yields the unique solution to Vw L(W) = 0:
BIi BIT
Wrot
(n 2) BBTT|
Finally, plugging this Wopt to (25) yields the following error:
L(WPL) = + |Brrl12
Lemma C.2 (Eigenvalues). Consider the optimal pre-trained W matrix, which minimizes the population loss over all possible lasks sampled from Eg in (1), Then, all eigenvalues of W = /////'/ lie in the interval [0, =
Proof. By Theorem of ei al. (2024), we write the following:
W - >/>=:/"'>; where W = ((1 + 1)I, + MA-I)= where M tr(A) + o² and = =x!// Σ 11/2
We know that the matrices Σ, and Σe are jointly diagonalizable with Σ, QA.QT and Σp = 2A where Q is an orthonormal matrix Q Raxdd and Ax. E Raddd are diagonal matrices with entries (2:) and (Bi). respectively. Then,
A = A (A Ap A 1/2 @T)
e
Define A4 1/2 Ag 1/2 which also diagonal. Hence
A- = == Q
23

---

Test-Time Training Provably Improves Transformers In-context Learners
where A diag (1/(1B))). Inside W", have
I M A 2 (n I)IQ" + M (2) =2 [ 1 M o². We know that Adiug
(n+ + + M A is also diagonal with diagonal entries (n
M
It follows that
Ap
A,B:
W = - + - - diag
e¹.
(n + 1)2, Bi
M
where M = tr(A) + o² =tr(AA/))= = o2 E AB, This concludes the proof.
Lemma C.3 (Covariance Shift). Consider the setting in Section where
The true labels are generated by y= = E with 2(0.2.). B~ N(O, Ep), and noise & N(O,
The prediction of linear attention model is of the form
$ x
where the context X in Ruxd collects previous samples x; as rows, y is the corresponding label vector, and e Rexd the model parameter matrix.
Then, the invertible transformation
(x, (x, (R. 8.B. where x
1-1/xx, 8 X B E B. W W WE
preserves the population risks in (1) and (3). More precisely, defining Σ(W) to be the population loss (3) evaluated under the transformed setting (x, X.B. W), we have
L(W) L(W).
Proof. Since N(O, Σx), multiplying by 2/ yields F == NO, 44). Likewise, because B NOO, have ==
E|/ B ~ N(O. 1l/ mp ) Next, we observe
x X
XB (XΣ, 1/2 ) (I!?2 B) XB.
ensuring that label vector corresponding to context data XB x XB & stays the same. Also, note that the label is preserved for the query token
x & F B &, as well. Under the map x F E- X H X = XΣ, 112 w the prediction of the linear attention model is also preserved: =-
ß B
1/2 B. W
- = 1/2
WX (E- x2 x 2 // (x-") T (Σ (w!/2
(X = w y.
As result, the emror in the labels remain numerically unchanged as
- x WX - y)
Hence, this implies that the population losses described in (1), (3) are preserved.
L(W) Z(W).
In particular, if W is the unique minimizer of the pre-training loss in (1) in the original coordinates, its counterpart in the transformed system precisely
W 1/2 W w
This completes the proof.
24

---

Test-Time Training Provably Improves Transformers In-context Learners
Theorem C.A. Let L/d O(1) and o² O. Suppose the covariance matrices are jointly diagonalizable by an orthogonal matrix Q Rand that Mx QA and Ep QA Let W be any matrix that's 's also jointly diagonalizable by Q such that if 2 W = QAQ". then all eigenvalues of A lie in the interval [0, r. Define B := 2 _ BYr. A := BTT (1 nA BrT, and B == x Then, the optimal step size that minimizes the population loss given in (3) after the test-time training update is
tr(A²).
A BBTll
n
2(k d) n²
(A+ B)
With this optimal step-size rf. the improvement in the loss due to test-time training and the initial loss are approximately
L(W) - Lm k A
W(W) ~ A
k+dA+B
Proof. We consider the general non-isotropic covariance scenario where Σn and Lg may be arbitrary covariance matrices. Our analysis will hold for any initial weight matrix W in Rabd such that if 1// W x// = QAQ", all eigenvalues of A are in [0, [+I] Now, we aim to calculate the expected gain in the loss with respect to (n + k) samples used during the test-time training process. That is, we will take the expectation of the loss function given by Lemma B.2 and compute L(W) (TT( = Ex,,. Ex [A]] where A& is given by:
AΣ L(W)- L(WTT)
= BTT'S n Ix (W - W) Ex + I2 (WITT WT) Ex + n(n + 1) (WT Σg W
WTT T E. WIT) Ln
+n (tr(WT Σ, WI,) - tr(Wrr" Σ, WIT 2x)) Zr| BrT.
(27)
Now, recall the test-time training update from Proposition 3.1 when the new task is Brr:
WIT W = 17 Xtrain "Xirain (I W comteex T X Bri BrT Xcontext Xonntexts
WIT - WT 2oneeot Xeooexx Bri BIT T (I -ooontexx Xoootext WT) X Xuain-
As in previous theorems, we will use the Gaussian approximation following Lemma B.1 by approximating Xoontexi XconeeaBr Byr nEx vn v g where g N(O, (1Brr?11 This non-isotropic form can be obtained by factoring out _ and reducing the problem to the isotropic case from Lemma B.1. Before going forward, let's apply the covariance shift discussed in Lemma C.3, which transforms:
F x->10 La). BTT Bf N(O. 1.0.E122,p x1101 and W =2 21/W2//2
Likewise, we define the transformed versions of training and context data as Xmain Xmain Σ, 1/2 and Xoontett Xcontext Mx 1/2 Now, suppose that W- QAQ and let AW := Wrr W E!!
Wrr 11/2
=!/2WE!! ww!! Then,
AW = 2n train & (BI: Q^Q 8T
context &ontes Bri) BTr 8T context &
z 2n train &urain (BIT nQaQ B vieAe 8) (nBTT Vngg)
= 2n main &train (MB BTT VAßTIgT - -ZQAQ" B BIT n VT TRQAQ BIT - N - =QAQT gg")
2n Irain &train
(M(1-nQAQ) BIT + (0II--QA"")
- -nveAA g - -QAQ gg
AWT = 2z (n BTr (1-nQA"))) + Vig (Q==Q V g TQqQ - -ngg (TQT))QT) 8 &
By plugging the above definitions of WrT, W into Equation (27), we encounter the following difference term in two of the expressions:
W W WT W WT = W ( +AW )W + AW)
-WWAWW- AWW-WWTH W AWTAW AW.
(28)
25

---

Test-Time Training Provably Improves Transformers In-context Learners
Hence, we need to calculate both the first-order expectations Ex.... [AW], Em...g [AWT] and the second-order expectation
Exmang [AWTAW] Calculating the first-order expectations gives:
Ex [AW] 2nykn (1-- -0n=nAA)) BIT 11Brill QAQ"). AWT 21/k (BIT BiT (1- (l--AAA)) -18TT11 QAQ").
(29)
Ex &uiio
(30)
Using (29) and (30), the first two terms of the loss difference are:
Xmmas PTT n (WIT- W) BTT] = 2nkk (nilBBrli (I mQAQ") Bm mImmm ( BrT). & BT n WTT - WT) BTT] 2rykn (nIlBBTII2
(31)
mQAQ") B - mI|BTTIB QNQ BIT)
(32)
Similar to proofs of previous theorems, by plugging M / in Lemma B.3, we know that Bx ((XTX)(XX ktr(E)E where X E Rexd has its rows drawn i.dd from N(O, Σ). Therefore,
X)(X X)] = k(k 1)
Ex Ex
8 Irain &trai & |mmin] k(k + d + 1)I.
train
Utilizing the above fact, we compute:
xx (AWTAN AW)
= =11² B& BIBB (1- (I== + VigBT (1- (1- - BrT 8 QAQ - ngg QAQT)
k(k + d + 1) (n-I BII BII
+ Vn(I-nQAQ ) Brr g"
n Vn VUAQ gBTr nQAQ" gg
= k + d + 1)B,, 1' (1- 11-nAAQ' (1- BIT - 12 BrT BTr [1 nQ^Q ( )QAQ gg
BT n gg QAQ QAQ gg' + + +ggBT (1 - -1--QQQ@1 ) 10---QAT ( g" -n'gg Q^Q (1 - 10-- BrI
+n'Bm g' Q^Q QAQ- gBT n² BT (1 (I- m@AQ Q^Q" gBT n BIT QAQ (I (
) Brig' g-1
Let's ' inspect each term inside the expectation in the above sum:
Bg BinTTTeAQQ' QAQT gBTT| = B' Bg g QA 8 B|BTT
=n'Br B.
J A2 BTr where z N (0. 11BrrBl 1)
= .BBBBBTT tr(A2) BIT PTr
Eg In' gB (I --n====x)))) g
n² B2 (BT 11 --A===ne="))n eare 8) BTr
B [ ((8 (1-
BIT) BTr
n'/BTTl QxQ (1 (1-nQAQ") (IQQAQ")BTTL BiT
Bg [n2Brin QTATTT (
n²Brr By [( QAQ (1-nQAQ)) ) g*|
""Brr Bg [PTr (1- (l-nQAQ)) (0QA")(QQ') 8)&² QAQ = BIT BTT 11-nnAA)) QAQ )
Bg IngBin (1- (1-=QAQ)((1-n0))) (1-10 B
=nBBg & (BTr (I 11-n010 ---AAA")) (1 -nQAQ"))) Bri) g"
= =nB (1 - 1--n2ee) (1-QAQ BTT BR [g8"
BIT
= n =11Brrll (BTT (1 11---AA') (1- BrT)
Ex (2288 [nrggs@QQ' (TTTTTT00O QAQ" gg
= n B gg eA-0" 88"
= =1IIBTII (r(())) 20)
Ex [n=gg" QAQ (1 nQne") BTT
= ~311BT11 (I (I- (I- ( Bir
By | BTv (1- (1-=QAQT)@N exQ' gg"
= = PTT BT (1- (ll=nAAQ))0m) eNe".
26

---

Test-Time Training Provably Improves Transformers In-context Learners
Hence, considering that the eigenvalues of A are smaller than _. above:
we drop lower order terms after combining the results n+1
AW AW|
4nk + d
BTB ~ ~ BirBir - Q 2n 1nrllB 2 (A
+n =111T111 (1- Q 1 (tr() 220=)+2222)))))
(
= knk + d trIBBTTI BBTTTTT B nA) Q BrT
+n |B (B Qu B
Thus, we reach the following result:
8x PHA AWBrT k(k d 1) tr(A3) + (n² +)) 1 -
(33)
Now, recall from Equation (28) that the second-order difference is in the form:
W W- WT WTT WTAW AW W AWTAN AW.
Therefore, also need calculate the expectations of the terms WT AW, AWT W. Using previous results (29) and (30), get:
8
WT AWW
112k (1 n2AQ" BTT B n11B 1 eAQe)
22nkm (RiA nA 1 BBTll e*²eT)
-AW W|
2nkn (BIT Bi-
nA - 18TTlh e*
x &
This gives us:
E Ex
BTT WT AWBTT = 2nphr (TBTTIBBT1Q(A - BTT - BTT)
(34)
E Ex..ng
AW WBIr = 2npkn (IBrTI1& (BBrIIBBBT. BT Q(A n 10 (TT)).
(35)
Exploiting (33), (34) and (35), we get
Ex 8 Bir n(n + 1) (WT W WT Wrr Brr]
kn + D)m(n 1) Br'llB tr(A?) (n² + m) 18 nA) QTT]
4n (1BrTI1 nA - BIT)
(36)
Expa Bir n (tr(W" W) - tr WrT)) BrT
k + + 1) n Br tr(A3 2) + n + d) [Br + nA) e BrT]
4nkn? 2T nd B.r [Brrll|3 tr(A3)]
(37)
Combining (31), (32), (36) and (37), and plugging into the Equation (27) gives the expected loss improvement:
L(W) Lrr(WTT) 4nkn² (1BrrIl? 2 BTT QI - nA) Q P.. 18 QAQ BrT)
(38)
4n + 1) 11BTT BTT @(A nA 10 n2 Brr 1BTT BT Q
4mkn2 |IBTTI| 1 BT Q(A n 1m 1-B--11BBl
1BTT tr
-472 + 1)ml n (n +
+ 1 (11 1)2 + +n+ d) BBTTI QQ- e BrT]
= 4nkkn"1/@TT16 BT, Q - -nn-nAA22 2TT - kr2 + d 1)m² (n' 1BTTI tr(A3) + (n² + +1) (1 nA)
QTBml]
(39)
27

---

Test-Time Training Provably Improves Transformers In-context Learners
This is quadratic expression in 77, and we can solve for the optimal rf' by setting the derivative This way, the optimal step size is approximately:
F BTT 2(k + 1)1BB / tr(A)) + (n² +
of
QTB
Plugging this optimal 1y' value into Equation (39), the population loss gain becomes:
( 11 Q QTBIT)
k k d |BrrIB tr(AA) + (1+ Qu Q
Thus, considering the definitions ß.T BTB and A = p.r (1 (1- ~->= BTT. B = ntr(A²) 1BTTT (notice that &2 norm is unitarily-invarian so that [βrrlll ||Brrll) and recalling n/d = 0(1) recovers the desired final expression. Also, by viewing the initial task as Brr, we match the diagonal covariance form described in Section 5. Therefore, we conclude our argument
A2
L(W) (
(Tr(WT)) = k+ ++dA+ B
For the second part of the proposition, let's write the initial loss by applying the covariance shift from Lemma C.3 so that the transformed features = -, === x have identity covariance and = x W' is the corresponding minimizer the transformed system. Again, we set Brr = 2 Brr in order to have diagonal covariances. Assuming or and plugging these in the population loss formula from Lemma B.2 with respect to the task Brr gives:
L(W) = PT r n1 W - nJWT + n(n 1)/W /W ntr(W IW1)] BrT + n'rtrWW IW1) +
1BrilK 0T Brr n(n 1) 1I e Br nt +ntr(A3)lBBIlll)
= IBrrl13 - 2n BrT + QA T B + n1r(13)1B11113
= (1-2(I-nA)) e' ß.g r
=BT (I-ni) (-= BIT (BTT
=A+ A+B B.
Hence, L(W) A B. The proof is complete.
Proposition 5.2. Let Σx 12 0 and suppose Lp is arbitrary diagonal covariance matrix with the first diagonal entries being non-zero. Then, the minimizer W of the pre-training loss (1) that has minimal Frobenius norm and its population loss are given by
W = (In ((n+1)I 1)f' + tx(E))E)) L(W) = A
where I' = diag(1, O2-r).
Proof. While Li et al. (2024) derives the solution form in the full-rank case, the same closed-fome expression can be adapted when Σp is low-rank. Indeed, if Σp has zeros in its last d - r diagonal entries, those coordinates do not appear in the pre-training loss. Consequently, there is set of minimizers W which differ only in those degenerate directions. Among these, the Sinimal-Frooenius-ee subspace would increase [|W]] Consequently,
criterion forces all degenerate coordinates to be zero as any nonzero component that
W ((1+1) + tr(Ex)
zeroes out precisely the degenerate directions and is the unique Probenius-minimal solution.
For the second part of the proposition, the argument directly follows from the proof of Theorem 5.3.
Corollary 5.4. Consider the setting in Theorem 5.3, Recall the definitions of = n//d and y = k/d and define ci
|1Brzll""
B
(C1 c2) (C1 (2)
C2
1BTl Then, there exists a phase transition point y*
such that y < <x" if and only if
C2 (20c + C2)
better to utilize the pre-trained W over the null initialization Oxxa-
28

---

Test-Time Training Provably Improves Transformers In-context Learners
Proof. Recall the population loss for W given by Lemma B.2:
L(W) = =Brr" Σx n2 x WT Σ. + n(n + W),, + ntr(w W2,,) + o
WI,,, + o².
First, plugging Oax yields the initial loss of 1/Brl Also, setting W Oxxdd in Theorem 5.3, we know that the improvement by eesttiim-training approximately:
k
1BTTII2
k d
At the same time, still by Theorem 5.3, we know that if our initial weight matrix is the pre-trained W the improvement is approximately:
k A²
k d B
Whereas, the corresponding initial loss is approximately L(W") = A B. Therefore, we check when it's better to use pre-trained matrix W over 0axd (null) matrix with the below inequality, which compares the losses after the test-time- sest-time-training
update:
A2 A B
k
< 1BBTTI2
k + ++dlBmrli
k+dA+R B
In lines with the proof of Corollary 4.6, denote B ked
Then, a series of algebraic manipulations give:
A
c
A + B
11pTTl13 BIBBTllB = c
<1
A B
C +2
oj
B
< +
CI +2
CI a
(C) + +
(C) + (2))
d
C2 (20 + 62)
>
k
C +2 - (Ci + +22
k
C + (9 +2)
=
C (2c) (2)
This completes our argument.
D. Further Experimental Results for Section 6
To illustrate the improvement more clearly, we also provide version of Figure with the x-axis on log scale, shown in Figure 4.
29

---

Test-Time Training Provably Improves Transformers In-context Learners
0.7
0.6
Accuracy
TabPFN
TabPFN + TTT
0.2
10a
10a
10a
10a
(Context length)
Figure 4. Accuracy of TabPEN model with and without test-timettraiiing as a function of number of in-context samples n with the x-axis log scale.
30


Reviews:
[{"review_id": "vIrOTwPHUC", "reviewer": "Reviewer_yFX9", "summary": "In-context learning and Test-time Training (TTT) are two ways of enhancing the predictive power of pretrained models on new tasks at test time. In-context learning involves incorporating demonstrations from the task into the prompt context. TTT involves light finetuning of the model on data related to the test task. Recent work has shown that TTT can naturally integrate with in-context learning: one can lightly finetune the model on provided demonstrations and subsequently use them as an in-context prompt with the finetuned model.\n\nThis paper develops a theoretical understanding of TTT to enhance the in-context learning ability of a model for a given task. Specifically, given a model trained for in-context learning of linear functions, the model is test-time trained on in-context prompts from a target linear function. The analysis focuses on single-layer linear attention models where a single gradient step is used for TTT. The paper characterizes the performance of the test-time trained model as a function of context length, the number of TTT samples, and the alignment between the pretrained model and the target task. High-level takeaways from the analysis include:\n\n- With TTT, one can reduce the number of in-context learning examples needed to perform well on a given task. Of course, this comes at the cost of additional training samples and extra time during the TTT phase.\n\n- Comparison of TTT from scratch vs. TTT on a model pretrained for in-context learning: (i) The relative advantage of TTT on top of a pretrained in-context learning model depends on how \"aligned\" the target task is with the tasks the model was trained on. (ii) As the number of TTT samples increases, the relative advantage of TTT on a pretrained model diminishes.\n\nSimulations demonstrate the theory, and some findings are also empirically validated on multi-layer non-linear Transformers.\n\nFinally, the paper takes a model pretrained for in-context learning of tabular data and shows that with TTT, one can substantially reduce the number of in-context examples needed to perform well on given tasks. This leads to efficient inference, especially given the quadratically growing cost of softmax attention.\n\n### update after rebuttal\n\nI thank the authors for their detailed response and answering my clarifying questions. My general view of strengths and weaknesses of the paper remains the same. Hence, I would stick to my original rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s positive assessment of our work and their detailed feedback. We now address each point below.\n\n> **Weakness.** While the theory presented concretely highlights the effect of various factors on TTT, most results do not seem particularly surprising. ... theory does not provide new insights that significantly change existing perspectives on TTT or in-context learning.\n\n**Response**: The main goal of our paper is to provide a rigorous theoretical framework that precisely quantifies the improvement by single-step gradient descent with TTT for in-context learning tasks. While it may seem intuitive that TTT would help or that transition points might emerge, to our knowledge, our work is the first formal analysis revealing why and how these phenomena occur. Achieving this required overcoming significant technical challenges, including computing high-order matrix moments using Gaussian approximations with a rigorous error analysis.\n\nParticularly, we characterize the improvement by TTT in terms of context length $(n)$, embedding dimension $(d)$, test-time training set size $(k)$, and the alignment between the pre-trained model and target task. This characterization allows us to show how the transition points depend on different regimes of test-time training set size $(k)$ and alignment between the model and target task, and we reveal the resulting non-monotonicities after TTT. Moreover, our theory and experiments show that one step of TTT is sufficient, which complements the existing practical research (Akyürek et al., 2024) stating that a few gradient steps offer significant test-time improvements. We hope that our results will serve as a foundation for future research on the interaction of TTT and ICL methods. We've provided a further discussion of insights in Q2.\n\n> Minor typo: The gradient seems to be missing from the gradient descent update equation in line 134.\n\n**Response:** Thanks; we will edit the revision accordingly.\n\n> **Q1.** What is the intuition behind non-monotonic trends in Figure 1a? \n\n**Response**: The non-monotonic behavior is observed as the result of two opposing effects, which are the initial loss and the improvement by TTT. \n\n* As $n$ grows against $d$ (i.e. $\\alpha$ increases), the pre-trained model does better initially and already has a lower loss before TTT, which makes it harder to be further reduced by TTT as the rank-1 update is unable to correct all directions. \n* On the other hand, when $d$ grows against $n$  ($\\alpha$ decreases), the initial loss is high, and there's more room (error) to be corrected by rank-1 update, and thus, the improvement by TTT is larger. This intuitively aligns with Theorem 4.2, which states that the TTT improvement scales as $(\\frac{d}{n+d})^3$. \n\nTogether, these two trends cause the non-monotonic behavior. Practically, this tells us that TTT through single-step gradient descent is more likely to be useful when $d = \\Omega (n)$. \n\n> **Q2.** Are there any insights from your theoretical analysis that you think would hold generally and add to existing perspectives on TTT or in-context learning? \n\n**Response:** We thank the reviewer for the valuable question. A key insight from our paper is that test-time training can serve as a lightweight, supervised complement to ICL. When TTT is combined with ICL, one might benefit from the examples at test time in a supervised fashion and do a quick parameter update by treating a subset of them as training data instead of simply using them all as context. Accordingly, we analyze the single-step gradient descent as it has minimal test-time compute overhead (which makes it close to realistic settings). This way, the TTT update can be done only in a single forward-backward pass with appropriate masking, as discussed in the Remark in Section 3. Through our experiments, we also observe that a single gradient step of TTT captures a large part of the benefit compared to multiple steps, and might be an appealing method for practical scenarios. This observation also holds in practice as Akyürek et al. (2024) find that only a small number of gradient steps on test-time training data is enough to achieve strong performance. We believe that our work might inspire future integrations of TTT with ICL.\n\n> **Q3.** The experiment with GPT-2 is not on top of a pretrained GPT-2 model but using a randomly initialized GPT-2 style architecture?\n\n**Response:** We thank the reviewer for bringing this detail. Between lines 428-431, we discuss how we obtain the pre-trained and scratch models. The scratch model is just initialized randomly with the GPT-2 architecture before the TTT process. For the pre-trained model, we obtain it by training it until convergence with the tasks sampled from the task covariance $\\boldsymbol{\\Sigma}_{\\boldsymbol{\\beta}}$, and before the training, it's also initialized *randomly* using the GPT-2 style architecture. We will add this missing detail and clarify the related paragraph further.", "claims_and_evidence": "Yes, the claims made in the submission are supported by clear evidence.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria make sense for the problem.", "theoretical_claims": "I did not verify the proofs but the claims made intuitive sense and were also corroborated with simulations.", "experimental_designs_or_analyses": "The experimental design and analysis seems sound.", "relation_to_broader_scientific_literature": "Discussed in summary and strengths and weaknesses sections.", "essential_references_not_discussed": "Prior work has been adequately discussed.", "other_strengths_and_weaknesses": "Strengths:\n\nThe intersection of TTT and in-context learning provides a valuable framework for understanding the interplay between in-context learning and in-weights learning. This paper advances this understanding by formally exploring the role of factors such as alignment between target and pretraining tasks, the number of TTT samples, and context length.\n\nI also liked the experiment with tabular data. Recently, Transformer models pretrained for in-context learning of tabular data have proven highly effective. However, inference with these models can be expensive for large tabular datasets due to the cost increasing qudratically with the sequence length. This work demonstrates that such costs can be significantly reduced using TTT, though at the expense of additional training time.\n\nOverall, I enjoyed reading the paper. It is clearly written, with simulations accompanying most theoretical results, making it easy to absorb the main claims.\n\nWeaknesses:\n\nWhile the theory presented concretely highlights the effect of various factors on TTT, most results do not seem particularly surprising. A possible criticism is that the theory does not provide new insights that significantly change existing perspectives on TTT or in-context learning.", "other_comments_or_suggestions": "Minor typo: The gradient seems to be missing from the gradient descent update equation in line 134.", "questions_for_authors": "Q1. What is the intuition behind non-monotonic trends in Figure 1a? This was probably the most non-intuitive result for me. It would be great if the authors could discuss the intuion and practical relevance of this result.\n\nQ2. Are there any insights from your theoretical analysis that you think would hold generally and add to existing perspectives on TTT or in-context learning? I know this question is a bit subjective but I would love to hear authors' thoughts on this.\n\nQ3. The experiment with GPT-2 is not on top of a pretrained GPT-2 model but using a randomly initialized GPT-2 style architecture? It would be good to clarify this to avoid confusion.", "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": "NBMS47WA9F", "reviewer": "Reviewer_KMYg", "summary": "This paper investigates how Test-Time Training (TTT) affects transformer models' in-context learning capabilities. The authors formulate this as a two-stage problem: a model is first trained on a pre-training dataset, then further trained on test data using single-step gradient descent. They provide theoretical proofs for performance gains through TTT, analyze the relationship between context length and test data size, and compare TTT with direct training on test data from scratch. Experimental results validate their theoretical findings.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed evaluation. We reply to each of the points below.\n\n> **Theoretical claims**. The mathematical proofs contain imprecise statements. ...\n\n**Response:** Thank you for pointing this out. The approximations in Theorems 4.2, 5.3, and Corollary 4.5 stem from ignoring lower order terms and using Gaussian Approximation in Lemma B.1, which rigorously bounds the error term. In the proportional regime ($n/d = \\Theta(1)$), all approximations yield an error on the order of $n^{-1}$, while the main loss terms are in the constant order. Consequently, any approximation error is small, which is evident in the close match between theoretical and empirical curves in Figures 1-2. We also wanted to kindly note that **after Theorems 4.2 and 5.3 in lines 190-192 and 290-291, we refer readers to Remark B.5 in Appendix B for discussion on the validity of the approximations.** We appreciate your feedback, and to clarify these points, we plan to move Remark B.5 into the main body.\n\n> **Methods And Evaluation Criteria**. Their approach fixes the context and uses k different mini-batch data ..., testing it using the same context as training.\n\n> **Weakness 1.** The paper's oversimplified TTT procedure. ... comparison between randomly initialized models and pretrained models lacks practical insights due to this oversimplified framework.\n\n**Response:** Thank you for your question. While our single-step TTT setup may not capture every real-world scenario, it's both theoretically interesting and practically informative. The main reason behind using the same context is to make gradient update tractable. For testing, however, we do not reuse the same context as in training but instead evaluate the model over *any* new context–query pairs by taking expectations (lines 140–142). Even under the current setup, there are major technical challenges, such as computing $\\mathbb{E}[\\mathbf{W}^\\top \\mathbf{W}]$, which requires the Gaussian Approximation technique and a rigorous error analysis. Without fixing the context in the training, it's practically infeasible to obtain a closed-form update (Proposition 3.1) and compute this expectation. Still, we believe that our framework presents two clear, theoretically verified practical insights:\n\n1) We characterize how alignment between the pre-trained model and target task affects TTT improvements and phase-transition points, which is illustrated in Figures 2a-2b and practically verified by GPT-2 experiments of Figure 2c.\n\n2) Both our theory and experiments indicate that one-step TTT can be sufficient, which complements the existing practical findings (Akyürek et al., 2024) that a few gradient steps yield significant test-time improvements.\n\nOverall, we believe our framework offers useful practical insights and can serve as a baseline theoretical reference for future research on TTT-ICL interactions.  \n\n> **Experimental Designs Or Analyses.** .... the real-world experiments in Section 6.2 have methodological issues. ...\n>\n> **Weakness 2.** The experiments in Section 6.2 need more detailed analysis and better clarification of the methodological choices.\n\n**Response:** Thank you for your feedback. While TabPFN+TTT indeed uses an extra $1000$ samples, our main goal is to demonstrate how TTT reduces the needed sample complexity for each query from the new task, not just one. The TTT update is performed only once per task to adapt the model, which helps improve the inference time efficiency as each subsequent evaluation needs much fewer in-context examples. We want to highlight that this is more evident near the zero-shot regime as TTT helps the model \"memorize\" new task dynamics and improves accuracy from 0.2 to 0.45. Akyürek et al. (2024) also argue that during TTT, for each task, the model is evaluated on a batch of test points after a small number of gradient steps on a test-time training set, specifically in Section 5 for BBH tasks.\n\nRegarding the dataset, we use the T4 benchmark from Gardner et al. (2024), which is a large set of real-world tabular tasks and consists of four million tables. Following the official TabPFN v2 implementation and our own setup, we select the datasets with at least 1,250 samples (1,000 for training, using an 80–20 split), at most 10 classes (selecting the 10 most frequent), and a 100-feature limit (chosen randomly). We also convert regression tasks into 10-class classifications based on quartiles to maintain consistency in training and evaluation. We appreciate your feedback and plan to detail our dataset choice and experimental methodology both in the main text and appendix.\n\n> **Weakness 3.** ... propositions and theorems are unclear and could be better formulated\n\n**Response:** We are planning to do a pass over the statements and relocate Remark B.5. into the main text for clarity and completeness of the approximations in the theorem statements. We welcome any additional suggestions to clarify our statements.", "claims_and_evidence": "Yes. The claims are supported by theoretical proofs and experimental validation.", "methods_and_evaluation_criteria": "The paper presents a simplified version of TTT compared to practical implementations. Their approach fixes the context and uses k different mini-batch data with MSE loss to train the model, testing it using the same context as training. This simplification differs significantly from TTT used in practice such as [1],[2], potentially limiting the practical insights this work can provide.\n\n[1] Akyürek, E., et al. \"The surprising effectiveness of test-time training for abstract reasoning.\"\n [2] Sun, Y., et al. \"Learning to (learn at test time): RNNs with expressive hidden states.\"", "theoretical_claims": "The mathematical proofs contain imprecise statements, particularly in the use of approximation symbols ($\\approx$) without specifying error bounds. For a paper focused on rigorous theoretical analysis, such imprecision is problematic, as seen in Theorem 4.2 and Corollary 4.5.", "experimental_designs_or_analyses": "While the theoretical validation experiments are sound, the real-world experiments in Section 6.2 have methodological issues. The comparison between TabPFN and TabPFN+TTT should account for the additional 1000 samples used in TTT as part of the context length. The comparison in Figure 3a is therefore problematic. Additionally, for the T4 benchmark experiments, more detailed information about test dataset size and sampling methodology would be beneficial.", "relation_to_broader_scientific_literature": "The paper's main contribution lies in proving that fine-tuning on labeled test data can improve performance. However, this empirical result is somewhat expected, limiting the paper's potential impact.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "**Strengths**:\n1. Comprehensive comparison of TTT performance with different initialization strategies (pre-trained vs. zero initialization)\n2. Thorough analysis of performance gains relative to context and test example numbers, with experimental results supporting theoretical findings\n\n**Weaknesses**:\n1. The paper's oversimplified TTT procedure. The main comparison between randomly initialized models and pretrained models lacks practical insights due to this oversimplified framework.\n2. The experiments in Section 6.2 need more detailed analysis and better clarification of the methodological choices.\n3. The expressions and claims in propositions and theorems are unclear and could be better formulated.", "other_comments_or_suggestions": "If there are any misunderstandings on my part, please point them out, and I will reconsider my evaluation of this work.", "questions_for_authors": "please refer to the weakness part in Other Strengths And Weaknesses, I don’t have more question, my main concern is the oversimplified problem setting. This simplification makes the approach significantly diverge from practical TTT applications.", "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": {"ethical_review_concerns": "NA"}}, {"review_id": "cgS0gHqzz5", "reviewer": "Reviewer_StUH", "summary": "This paper examines the impact of test-time training (TTT) on the performance of a single-layer linear attention model (without an MLP) after a single gradient step during test-time fine-tuning. The authors focus on the problem of in-context linear regression and characterize the performance improvement achieved through one step of gradient descent for a linear attention model in terms of the context length (\\(n\\)) and the ambient dimension (\\(d\\)).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed feedback and positive evaluation of our work. We address the weaknesses and questions brought up below.\n\n> **Experimental Designs Or Analyses.** The experiments are reasonable; however, to support the claim that this theory provides insight into broader in-context learning (ICL) problems, it would be beneficial to include additional ICL tasks, such as a two-layer ReLU neural network, or expand the evaluation with more datasets for the TabPFN experiments (currently it is only tested on 1 data set).\n\n**Response:** Thank you for highlighting this important detail. In lines 379–381, we mistakenly referred to a “single dataset,” but we intended to say “each dataset” in the The Tremendous TabLib Trawl (T4) benchmark (Gardner et al., 2024). T4 is a large-scale, high-quality collection of public tabular benchmarks containing approximately 4 million tables. Our choice to use T4 was indeed to make evaluations more comprehensive by covering a broad range of data. As stated in lines 412-413, we report the average performance across all datasets in T4 with sufficient samples in Figure 3-a.\n\n> **Weakness.** The main limitation of the paper is its focus on a single-step gradient descent and a simplified model.\n\n**Response:** Thank you for your feedback. We want to highlight that we focus on single‐step gradient descent because it is both practically motivated due to its efficiency and theoretically interesting. The TTT update can be implemented in a single forward-backward pass with appropriate masking, as described by the Remark in Section 3, which makes it a meaningful choice with minimal test-time compute overhead. Moreover, our theory and experiments demonstrate that single-step gradient descent can be sufficient, which aligns with the existing practical works (Akyürek et al., 2024) that just a few gradient steps offer significant test-time improvements. Also, even with the linear attention model, the analysis poses significant technical challenges: we must calculate the 8th moments of matrices, and since this is practically intractable, we rely on Gaussian approximation techniques presented in Lemma B.1 in Appendix B. \n\n> **Q1.** How would performing additional gradient descent steps at test time affect the results? Specifically, how would it impact the loss improvement described in the paper's theorem?\n\n**Response:**  Thank you for your thoughtful question. Theoretically, there are technical difficulties in analyzing multiple steps of gradient descent as each subsequent update introduces higher-order dependencies and makes it hard to compute the expectation $\\mathbb{E}[\\mathbf{W}^\\top \\mathbf{W}]$. However, as we discussed in our multi-gradient step experiment in Figure 2c, one would expect the improvement obtained from the next steps to diminish. This can also be intuitively seen in Section 5, as the improvement is given by $A^2/(A+B)$ where $A$ is the misalignment and $B$ is the magnitude term. Each gradient step reduces $A$ as we decrease the misalignment between $\\mathbf{W}$ and $\\boldsymbol{\\beta}_{TT}$, and that $B$ remains relatively stable with the updates. Then, it can be argued that the loss improvement decreases in the next gradient steps.\n\n> **Q2.** Can this approach be extended to a one-layer transformer with softmax attention and an MLP? How challenging would this extension be, and what factors contribute to the difficulty?\n\n**Response:** We thank the reviewer for raising this point. For a single-layer transformer with softmax attention and possibly an MLP layer at the end, the gradient involves calculating terms like $\\nabla_{\\mathbf{W} } \\mathrm{softmax}(\\mathbf{X} \\mathbf{W} \\mathbf{X}^\\top)$ where $\\mathbf{W}_Q \\mathbf{W}_K^\\top = \\mathbf{W}$. For example, considering one row of this resulting matrix $\\mathbf{s} = \\begin{bmatrix} s_1 ~ \\cdots ~ s_n \\end{bmatrix}^\\top = \\mathrm{softmax}(\\mathbf{X}\\mathbf{W}\\mathbf{z}) \\in \\mathbb{R}^{n}$ and taking its derivative w.r.t $\\mathbf{W}$ gives the following form:\n\n$$\n\\qquad \\frac{\\partial s_i}{\\partial \\mathbf{W}} = s_i  (\\mathbf{x_i} - \\sum_{j} s_j \\mathbf{x_j}) \\mathbf{z}^{\\top}.\n$$\n\nAs we analyze the performance of the model after the gradient descent update, we encounter expressions involving high-order moments of data matrices while calculating $\\mathbb{E}[\\mathbf{W}^\\top \\mathbf{W}]$. With this much complicated gradient form, analyzing the update and especially calculating high-order moments of the data matrix $\\mathbf{X}$ becomes far more challenging and intractable. Also, adding an MLP layer might further complicate the backward pass by adding another layer of non-linear transformation.", "claims_and_evidence": "Yes, the claim well supported with empirical experiment on synthetic data.", "methods_and_evaluation_criteria": "Yes, they all make sense.", "theoretical_claims": "The theoretical framework is sound; however, it is highly limited to a single attention layer (without an MLP) and applies only to the linear regression problem.", "experimental_designs_or_analyses": "The experiments are reasonable; however, to support the claim that this theory provides insight into broader in-context learning (ICL) problems, it would be beneficial to include additional ICL tasks, such as a two-layer ReLU neural network, or expand the evaluation with more datasets for the TabPFN experiments(currently it is only tested on 1 data set).", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "I found the paper interesting, and the theoretical insights were valuable. The experiment with real TabPFN was also compelling. The main limitation of the paper is its focus on a single-step gradient descent and a simplified model. However, despite this constraint, the paper still provides meaningful insights.", "other_comments_or_suggestions": "trying on more data set for TabPFN would be an interesting addition helps with completeness.", "questions_for_authors": "1. How would performing additional gradient descent steps at test time affect the results? Specifically, how would it impact the loss improvement described in the paper's theorem?  \n\n2. Can this approach be extended to a one-layer transformer with softmax attention and an MLP? How challenging would this extension be, and what factors contribute to the difficulty?", "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": "ddMPPys0mA", "reviewer": "Reviewer_XqSw", "summary": "The paper investigates the theoretical and empirical advantages of Test-Time Training (TTT) for improving transformers as in-context learners. Authors develop a theoretical framework to characterize how a single-step gradient update at test time enhances in-context learning. The authors provide a rigorous theoretical proof. The analysis covers the role of alignment between pretraining distribution and target tasks, the ability of TTT to mitigate distribution shift, and its sample complexity benefits. Empirical evaluations support the theoretical findings and demonstrate that TTT significantly reduces required sample sizes while maintaining performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their favorable evaluation of our work and for the helpful feedback. We now reply to each of the points raised below.\n\n> **Experimental Designs Or Analyses.** Yes, the experimental design is valid ... Potential limitations are the narrow focus on TabPFN and the need for additional benchmarks to validate broader applicability.\n\n**Response:** Thank you for your feedback. We chose to conduct our experiments with the TabPFN model because it's well-aligned with our theoretical setting with similar token encodings but different prior distributions. This favorable agreement made TabPFN a particularly good setting to conduct large-scale experiments on how TTT reduces the context length needed. For the benchmarks used while evaluating TabPFN, we realized that we did not provide enough contextual information in the paper for our choice of the T4 dataset (Gardner et al., 2024), which is an extensive collection of public benchmarks containing diverse real tabular tasks. In the revision, we will include how we filtered the datasets in T4 according to the requirements of TabPFN and our problem setup, as well as how we constructed the experimental pipeline. Still, we agree that additional domains beyond tabular data would help demonstrate the broader applicability of our theoretical claims.\n\n\n\n> **Weakness.** Limited to simple or toy linear transformer architectures and basic TabPFN tasks, restricting broader applicability. \n\nWe thank the reviewer for the feedback. The TabPFN v2 model itself has been pre-trained on a wide range of tabular distributions, and we evaluate it on the T4 dataset (Gardner et al., 2024) through many diverse real-life tabular tasks. So, we actually demonstrate how TTT reduces the number of in-context examples across a broad range of realistic classification tasks by plotting the average performance across all tasks in T4. This extensive coverage enables us to. For the discussion on model architecture, we refer the reviewer to our related response for the weakness mentioned by Reviewer StUH. \n\n> **Q1.** Would it be necessary to analyze the effects of different optimizers separately?\n\n**Response:**\n\nWe appreciate this question. In our TabPFN experiments, we initially used AdamW as the optimizer, as it's known to work well with transformer architectures. Based on your question, we tried the same experiment using the standard SGD with momentum and observed minimal differences in final performance. The results are available [here](https://anonymous.4open.science/r/TestTimeTrainingProvablyImprovesTransformersAsInContextLearners-F845/TabPFN_SGD.pdf).\n\n\n\n> **Q2.** Do you know about work that explains in-context learning using gradient updates? Since TTT also relies on gradient updates, what are the deeper differences and connections between the two when viewed from the perspective of gradient updates?\n\n**Response:**\n\nThank you for highlighting this connection. The existing works (e.g., Ahn *et al.*, 2023; Mahankali *et al.*, 2023; Zhang *et al.*, 2023; Li *et al.*, 2024) show that in-context learning can be viewed as implicit gradient descent within a transformer's forward pass which effectively simulates a preconditioned GD governed by the task prior. In contrast, our TTT approach explicitly updates the model parameters at inference and provides task-specific adaptation rather than one governed only by the prior. In ICL, any gradient descent occurs internally without updating the weights; as a result, the model can only implement the preconditioned GD it meta‐learned from pretraining. Instead, TTT uses a real gradient step on test data and handles the distribution shift with a low computational cost. Finally, while most ICL analyses operate at the population level (i.e., with infinite samples), our setting provides finite‐sample guarantees on performance improvements on how a single‐step update can reduce the needed context length.\n\n\nAhn, K., Cheng, X., Daneshmand, H., and Sra, S. Transformers learn to implement preconditioned gradient descent for in‐context learning, 2023. URL https://arxiv.org/abs/2306.00297\n\nZhang, R., Frei, S., and Bartlett, P. L. Trained transformers learn linear models in-context, 2023. URL https://arxiv.org/abs/2306.09927.\n\nLi, Y., Rawat, and Oymak, S. Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond, 2024. URL https://arxiv.org/abs/2407.10005.\n\nMahankali, A. V., Hashimoto, T., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention, 2024. URL https://arxiv.org/abs/2307.03576", "claims_and_evidence": "The paper provides sound theoretical backing and empirical results to support the claims. \nHowever, the reliance on linear transformers may limit generalizability to more complex architectures.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes, the paper presents sound and rigorous theoretical derivations.\nI've read and checked the Analysis for Isotropic Covariances and  Analysis for General Covariance parts for the theoretical claims and there's no obvious issues.", "experimental_designs_or_analyses": "Yes, the experimental design is valid and directly supports the theoretical claims, including the experiments on TabPFN (demonstrating efficiency improvements and sample size reductions) and the simulations with GPT-2 (showing TTT for handling distribution shift).\n\nPotential limitations are the narrow focus on TabPFN and the need for additional benchmarks to validate broader applicability.", "relation_to_broader_scientific_literature": "Yes, I think this paper can help the broad community to understand TTT better and provide new insights for utilizing TTT. Also, the experiments on TabPFN shows new opportunities to reduce test time scaling cost (when using ICL).", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strength:\nVery novel and natural theoretical explanation for TTT and ICL.\nWeaknesses:\nLimited to simple or toy linear transformer architectures and basic TabPFN tasks, restricting broader applicability.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1. Would it be necessary to analyze the effects of different optimizers separately?\n2. Do you know about work that explains in-context learning using gradient updates? Since TTT also relies on gradient updates, what are the deeper differences and connections between the two when viewed from the perspective of gradient updates?", "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:

Meta-review: This paper investigates the impact of test-time training on in-context learning for linear models. Within a rigorous theoretical framework, the authors quantify its benefits, such as reducing the number of in-context examples required for downstream tasks. A nice balance of theoretical analysis with experimental observation distinguishes this work from other studies on in-context learning which often focus on either theoretical analyses or experimental observations.
Decision: Accept (poster)