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

TL;DR

This paper analyzes how Sign Gradient Descent optimizes two-layer transformers, revealing four training stages, convergence properties, and generalization limitations, with insights applicable to Adam and real-world data.

Contribution

It provides a theoretical analysis of SignGD's optimization dynamics and generalization behavior on transformers, connecting it to Adam and highlighting data quality importance.

Findings

SignGD exhibits four distinct training stages.

The transformer converges quickly but generalizes poorly on noisy data.

Experiments support the theoretical analysis and real-world relevance.

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.