Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Deconstructing What Makes a Good Optimizer for Language Models

Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, Sham Kakade

TL;DR

This paper challenges the conventional wisdom that Adam is the sole best optimizer for large-scale language modeling by conducting a broad, controlled comparison of SGD, Adafactor, Lion, Signum, and Sophia across model sizes and architectures. It shows that, once hyperparameters are well tuned, most adaptive, diagonal preconditioners perform similarly to Adam, with SGD lagging behind. The authors further dissect adaptivity, showing Signum can match Adam when $β_1=β_2$ and introducing Adalayer, a layerwise preconditioner whose effectiveness hinges on adaptive treatment of the last layer and LayerNorm parameters. Practically, these findings imply that optimizer choice may be driven more by memory and implementation considerations, while the key to stable and high-performing training lies in targeted adaptivity for specific network components.

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.

Deconstructing What Makes a Good Optimizer for Language Models

TL;DR

This paper challenges the conventional wisdom that Adam is the sole best optimizer for large-scale language modeling by conducting a broad, controlled comparison of SGD, Adafactor, Lion, Signum, and Sophia across model sizes and architectures. It shows that, once hyperparameters are well tuned, most adaptive, diagonal preconditioners perform similarly to Adam, with SGD lagging behind. The authors further dissect adaptivity, showing Signum can match Adam when and introducing Adalayer, a layerwise preconditioner whose effectiveness hinges on adaptive treatment of the last layer and LayerNorm parameters. Practically, these findings imply that optimizer choice may be driven more by memory and implementation considerations, while the key to stable and high-performing training lies in targeted adaptivity for specific network components.

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.
Paper Structure (26 sections, 1 theorem, 1 equation, 21 figures, 1 table, 1 algorithm)

This paper contains 26 sections, 1 theorem, 1 equation, 21 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Comparing Optimizers Across Hyperparameters, Architectures and Scale
  3. Methodology
  4. Setup
  5. Algorithms.
  6. Models.
  7. Training variants.
  8. Token counts.
  9. Other hyperparameters.
  10. Sweeping learning rates
  11. Takeaway:
  12. Sweeping momentum
  13. Takeaway:
  14. Additional hyperparameter sweeps
  15. Takeaway:
  16. ...and 11 more sections

Key Result

Lemma 1

Consider a parameter with a history of gradients $g_t, g_{t-1}, \ldots$. Let $m$ be the random variable that is equal to $g_{t-\tau}$ with probability $(1-\beta_1)\beta_1^{\tau}$ and $v$ be the random variable that is equal to $g_{t-\tau}$ with probability $(1-\beta_2)\beta_2^{\tau}$. The Adam updat

Figures (21)

  • Figure 1: Final validation loss when training language models with 150m, 300m, 600m, and 1.2b parameters, sweeping across learning rates for five standard optimizers (SGD, Adam, Adafactor, Lion, and Signum). Plots have been shifted to align the optimal learning rates for each optimizer. Except for SGD, other optimizers seem comparable in their optimal performance and stability with respect to learning rate tuning.
  • Figure 2: Sweeping learning rate without QK norm or z-loss for (Left) the 150m model, and (Right) the 300m model. These models are less stable than the standard model, but the same general trend across algorithms hold here.
  • Figure 3: Sweeping learning rate on 150m models trained for 4x longer (100k steps) than in the base runs for (Left) the standard model, and (Right) the model without QK norm or z-loss. Compared to the shorter runs, these models achieve better performance and increased stability across learning rates.
  • Figure 4: Sweeping momentum for fixed learning rate across three settings: (Left) 150m standard, (Middle) 150m with no QK norm or z-loss, (Right) 300m standard. Adam and Adafactor are similarly robust to $\beta_1$, while Lion and Signum are slightly more sensitive to low values and SGD is substantially more sensitive.
  • Figure 5: Sweeps over other hyperparameters. Top: weight decay, warmup duration, and batch size. Bottom: $\epsilon$ and $\beta_2$. We generally find little effect for the non-SGD algorithms, however there are parameters that differ from our defaults that can offer up to 0.02 improvements in validation loss
  • ...and 16 more figures

Theorems & Definitions (1)

  • Lemma 1: balles2018dissecting