Efficient and Minimax-optimal In-context Nonparametric Regression with Transformers

TL;DR

This work addresses in-context learning for nonparametric regression with $α$-Hölder functions in $d$ dimensions, and proves that a pretrained transformer with $Θ(\log n)$ parameters and $Ω(n^{2α/(2α+d)}\log^3 n)$ pretraining sequences can achieve the minimax rate $O(n^{-2α/(2α+d)})$ in mean squared error. The authors show how a transformer can efficiently approximate local polynomial estimators by implementing a kernel-weighted polynomial basis and then performing gradient descent, leveraging linear attention to realize the optimization steps. The main contributions include a transformer approximation theorem with an $O(1/n)$ excess risk relative to a local polynomial estimator and a global risk bound that matches the minimax rate under mild assumptions, all with significantly fewer parameters and pretraining resources than prior work. This advances the theoretical understanding of ICL for nonparametric regression and suggests practical pathways for building data-efficient, in-context learners using transformer architectures.

Abstract

We study in-context learning for nonparametric regression with $α$-Hölder smooth regression functions, for some $α>0$. We prove that, with $n$ in-context examples and $d$-dimensional regression covariates, a pretrained transformer with $Θ(\log n)$ parameters and $Ω\bigl(n^{2α/(2α+d)}\log^3 n\bigr)$ pretraining sequences can achieve the minimax-optimal rate of convergence $O\bigl(n^{-2α/(2α+d)}\bigr)$ in mean squared error. Our result requires substantially fewer transformer parameters and pretraining sequences than previous results in the literature. This is achieved by showing that transformers are able to approximate local polynomial estimators efficiently by implementing a kernel-weighted polynomial basis and then running gradient descent.

Figures (6)

• Figure 1: What is the relationship between these two concepts?
• Figure 2: Example of ICL in large language models. Here, there are $n=5$ in-context examples and one query point in the prompt.
• Figure 3: Plate diagram showing the data generating mechanism of our in-context nonparametric regression problem. For each pretraining sequence index $\gamma \in [\Gamma]$, a regression function is drawn from $P_\mathcal{H}$. Then, for each $i \in [n+1]$, i.i.d. covariates $X_i^{(\gamma)}$ and errors $\varepsilon_i^{(\gamma)}$ are sampled from $P_{X, \varepsilon}$. Finally, responses are generated as $Y_i^{(\gamma)} = m^{(\gamma)}\bigl(X_i^{(\gamma)}\bigr) + \varepsilon_i^{(\gamma)}$.
• Figure 4: A single block in the transformer architecture, consisting of a single head of linear attention followed by a one-layer feed-forward neural network applied identically to every row.
• Figure 5: Key steps in the construction of the approximating transformer $f_\mathrm{TF}$; unchanged quantities are omitted for clarity. The first three blocks compute the centred and scaled covariates $\bm{X}$, relative to the test point $X_{n+1}$, along with the kernel matrix $\bm{K}_h(X_{n+1})$. Next, $\Theta(\log n)$ blocks are used to repeatedly multiply these, producing approximations of the kernel-weighted monomial basis $\tilde{\bm{X}}$ and responses $\tilde{Y}$. Finally, $\Theta(\log n)$ steps of gradient descent are applied to arrive at an approximation of the optimal point $w_*$.

Theorems & Definitions (53)

• Definition 1.1: Hölder-ball
• Definition 2.1: Linear attention layer
• Definition 2.2: Feed-forward network layer
• Definition 2.3: Transformer
• Definition 2.4: Transformer-based estimator
• Theorem 2.5
• Theorem 3.1
• Theorem 3.2
• Lemma A.1
• proof : Proof of Lemma \ref{['lem:locpol_H']}