Understanding Why Adam Outperforms SGD: Gradient Heterogeneity in Transformers
This work addresses a fundamental optimization problem for researchers and practitioners training transformers, though it is incremental in explaining known performance differences.
The paper investigates why Adam outperforms SGD in optimizing transformers, attributing it to gradient heterogeneity—disparities in gradient norms among parameters—which hinders SGD but affects sign-based optimization (a simplified Adam variant) less. Experiments in NLP and vision domains validate that gradient heterogeneity is influenced by layer normalization placement.
Transformers are challenging to optimize with SGD and typically require adaptive optimizers such as Adam. However, the reasons behind the superior performance of Adam over SGD remain unclear. In this study, we investigate the optimization of transformers by focusing on gradient heterogeneity, defined as the disparity in gradient norms among parameters. Our analysis shows that gradient heterogeneity hinders gradient-based optimization, including SGD, while sign-based optimization, a simplified variant of Adam, is less affected. We further examine gradient heterogeneity in transformers and show that it is influenced by the placement of layer normalization. Experimental results from fine-tuning transformers in both NLP and vision domains validate our theoretical analyses. This study provides insights into the optimization challenges of transformers and offers guidance for designing future optimization algorithms. Code is available at https://github.com/tom4649/gradient-heterogeneity.