Transformers are Universal In-context Learners
Takashi Furuya, Maarten V. de Hoop, Gabriel Peyré
TL;DR
The paper develops a rigorous, measure-theoretic framework for transformers operating on arbitrarily long token contexts by treating contexts as probability distributions and using Wasserstein-based smoothness. It proves universal approximation results for both unmasked and masked (causal) transformers with fixed embedding dimensions and a number of heads that do not grow with precision, enabling a single transformer to approximate any continuous in-context mapping uniformly over compact token domains. The approach relies on expressing in-context mappings as compositions of context-dependent (attention) and context-free (MLP) operators, then leveraging dense families of cylindrical functions and Stone–Weierstrass arguments to show universal approximation, including a space-time lifting to handle causality. The masked-case result additionally requires Lipschitz regularity of the time-evolving context and identifiability to guarantee the reduction to a two-argument representation amenable to approximation. Overall, the work provides a foundational, general theory for the expressive power of Transformers as universal in-context learners operating on distributions, with implications for understanding long-context reasoning and mean-field limits in deep architectures.
Abstract
Transformers are deep architectures that define "in-context mappings" which enable predicting new tokens based on a given set of tokens (such as a prompt in NLP applications or a set of patches for a vision transformer). In this work, we study in particular the ability of these architectures to handle an arbitrarily large number of context tokens. To mathematically, uniformly address their expressivity, we consider the case that the mappings are conditioned on a context represented by a probability distribution of tokens which becomes discrete for a finite number of these. The relevant notion of smoothness then corresponds to continuity in terms of the Wasserstein distance between these contexts. We demonstrate that deep transformers are universal and can approximate continuous in-context mappings to arbitrary precision, uniformly over compact token domains. A key aspect of our results, compared to existing findings, is that for a fixed precision, a single transformer can operate on an arbitrary (even infinite) number of tokens. Additionally, it operates with a fixed embedding dimension of tokens (this dimension does not increase with precision) and a fixed number of heads (proportional to the dimension). The use of MLPs between multi-head attention layers is also explicitly controlled. We consider both unmasked attentions (as used for the vision transformer) and masked causal attentions (as used for NLP and time series applications). We tackle the causal setting leveraging a space-time lifting to analyze causal attention as a mapping over probability distributions of tokens.