# Test time training enhances in-context learning of nonlinear functions

arXiv: https://arxiv.org/abs/2509.25741 · HTML: https://mlsanity.com/paper/2509.25741/

## 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 (most cited first, 30 of 36)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Many-Shot In-Context Learning](https://mlsanity.com/paper/2404.11018/) (2024)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [Linearized two-layers neural networks in high dimension](https://mlsanity.com/paper/1904.12191/) (2019)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [Test-Time Training Done Right](https://mlsanity.com/paper/2505.23884/) (2025)
- [Online stochastic gradient descent on non-convex losses from high-dimensional inference](https://mlsanity.com/paper/2003.10409/) (2020)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [The Surprising Effectiveness of Test-Time Training for Few-Shot Learning](https://mlsanity.com/paper/2411.07279/) (2024)
- [Learning Single-Index Models with Shallow Neural Networks](https://mlsanity.com/paper/2210.15651/) (2022)
- [Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networks](https://mlsanity.com/paper/1910.01619/) (2019)
- [How Two-Layer Neural Networks Learn, One (Giant) Step at a Time](https://mlsanity.com/paper/2305.18270/) (2023)
- [Many-Shot In-Context Learning in Multimodal Foundation Models](https://mlsanity.com/paper/2405.09798/) (2024)
- [Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?](https://mlsanity.com/paper/2410.08292/) (2024)
- [Neural Networks Efficiently Learn Low-Dimensional Representations with SGD](https://mlsanity.com/paper/2209.14863/) (2022)
- [Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context](https://mlsanity.com/paper/2312.06528/) (2023)
- [Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models](https://mlsanity.com/paper/2305.10633/) (2023)
- [Learning time-scales in two-layers neural networks](https://mlsanity.com/paper/2303.00055/) (2023)
- [Transformers Provably Solve Parity Efficiently with Chain of Thought](https://mlsanity.com/paper/2410.08633/) (2024)
- [Test-Time Learning for Large Language Models](https://mlsanity.com/paper/2505.20633/) (2025)
- [Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit](https://mlsanity.com/paper/2406.01581/) (2024)
- [Computational-Statistical Gaps in Gaussian Single-Index Models](https://mlsanity.com/paper/2403.05529/) (2024)
- [Learning without training: The implicit dynamics of in-context learning](https://mlsanity.com/paper/2507.16003/) (2025)
- [Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions](https://mlsanity.com/paper/2405.15459/) (2024)
- [In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization](https://mlsanity.com/paper/2402.14951/) (2024)

## Built on by (most cited first, 3 of 3)

- [Efficient and Minimax Optimal In-context Nonparametric Regression with Transformers](https://mlsanity.com/paper/2601.15014/) (2026)
- [Optimal In-context Adaptivity and Distributional Robustness of Transformers](https://mlsanity.com/paper/2510.23254/) (2025)
- [A Decision-Theoretic View of Test-Time Training: When, How Far, and Which Directions to Adapt](https://mlsanity.com/paper/2606.15569/) (2026)

## 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.