A Decision-Theoretic View of Test-Time Training: When, How Far, and Which Directions to Adapt
cs.LG, math.ST, stat.ML, stat.TH
Submitted: 2026-06-14
Updated: 2026-10-19
Journal ref: Proceedings of Machine Learning Research, 306:124279-124317, 2026 (ICML 2026)
Project page: http://probml.github.io/book1
License: http://creativecommons.org/licenses/by/4.0/
The gist: Test-time training (TTT) adapts a pretrained model to each prompt via parameter updates, improving accuracy under pretraining-to-test distribution shifts.
Terminology
Abstract
Test-time training (TTT) adapts a pretrained model to each prompt via parameter updates, improving accuracy under pretraining-to-test distribution shifts. Yet, its performance often suffers from instability and sensitivity to hyperparameters such as update steps and subspace. We explain this behavior through a decision-theoretic lens, treating TTT as implicit Bayesian inference in the kernel regime. Under a Gaussian process benchmark, we show that TTT reduces prediction error when updates are spectrally matched to the prompt's signal-to-noise ratio and aligned with query-relevant eigen-directions. This perspective underpins the following results: (1) we show when fixed update steps and subspaces fail under distribution shifts, motivating adaptive strategies; (2) we prove that selecting update steps via prompt evidence admits a PAC-Bayes guarantee against overfitting; and (3) we characterize the Bayes-optimal update subspace under a linear-Gaussian correction model, yielding a scoring rule for selecting Transformer blocks and heads. Our theory helps explain the empirical instability of TTT, taking a step toward principled guidance for when, how far, and which directions to adapt.
Sources
- Pac-Bayesian Supervised Classification: The Thermodynamics of Statistical Learning
- Specialization after Generalization: Towards Understanding Test-Time Training in Foundation Models
- Test time training enhances in-context learning of nonlinear functions
- Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent
- Amortized In-Context Bayesian Posterior Estimation
- In-Context Learning Is Provably Bayesian Inference: A Generalization Theory for Meta-Learning
- Test-Time Training Done Right
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks