Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer
arXiv: 2605.05176 Authors: Alexander Hsu, Zhaiming Shen, Wenjing Liao, Rongjie Lai Field: Machine Learning
Abstract
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, this paper studies ICL in the nonlinear regression setting.
Through the interaction mechanism in attention, the authors 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, they establish a framework to analyze end-to-end in-context nonlinear regression with the constructed features.
Key Contributions
- Nonlinear featurization via attention: Explicit construction of transformer networks whose attention mechanisms realize nonlinear feature bases (e.g., polynomials, splines), extending beyond linear-model analyses of ICL.
- Generalization framework: A theoretical framework for end-to-end in-context nonlinear regression using the constructed features.
- Finite-sample error bounds: Generalization error bounds expressed in terms of context length and training set size.
- Empirical validation: Numerical experiments on synthetic regression tasks that corroborate the theoretical results.