Is your batch size the problem? Revisiting the Adam-SGD gap in language modeling
This work addresses the optimizer selection problem for language model practitioners, offering incremental insights into training dynamics.
The paper revisits the performance gap between Adam and SGD in language modeling, finding that SGD with momentum can match Adam in small-batch settings with proper tuning, and provides new insights into the role of batch size using stochastic differential equation models.
Adam is known to perform significantly better than Stochastic Gradient Descent (SGD) in language models, a phenomenon for which a number of explanations have been proposed. In this work, we revisit this "optimizer gap" through a series of comprehensively tuned baseline training runs for language modeling with Transformers. We exhaustively study how momentum, gradient clipping, and batch size affect the gap between SGD and Adam. Our empirical findings show that SGD with momentum can actually perform similarly to Adam in small-batch settings, if tuned correctly. We revisit existing explanations for Adam's advantage, including heavy-tailed class imbalance, directional sharpness, and Hessian heterogeneity, which struggle to directly explain this phenomenon. Towards bridging this gap in our understanding, by analyzing our Transformer training runs and simple quadratic settings inspired by the literature, we provide new insights, driven by stochastic differential equation models, into the role of batch size on the training dynamics.