Is your batch size the problem? Revisiting the Adam-SGD gap in language modeling

TL;DR

This paper investigates the optimizer gap between Adam and SGD in language modeling, revealing that with proper tuning, SGD with momentum can match Adam's performance in small-batch scenarios, and provides new insights into training dynamics.

Contribution

It demonstrates that SGD with momentum can perform comparably to Adam in small-batch settings when properly tuned, challenging previous explanations for Adam's superiority.

Findings

SGD with momentum matches Adam in small-batch training when tuned

Existing explanations like class imbalance and sharpness do not fully explain the gap

Batch size significantly influences training dynamics as shown by SDE models

Abstract

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.