In-Context Learning in Linear vs. Quadratic Attention Models: An Empirical Study on Regression Tasks

Ayush Goel; Arjun Kohli; Sarvagya Somvanshi

cs.LGFeb 19, 2026

In-Context Learning in Linear vs. Quadratic Attention Models: An Empirical Study on Regression Tasks

Authors: Ayush Goel, Arjun Kohli, Sarvagya Somvanshi

Organizations: University of California, Berkeley

Abstract

Recent work has demonstrated that transformers and linear attention models can perform in-context learning (ICL) on simple function classes, such as linear regression. In this paper, we empirically study how these two attention mechanisms differ in their ICL behavior on the canonical linear-regression task of Garg et al. We evaluate learning quality (MSE), convergence, and generalization behavior of each architecture. We also analyze how increasing model depth affects ICL performance. Our results illustrate both the similarities and limitations of linear attention relative to quadratic attention in this setting.

Figures & tables

Config	Params	Train Loss	Test Loss	Anisotropic Test Loss
1-layer	13.7 M	0.7976 ± 0.0543	0.8183 ± 0.0688	0.8446 ± 0.0656
3-layer	15.2 M	0.0502 ± 0.0063	0.0894 ± 0.0097	0.0936 ± 0.0112
6-layer	17.6 M	0.0203 ± 0.0016	0.0365 ± 0.0041	0.0398 ± 0.0059

Table 1: Quadratic Attention Transformer performance

Panel of Figure 1 — Figure 1: Training and Testing loss. Top row shows shallower networks (1 and 3 layers), while the bottom row shows the deeper 6-layer model.

Config	Params	Train Loss	Test Loss	Anisotropic Test Loss
1-layer	0.8 M	0.5946 ± 0.0483	0.7430 ± 0.0671	0.7867 ± 0.0488
3-layer	2.38 M	0.0366 ± 0.0021	0.0959 ± 0.0026	0.1080 ± 0.0035
6-layer	4.75 M	0.0102 ± 0.0012	0.0302 ± 0.0034	0.0328 ± 0.0030

Table 2: Linear Attention Transformer performance

Panel of Figure 2 — Figure 2: Training and Testing loss for the Linear Attention Transformer

Panel of Figure 3 — Figure 3: Isotropic vs. anisotropic performance across depths.

Model	Batch Size	Steps to 90%	Samples to 90%
Quadratic 1 Layer	32	7000	224,000
Quadratic 3 Layer	32	21,500	688,000
Quadratic 6 Layer	32	25,000	800,000
Linear 1 Layer	64	4000	256,000
Linear 3 Layer	64	4000	256,000
Linear 6 Layer	64	7500	480,000

Table 3: Steps and equivalent samples seen to reach 90% convergence for each architecture and depth.

Explore similar work

Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer
May 6, 2026Alexander Hsu, Zhaiming Shen, Wenjing Liao +1Transformer Attention In-Context Learning
Department of Mathematics, Purdue University · School of Mathematics, Georgia Institute of Technology
What can linear attention learn from nonlinear teachers in-context?
Oct 7, 2026Mary Letey, Arman Rysmakhanov, Yue M. Lu +2In-Context Learning Linear Attention
Applied Mathematics, Harvard University Kempner Institute, Harvard University · Williams College · Center for Brain Science, Harvard University Society of Fellows, Harvard University
In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention
Jul 17, 2026Katsuyuki HagiwaraIn-Context Learning Linear Regression
Faculty of Education, Mie University · Kurima-Machiya-cho, Tsu, 514-8507, Japan

May 6, 2026J/K move · Enter open · S save

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

Alexander Hsu, Zhaiming Shen, Wenjing Liao, Rongjie Lai

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

Transformer Attention In-Context Learning Statistical Learning Theory

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

What can linear attention learn from nonlinear teachers in-context?

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