Provable test-time adaptivity and distributional robustness of in-context learning
Tianyi Ma, Tengyao Wang, Richard J. Samworth
TL;DR
The paper develops a theory for provable test-time adaptivity and distributional robustness of in-context learning with Transformers pretrained on a mixture of tasks of varying difficulty. It proves a decomposition of the ICL excess risk under test-time shifts and establishes that large, data-rich Transformers achieve optimal convergence rates corresponding to the test task difficulty, uniformly over divergence-bounded test distributions. By leveraging posterior contraction theory, Besov and multi-index-model analyses yield adaptive, minimax-optimal rates and robustness to distribution shift, complemented by Bayes-risk lower bounds showing no estimator can surpass these rates even with knowledge of the test distribution. These results provide a principled, information-theoretic justification for ICL reliability in the presence of distributional shifts and varying task complexities, with implications for designing and pretraining Transformers for robust downstream use.
Abstract
We study in-context learning problems where a Transformer is pretrained on tasks drawn from a mixture distribution $π=\sum_{α\in\mathcal{A}} λ_α π_α$, called the pretraining prior, in which each mixture component $π_α$ is a distribution on tasks of a specific difficulty level indexed by $α$. Our goal is to understand the performance of the pretrained Transformer when evaluated on a different test distribution $μ$, consisting of tasks of fixed difficulty $β\in\mathcal{A}$, and with potential distribution shift relative to $π_β$, subject to the chi-squared divergence $χ^2(μ,π_β)$ being at most $κ$. In particular, we consider nonparametric regression problems with random smoothness, and multi-index models with random smoothness as well as random effective dimension. We prove that a large Transformer pretrained on sufficient data achieves the optimal rate of convergence corresponding to the difficulty level $β$, uniformly over test distributions $μ$ in the chi-squared divergence ball. Thus, the pretrained Transformer is able to achieve faster rates of convergence on easier tasks and is robust to distribution shift at test time. Finally, we prove that even if an estimator had access to the test distribution $μ$, the convergence rate of its expected risk over $μ$ could not be faster than that of our pretrained Transformers, thereby providing a more appropriate optimality guarantee than minimax lower bounds.