ADOPT: Modified Adam Can Converge with Any $\beta_2$ with the Optimal Rate

Thank you for your insightful comments. We will answer your questions to address your concerns. Please also refer to the general response and the PDF attached with it.

The non-convergence of Adam is not an issue in practice and I am unsure about the practical implications of ADOPT over Adam. The non-convergence issue arises near the minima and most modern deep learning models are not trained to convergence.

In realistic settings, where is the marginal convergence speed in ADOPT coming from? My understanding is that non-convergence of Adam would matter near the minima only. Faster training way before reaching the minima is not expected. In particular, in language modeling pretraining, convergence is usually not achieved.

The paper does not thoroughly investigate the sensitivity of ADOPT optimizer parameters in practical scenarios. A more comprehensive study on the optimizer parameters and comparison with Adam is required to understand its robustness and practical utility.

I agree with you that in modern deep learning models, such as large language models, training is rarely run until full convergence is achieved. However, even in such cases, Adam's non-convergence can be a problem. For example, when training with a large dataset or small batch size, the constant term in Theorem 3.1 can become non-negligible due to the large gradient noise. In such cases, Adam may become unstable even in the early stages of training. To confirm that this happens empirically, we have added a new experiment of pre-training of GPT-2. See the figure in the PDF attached to the general response for the results. We observed that ADOPT is always stable in this experiment, whereas Adam actually diverges in the early phase of training due to loss spikes when the batch size is small. We also observed that ADOPT performed slightly better even when the batch size was large. We believe that these results clearly demonstrate the practical effectiveness of ADOPT in modern deep learning.

Does the difference in the recommendation of the epsilon parameter for ADOPT simply arise because it is added inside the square root instead of outside?

This is a good question. In fact, the reason that ADOPT requires larger $\epsilon$ than Adam is not simply due to the order of square root and addition. This can be explained by theory: as can be seen from Theorem 3.1, the convergence bounds for Adam are of the order of $\log(\epsilon^{-2})$, while the convergence bounds for ADOPT in Eq. (33) are of the order of $\epsilon^{-2}$. In other words, when$\epsilon$ is small, the convergence bound of ADOPT becomes loose more rapidly than those of Adam. For this reason, it is safer for ADOPT to use slightly larger $\epsilon$ than Adam. However, our default settings have been found to work robustly in all experiments, so tuning of the eps is rarely needed, even in practice.

We would be glad to respond to any further questions and comments that you may have.

Thanks.

References

[1] github.com/karpathy/nanoGPT

I thank the reviewers for their responses.

...when training with a large dataset or small batch size, the constant term in Theorem 3.1 can become non-negligible due to the large gradient noise. In such cases, Adam may become unstable even in the early stages of training. To confirm that this happens empirically, we have added a new experiment of pre-training of GPT-2. See the figure in the PDF attached to the general response for the results. We observed that ADOPT is always stable in this experiment, whereas Adam actually diverges in the early phase of training due to loss spikes when the batch size is small....

I am unsure if the statements from Theorem 3.1 generalize to realistic scenarios:

1. First, the Theorem corresponds to RMSprop and there are various assumptions 2.1-2.3 and 2.5, which can become invalid in practical scenarios.
2. Regarding the GPT pretaining result, as the optimal learning rate range for the two optimizers can be very different, it is unclear if the learning rate for Adam was just large whereas it wasn't the case for ADOPT. A compelling result would include performance vs. learning rate for the different batch sizes. I do not expect the authors to do these experiments in this short rebuttal. If the authors have already scanned the learning rates, I would request them to point me to the result.

To explicitly show that ADOPT is more stable during early training, the authors can track the top eigenvalue of the pre-conditioned Hessian [1], which determines the stability of adaptive optimizers for large batches. An alternative experiment would involve plotting heatmaps of performance against learning rate and batch size. In small-scale experiments, if the authors can show that ADOPT either (i) has a smaller pre-conditioned sharpness during early training or (ii) the optimal performance for optimal covers a large range of hyperparameters such as learning rate and batch size. Then, it would be convincing that ADOPT works reasonably well compared to Adam.

[1] Adaptive Gradient Methods at the Edge of Stability, https://arxiv.org/abs/2207.14484

I would like to reiterate that I do not expect the authors to perform these experiments in this discussion period.

My concerns regarding the practical implications of the non-convergence issues remain. Therefore, I am keeping my current score.

Thank you for your comment.

For the optimal hyperparameters, such as learning rate, does ADOPT always perform better than Adam?

We have summarized the result regarding the optimal hyperparameter settings in terms of the batch size and the learning rate in the table below. The result shows the test loss after running 100K training iterations. The best results are shown in bold. When comparing the best results, ADOPT shows a slightly better result than Adam.

Adam

ADOPT

Does ADOPT require a smaller warmup (or no warmup)? Is warmup used in the GPT experiments?

In the GPT-2 experiment, following the default setting of the nanoGPT code base, the linear warmup is used for the first 2000 iterations. Please refer to the original code provided at github.com/karpathy/nanoGPT/blob/master/train.py for more detailed settings. We did not change the experimental settings except for choices of the optimizer, the batch size, and the learning rate.

Is the range of optimal learning rates longer for ADOPT compared to Adam? Or the size of the range is the same but it is shifted?

As can be seen from the table above, ADOPT seems to work more robustly to the choice of the learning rate compared to Adam.

If you have any further comments or questions, we would be glad to respond them.

Thanks!

Thank you for your insightful comments. We will answer your questions to address your concerns. Please also refer to the general response and the attached PDF.

The analysis still relies on the assumption that the second moment of the stochastic gradient is uniformly bounded, which might not always hold in practice. How can the assumption on the second moment being uniformly bounded be relaxed? Are there any potential methods or future work suggested to address this?

As you mentioned, our analysis still relies on the bounded second moment assumption, which is sometimes violated in practice, although it is milder than the bounded stochastic gradient assumption used in the previous works. A promising direction to further relax this assumption is to use the bounded variance assumption, where $\mathbb{E} [ \|| g - \nabla f \||^2 ]$ is assumed to be bounded. We mentioned it in the third paragraph of Section 6 as a limitation and future work.

The paper could have included comparisons with more recent optimizers beyond Adam to strengthen its empirical claims, like the recent work CO2 ("CO2: Efficient distributed training with full communication-computation overlap." ICLR (2024).).

Thank you for your suggestion. We think that CO2 is specifically designed for efficient distributed training, so their contributions are orthogonal to ours. In our experiment, we compare other optimizers which share the motivation of addressing the non-convergence issue of Adam (e.g., AMSGrad). Of course, combining the techniques of other works (e.g., CO2) with our ADOPT could be a promising direction to further improve the performance, but we leave it for future work.

The paper does not thoroughly discuss the potential computational overhead introduced by ADOPT compared to Adam and other optimizers. What is the computational cost of ADOPT compared to Adam and AMSGrad? Is there a significant overhead that might affect its practicality in large-scale applications?

Thank you for your comment. Let us clarify the computational cost of ADOPT compared to Adam and AMSGrad. The computational cost of ADOPT is equal to that of Adam and less than that of AMSGrad. Since both Adam and ADOPT need to store momentum $m_t$ and second-order moments $v_t$, their memory costs are about the same. In addition, Adam computes bias corrections when updating parameters, while our implementation of ADOPT omits them, making ADOPT slightly less computationally expensive than Adam. On the other hand, AMSGrad requires storing $\hat{v}_t$ in addition to $m_t$ and $v_t$, so its memory cost is larger than ADOPT (and Adam).

While ADOPT performs well on a variety of tasks, how does it perform on tasks not covered in the experiments? Are there specific types of problems where ADOPT might not be as effective?

As far as we have experimented, ADOPT has always shown better results than Adam. To further strengthen our experimental results, we also performed pre-training experiments on GPT-2 and found that ADOPT performed better than Adam in those experiments as well. In particular, we observed that Adam suffers from loss spikes when the batch size is small and the gradient noise is large, while ADOPT always performs well stably. See the general response for more details.

We would be glad to respond to any further questions and comments that you may have.

Thanks.

Thank you again for your efforts in reviewing our paper and your constructive comments. The discussion period will end soon, so please let us know if you have further comments about our reply to your feedback.

Thanks.

Thank you for your insightful comments. We will answer your questions to address your concerns. Please also refer to the general response and the PDF attached with it.

In Theorem 4.1, please specify the choice of hyperparameters (beta_1, beta_2, \eps). I could not find this even in the appendix, even though the main text said it can be found there. In particular, does the choice found by your analysis similar to what people use in practice?

If the choice of beta's do not match with your theoretical prediction, I think the claim that the new optimizer AdaOPT alleviates the parameter-tuning is not an overclaim. At the end of the day, it seems that the authors have to tune those parameters depending on the settings. Hence, practically speaking, there's no reason for practitioners to use the proposed algorithm over the original Adam.

The concrete convergence bound is provided in Eq. (33) in Appendix E, which shows that the convergence rate is $O ( 1 / \sqrt{T} )$ with any $(\beta_1, \beta_2, \epsilon)$. The bound gets tighter when $\beta_2$ is chosen to be close to 1, which corresponds to practical choices. In fact, as shown in Figure 1, the training with ADOPT tends to be stable when $\beta_2$ is close to 1, although the convergence is achieved even with small $\beta_2$.

In terms of $\beta_1$, there is a gap between theory and practice. In theory, the bound will be tighter when $\beta_1$ is small, whereas $\beta_1 = 0.9$ is used in practice. This gap is consistently observed in the literature on the convergence analysis of Adam (e.g., [2, 3]). To the best of our knowledge, the effectiveness of momentum in Adam-type optimizers is still an open question.

It seems that for your algorithm, needs to be chosen much larger than that of Adam. How sensitive is this choice? Why is this the case? really should be for numerical stability in the original Adam. Does your analysis suggest such a large value for ?

You are correct that ADOPT requires a larger $\epsilon$ than Adam. This can be explained by theory: as can be seen from Theorem 3.1, the convergence bounds for Adam are of the order of $\log(\epsilon^{-2})$, while the convergence bounds for ADOPT in Eq. (33) are of the order of $\epsilon^{-2}$. In other words, when $\epsilon$ gets small, the convergence bounds of ADOPT become loose more rapidly than those of Adam. For this reason, it is safer for ADOPT to use a slightly larger value for $\epsilon$ than Adam. Our default settings have been found to work robustly in all experiments, so tuning of the $\epsilon$ is rarely needed, even in practice.

Overall, the main contribution in this work seems quite marginal. For a submission of this type, I would support a clear acceptance if

(i) the theoretical guarantees have a noticeable innovation over the previous ones.

(ii) or the proposed method works much better than the previous ones.

It seems that the theoretical improvement is the removal of the uniformly upper bounded assumption of stochastic gradients. To me, this looks like a minor improvement over the previous work. Moreover, the resulting algorithm does not seem to have major advantages over Adam. In particular, unlike the claim made in this paper, I don't think the resulting algorithm alleviates the hyper-parameter tuning. I don't see why practitioners should choose this algorithm over Adam.

Thank you for your suggestions. To demonstrate the practical effectiveness of our ADOPT more clearly, we have added an experiment in pre-training language models, in which Adam tends to suffer from optimization difficulties like loss spikes. In this experiment, we used the nanoGPT [1] code base to run the GPT-2 pre-training in OpenWebText. We observed that Adam suffered from loss spikes and completely failed the training when the batch size was small, while ADOPT was able to train stably even with small batch sizes. We also observed that ADOPT performed slightly better than Adam even for large batch sizes. This result is consistent with theory, since, when the batch size is small, the gradient noise is larger and the constant term of Adam's convergence bound in Theorem 3.1 also gets larger. Thus, the experimental results confirm that ADOPT has an advantage over Adam even in practical cases such as pre-training language models. We hope these results will address your concerns.

We would be glad to respond to any further questions and comments that you may have.

Thanks.

References

[1] github.com/karpathy/nanoGPT

[2] Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.

[3] Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen. Closing the gap between the upper bound and lower bound of adam’s iteration complexity. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.

Thank you for your additional comments.

On a side note, it seems that the benefits of momentum seems to be theoretically justified in the nonsmooth and nonconvex setting (as shown by [1]), and the authors might find it helpful to further clarify the main contributions of this work.

We were not aware of the result of the paper you mentioned, so thank you for letting us know. As you pointed out, that paper does indeed partially explain the role of momentum in Adam, but there appear to be some limitations: first, their analysis is limited to the case $\beta_2=\beta_1^2$, which deviates from the practical choice. In addition, their analysis can only be applied to the case where Adam's update is clipped, which also deviates from the practical algorithm. Thus, the role of momentum in Adam is not yet fully understood and seems to be an open question. We promise to cite that paper and describe this point clearly in the final version.

If you have any further comments or questions, we would be glad to respond them.

Thanks.

Thank you for your insightful comments. We will answer your questions to address your concerns. Please also refer to “Author Rebuttal by Authors” and the attached PDF.

The convergence of a modified version of Adam is not significant from theoretical sense unless the ADOPT can beat the performance of Adam in practice given existing convergence proofs of Adam.

From the empirical results, the performance of ADOPT is not superior over Adam very much. People may be reluctant to use ADOPT in practice.

Your point is that ADOPT does not seem to perform so well compared to Adam in the practical experiments in Section 5 except for the toy examples. In fact, what our theory shows is that Adam fails catastrophically when the constant term of Adam's convergence bound in Theorem 3.1 becomes large; hence, in a correctly tuned situation in practice (e.g., with a sufficiently large $\beta_2$), it is natural that ADOPT will not improve significantly relative to Adam.

However, to show that even in practice Adam can make catastrophic failures and ADOPT can avoid them, we have added an experiment in pre-training language models. In this experiment, we used the nanoGPT [1] code base to run the GPT-2 pre-training with OpenWebText. See the general response and the attached PDF for more details. We observed that Adam suffered from loss spikes and completely failed the training when the batch size was small, while ADOPT was able to train stably even with small batch sizes. We also observed that ADOPT performed slightly better than Adam even for large batch sizes. This result is consistent with theory, since the gradient noise is larger and the constant term of Adam's convergence bound is larger when the batch size is small. Thus, the experimental results confirm that ADOPT has an advantage over Adam even in practical cases such as pre-training language models. We hope these results will address your concerns.

More importantly, the algorithm 1 (ADOPT) seems to require storage of three parts: g_t, m_t, v_t, which is more than what the standard Adam requires. This is quite significant drawback of ADOPT if this cannot be optimized.

I respectfully point out that this is a misunderstanding: the memory cost of ADOPT is equivalent to that of Adam. Only $m_t$ and $v_t$ need to be stored in ADOPT, and $g_t$ can be discarded once those updates are done. Presumably, this misunderstanding arises from the fact that $g_{t+1}$ is used to update $m_{t+1}$ in Algorithm 1, but this does not mean that $g_t$ needs to be stored. To clarify this, an equivalent alternative representation is given in Algorithm 2 in the attached PDF of the general response. In fact, experimental results confirm that the memory costs of ADOPT and Adam are equivalent. For example, the memory cost for pre-training GPT-2 is approximately 18 GB per GPU for both.

We would be glad to respond to any further questions and comments that you may have.

Thanks.

References

[1] github.com/karpathy/nanoGPT

Thank you again for your efforts in reviewing our paper and your constructive comments. The discussion period will end soon, so please let us know if you have further comments about our reply to your feedback.

Thanks.