Reproduce “Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer” on OpenResearch

About this paper

Alexander Hsu, Zhaiming Shen, Wenjing Liao, Rongjie Lai

Researchers at Purdue University and Georgia Institute of Technology constructed an end-to-end transformer architecture where ReLU-activated attention explicitly generates nonlinear features for polynomial and spline regression, then solves the in-context least squares problem. This framework provides generalization error bounds for nonlinear in-context learning, demonstrating dependence on context length and training set size.

Key insights

• ReLU-activated attention can serve as a powerful featurizer, capable of performing precise in-context arithmetic operations to construct complex nonlinear basis functions.
• The entire nonlinear regression process, from feature construction to coefficient estimation, can be unified within a single transformer architecture.
• Generalization error in nonlinear ICL depends predictably on context length (n) and pre-training set size (L), with specific architectural scaling requirements for depth and width.

Summary by alphaXiv.