Transformers are almost optimal metalearners for linear classification

TL;DR

This paper provides a theoretical analysis demonstrating that simplified transformers trained via gradient descent can serve as near-optimal metalearners for linear classification tasks, efficiently generalizing with few in-context examples.

Contribution

It introduces the first formal analysis showing transformers as effective metalearners in a linear classification setting, with bounds independent of ambient dimension.

Findings

Transformers trained with gradient descent can act as near-optimal metalearners.

The number of in-context examples needed scales with the subspace dimension, not the ambient dimension.

Transformers outperform learners limited to in-context data, requiring fewer examples to generalize.

Abstract

Transformers have demonstrated impressive in-context learning (ICL) capabilities, raising the question of whether they can serve as metalearners that adapt to new tasks using only a small number of in-context examples, without any further training. While recent theoretical work has studied transformers' ability to perform ICL, most of these analyses do not address the formal metalearning setting, where the objective is to solve a collection of related tasks more efficiently than would be possible by solving each task individually. In this paper, we provide the first theoretical analysis showing that a simplified transformer architecture trained via gradient descent can act as a near-optimal metalearner in a linear classification setting. We consider a natural family of tasks where each task corresponds to a class-conditional Gaussian mixture model, with the mean vectors lying in a shared $k$-dimensional subspace of $R^d$. After training on a sufficient number of such tasks, we show that the transformer can generalize to a new task using only $O(k / R^4)$ in-context examples, where $R$ denotes the signal strength at test time. This performance (almost) matches that of an optimal learner that knows exactly the shared subspace and significantly outperforms any learner that only has access to the in-context data, which requires $\Omega(d / R^4)$ examples to generalize. Importantly, our bounds on the number of training tasks and examples per task needed to achieve this result are independent of the ambient dimension $d$.