Deconstructing What Makes a Good Optimizer for Autoregressive Language Models

We thank the reviewer for their review and recognizing that our results challenge the prevailing notion of Adam being uniquely suited for language model training. We’d like to address the issues brought up by the reviewer, as well as answer their specific questions.

Weaknesses:

1. Limited contribution: We would like to reiterate the main contributions of our work and highlight their broader implications. Our paper provides a comprehensive empirical evaluation of optimizers for autoregressive language model pretraining, including a wide range of hyperparameters and model scales. One key finding is that no single algorithm consistently outperforms others in terms of both performance and hyperparameter stability, challenging the prevailing notion of Adam as the default optimizer for language models. In previous work, as we have noted in the manuscript, empirical studies of optimizers across hyperparameters have been either outside of the autoregressive language modeling domain, limited to Adam only, or limited to a fixed scale (see, eg. [1,2,3]). Additionally, as you noted, we show that adaptivity in both the last layer and LayerNorm parameters is essential for maintaining performance and learning rate stability. This insight not only informs optimization design but also suggests pathways for developing more memory-efficient optimizers—a direction that concurrent work has already started to explore. We believe our findings can encourage groups with more computational resources or those working in other domains to replicate and extend our results, scaling to larger models and exploring new designs. We would appreciate further clarification on why the reviewer views our contribution as limited, and are happy to discuss this point further.
2. Gaining more insight about choosing optimizers: Indeed, one aim of our work is to provide a guide for practitioners to refer to when comparing across various optimizers and hyperparameters. Regarding your comment “...I didn’t feel I gained a clear insight to help decide on a better optimizer…”, this actually touches on the primary takeaway from our ablations: across the optimizers we tested, there isn’t any clear winner based on both optimal performance and hyperparameter stability. The only differing axis for practitioners to consider is thus memory constraints, but based on our results, they can be reassured that memory considerations for optimizer choice should not significantly affect the performance of their training run, at least across the optimizers tested (i.e. optimizers using diagonal preconditioning). For more concrete takeaways regarding our hyperparameter sweeps, we reported the optimal learning rates found in our sweeps in Appendix A as a starting point for practitioners in tuning (eg. tuning optimizers like Signum and Lion seem to require learning rates an order of magnitude smaller compared to Adam/Adafactor). In terms of other tuning practices, we observed most other parameters are very stable for the optimal learning rate (with exception to momentum) and thus we believe tuning learning rate and momentum are most critical. We agree that the main takeaways from our results could have been consolidated and communicated more clearly; we will be sure to add everything that we have written above in a paragraph in the Discussion section in the revised pdf (please also see our response to Reviewer 4oK2 for a similar discussion). Regarding your last point (“...the paper has potential but feels incomplete”), if there’s anything in addition to our clarifications above that could provide a more complete message to readers, we would appreciate hearing any concrete suggestions and would look into implementing those.

[1] Kunstner, Frederik, et al. "Heavy-tailed class imbalance and why adam outperforms gradient descent on language models." (2024).

[2] Zohrevand, Abbas, and Zahra Imani. "An empirical study of the performance of different optimizers in the deep neural networks." 2022 International Conference on Machine Vision and Image Processing (MVIP). IEEE, 2022.

[3] Wortsman, Mitchell, et al. "Small-scale proxies for large-scale transformer training instabilities."(2023).

We thank the reviewer for their review and for recognizing our work’s findings as interesting. We’d like to address the issues brought up by the reviewer, as well as answer their specific questions.

Weaknesses:

1. Limited Contribution: We acknowledge that our paper is an empirical comparison, and we appreciate your recognition of its value as a resource providing a comprehensive comparison across optimizers in a standardized codebase. While our scope is indeed limited to autoregressive language models, we believe the paper offers key takeaways for practitioners, such as the importance of adaptivity in the last layer and layer norm parameters for maintaining performance and stability, particularly relative to Adam. If the title appears misleading, we are very open to suggestions on how to make it more specific and reflective of these findings. Regarding the relationship between the two parts of the paper, our goal was to bridge the observed performance gap between SGD and other diagonal optimizers identified in the first section. We approached this by empirically examining the parameters requiring adaptivity, offering insights into why algorithms like SGD fail—namely, its inability to assign appropriate learning rate scales to these critical parameters. We will work to clarify the connection between these two sections in the revised version.
2. Main takeaways from our work: Thank you for bringing this up; indeed, one aim of our work is to provide a source for practitioners to refer to when comparing across various standard optimizers and hyperparameters. We’d like to provide an outline below addressing each point you have mentioned (eg. best practices for practitioners, conceptual/theoretical insights, and most fruitful further directions in algorithm design). We agree that the main takeaways from our results could have been consolidated and communicated more clearly; we will add a paragraph in the Discussion section of the revised pdf including these points below.

The primary takeaway from our ablations is that across the optimizers we tested, there is no ‘better optimizer’ based on both optimal performance and hyperparameter stability. The only differing axis for practitioners to consider is thus memory constraints, but based on our results, they can be reassured that memory considerations for optimizer choice should not affect the performance of their training run, at least across the optimizers tested (i.e. optimizers using diagonal preconditioning). For more concrete takeaways regarding our hyperparameter sweeps, we reported the optimal learning rates found in our sweeps as a starting point for practitioners in tuning (eg. tuning optimizers like Signum and Lion seem to require learning rates an order of magnitude smaller compared to Adam/Adafactor). In terms of other tuning practices, we observed most other parameters are very stable for the optimal learning rate (with exception to momentum) and thus we believe tuning learning rate and momentum should be prioritized if one has a limited compute budget for hyperparameter sweeping.

Our work highlights the following research directions as the most promising for developing optimizers for language model training. Firstly, our results suggest that within the realm of diagonal preconditioning optimizers, there may be room for more memory-efficient alternatives to AdamW to be designed (and has already been explored in concurrent work [1]) for domains where one is compute-constrained; we have provided one direction for cutting down on memory, where we have shown that the need for adaptivity is distributed unevenly across different parameters of the network. Secondly, in terms of improving the performance of these models, our work suggests that optimizer choice is unlikely to be a promising approach– if we are limited to diagonal preconditioning optimizers. Thus, it is likely that we need to turn to non-diagonal methods to improve on our current methods; this is indeed where previous work has demonstrated improved performance to Adam [2].

3. Related work on optimization comparisons: Thank you for mentioning this, we have revised the manuscript to include prior literature conducting empirical optimizer comparisons across various settings.

[1] Zhang, Yushun, et al. "Adam-mini: Use fewer learning rates to gain more." (2024).

[2] Gupta, Vineet, Tomer Koren, and Yoram Singer. "Shampoo: Preconditioned stochastic tensor optimization." (2018).

Thanks for the insightful comments. We are happy to see that the reviewer finds our work interesting, and appreciates the focus on hyperparameter stability. The specific comments are addressed below:

1. Sweeping min cosine learning rate - We did sweep this hyperparameter with Adam and used 0.1 as it seemed to perform the best for Adam in our setup. We will explicitly add the sweeps for this parameter in future revisions.
2. Attention heads disentangling in Adafactor and Adalayer - Indeed, from the Hessian structure perspective, Adam-mini [1] suggests having an adaptive learning rate assigned to each attention head’s query and key matrix could improve Adalayer* performance further. This is an interesting direction of future work and we would be happy to add these variations in a future revision.
3. Shampoo - We explicitly wanted to limit our study to the diagonal preconditioning optimizers. This is the reason we have not added Shampoo to the study.
4. Optimality of hyperparameters with scale - Most of our sweeps were done at 150m scale,and thus we cannot provide the scaling of optimal hyperparameters with scale given our compute budget. Regarding learning rate, keeping other hyperparameters fixed, we found the optimal learning rate to be reasonably stable from 150m to 1.2b scale (as shown in Table 1 in the Appendix).
5. Last layer adaptivity - This is required as different tokens have varying frequencies in the training data, thus, without adaptivity, the infrequent tokens won’t even be trained. For this reason, we need adaptivity in the last layer.

We hope that we adequately addressed the reviewer’s concerns and questions, and are happy to provide any additional clarifications or engage in further discussion.

[1] Zhang, Yushun, et al. "Adam-mini: Use fewer learning rates to gain more." (2024).

We thank the reviewer for their insightful comments. We are happy to see that the reviewer finds the paper well-written and easy to understand. The specific comments are addressed below:

1. Comparable performance of optimizers - As discussed in the introduction, the primary motivation for our initial experiments was to empirically verify the community’s preconceptions about Adam being superior to other optimizers; this implies an underlying assumption of a substantial gap existing between Adam and these other optimizers. When we have written these optimizers are "comparable in performance," we mean that the actual performance differences between them is much smaller than what might have previously been perceived, especially when we consider this gap relative to the gap of the outlier performance of SGD. As seen in Figure 1, compared to SGD, other algorithms are relatively stable when tuning learning rate. This does not mean that practitioners should not tune learning rate as a hyperparameter, but our intended main takeaway is that at scale–where it may only be possible to launch a single run– optimizers apart from SGD will not be very far from their optimal performance. This is an essential property that an optimizer needs to satisfy to be used in large-scale pre-training runs. This is indeed a comparison on a more crude scale, but we believe it is justified to refrain from defining a precise quantitative gap because such a definition may not be robust to other design choices in our setup, such as architectural variations. If the reviewer disagrees with our reasoning, we are keen to hear their thoughts to discuss this point further.
2. Stability of other hyperparameters except momentum and learning rate - Yes, we agree that our takeaway in line 327 is based on the relative ‘visual’ stability of the 1D plots. We would be happy to make it quantitative by providing the range for each hyperparameter which incurs a maximum deviation of a specified quantity (may be 0.01) from the optimal along the 1D sweep. If the reviewer considers this information important for their evaluation of the paper, please make note and we would be glad to include it.
3. Sweep of $\beta_2$ and warmup - We would be happy to add more extreme values for the two parameters in the sweeps and will include this in the revision.
4. Sweep of batch size - As pointed out by the reviewer, when changing batch size, tuning learning rate and momentum is known to be essential for obtaining optimal performance. This was the reason we did not do large sweeps for this hyperparameter. We would be happy to add more points in the 1-D sweep for batch size as well.
5. Section 3.6 finding - We are sorry for the confusion. We did not mean to refer to Lemma 1 of Balles and Hannig 2018, and simply wanted to point out that this observation has been noted in this previous work. Also, thank you for the correction– we indeed need |E[m]| in the numerator. Regarding the discussion when $\beta_1 = \beta_2$, we only mean to say that as both the numerator and denominator average over the same window size, it is more likely for Adam and Signum to behave similarly. Indeed, for this to hold mathematically, we will need the ratio of the mean to the square root of the second moment to be similar for all the coordinates, which is unlikely to happen. However, as shown in Figure 6, we do get similar performance as Signum for Adam with $\beta_1 = \beta_2$.
6. Loss curve for Signum in Figure 4 and Figure 5 - Indeed, in the $\beta_2$ sweep in Figure 5, we don’t have a curve for Signum, and that is precisely because Signum is Lion with $\beta_1 = \beta_2$.
7. AdaSGD - We would be very happy to include AdaSGD in subsequent revisions.
8. Typos - Thanks for catching them. We will fix them.

We hope that we adequately addressed the reviewer’s concerns and questions, and are happy to provide any additional clarifications or engage in further discussion.