Research map / Architectures and efficiency / Attention and transformers
In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows theoretically and empirically that softmax attention pretrained on regression ICL tasks learns an attention-window scale tied to task Lipschitzness and noise, and a projection onto the low-rank subspace of linear tasks, which linear attention cannot do.
Problem
How softmax attention learns during pretraining to perform in-context learning on regression tasks, i.e., what shared structure of pretraining tasks it recovers that enables ICL on downstream tasks.
Approach
Views pretrained softmax attention as a Nadaraya-Watson nearest-neighbor regressor whose bandwidth (attention window) is meta-learned; proves that the optimal window scale grows inversely with task Lipschitzness and with noise, and that for low-rank linear tasks the optimal key-query matrix projects onto the task subspace, while linear attention provably fails.
Key results
- Pretrained softmax attention's key-query norm ||M|| increases with task Lipschitzness L, decreases with context size n, and increases more slowly with noise sigma, matching the theoretical bounds (d=5, L=1, sigma=0.01, n=20).
- Pretraining on 1-Lipschitz affine tasks gives downstream ICL performance on cosine tasks similar to pretraining on cosine tasks of the same Lipschitzness, while mismatched Lipschitzness (0.1 or 10) degrades performance sharply.
- Lower bound: linear attention's pretraining loss is Omega(L^2) for ReLU tasks, exceeding the upper bound on the softmax loss up to factors in d.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Pretrain softmax attention on tasks with Lipschitzness L | Softmax attention · D(Faff_L), D(F+_L) | attention scale |wKQ|: Case 1: Ω(Λ^α) ≤ | w_{KQ} | ≤ O(Λ^{\frac{2α}{1-β} }); Case 2: Ω(Λ^β) ≤ |w_{KQ}|≤ O(Λ^{2β}) | Helped |
| Pretrain on tasks with Lipschitzness matching the downstream task | Softmax attention · Downstream L-Lipschitz tasks | ICL loss: \uL ≤ \uO(\frac{L^2}{\Lambda^{\beta}}) | Helped |
| Use downstream tasks with higher Lipschitzness than pretraining tasks | Softmax attention pretrained on D(F+_L) · D(F+_L') with L' > L | ICL loss: \uL\ge\Omega(\frac{L'^2}{\Lambda^{2\beta}}) | Hurt |
| Use downstream tasks with lower Lipschitzness than pretraining tasks | Softmax attention pretrained on D(F+_L) · D(F+_L') with L' < L | ICL loss: \uL\ge\Omega(\frac{\Lambda^{\beta d/2}}{n}) | Hurt |
| Replace softmax activation with identity | Attention for tasks from D(F+_L) · D(F+_L), covariates from U^d | pretraining loss: \mathcal{L}_{LA}(\mathbf{M}) = \Omega({L^2}) | Hurt |
| Increase task Lipschitzness L (ablation) | Softmax attention pretrained on D(F+_L) or D(Fcos_L) · Lipschitzness ablation | attention matrix spectral norm ||M|| | Helped |
| Increase noise standard deviation σ (ablation) | Softmax attention · Noise ablation | attention matrix spectral norm ||M|| | Hurt |
| Increase number of context samples n (ablation) | Softmax attention · Context-size ablation | attention matrix spectral norm ||M|| | Helped |
| Pretrain on a different function class with the same Lipschitzness (ablation) | Softmax attention · 1-Lipschitz affine, ReLU, and cosine tasks | test ICL error | No clear effect |
| Pretrain on cosine tasks with Lipschitzness 0.1 or 10 instead of 1 (ablation) | Softmax attention evaluated on D(Fcos_1) · D(Fcos_1) test tasks | test ICL error | Hurt |
| Train on low-rank affine, quadratic, or cosine tasks | Softmax attention · Low-rank affine, quadratic, and cosine tasks | test error | Helped |
| Replace softmax with linear attention | Attention on low-rank affine, quadratic, and cosine tasks · Low-rank affine, quadratic, and cosine tasks | recovery of low-rank structure | Hurt |
Surprising findings
- Larger task Lipschitzness yields larger pretrained key-query norm even though gradient norms with respect to the key/query matrices are not obviously larger, which the authors note is surprising.
Limitations
- Analysis covers only a single softmax attention layer, not multi-layer transformers.
- Does not model the autoregressive next-token pretraining of LLMs; the model is an abstraction of LLM few-shot priming.
- Theory for direction recovery requires either sigma=0 or n=2.
Experimental setup
Datasets: synthetic regression tasks (affine, ReLU, cosine, quadratic GLM function classes)
Benchmarks: synthetic ICL regression tasks
Metrics: ICL test squared error attention window scale ||M|| subspace recovery metric rho(M,B)
Baselines: linear attention
Builds on: linear attention ICL analyses Nadaraya-Watson estimator
Closest papers by meaning
- The Closeness of In-Context Learning and Weight Shifting for Softmax Regression (NeurIPS 2024 · similarity 0.92)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.91)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.91)
- How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? (ICLR 2024 · similarity 0.91)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.91)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.