# On the Optimization and Generalization of Two-layer Transformers with Sign Gradient Descent

arXiv: https://arxiv.org/abs/2410.04870 · HTML: https://mlsanity.com/paper/2410.04870/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 7 Oct 2024.

## Built on by (most cited first, 12 of 12)

- [Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks](https://mlsanity.com/paper/2506.02651/) (2025)
- [NIRVANA: Structured Pruning Reimagined for Large Language Model Compression](https://mlsanity.com/paper/2509.14230/) (2025)
- [LoTA-QAF: Lossless Ternary Adaptation for Quantization-Aware Fine-Tuning](https://mlsanity.com/paper/2505.18724/) (2025)
- [How Does Label Noise Gradient Descent Improve Generalization in the Low SNR Regime?](https://mlsanity.com/paper/2510.17526/) (2025)
- [Never Saddle for Reparameterized Steepest Descent as Mirror Flow](https://mlsanity.com/paper/2603.02064/) (2026)
- [Understanding Generalization in Transformers: Error Bounds and Training Dynamics Under Benign and Harmful Overfitting](https://mlsanity.com/paper/2502.12508/) (2025)
- [Unveiling m-Sharpness Through the Structure of Stochastic Gradient Noise](https://mlsanity.com/paper/2509.18001/) (2025)
- [Adaptive Methods Are Preferable in High Privacy Settings: An SDE Perspective](https://mlsanity.com/paper/2603.03226/) (2026)
- [Looped Transformers with Layer Normalization Provably Learn the Power Method](https://mlsanity.com/paper/2606.00605/) (2026)
- [Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent](https://mlsanity.com/paper/2605.06609/) (2026)
- [Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models](https://mlsanity.com/paper/2603.22801/) (2026)
- [From Parallel to Sequential: Small-Batch Noise Restores Simplicity Bias of Muon](https://mlsanity.com/paper/2610.16868/) (2026)

## Abstract (arXiv)

The Adam optimizer is widely used for transformer optimization in practice, which makes understanding the underlying optimization mechanisms an important problem. However, due to the Adam's complexity, theoretical analysis of how it optimizes transformers remains a challenging task. Fortunately, Sign Gradient Descent (SignGD) serves as an effective surrogate for Adam. Despite its simplicity, theoretical understanding of how SignGD optimizes transformers still lags behind. In this work, we study how SignGD optimizes a two-layer transformer -- consisting of a softmax attention layer with trainable query-key parameterization followed by a linear layer -- on a linearly separable noisy dataset. We identify four stages in the training dynamics, each exhibiting intriguing behaviors. Based on the training dynamics, we prove the fast convergence but poor generalization of the learned transformer on the noisy dataset. We also show that Adam behaves similarly to SignGD in terms of both optimization and generalization in this setting. Additionally, we find that the poor generalization of SignGD is not solely due to data noise, suggesting that both SignGD and Adam requires high-quality data for real-world tasks. Finally, experiments on synthetic and real-world datasets empirically support our theoretical results.