Transformers are Minimax Optimal Nonparametric In-Context Learners

TL;DR

This paper provides a theoretical analysis of in-context learning with transformers, showing they can achieve minimax optimal estimation risk by encoding relevant basis representations, and explores the roles of task diversity and representation learning.

Contribution

It develops approximation and generalization error bounds for transformers in nonparametric regression, establishing their optimality and analyzing the effects of pretraining and in-context learning.

Findings

Transformers can achieve minimax optimal estimation risk in ICL.

Pretraining encodes relevant basis representations improving in-context learning.

Theoretical bounds and lower bounds clarify roles of task diversity and representation learning.

Abstract

In-context learning (ICL) of large language models has proven to be a surprisingly effective method of learning a new task from only a few demonstrative examples. In this paper, we study the efficacy of ICL from the viewpoint of statistical learning theory. We develop approximation and generalization error bounds for a transformer composed of a deep neural network and one linear attention layer, pretrained on nonparametric regression tasks sampled from general function spaces including the Besov space and piecewise $\gamma$-smooth class. We show that sufficiently trained transformers can achieve -- and even improve upon -- the minimax optimal estimation risk in context by encoding the most relevant basis representations during pretraining. Our analysis extends to high-dimensional or sequential data and distinguishes the \emph{pretraining} and \emph{in-context} generalization gaps. Furthermore, we establish information-theoretic lower bounds for meta-learners w.r.t. both the number of tasks and in-context examples. These findings shed light on the roles of task diversity and representation learning for ICL.