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)

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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:Rd→RDh:\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 NN . We fix N0=24N_{0}=24 and L=L~=1024L=\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(Ekernel/Efeature)\log_{2}(\mathcal{E}^{\mathrm{kernel}}/\mathcal{E}^{\mathrm{feature}}) . (a,b) Memorization and generalization as N/d2N/d^{2} and N0/dN_{0}/d vary, with α=α~=32\alpha=\widetilde{\alpha}=32 ; vertical dashed lines mark N=d2N=d^{2} and N=D2N=D^{2} . (c,d) Memorization and generalization as the training length α=L/d\alpha=L/d and inference length α~=L~/d\widetilde{\alpha}=\widetilde{L}/d vary, with N0=24N_{0}=24 and N=12288N=12288 ; the diagonal dotted line marks equal lengths, and the horizontal dashed line marks α~=32\widetilde{\alpha}=32 . Gold stars mark the shared parameter setting ( N=12288,N0=24,α=α~=32N=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⁡ds=\log L/\log d , with equal training and inference lengths L=L~=dsL=\widetilde{L}=d^{s} . The teacher is ( 24 ). Circles show experiments at d=16d=16 , N0=8192N_{0}=8192 , N=1200000N=1200000 , Δ=0\Delta=0 , and λ=0\lambda=0 . Solid curves are Theorem 1 (Eqs. ( 20 ) and ( 21 )) evaluated at d=16d=16 . Dashed curves are the same formulas evaluated at d=106d=10^{6} and display the separation between the scales L≍dL\asymp d and L≍d2L\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 NN at d=32d=32 , and the lower row varies N0N_{0} at d=24d=24 , N=1200000N=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,qˉ,rq,m,\bar{q},r , respectively, for the learner in that column.
Figure 6: Memorization (left) and generalization (right) errors as a function of N0N_{0} . We fix d=24d=24 , D=323D=323 , N=1200000N=1200000 , L=L~=768L=\widetilde{L}=768 , Δ=0.01\Delta=0.01 , and λ=10−3\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 N0N_{0} , with 3276832768 inference prompts per error estimate.

Explore similar work

CardsList

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

    May 6, 2026Alexander Hsu, Zhaiming Shen, Wenjing Liao +1In-Context LearningTransformer Architectures

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

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

    Sep 24, 2026Zhongjie Shi, Rongjie Lai, Alexander Cloninger +1In-Context LearningTransformer Architectures

    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

  3. Test time training enhances in-context learning of nonlinear functions

    Sep 30, 2025Kento Kuwataka, Taiji SuzukiIn-Context LearningTest-Time Training

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