Research map / Optimization / Implicit bias and generalization
Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression
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 that Transformers trained on in-context linear regression converge at a rate matching Iterative Newton's method (second-order) rather than gradient descent, and constructs a Transformer implementing k Newton iterations with k+O(1) layers.
Problem
How Transformers perform in-context learning for linear regression: whether they implicitly run first-order gradient descent or a faster second-order method, and what convergence rate their layers achieve.
Approach
Compares layer-wise Transformer predictions and induced weights against iterative algorithms (GD, Iterative Newton, OGD, BFGS, CG) via best-matching steps; finds a linear layer-to-Newton-iteration correspondence and exponential mismatch with GD, and proves one attention+MLP layer can compute one Newton iteration with O(d) hidden size.
Key results
- Transformer layers 3 to 9 match Iterative Newton iterations linearly (each layer ~3 Newton iterations) while GD matches exponentially, indicating a convergence rate of O(log log(1/eps)) for Transformers.
- On ill-conditioned data (kappa=100), the Transformer still closely matches Iterative Newton with 21 iterations, whereas GD would need about 2,000 steps.
- Theorem: Transformers can implement k Newton iterations with k+8 layers and O(d) hidden dimension; the Transformer error similarity with Newton is 0.991 vs 0.957 with GD (LSTM: 0.920 Newton, 0.916 GD).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train a Transformer for in-context linear regression | GPT-2 · In-context linear regression | convergence rate relative to Iterative Newton | No clear effect |
| Train a Transformer for in-context linear regression | GPT-2 · In-context linear regression | convergence rate relative to Gradient Descent | Helped |
| Use a Transformer instead of an LSTM | In-context linear regression · In-context linear regression | error similarity to OGD: 0.954 (LSTM) vs 0.806 (Transformers) | Mixed |
| Use a Transformer instead of an LSTM | In-context linear regression · In-context linear regression | mean-squared error | Helped |
| Increase hidden size to 32 or 64 (ablation) | 12-layer, 1-head Transformer on d=20 · In-context linear regression | ability to mimic OLS | Helped |
| Reduce hidden size below 32 (ablation) | 12-layer, 1-head Transformer on d=20 · In-context linear regression | ability to mimic OLS | Hurt |
| Use ill-conditioned covariance with condition number 100 (ablation) | Transformer · Ill-conditioned linear regression | convergence relative to Iterative Newton | No clear effect |
| Use ill-conditioned covariance with condition number 100 (ablation) | Gradient Descent · Ill-conditioned linear regression | steps required for convergence: 2,000 steps | Hurt |
| Use one attention head per layer instead of eight (ablation) | 12-layer GPT-2 Transformer · In-context linear regression | convergence-rate similarity to Iterative Newton | No clear effect |
| Increase depth to 24 layers (ablation) | GPT-2 Transformer · In-context linear regression | convergence rate | No clear effect |
| Use a noisy regression target with σ = 0.1 (ablation) | Transformer · Noisy linear regression | convergence rate | No clear effect |
| Use a 2-layer ReLU neural-network function class instead of linear regression (ablation) | Transformer · 2-layer neural-network regression with ReLU | convergence rate relative to Gradient Descent | Helped |
Surprising findings
- LSTMs do not improve in predictions across layers and behave more like online gradient descent than Transformers.
- Transformers' similarity to GD in cosine terms is comparable at later layers, yet their convergence correspondence is exponential rather than linear.
Limitations
- Evidence is for a specific synthetic setup (Gaussian inputs, linear regression, d=20, n up to 40) and the paper does not claim Transformers implement Newton exactly.
- Matching relies on best-matching steps and cosine similarity; the rate claims are inferred from correspondence trends rather than a direct convergence proof for trained models.
- Theory uses full attention and ReLU activations rather than the causal softmax attention used in experiments.
Experimental setup
Datasets: synthetic linear regression prompts (Gaussian x, w*)
Benchmarks: synthetic linear regression (isotropic and ill-conditioned covariance, noisy, 2-layer MLP)
Metrics: similarity of errors (cosine) similarity of induced weights (cosine) normalized MSE best-matching steps
Baselines: Gradient Descent Online Gradient Descent Iterative Newton BFGS L-BFGS Conjugate Gradient LSTM OLS
Builds on: Akyürek et al. 2022 (what learning algorithm do transformers implement) von Oswald et al. 2023 (transformers learn GD in context) Bai et al. 2023 (transformers as algorithms) Giannou et al. 2023 (looped transformers for Newton)
Closest papers by meaning
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.93)
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.92)
- In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization (NeurIPS 2024 · similarity 0.91)
- On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery (ICLR 2025 · similarity 0.91)
- Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context (NeurIPS 2024 · similarity 0.91)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.89)
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.