Deconstructing What Makes a Good Optimizer for Language Models

TL;DR

This paper compares various optimization algorithms for language model training, finding that most perform similarly and that optimizer choice can be based on practical factors rather than performance differences.

Contribution

The study provides a comprehensive comparison of optimizers like Adam, SGD, Adafactor, Lion, and Sophia across multiple models and hyperparameters, and introduces simplified variants of Adam.

Findings

Most optimizers perform comparably in optimal and hyperparameter-robust settings.

No single optimizer is clearly superior in performance or stability.

Adaptive methods like Adam and its variants are crucial for maintaining performance and stability.

Abstract

Training language models becomes increasingly expensive with scale, prompting numerous attempts to improve optimization efficiency. Despite these efforts, the Adam optimizer remains the most widely used, due to a prevailing view that it is the most effective approach. We aim to compare several optimization algorithms, including SGD, Adafactor, Adam, Lion, and Sophia in the context of autoregressive language modeling across a range of model sizes, hyperparameters, and architecture variants. Our findings indicate that, except for SGD, these algorithms all perform comparably both in their optimal performance and also in terms of how they fare across a wide range of hyperparameter choices. Our results suggest to practitioners that the choice of optimizer can be guided by practical considerations like memory constraints and ease of implementation, as no single algorithm emerged as a clear winner in terms of performance or stability to hyperparameter misspecification. Given our findings, we further dissect these approaches, examining two simplified versions of Adam: a) signed momentum (Signum) which we see recovers both the performance and hyperparameter stability of Adam and b) Adalayer, a layerwise variant of Adam which we introduce to study the impact on Adam's preconditioning for different layers of the network. Examining Adalayer leads us to the conclusion that, perhaps surprisingly, adaptivity on both the last layer and LayerNorm parameters in particular are necessary for retaining performance and stability to learning rate.