In-context Learning of Single-index Targets: Comparing Kernel and Feature Learners

Haotian Gu; Yizhou Xu; Lenka Zdeborová

cs.LGOct 1, 2026

In-context Learning of Single-index Targets: Comparing Kernel and Feature Learners

Authors: Haotian Gu, Yizhou Xu, Lenka Zdeborová

Organizations: Statistical Physics of Computation Laboratory, École Polytechnique Fédérale de Lausanne (EPFL) · University of Chinese Academy of Sciences · Information, Learning and Physics Laboratory, École Polytechnique Fédérale de Lausanne (EPFL)

Abstract

In-context learning (ICL) enables a pretrained model to infer a task from demonstrations without updating its parameters. While much of the existing theory focuses on linear target functions, in this paper we study nonlinear cases by comparing two one-layer attention architectures on the same family of single-index tasks. A kernel learner first maps inputs through a fixed nonlinear feature map and then applies linear attention, whereas a feature learner applies attention to the original input, followed by a learned nonlinear readout. We derive predictions for their memorization and generalization errors using the replica method, retaining the effects of pretraining size, task-pool diversity, and training and inference context lengths. The resulting predictions closely match numerical experiments across a broad range of regimes. Our analysis yields phase diagrams that characterize when each architecture is advantageous as the amount of pretraining data, task diversity, and context lengths vary. We further identify qualitatively different context-length scalings for the two learners. Together, these results clarify how architectural choices interact with the dataset and govern nonlinear in-context learning.

Figures & tables

Figure 1: The kernel and feature learners. The kernel learner applies a fixed map $h:\mathbb{R}^{d}\to\mathbb{R}^{D}$ before self-attention, whereas the feature learner applies a trainable readout $\widehat{\sigma}$ after self-attention.

Figure 2: Memorization (left) and generalization (right) errors as a function of

N

. We fix

N_{0}=24

and

L=\widetilde{L}=1024

. Red, orange, and blue denote the feature learner, linear attention, and the kernel learner, respectively. Solid lines show replica predictions, and circles show numerical results from a single training run. The three learners use the same training and inference prompts.

Figure 3: Replica predictions for the relative errors of the two learners. In all panels, the color shows $\log_{2}(\mathcal{E}^{\mathrm{kernel}}/\mathcal{E}^{\mathrm{feature}})$ . (a,b) Memorization and generalization as $N/d^{2}$ and $N_{0}/d$ vary, with $\alpha=\widetilde{\alpha}=32$ ; vertical dashed lines mark $N=d^{2}$ and $N=D^{2}$ . (c,d) Memorization and generalization as the training length $\alpha=L/d$ and inference length $\widetilde{\alpha}=\widetilde{L}/d$ vary, with $N_{0}=24$ and $N=12288$ ; the diagonal dotted line marks equal lengths, and the horizontal dashed line marks $\widetilde{\alpha}=32$ . Gold stars mark the shared parameter setting ( $N=12288,N_{0}=24,\alpha=\widetilde{\alpha}=32$ ). Blue favors the kernel learner, red favors the feature learner, and black curves indicate equal errors.

Figure 4: Learning different harmonic degrees. The horizontal coordinate is

s=\log L/\log d

, with equal training and inference lengths

L=\widetilde{L}=d^{s}

. The teacher is ( 24 ). Circles show experiments at

d=16

N_{0}=8192

N=1200000

\Delta=0

, and

\lambda=0

. Solid curves are Theorem 1 (Eqs. ( 20 ) and ( 21 )) evaluated at

d=16

. Dashed curves are the same formulas evaluated at

d=10^{6}

and display the separation between the scales

L\asymp d

and

L\asymp d^{2}

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Consistency of the replica method with numerical experiments. Left: feature learner; right: kernel learner. The upper row varies

N

d=32

, and the lower row varies

N_{0}

d=24

N=1200000

. Solid lines are replica predictions, and stars show the order parameters measured from the learned attention matrices in one run. Red, teal, green, and purple denote

q,m,\bar{q},r

, respectively, for the learner in that column.

Figure 6: Memorization (left) and generalization (right) errors as a function of

N_{0}

. We fix

d=24

D=323

N=1200000

L=\widetilde{L}=768

\Delta=0.01

, and

\lambda=10^{-3}

. Red, orange, and blue denote the feature learner, linear attention, and the kernel learner, respectively. Solid lines show replica predictions, and circles show numerical results. The three learners use the same training and inference prompts at each

N_{0}

, with

32768

inference prompts per error estimate.

Explore similar work

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

Sep 24, 2026stat.ML

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

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.

Zhongjie Shi, Rongjie Lai, Alexander Cloninger +1

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

In-Context Learning Transformer Architectures Cross-Task Interference

Ask PDF

Sep 30, 2025stat.ML

Test time training enhances in-context learning of nonlinear functions

Test-time training (TTT) enhances model performance by explicitly updating designated parameters prior to each prediction to adapt to the test data. While TTT has demonstrated considerable empirical success, its theoretical underpinnings remain limited, particularly for nonlinear models. In this paper, we investigate the combination of TTT with in-context learning (ICL), where the model is given a few examples from the target distribution at inference time. We analyze this framework in the setting of single-index models, where the feature vector is drawn from a hidden low-dimensional subspace. For single-layer transformers trained with gradient-based algorithms and adopting TTT, we establish an upper bound on the prediction risk. Our theory reveals that TTT enables the single-layer transformers to adapt to both the feature vector and the link function, which vary across tasks. This creates a sharp contrast with ICL alone, which is theoretically difficult to adapt to shifts in the link function. Moreover, we provide the convergence rate with respect to the data length, showing the predictive error can be driven arbitrarily close to the noise level as the context size and the network width grow.

Kento Kuwataka, Taiji Suzuki

Department of Mathematical Engineering and Information Physics The University of Tokyo Hongo 7-3-1, Bunkyo-ku, Tokyo

In-Context Learning Test-Time Training Transformer Architectures Inference-Time Nonlinearity

Ask PDF