mlsanityData to 14 Oct 2026

Test-Time Training Provably Improves Transformers as In-context Learners

arXiv 2503.11842 ↗First posted 14 Mar 2025cs.LG stat.MLcited by 16

Not on the map. Why?

Abstract (arXiv)

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.

Key facts. Cited by 16 papers in the map (data through 14 Oct 2026). Builds on 24 of its 35 references that are in the map. First posted on arXiv 14 Mar 2025.

Builds on

24 of 35 references.

References · 24

  1. 2020Language Models are Few-Shot Learnerscited by 32,419
  2. 2017Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networkscited by 4,897
  3. 2022Large Language Models are Zero-Shot Reasonerscited by 4,374
  4. 2021Finetuned Language Models Are Zero-Shot Learnerscited by 3,780
  5. 2021Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processingcited by 2,234
  6. 2024Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameterscited by 1,986
  7. 2022Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?cited by 1,434
  8. 2021An Explanation of In-context Learning as Implicit Bayesian Inferencecited by 938
  9. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 715
  10. 2022TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Secondcited by 674
  11. 2022Efficient Test-Time Model Adaptation without Forgettingcited by 469
  12. 2024Learning to (Learn at Test Time): RNNs with Expressive Hidden Statescited by 355
12 more
  1. 2023Trained Transformers Learn Linear Models In-Contextcited by 337
  2. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 279
  3. 2023Transformers as Algorithms: Generalization and Stability in In-context Learningcited by 222
  4. 2022Test-Time Training with Masked Autoencoderscited by 212
  5. 2024Many-Shot In-Context Learningcited by 206
  6. 2023Robust Test-Time Adaptation in Dynamic Scenarioscited by 189
  7. 2023One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attentioncited by 170
  8. 2023How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?cited by 96
  9. 2023Test-Time Training on Nearest Neighbors for Large Language Modelscited by 88
  10. 2024Large Scale Transfer Learning for Tabular Data via Language Modelingcited by 76
  11. 2024Retrieval & Fine-Tuning for In-Context Tabular Modelscited by 59
  12. 2024Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyondcited by 24

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 6; 2026: 10.

16 papers.

Sort:

By area:

  1. 2025Making Small Language Models Efficient Reasoners: Intervention, Supervision, Reinforcementcited by 14
  2. 2025Pretrain-Test Task Alignment Governs Generalization in In-Context Learningcited by 8
  3. 2025Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspectivecited by 6
  4. 2025Retrieval-Augmented Generation as Noisy In-Context Learning: A Unified Theory and Risk Boundscited by 4
  5. 2025Test time training enhances in-context learning of nonlinear functionscited by 3
  6. 2026Evolutionary Multi-Task Optimization for LLM-Guided Program Discoverycited by 3
  7. 2026Phantom Gains: Auditing Self-Improvement Against a Measured Nullcited by 2
  8. 2025Attention with Trained Embeddings Provably Selects Important Tokenscited by 1
  9. 2026Fine-Tuning Without Forgetting In-Context Learning: A Theoretical Analysis of Linear Attention Modelscited by 1
  10. 2026The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspacescited by 1
  11. 2026Sequential Correlations Change In-Context Learning: Effective Context Length and Architectural Mismatchcited by 1
  12. 2026Covariance-Aware Transformers for Quadratic Programming and Decision Makingcited by 0
4 more
  1. 2026A Decision-Theoretic View of Test-Time Training: When, How Far, and Which Directions to Adaptcited by 0
  2. 2026Complementary Roles of Activation and Parametric Memory in Few-Shot Learningcited by 0
  3. 2026GradLev: Token-Parallel Test-Time Training Via Costate Predictioncited by 0
  4. 2026Retrieval-Based In-Context Learning: A Domain Adaptation Frameworkcited by 0