# Why Transformers Need Adam: A Hessian Perspective

**Authors:** Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, Zhi-Quan Luo  
**arXiv:** [2402.16788](https://arxiv.org/abs/2402.16788)  
**Date:** 2024-02-26  
**Summary page:** https://www.summarizepaper.com/en/paper/2402.16788v1/

## Key points

- The authors explore the performance difference between Stochastic Gradient Descent (SGD) and Adam optimization algorithms on Transformers
- Transformers exhibit a unique characteristic called "block heterogeneity," where the Hessian spectrum varies significantly across parameter blocks
- SGD struggles to navigate problems with block heterogeneity due to its uniform learning rate application across all parameter blocks
- Adam's adaptive learning rates tailored to individual blocks enable it to effectively address block heterogeneity and outperform SGD on Transformer tasks
- Experimental results show that Adam consistently outperforms SGD on Transformer-based tasks, while both algorithms yield comparable results for CNN-based tasks
- The study emphasizes the critical role of adaptive learning rates in addressing block heterogeneity challenges in Transformer models

## Summary

In their work titled "Why Transformers Need Adam: A Hessian Perspective," Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo delve into the performance discrepancy between Stochastic Gradient Descent (SGD) and Adam optimization algorithms on Transformers. The authors provide an insightful explanation for SGD's inferior performance on Transformers by examining the concept of block heterogeneity in the context of the Hessian matrix. They highlight that Transformers exhibit a unique characteristic termed "block heterogeneity," where the Hessian spectrum varies significantly across parameter blocks. This phenomenon poses a challenge for SGD optimization as it struggles to effectively navigate problems with such block heterogeneity. To validate this assertion, the team conducts experiments on various models including Transformers, Convolutional Neural Networks (CNNs), Multi-Layer Perceptrons (MLPs), and quadratic problems. Their findings indicate that while SGD performs well on tasks without block heterogeneity, it falters when faced with heterogeneous blocks. Through theoretical analysis, the authors propose that SGD's failure can be attributed to its uniform application of a single learning rate across all parameter blocks. This approach proves inadequate in handling the diverse characteristics of different blocks within Transformers. In contrast, Adam's adaptive learning rates tailored to individual blocks enable it to effectively address block heterogeneity and outperform SGD on Transformer tasks. Expanding on their research, the team applies Second-Order Langevin Dynamics (SLQ) to examine various CNNs and Transformer models across different datasets. They explore models such as ResNet18 and VGG16 for CNNs on ImageNet, as well as ViT-base, BERT, GPT2, and GPT2-nano2 for Transformers on various corpora. Experimental results demonstrate that while Adam consistently outperforms SGD on Transformer-based tasks due to block heterogeneity issues, both optimization algorithms yield comparable results for CNN-based tasks. Furthermore, the researchers analyze the Hessian spectrum of both full Hessian matrices and individual parameter blocks within these models. By dissecting parameters based on PyTorch default partitions like Embedding layers and attention components in Transformers, they gain deeper insights into how block heterogeneity impacts optimization performance. In conclusion, this study sheds light on the critical role of adaptive learning rates in addressing block heterogeneity challenges in Transformer models. By highlighting the limitations of SGD and showcasing Adam's effectiveness in handling heterogeneous blocks through tailored learning rates, the research underscores the importance of algorithm selection in optimizing complex neural network architectures like Transformers.