Optimal In-context Adaptivity and Distributional Robustness of Transformers

TL;DR

This paper analyzes how pretrained Transformers perform on tasks with distribution shifts, demonstrating they adapt optimally to task difficulty and are robust to distributional changes, with theoretical guarantees.

Contribution

It provides a theoretical framework showing pretrained Transformers achieve optimal convergence rates under distribution shifts, outperforming traditional minimax bounds.

Findings

Transformers pretrained on sufficient data adapt to task difficulty levels.

They maintain optimal convergence rates within chi-squared divergence bounds.

Pretrained Transformers outperform estimators with access to test distributions.

Abstract

We study in-context learning problems where a Transformer is pretrained on tasks drawn from a mixture distribution $\pi=\sum_{\alpha\in\mathcal{A}} \lambda_{\alpha} \pi_{\alpha}$, called the pretraining prior, in which each mixture component $\pi_{\alpha}$ is a distribution on tasks of a specific difficulty level indexed by $\alpha$. Our goal is to understand the performance of the pretrained Transformer when evaluated on a different test distribution $\mu$, consisting of tasks of fixed difficulty $\beta\in\mathcal{A}$, and with potential distribution shift relative to $\pi_\beta$, subject to the chi-squared divergence $\chi^2(\mu,\pi_{\beta})$ being at most $\kappa$. In particular, we consider nonparametric regression problems with random smoothness, and multi-index models with both random smoothness and random effective dimension. We prove that a large Transformer pretrained on sufficient data achieves the optimal rate of convergence corresponding to the difficulty level $\beta$, uniformly over test distributions $\mu$ 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 $\mu$, the convergence rate of its expected risk over $\mu$ could not be faster than that of our pretrained Transformers, thereby providing a more appropriate optimality guarantee than minimax lower bounds.