Research map / Optimization / Implicit bias and generalization
Linear Transformers are Versatile In-Context Learners
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 linear transformers maintain an implicit linear regression model per layer and, on mixed-noise-variance regression, learn a momentum- and noise-adaptive optimizer that matches or beats ridge-regression baselines.
Problem
In-context learning with linear transformers on linear regression, especially the harder mixed noise-variance setting where the noise level varies across sequences and must be inferred from context.
Approach
Proves each layer of a linear transformer maintains a latent linear regression weight vector (interpreted as preconditioned GD with momentum-like terms), shows GD++ is second-order, and finds that trained linear transformers (full or diagonal) learn noise-adaptive step sizes and rescaling that outperform closed-form ridge baselines under mixed noise.
Key results
- On uniform noise sigma~U(0,sigma_max), 5+ layer linear transformers match or exceed TunedRR; e.g., at sigma_max=7 7-layer Full reaches 0.057 vs TunedRR 0.068.
- With 4 layers, linear transformers outperform AdaRR for sigma_max>=4.
- GD++ plateaus near ConstRR (e.g., 0.530 at sigma_max=7) and fails to adapt to mixed noise, while Diag and Full models keep improving with depth.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train a 7-layer Full linear transformer instead of using TunedRR | noisy linear regression · Uniform σ_τ∼U(0,5) | adjusted evaluation loss: 0.049 -> 0.047 | Helped |
| Use 7-layer Full instead of GD++ | noisy linear regression · Uniform σ_τ∼U(0,5) | adjusted evaluation loss: 0.344 -> 0.047 | Helped |
| Constrain attention matrices to be diagonal (ablation) | 7-layer linear transformer · Uniform σ_τ∼U(0,5) | adjusted evaluation loss: 0.047 -> 0.047 | No clear effect |
| Use 2 layers instead of 7 (ablation) | Full linear transformer · Uniform σ_τ∼U(0,5) | adjusted evaluation loss: 0.047 -> 0.245 | Hurt |
| Use 2 layers instead of 7 (ablation) | Diag linear transformer · Uniform σ_τ∼U(0,5) | adjusted evaluation loss: 0.047 -> 0.242 | Hurt |
| Use 4 layers instead of 7 (ablation) | Full linear transformer · Uniform σ_τ∼U(0,5) | adjusted evaluation loss: 0.047 -> 0.065 | Hurt |
| Use 4 layers instead of 7 (ablation) | Diag linear transformer · Uniform σ_τ∼U(0,5) | adjusted evaluation loss: 0.047 -> 0.059 | Hurt |
| Use Full instead of Diag attention matrices (ablation) | 7-layer linear transformer · Categorical σ_τ∈{1,3} | adjusted evaluation loss: 0.018 -> 0.010 | Helped |
| Use Full instead of Diag attention matrices (ablation) | 7-layer linear transformer · Categorical σ_τ∈{1,3,5} | adjusted evaluation loss: 0.046 -> 0.035 | Helped |
| Train a 7-layer Full linear transformer instead of using TunedRR | noisy linear regression · Categorical σ_τ∈{1,3} | adjusted evaluation loss: 0.021 -> 0.010 | Helped |
| Train a 7-layer Full linear transformer instead of using TunedRR | noisy linear regression · Categorical σ_τ∈{1,3,5} | adjusted evaluation loss: 0.054 -> 0.035 | Helped |
| Train a 7-layer Diag linear transformer instead of using TunedRR | noisy linear regression · Uniform σ_τ∼U(0,5) | adjusted evaluation loss: 0.049 -> 0.047 | Helped |
Surprising findings
- Diagonal-parameterized linear transformers perform comparably to full-matrix ones on mixed-noise regression.
- Trained linear transformers outperform tuned closed-form ridge regression baselines despite lacking matrix inversion.
Limitations
- Focus on simplified linear models with no MLP or LayerNorm and mostly diagonal attention matrices.
- Unclear whether discovered algorithm is optimal; Full and Diag may differ in ways not fully characterized.
- Only tested on synthetic regression; generalization to complex function classes and larger datasets is not shown.
Experimental setup
Datasets: synthetic noisy linear regression sequences
Benchmarks: synthetic noisy linear regression (uniform and categorical noise variance)
Metrics: adjusted evaluation loss
Baselines: GD++ ConstRR AdaRR TunedRR
Builds on: GD++ Transformers learn in-context by gradient descent Ridge regression
Closest papers by meaning
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.93)
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.93)
- How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression (NeurIPS 2024 · similarity 0.90)
- One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention (ICLR 2024 · similarity 0.90)
- In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization (NeurIPS 2024 · similarity 0.90)
- Unraveling the Gradient Descent Dynamics of Transformers (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.