Research map / Optimization / Implicit bias and generalization
Unraveling the Gradient Descent Dynamics of Transformers
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.
Analyzes gradient descent on a single multi-head Transformer layer with Softmax or Gaussian attention, proving global convergence under large embedding dimension and careful initialization, and showing Softmax can get stuck at non-optimal stationary points while Gaussian kernel attention converges.
Problem
Which Transformer architectures and initial conditions let gradient descent provably reach global optima of training loss, given the non-convexity introduced by softmax attention.
Approach
Derives GD convergence for a one-layer multi-head attention regression model (no layer norm, no MLP) with Softmax or Gaussian kernels; shows that updating only W^V gives linear convergence, and that with only W^Q updated, Gaussian kernel attention keeps a PL-type guarantee while Softmax admits vanishing-gradient stationary points.
Key results
- Under large embedding dimension and proper initialization, GD on W^V alone converges linearly to a global optimum with rate (1-eta*alpha)^t, where alpha = ||W^O||^2 (lambda^B)^2.
- With only W^Q updated, the Gaussian kernel admits a linear convergence guarantee, while for Softmax the authors give an example where the gradient vanishes at a non-optimal initialization.
- On IMDb and Pathfinder, Gaussian-kernel Transformers show faster convergence and higher test accuracy than Softmax ones; worst Pathfinder test accuracy after 20,000 epochs is about 0.58 (Softmax) vs 0.62 (Gaussian).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Gaussian kernel attention instead of Softmax attention | 2-layer Transformer · IMDb review text classification and Pathfinder | test accuracy | Helped |
| Use Gaussian kernel attention instead of Softmax attention | 2-layer Transformer · Pathfinder test, after 20,000 epochs | test accuracy: around 0.58 -> 0.62 | Helped |
| Use Softmax attention instead of Gaussian kernel attention | 2-layer Transformer · Pathfinder | test accuracy variance | Hurt |
| Use Softmax attention instead of Gaussian kernel attention | 2-layer Transformer · IMDb review text classification and Pathfinder | convergence speed | Hurt |
| Use Softmax attention instead of Gaussian kernel attention | 2-layer Transformer · IMDb review text classification and Pathfinder | optimization landscape complexity | Hurt |
Surprising findings
- Softmax attention can produce vanishing gradients at a non-optimal initialization even when the initialization satisfies the convergence conditions, so GD does not reach a global optimum.
Limitations
- Analysis restricted to one layer without layer normalization, MLP, or mean pooling; the MLP extension is only a corollary.
- Requires large embedding size (HD >= Nn or Dd >= Nn^2) and strong initialization conditions.
- Softmax failure is shown via a specific constructed example rather than generically.
Experimental setup
Datasets: IMDb Pathfinder
Benchmarks: IMDb Pathfinder
Metrics: test accuracy test loss loss landscape
Baselines: Softmax attention
Builds on: Softmax attention Transformer over-parameterized neural network convergence analysis
Closest papers by meaning
- Linear attention is (maybe) all you need (to understand Transformer optimization) (ICLR 2024 · similarity 0.91)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.90)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.90)
- What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis (ICLR 2025 · similarity 0.89)
- Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency (NeurIPS 2025 · similarity 0.89)
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (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.