Transformers Meet In-Context Learning: A Universal Approximation Theory

TL;DR

This paper presents a universal approximation theory for transformers, explaining how they enable in-context learning for a broad class of functions without weight updates, extending beyond convex optimization problems.

Contribution

It introduces a new theoretical framework combining Barron's approximation theory with the algorithm mimicking view to explain transformers' in-context learning capabilities.

Findings

Transformers can approximate a wide class of functions with small risk.

A transformer can be constructed to find linear representations similar to Lasso.

The theory extends beyond convex problems like linear regression.

Abstract

Large language models are capable of in-context learning, the ability to perform new tasks at test time using a handful of input-output examples, without parameter updates. We develop a universal approximation theory to elucidate how transformers enable in-context learning. For a general class of functions (each representing a distinct task), we demonstrate how to construct a transformer that, without any further weight updates, can predict based on a few noisy in-context examples with vanishingly small risk. Unlike prior work that frames transformers as approximators of optimization algorithms (e.g., gradient descent) for statistical learning tasks, we integrate Barron's universal function approximation theory with the algorithm approximator viewpoint. Our approach yields approximation guarantees that are not constrained by the effectiveness of the optimization algorithms being mimicked, extending far beyond convex problems like linear regression. The key is to show that (i) any target function can be nearly linearly represented, with small $\ell_1$-norm, over a set of universal features, and (ii) a transformer can be constructed to find the linear representation -- akin to solving Lasso -- at test time.