Research map / Architectures and efficiency / Attention and transformers
Linear attention is (maybe) all you need (to understand Transformer optimization)
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 a shallow linear Transformer trained on random linear regression reproduces the known distinctive features of full Transformer optimization (Adam>SGD gap, heavy-tailed gradient noise, ill-conditioned robust condition number, directional smoothness, generalized smoothness) and uses it to probe effects of covariate tails and depth.
Problem
Transformer training is notoriously difficult and full-scale Transformers are too complex to analyze, so the paper asks whether a simpler model can capture the distinctive features of Transformer optimization.
Approach
Replace full Transformers with a linearized multi-layer linear self-attention model (no softmax, no feedforward layers, single key-query matrix Q) trained on random linear regression, and show it reproduces the known Transformer optimization phenomena.
Key results
- Linear 3-layer Transformers trained with SGDM and Adam reproduce the Adam>SGD gap, heavy-tailed gradient noise, lower robust condition number for Adam, and directional smoothness trends seen in full Transformers.
- Heavier-tailed covariates (Gamma k=0.1, theta=10) make the stochastic gradient noise more heavy-tailed and enlarge the robust condition number gap between SGD and Adam.
- Increasing depth L from 2 to 8 makes the Adam-SGD loss gap more pronounced and the gradient noise more heavy-tailed.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Adam instead of SGD with momentum | 3-layer linear Transformer · Random linear regression, Settings 1–3 | training loss / optimization speed | Helped |
| Use spherical (light-tailed) covariates instead of Gamma-scaled heavy-tailed covariates (ablation) | 3-layer linear Transformer · Random linear regression | optimization speed | Helped |
| Use spherical (light-tailed) covariates instead of Gamma-scaled heavy-tailed covariates (ablation) | 3-layer linear Transformer · Random linear regression | relative optimization-speed gap between Adam and SGD | No clear effect |
| Increase Transformer depth from 2 layers to 4, 6, or 8 layers (ablation) | Linear Transformer · Random linear regression, Setting 1 | Adam–SGD training-loss gap | Mixed |
| Increase Transformer depth (ablation) | Linear Transformer trained with SGD or Adam · Random linear regression, Setting 1 | absolute training loss | Helped |
| Increase Transformer depth from 2 layers to 4, 6, or 8 layers (ablation) | Linear Transformer · Random linear regression, Setting 1 | SGD–Adam robust-condition-number gap | Mixed |
Surprising findings
- The relative Adam-over-SGD speedup does not noticeably improve under light-tailed covariates, suggesting heavy-tailed gradient noise alone does not explain Adam's advantage.
Limitations
- Lacks a solid theoretical foundation; findings are empirical.
- Linear regression setting may not fully capture features of language data.
- Only shallow/small settings (d=5, L up to 8) studied.
Experimental setup
Datasets: synthetic random linear regression (Gaussian covariates) heavy-tailed Gamma-scaled spherical covariates MLP-distorted nonlinear regression
Benchmarks: synthetic linear regression training loss (Settings 1-3)
Metrics: training loss (log loss) stochastic gradient noise tail (q-q plot) robust condition number directional smoothness generalized smoothness
Baselines: SGD with momentum Adam
Builds on: linear Transformer in-context learning (von Oswald et al. 2023) Ahn et al. 2023 linear attention Adam vs SGD for attention models (L0,L1)-smoothness directional smoothness
Closest papers by meaning
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.91)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.89)
- Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency (NeurIPS 2025 · similarity 0.88)
- Taming Transformer Without Using Learning Rate Warmup (ICLR 2025 · similarity 0.88)
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.86)
- Unveil Benign Overfitting for Transformer in Vision: Training Dynamics, Convergence, and Generalization (NeurIPS 2024 · similarity 0.86)
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.