Attention models differ in learning nonlinear tasks from context

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

Machine Learning

Summary

Understanding how AI models learn from examples without changing their inner setup is important for making smarter systems. This paper compares two simple attention-based AI designs that try to learn tasks where the output depends in a complex way on input data. One model uses a fixed way to transform inputs before focusing attention, while the other learns to pick and combine input features directly. The authors develop formulas predicting how well each model memorizes what it sees and how well it applies that to new examples. Their results help explain when one model design works better depending on data size and variety, as well as the amount of example context the model gets.

What this means in practice

  • For machine learning engineers: Choose attention architectures that better fit dataset variety and context length to improve nonlinear task learning without model retraining.
  • For ai system architects: Design pretrained models by balancing kernel and feature learners to optimize in-context learning effectiveness for complex tasks.

Authors

Haotian Gu, Yizhou Xu, Lenka Zdeborová

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.
View PDFOpen arXiv