Research map / Optimization / Implicit bias and generalization
On the Optimization and Generalization of Two-layer Transformers with Sign Gradient Descent
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 the training dynamics of SignGD (a proxy for Adam) on a two-layer transformer with softmax query-key attention and linear head on a linearly separable noisy dataset, showing fast convergence but poor generalization.
Problem
How sign gradient descent (a surrogate for Adam) optimizes and generalizes when training a two-layer transformer with trainable query-key attention on noisy, linearly separable data.
Approach
Characterizes the full SignGD training dynamics of a softmax attention layer with trainable query-key parameters in four stages, proving linear convergence of training loss but constant test loss ("destructive convergence"), and shows Adam behaves similarly.
Key results
- SignGD reaches training loss below epsilon in T = O(log(1/epsilon)) steps (linear convergence) while test loss remains Theta(1), i.e., fast convergence but poor generalization.
- Adam follows similar four-stage dynamics and shows high test loss on noisy MNIST, especially under strong noise, while GD generalizes better across noise levels.
- Noise-signal softmax outputs decay exponentially fast in Stage IV, making attention concentrate on the noise patch (noise memorization).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train the two-layer transformer with SignGD | Two-layer transformer · Noisy synthetic dataset | training loss: L_S(W^(T)) ≤ ε | Helped |
| Train the two-layer transformer with SignGD on noisy data | Two-layer transformer · Noisy synthetic dataset | test loss: L_D(W^(T)) = Θ(1) | Hurt |
| Use Adam instead of GD | Two-layer transformer · Noisy MNIST, classes 3 and 7 | test loss | Hurt |
| Use SignGD instead of GD | Two-layer transformer · Synthetic data, data setting (d) | softmax output convergence speed | Helped |
| Use GD instead of SignGD | Two-layer transformer · Noisy data at different noise levels | test loss | Helped |
| Use Adam with β1 and β2 close to 1 | Two-layer transformer · Synthetic data, Stage IV | query-noise sign alignment | Hurt |
| Increase context length to L = 10 (ablation) | Two-layer transformer trained with SignGD · Synthetic data, context length L = 10 | attention concentration on noise patches | No clear effect |
| Add four attention heads (ablation) | Simplified transformer trained with SignGD · Synthetic data, four-head model | attention dynamics | No clear effect |
| Add residual connections to deeper transformers (ablation) | Deeper transformer trained with SignGD · Synthetic data, deeper transformer with residual connections | theoretical stage behaviors | Helped |
| Remove residual connections from deeper transformers (ablation) | Deeper transformer trained with SignGD · Synthetic data, deeper transformer without residual connections | theoretical stage behaviors | Hurt |
| Add MLP layers to transformers with residual connections (ablation) | Deeper transformer trained with SignGD · Synthetic data, deeper transformer | completion of Stage IV sign alignment | Hurt |
| Use non-orthogonal and/or non-sparse data instead of the default data setting (ablation) | Two-layer transformer trained with SignGD · Synthetic data, non-sparse and/or non-orthogonal variants | agreement with theoretical dynamics | No clear effect |
Surprising findings
- Fast training does not imply better generalization: SignGD converges fast yet generalizes poorly, contradicting the 'train faster, generalize better' argument from algorithmic stability.
- Poor generalization of SignGD and Adam is not solely due to data noise; GD generalizes better under the same noise levels.
Limitations
- Analysis restricted to linearly separable data, sparse and orthogonal signal/noise, context length 2, single-head attention, and a fixed linear head.
- Adam with beta1 close to 1 does not always match SignGD; in Stage IV Adam did not show query-noise sign alignment, attributed to momentum.
- Language data and deeper or multi-head transformers are not covered by the theory.
Experimental setup
Datasets: synthetic signal-noise data (Definition 1) MNIST (classes 3 and 7, noised)
Benchmarks: MNIST (noisy, classes 3 and 7)
Metrics: training loss test loss softmax output dynamics 0-1 test loss
Baselines: gradient descent Adam
Builds on: feature learning framework Adam
Closest papers by meaning
- Exact risk curves of signSGD in High-Dimensions: quantifying preconditioning and noise-compression effects (ICML 2025 · similarity 0.90)
- In Search of Adam’s Secret Sauce (NeurIPS 2025 · similarity 0.88)
- Adaptive Methods through the Lens of SDEs: Theoretical Insights on the Role of Noise (ICLR 2025 · similarity 0.88)
- The Hidden Ratio in Adam: Stable Structure, Compression, and Sign Dynamics (NeurIPS 2026 · similarity 0.88)
- Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold (NeurIPS 2025 · similarity 0.87)
- The Rich and the Simple: On the Implicit Bias of Adam and SGD (NeurIPS 2025 · similarity 0.87)
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.