A Decision-Theoretic View of Test-Time Training: When, How Far, and Which Directions to Adapt

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.

Large AI models are usually trained before deployment and then asked to handle many different tasks. But real questions often differ from the data and situations seen during training, which can hurt accuracy. Test-time training (TTT) tries to fix this by making small updates using the prompt itself before the model answers. In practice, however, TTT is fragile: too many updates can overfit the prompt, and updating the wrong parts of the model can make predictions worse. We show that TTT can be understood as a form of Bayesian reasoning: the prompt provides evidence that should update the model’s prior beliefs, but only by the right amount and in the right directions. This view yields three practical principles: a test for when adaptation should help, a way to choose the number of update steps from the prompt with a guarantee against overfitting, and a query-aware score for selecting which Transformer attention components to update. These results explain why TTT is often unstable and provide principled recipes for making on-the-fly adaptation more reliable.