In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness

TL;DR

The paper investigates why softmax-attention in transformers supports in-context learning (ICL) beyond prior linear-attention analyses. It shows that pretraining calibrates the attention window by learning a bandwidth parameter $w_{KQ}$ such that the attention acts like a Nadaraya–Watson estimator with bandwidth $w_{KQ}$, scaling with the Lipschitz constant $L$, label noise $\sigma^2$, and context size $n$; hence the window widens for smaller $L$ and larger noise and narrows with larger $n$. It further demonstrates directional adaptivity in low-rank settings: for function classes with nonzero Lipschitzness only in a subspace spanned by $\mathbf{B}$, the pretrained $\mathbf{M}$ aligns with $\mathbf{B}\mathbf{B}^\top$, effectively projecting inputs onto the informative subspace. The work shows softmax is essential for this adaptivity, as linear attention cannot replicate the same ICL performance, and provides generalization bounds linking downstream Lipschitzness to pretraining. Together, these results illuminate how softmax-attention captures task structure during pretraining and enables robust ICL across related downstream tasks.

Abstract

A striking property of transformers is their ability to perform in-context learning (ICL), a machine learning framework in which the learner is presented with a novel context during inference implicitly through some data, and tasked with making a prediction in that context. As such, that learner must adapt to the context without additional training. We explore the role of softmax attention in an ICL setting where each context encodes a regression task. We show that an attention unit learns a window that it uses to implement a nearest-neighbors predictor adapted to the landscape of the pretraining tasks. Specifically, we show that this window widens with decreasing Lipschitzness and increasing label noise in the pretraining tasks. We also show that on low-rank, linear problems, the attention unit learns to project onto the appropriate subspace before inference. Further, we show that this adaptivity relies crucially on the softmax activation and thus cannot be replicated by the linear activation often studied in prior theoretical analyses.

Figures (10)

• Figure 1: Top Row: The black line denotes the target function over a domain (horizontal axis). The gray dots are noisy training data, and the white dot is a query. From left to right, the Lipschitzness of the target function grows and the optimal softmax attention window (shaded blue) shrinks. Middle Row: Attention weights -- which determine the attention window -- as a function of the relative position from the query for softmax and linear attention. The softmax weights adjust to the Lipschitzness. Bottom Row: ICL error versus number of context samples for the three settings. Adapting to function Lipschitzness leads softmax attention to achieve small error. Please see Remark \ref{['rem:scaling']} and Appendix \ref{['appendix:sims']} for further discussion and details.
• Figure 2: From left to right, as we shrink the attention window (shaded in blue), the estimator has lower bias (the expected value of the estimate, depicted in purple, is closer to the ground-truth label, depicted by the white circle) but larger variance (shaded in tan).
• Figure 3: Spectral norm of $\mathbf{M}$ during pretraining with varying $L$. Each plot shows results for different task and covariate distributions, with (tasks, covariates) drawn from (Left) ($D(\mathcal{F}_L^+), \mathcal{U}^d$), (Middle-Left) ($D(\mathcal{F}_L^+), \tilde{\mathcal{U}}^d$), (Middle-Right) ($D(\mathcal{F}_L^{\cos}), \mathcal{U}^d$), (Right) ($D(\mathcal{F}_L^{\cos}), \tilde{\mathcal{U}}^d$), where $\tilde{\mathcal{U}}^d$ is a non-isotropic distribution on $\mathbb{S}^{d-1}$ (see Section \ref{['sec:general-exps']} for its definition).
• Figure 4: Spectral norm of $\mathbf{M}$ during pretraining on tasks drawn from $D(\mathcal{F}^+_{1})$ in Left, Middle-Right and $D(\mathcal{F}^{\cos}_{1})$ in Middle-Left, Right. Left, Middle-Left show ablations over the noise standard deviation $\sigma$ and Middle-Right, Right show ablations over the number of context samples $n$.
• Figure 5: Left, Middle-Left, Middle-Right: The test error for softmax attention as it is trained on the distributions over 1-Lipschitz affine, ReLU, and cosine function ($D(\mathcal{F}^{\mathop{\mathrm{\text{aff}}}\nolimits}_1)$, $D(\mathcal{F}^{+}_1)$, and $D(\mathcal{F}^{\cos}_1)$, respectively), where the error is evaluated at each pretraining iteration on 5 tasks drawn from the distributions over the 1-Lipschitz (affine, ReLU, cosine) function classes in (Left, Middle-Left, Middle-Right), respectively. Right: The test error evaluated on tasks drawn from $D(\mathcal{F}^{\cos}_1)$ for three softmax attention trained on tasks drawn from $D(\mathcal{F}^{\mathop{\mathrm{\text{aff}}}\nolimits}_1)$, $D(\mathcal{F}^{\cos}_{0.1})$, and $D(\mathcal{F}^{\cos}_{10})$, respectively.

Theorems & Definitions (71)

• Remark 2.1: Extreme cases
• Remark 2.2: Softmax advantage
• Definition 3.1: Lipschitzness
• Definition 3.2: Affine and ReLU Function Classes
• Theorem 3.4
• Theorem 3.5
• Theorem 3.6: Lower Bound for Linear Attention
• Definition 4.1: Direction-wise Lipschitzness of Function Class
• Definition 4.2: Low-rank Linear Function Class
• Theorem 4.4