Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer
For the theoretical understanding of in-context learning, this work extends analysis from linear to nonlinear regression, offering a framework and bounds that are currently missing.
The paper studies in-context learning for nonlinear regression with transformers, constructing attention-based networks that realize nonlinear features like polynomial or spline bases, and provides finite-sample generalization error bounds in terms of context length and training set size.
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.