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

Bingrui Li, Wei Huang, Andi Han, Zhanpeng Zhou, Taiji Suzuki, Jun Zhu, Jianfei Chen

Keywords: Sign Gradient Descent; Transformer; Training Dynamics; Theory

Abstract:

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.

Primary Area: learning theory

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 10289