Reproduce “Revisiting the Adam-SGD Gap in LLM Pre-Training: The Role of Large Effective Learning Rates” on OpenResearch

About this paper

Athanasios Glentis, Dawei Li, Chung-Yiu Yau, Mingyi Hong

The authors investigate why SGD trails Adam in pre-training and test targeted clipping to stabilize larger learning rates. They link SGD’s small updates to small stochastic gradients and identify layer-wise spikes and uneven output-class gradients as constraints. On 1B-parameter LLaMA trained on C4, SGD-LL lowers validation loss from 3.95 to 2.60, near Adam’s 2.51, but leaves a gap.

Key insights

• Adam's 'effective momentum learning rates' are consistently 2-4 orders of magnitude larger than SGD's, resulting in significantly larger update norms and parameter space exploration.
• LLM pre-training exhibits unusually high weight-to-stochastic-gradient norm ratios (10^3 to 10^5), necessitating large effective learning rates for meaningful parameter updates, a phenomenon exacerbated by large batch sizes.
• Gradient irregularities, specifically layer-wise spikes and output layer per-token-class disparity (scaling with token frequency), are the primary factors preventing standard SGD from stably utilizing larger learning rates.

Summary by alphaXiv.