Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper analyzes nonlinear in-context learning for single-index tasks by comparing kernel and feature learners, two one-layer attention architectures. Using the replica method, it derives theoretical predictions for memorization and generalization errors, accounting for pretraining size, task diversity, and context lengths, which are validated by numerical experiments. The analysis yields phase diagrams indicating when each architecture is superior and reveals distinct context-length scalings, clarifying architectural impacts on nonlinear ICL.
Kernel learners outperform feature learners for nonlinear in-context learning when task diversity is high and context length is short, but the opposite holds for low diversity and long contexts.
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.