Test time training enhances in-context learning of nonlinear functions
Not on the map. Why?
Abstract (arXiv)
Test-time training (TTT) enhances model performance by explicitly updating designated parameters prior to each prediction to adapt to the test data. While TTT has demonstrated considerable empirical success, its theoretical underpinnings remain limited, particularly for nonlinear models. In this paper, we investigate the combination of TTT with in-context learning (ICL), where the model is given a few examples from the target distribution at inference time. We analyze this framework in the setting of single-index models, where the feature vector is drawn from a hidden low-dimensional subspace. For single-layer transformers trained with gradient-based algorithms and adopting TTT, we establish an upper bound on the prediction risk. Our theory reveals that TTT enables the single-layer transformers to adapt to both the feature vector and the link function, which vary across tasks. This creates a sharp contrast with ICL alone, which is theoretically difficult to adapt to shifts in the link function. Moreover, we provide the convergence rate with respect to the data length, showing the predictive error can be driven arbitrarily close to the noise level as the context size and the network width grow.
Key facts. Cited by 3 papers in the map (data through 14 Oct 2026). Builds on 36 of its 46 references that are in the map. First posted on arXiv 30 Sep 2025.
Builds on
36 of 46 references.
References · 36
- 2020Language Models are Few-Shot Learnerscited by 32,419
- 2022Transformers learn in-context by gradient descentcited by 732
- 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 715
- 2023Trained Transformers Learn Linear Models In-Contextcited by 337
- 2024Many-Shot In-Context Learningcited by 206
- 2022Neural Networks can Learn Representations with Gradient Descentcited by 161
- 2019Linearized two-layers neural networks in high dimensioncited by 155
- 2022High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representationcited by 150
- 2023SGD learning on neural networks: leap complexity and saddle-to-saddle dynamicscited by 143
- 2025Test-Time Training Done Rightcited by 130
- 2020Online stochastic gradient descent on non-convex losses from high-dimensional inferencecited by 119
- 2024How Transformers Learn Causal Structure with Gradient Descentcited by 118
24 more
- 2023In-Context Convergence of Transformerscited by 115
- 2024The Surprising Effectiveness of Test-Time Training for Few-Shot Learningcited by 96
- 2022Learning Single-Index Models with Shallow Neural Networkscited by 92
- 2019Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networkscited by 79
- 2023How Two-Layer Neural Networks Learn, One (Giant) Step at a Timecited by 68
- 2024Many-Shot In-Context Learning in Multimodal Foundation Modelscited by 67
- 2024Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?cited by 65
- 2022Neural Networks Efficiently Learn Low-Dimensional Representations with SGDcited by 63
- 2023Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Contextcited by 63
- 2023Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Modelscited by 61
- 2023Learning time-scales in two-layers neural networkscited by 57
- 2024Transformers Provably Solve Parity Efficiently with Chain of Thoughtcited by 50
- 2025Test-Time Learning for Large Language Modelscited by 50
- 2024Neural network learns low-dimensional polynomials with SGD near the information-theoretic limitcited by 46
- 2024Computational-Statistical Gaps in Gaussian Single-Index Modelscited by 44
- 2025Learning without training: The implicit dynamics of in-context learningcited by 43
- 2024Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functionscited by 37
- 2024In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initializationcited by 33
- 2021How rotational invariance of common kernels prevents generalization in high dimensionscited by 32
- 2024Pretrained transformer efficiently learns low-dimensional target functions in-contextcited by 31
- 2025Test-Time Training Provably Improves Transformers as In-context Learnerscited by 16
- 2024On the Complexity of Learning Sparse Functions with Statistical and Gradient Queriescited by 15
- 2025Analyzing limits for in-context learningcited by 6
- 2025The Inherent Limits of Pretrained LLMs: The Unexpected Convergence of Instruction Tuning and In-Context Learning Capabilitiescited by 5
Built on by
Citing papers per year: 2025: 1; 2026: 2.
3 papers.