Transformers are Universal In-context Learners

We appreciate the detailed suggestions, criticisms and endorsement of the reviewer. We address these below:

(1) motivating the approach by potential (even non-immediate) applications
1 (Question). Could you provide one or two concrete examples ... accessibility of the paper to a wider audience.

We agree that we could have done a better job at giving some concrete example to illustrate the scope of our results. The first (and most important) example we have added to the manuscript is the one of a fixed number of tokens. This clearly show why discrete, fixed token length, is a special case of our theory. Namely, if we consider a fixed $n$ and an in-context map $G(X,x)$, with $X \in \mathbb{R}^{d_{\text{tok}} \times n}$ which is continuous for $\ell^2$ and permutation equivariant with respect to the token, then this defines a map on discrete probability measure

$$\Gamma(\frac{1}{n}\sum_i \delta_{x_i}, x) := G( (x_i)_i, x )$$

which is continuous for the weak$^*$ topology on the set $\mathcal{P}_n(\Omega) \subset \mathcal{P}(\Omega)$ of $n$-point measure (uniform distribution supported on $n$ points). The map $\Gamma$ is continuous on $\mathcal{P}_n(\Omega)$ (because on point sets, the weak$^*$ topology coincides with the $\ell^2$-topology up to permutations), and $\mathcal{P}_n(\Omega)$ is compact (because it is a closed subset of a compact set $\mathcal{P}(\Omega)$). Hence, we can use our theorem on $\mathcal{P}_n(\Omega)$, and obtain that $\Gamma$ can be approximated by a transformer on this space. This implies the approximation of $G$ by a transformer.

The second example we now provide is a regression task associated with the “in-context learning” phenomenon, which we detail below to address your next question.

(2) by providing graphical illustrations ... does not provide a single illustration).

We appreciate this suggestion. We will include a figure in the revised version.

The title "Transformers are Universal In-Context Learners" could be misleading to many readers ... clarify whether they see a connection between their work and the usual notion of ICL.

We do somewhat disagree; we are using "in context" in the same way as in the recent theoretical literature on ICL. We, however, acknowledge that we do not study "learning'' mechanisms (we do not study the optimization of the $(Q,K,V)$) but rather state a "possibility'' (i.e. universality) result: it is possible to learn I-C mapping. To make this connection with the literature more concrete, we have updated the manuscript with the example of the I-C linear regression task studied in [Johannes von Oswald et al., Transformers Learn In-Context by Gradient Descent, 2022]. In the discrete case, tokens are assumed to be of the form $x_i=(u_i,v_i)$ where $u_i$ are feature and $v_i$ labels to be predicted. Then simplified (linear attention) transformers are shown to learn in context a linear relation $v_i \approx W u_i$ and the in-context "prediction'' then maps some $(u,v)$ to $(u, W u)$ (the value of $v$ is discarded). Adding a ridge penalty $\lambda$ to makes the problem well-posed, this corresponds to the I-C map

$$G(X,(u,v)) := (u, W(X) u) \quad\text{where} \quad W(X) := \mathrm{argmin}_{W} \sum_i \|W u_i - v_i \|^2 + \lambda \|W\|^2$$

(the authors of the paper consider in fact a single attention layer and replace this minimization with a single step of gradient descent for simplicity, but this is just a modification of the IC map). Thanks to our framework which operates over measure, this can be written for any $n$ by considering a data distribution $\mu$ over the space $(u,v)$ of (feature, labels), and then defining the more general I-C map

$$\Gamma(\mu,(u,v)) := (u, W(\mu) u) \quad\text{where}\quad W(\mu) := \mathrm{argmin}_W \int \|W u - v\|^2 \mathrm{d} \mu(u,v).$$

This map has a closed-form

$$W(\mu) = \Big[ \int uu^\top \mathrm{d} \mu(u,v) + \lambda \mathrm{Id} \Big]^{-1} \Big[ \int vu^\top \mathrm{d} \mu(u,v) \Big],$$

and it is weak$^*$ continuous as long as $\lambda>0$, so our theorem states that it can be learned in context.

We appreciate the detailed suggestions, criticisms and endorsement of the reviewer. We address all of these below:

I'm not sure if people actually use the unmasked ... regularity constraints on the token distribution.

Classical transformers, such as BERT or ViT, indeed apply encoding to the tokens before passing them through the transformer layers. Therefore, they fit within our model by appropriately defining the $x_i$ to incorporate the embedding information. We have updated the manuscript to include this clarification. What remains open for future work is extending our approach to more recent positional encodings, such as Rotary Positional Embedding (RoPE) [Kazemnejad et al., The Impact of Positional Encoding on Length Generalization in Transformers, 2024]. RoPE modifies all attention layers to account for relative positional information, which would require slight adjustments to the proof to accommodate the different formulas.

We appreciate the detailed suggestions, criticisms and endorsement of the reviewer. We address all of these below:

Hard to think of any beyond those ... if any bearing this has on learning.

We acknowledge that our results do not directly translate into conclusions about the learning capabilities of transformers. Our techniques, particularly the use of the Weierstrass theorem combined with Optimal Transport, bear similarities to methods employed in analyzing the training dynamics of two-layer MLPs, such as in [Chizat, Bach, On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport, 2018]. Consequently, we believe that future research could extend our approach to explore convergence results for transformers. We have added further comments to elaborate on this point. That said, unlike classical MLPs with cosine activations, shallow architectures lack an algebraic structure (e.g., they cannot be multiplied), which makes the proof technique itself significant. To address this challenge, we relied on depth to overcome the limitation. We have also included additional remarks on the broader implications of the proof technique.