Test time training enhances in-context learning of nonlinear functions
This work addresses a theoretical gap in adapting nonlinear models during inference, which is incremental but important for improving robustness in few-shot learning scenarios.
The paper tackles the theoretical limitations of test-time training (TTT) combined with in-context learning (ICL) for nonlinear single-index models, showing that TTT enables single-layer transformers to adapt to task-specific features and link functions, reducing prediction risk to near-noise levels as context size and network width increase.
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 $y=σ_*(\langle β, \mathbf{x} \rangle)$, 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.