Research map / Optimization / Implicit bias and generalization
One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention
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.
Proves that the global minimizer of the pre-training loss for a one-layer linear self-attention transformer on noisy linear regression implements one step of GD (preconditioned GD under non-isotropic covariates, and still GD on a linear objective under nonlinear targets).
Problem
Understanding theoretically which algorithm a one-layer linear self-attention transformer learns for in-context linear regression, since prior work only observed this empirically.
Approach
Shows via a rotation-invariance and gradient-matching argument that the pre-training loss's global minimizer over (W_K, W_Q, W_V, h) reduces to matching the effective linear predictor to eta * X^T y, i.e., one GD step; extends to preconditioned GD for non-isotropic covariates and to nonlinear rotation-invariant target families.
Key results
- Global minimizer of the pre-training loss for isotropic Gaussian covariates implements one step of GD on least squares with learning rate eta given in closed form.
- With covariance Sigma and w ~ N(0, Sigma^{-1}), the global minimizer implements one step of GD preconditioned by Sigma^{-1}.
- With nonlinear rotation-invariant, sign-symmetric target functions, the global minimizer still implements one step of GD on a linear regression objective.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train on isotropic Gaussian covariates for noisy linear regression | One-layer linear self-attention · Noisy linear regression, isotropic Gaussian covariates | learned algorithm | Helped |
| Change covariates to a non-isotropic Gaussian with covariance Σ and weights distributed as N(0, Σ⁻¹) | One-layer linear self-attention · Noisy linear regression, non-isotropic Gaussian covariates | learned algorithm | Mixed |
| Change the response-generating function to a rotation-invariant, sign-symmetric nonlinear target | One-layer linear self-attention · Nonlinear target functions with isotropic Gaussian covariates | learned algorithm | No clear effect |
Surprising findings
- Changing the response distribution to nonlinear targets does not change the learned algorithm, whereas changing covariate covariance does.
Limitations
- Analysis restricted to a single-head, one-layer linear self-attention model without MLP or softmax.
- Learning rate eta is a global-minimizer value that depends on the data distribution, not a general fixed rate.
- Only global minima of pre-training loss are characterized; training dynamics are not analyzed.
Experimental setup
Builds on: Akyurek et al. 2023 (transformers learn ridge regression in context) von Oswald et al. 2023 (transformers learn GD in context) Ahn et al. 2023 (linear transformers implement preconditioned GD)
Closest papers by meaning
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.90)
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.89)
- How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression (NeurIPS 2024 · similarity 0.89)
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.89)
- In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization (NeurIPS 2024 · similarity 0.88)
- How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? (ICLR 2024 · similarity 0.88)
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.