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

Alexander Hsu; Zhaiming Shen; Wenjing Liao; Rongjie Lai

cs.LGMay 6, 2026

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

Authors: Alexander Hsu, Zhaiming Shen, Wenjing Liao, Rongjie Lai

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

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, 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.

Explore similar work

In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention
Jul 17, 2026Katsuyuki HagiwaraIn-Context Learning Transformer Architectures
Test time training enhances in-context learning of nonlinear functions
Sep 30, 2025Kento Kuwataka, Taiji SuzukiIn-Context Learning Transformer Layers
Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent
May 7, 2026Chenyang Zhang, Yuan CaoIn-Context Learning Transformer Architectures

Jul 17, 2026 · cs.LGJ/K move · Enter open · S save

In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention

Katsuyuki Hagiwara

Faculty of Education, Mie University · Kurima-Machiya-cho, Tsu, 514-8507, Japan

In-context learning is a remarkable property of transformers and has recently received a lot of interest. In many studies of in-context learning, it has been shown that transformers are capable of implementing solver for linear and non-linear regression problems, in which the most of them implement gradient descent algorithm. However, it is still unclear whether those implementations have actually been acquired through training. In this paper, we construct a transformer with linear self-attention, which in-context learns the least squares estimate in a simple regression task. The point here is that the closed form (analytical) solution is approximately obtained by using layer normalization rather than an approximate solution based on gradient descent algorithm. Then, we show an experimental example, in which our implementation is mainly used in the transformer trained with l1 regularization when the target output is the least squares estimate.

In-Context Learning Transformer Architectures Softmax Attention Closed-Form Batch Normalization

Open paper Ask about it PDF

Abstract

Explore similar work

In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention

Test time training enhances in-context learning of nonlinear functions

Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent