Transformers as Cross-Task Learners: Shared Structure Drives Sample Efficiency in In-Context Learning

Zhongjie Shi; Rongjie Lai; Alexander Cloninger; Wenjing Liao

stat.MLSep 24, 2026

Transformers as Cross-Task Learners: Shared Structure Drives Sample Efficiency in In-Context Learning

Authors: Zhongjie Shi, Rongjie Lai, Alexander Cloninger, Wenjing Liao

Organizations: School of Mathematics, Georgia Institute of Technology, Atlanta, GA 30332, United States · Department of Mathematics, Purdue University, West Lafayette, IN 47907, United States · Department of Mathematics and Halıcıoğlu Data Science Institute, University of California, San Diego, La Jolla, CA 92093, United States

Abstract

Transformers achieve remarkable performance by jointly learning broad families of tasks during pretraining and adapting to unseen tasks from only a short prompt. Yet a rigorous mathematical and statistical understanding of this phenomenon remains limited. This paper aims to study how Transformers exploit shared cross-task structure and how this structure affects the sample complexity of in-context learning (ICL). Specifically, we characterize task-space complexity through covering numbers under a prescribed metric, thereby quantifying the low-dimensional cross-task structure without requiring an explicit parametric representation. The resulting cover provides a set of anchor functions, which we use to introduce a task-identification-and-evaluation procedure: context observations localize an unseen task among the anchor functions, and the response at a query is predicted by aggregating the corresponding anchor function query evaluations. For approximation, we explicitly construct a Transformer with Softmax attention to approximate this procedure. For generalization, we derive an error bound that separates the effects of the number of pretraining tasks and the prompt length. The scaling with respect to the number of pretraining tasks is governed by the intrinsic dimensions of the task space and input domain; once sufficiently many tasks are available, the dependence on the prompt context length becomes dimension-free. To the best of our knowledge, this is the first work to quantify cross-task complexity for general nonlinear task families and explicitly construct a Transformer that exploits their low-dimensional structure to perform ICL. Our theory provides a quantitative explanation of how joint pretraining across related tasks improves in-context generalization.

Figures & tables

$n$	ICL on full family	ICL on 2D subspace	21D ridge regression for 2D subspace
4	0.66663	0.07832	0.67697
8	0.50026	0.03092	0.47792
16	0.35547	0.01718	0.16517
20	0.26714	0.01623	0.03750

Table 1: Test MSE on a common set of prompts from the hidden two-dimensional subspace. The two Transformers are pretrained on the full 21-dimensional quadratic family and the hidden subspace, respectively; ridge regression uses the full 21-dimensional polynomial basis and is fitted independently to each prompt.

Work	Task-space assumptions	Input and observation assumptions	Generalization error bound
This work	A uniformly bounded, uniformly $\alpha$ -Hölder task space $\mathcal{M}_{f}$ with $\alpha\in(0,1]$ , covering complexity $\mathcal{N}(r,\mathcal{M}_{f},L^{2}(\rho_{\mathbf{x}}))=\mathcal{O}(r^{-d_{f}})$ .	Inputs sampled i.i.d. from an arbitrary distribution supported on a compact $\mathcal{M}_{\mathbf{x}}\subset[0,1]^{d}$ with $\mathcal{N}(r,\mathcal{M}_{\mathbf{x}},\\|\cdot\\|_{2})=\mathcal{O}(r^{-d_{x}})$ ; noiseless responses.	$\displaystyle\begin{gathered}\widetilde{\mathcal{O}}\!\left(n^{-\frac{1}{3}}+\left(\frac{n^{2}}{\Gamma}\right)^{\frac{2}{q+2}}\right)\\[2.0pt] q=d_{f}+\frac{3d_{x}}{\alpha}\end{gathered}$
Kim et al. ( 27 )	A Besov ball $B_{p_{0},q_{0}}^{\alpha}([0,1]^{d})$ with $\alpha>d/p_{0},p_{0}\geq 2$ ; under a B-spline wavelet expansion, the task coefficients are centered, independent, and satisfy a prescribed scale-dependent variance decay.	Inputs sampled i.i.d. from a distribution with density bounded above and below on $[0,1]^{d}$ ; bounded, mean-zero observation noise.	$\displaystyle\widetilde{\mathcal{O}}\!\left(n^{-\frac{2\alpha}{2\alpha+d}}+\frac{n^{\frac{2d}{2\alpha+d}}}{\Gamma}\right)$
Shen et al. ( 42 )	Uniformly bounded $\alpha$ -Hölder functions with $\alpha\in(0,1]$ , defined on a compact $d$ -dimensional Riemannian manifold with positive reach.	Inputs sampled i.i.d. from the uniform distribution on the manifold; noiseless responses.	$\displaystyle\widetilde{\mathcal{O}}\!\left(n^{-\frac{2\alpha}{2\alpha+d}}+\frac{n}{\sqrt{\Gamma}}\right)$
Ching et al. ( 10 )	Tasks drawn from a distribution supported on a uniformly bounded $\alpha$ -Hölder ball on $[0,1]^{d}$ , with a common $\alpha>0$ .	Inputs sampled i.i.d. from a distribution with density bounded above and below on $[0,1]^{d}$ ; bounded observation noise satisfying $\mathbb{E}[\varepsilon\mid X]=0$ .	$\displaystyle\widetilde{\mathcal{O}}\!\left(n^{-\frac{2\alpha}{2\alpha+d}}+\frac{1}{\Gamma}\right)$
Hsu et al. ( 24 )	A uniformly bounded function class within $L^{\infty}$ distance $\delta$ of a fixed finite-dimensional polynomial space with bounded coefficients.	Inputs sampled i.i.d. from a distribution supported on a bounded interval, with uniformly well-conditioned feature covariance; noiseless responses.	$\displaystyle\widetilde{\mathcal{O}}\!\left(\frac{1}{n}+\delta^{2}+\sqrt{\frac{n}{\Gamma}}\right)$

Table 2: Comparison of generalization error bounds for ICL studied in our work and related works.

Explore similar work

Jun 5, 2026stat.ML

The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspaces

The transformer's emergent ability to perform in-context learning (ICL) has sparked a wide range of studies designed to understand its underlying mechanisms. Existing works often study how training task diversity, defined either as the number of ICL training task vectors or as the number of function classes from which the task vectors are drawn, shapes both the learning dynamics and generalization capabilities of ICL. While both definitions have uncovered many interesting phenomena, many observations under the latter definition remain theoretically unexplained. This paper presents a minimal analytical model under which these phenomena provably emerge from the properties of the training data. By modeling the training task vectors as a mixture of low-rank Gaussians, we show how training task diversity, defined by the number of non-overlapping columns between subspaces that parameterize the covariance matrices, improves both the generalization and optimization trajectory of ICL with linear attention. In particular, we show that our model can explain (i) why training with task diversity shortens the ICL plateau and (ii) why ICL appears to achieve out-of-distribution generalization. We conclude by empirically demonstrating how our results extend to nonlinear transformers and nonlinear function classes. Overall, our work presents a tractable framework to unify existing observations.

Soo Min Kwon, Alec S. Xu, Can Yaras +3

Department of Electrical and Computer Engineering, University of Michigan · Department of Statistics, University of California, Davis

In-Context Learning Equivariant Learning Subspace Out-Of-Distribution Energy-Based Transformer Impact Analysis

Ask PDF

Apr 28, 2026cs.LG

Investigation into In-Context Learning Capabilities of Transformers

Transformers have demonstrated a strong ability for in-context learning (ICL), enabling models to solve previously unseen tasks using only example input output pairs provided at inference time. While prior theoretical work has established conditions under which transformers can perform linear classification in-context, the empirical scaling behavior governing when this mechanism succeeds remains insufficiently characterized. In this paper, we conduct a systematic empirical study of in-context learning for Gaussian-mixture binary classification tasks. Building on the theoretical framework of Frei and Vardi (2024), we analyze how in-context test accuracy depends on three fundamental factors: the input dimension, the number of in-context examples, and the number of pre-training tasks. Using a controlled synthetic setup and a linear in-context classifier formulation, we isolate the geometric conditions under which models successfully infer task structure from context alone. We additionally investigate the emergence of benign overfitting, where models memorize noisy in-context labels while still achieving strong generalization performance on clean test data. Through extensive sweeps across dimensionality, sequence length, task diversity, and signal-to-noise regimes, we identify the parameter regions in which this phenomenon arises and characterize how it depends on data geometry and training exposure. Our results provide a comprehensive empirical map of scaling behavior in in-context classification, highlighting the critical role of dimensionality, signal strength, and contextual information in determining when in-context learning succeeds and when it fails.

Rushil Chandrupatla, Leo Bangayan, Sebastian Leng

In-Context Learning In-Context Transformer Architectures Intrinsic Dimensionality Improved Generalization Gaussian Mixture Models

Ask PDF

May 6, 2026cs.LG

Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer

Pre-trained transformers are able to learn from examples provided as part of the prompt without any weight updates, a remarkable ability known as in-context learning (ICL). Despite its demonstrated efficacy across various domains, the theoretical understanding of ICL is still developing. Whereas most existing theory has focused on linear models, we study ICL in the nonlinear regression setting. Through the interaction mechanism in attention, we explicitly construct transformer networks to realize nonlinear features, such as polynomial or spline bases, which span a wide class of functions. Based on this construction, we establish a framework to analyze end-to-end in-context nonlinear regression with the constructed features. Our theory provides finite-sample generalization error bounds in terms of context length and training set size. We numerically validate the theory on synthetic regression tasks.

Alexander Hsu, Zhaiming Shen, Wenjing Liao +1

Department of Mathematics, Purdue University · School of Mathematics, Georgia Institute of Technology

In-Context Learning Transformer Architectures Splines Feature Extraction

Ask PDF