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

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 14:59:52

Abstract: Adam is one of the most popular optimization algorithms in deep learning. However, it is known that Adam does not converge in theory unless choosing a hyperparameter, i.e., $\beta_2$, in a problem-dependent manner. There have been many attempts to fix the non-convergence (e.g., AMSGrad), but they require an impractical assumption that the gradient noise is uniformly bounded. In this paper, we propose a new adaptive gradient method named ADOPT, which achieves the optimal convergence rate of $\mathcal{O} ( 1 / \sqrt{T} )$ with any choice of $\beta_2$ without depending on the bounded noise assumption. ADOPT addresses the non-convergence issue of Adam by removing the current gradient from the second moment estimate and changing the order of the momentum update and the normalization by the second moment estimate. We also conduct intensive numerical experiments, and verify that our ADOPT achieves superior results compared to Adam and its variants across a wide range of tasks, including image classification, generative modeling, natural language processing, and deep reinforcement learning. The implementation is available at https://github.com/iShohei220/adopt.

Main paper: 
  
 Appendix:  


Decision: The paper introduces a new adaptive gradient method named ADOPT, which is designed to address the non-convergence issues of Adam, particularly when the $\beta_2$ parameter is not chosen in a problem-dependent manner. ADOPT modifies Adam in the computation of the second moment estimates and the order of momentum updates. It is claimed that ADOPT achieves optimal convergence rates without relying on the assumption that the gradient noise is uniformly bounded.

The reviewers acknowledged the theoretical contribution of ADOPT, though the contribution is marginal; for example, the convergence still relies on the assumption that the second moment of the stochastic gradient is uniformly bounded. Besides, the reviewers also agreed that ADOPT could be beneficial in specific scenarios, as shown by experiments. Overall, the reviewers are mostly supportive of this paper, and I think this paper could be accepted.

Reviews: review: paper_summary: paper_summary: The paper proposes a new adaptive gradient method called ADOPT, which addresses the non-convergence issue of popular methods like Adam and RMSprop.  The method modifies the calculation of second moment estimates and the order of momentum calculation and scaling operations. Extensive numerical experiments demonstrate that ADOPT achieves competitive or superior results compared to existing methods across various tasks.
strength_weakness: strength_weakness: 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.

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.

------------------------------------------------------------------------------------------------- comments after rebuttal -------------------------------------------------------------The authors' response clearly resolved the memory cost concerns. I would like to increase the score to 6.
questions: questions: see weakness.
limitations: limitations: Yes, The authors have adequately addressed the limitations.
paper_summary: paper_summary: Motivated by the counterexample due to Reddi et al., this work designs an adaptive optimizer that can converge for the choice of beta_2 that is independent of the problem instance. Their analysis works for a more general condition than the previous works.
strength_weakness: strength_weakness: The theoretical results look sound. I have some questions below.
questions: questions: - 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.

- It seems that for your algorithm, $\epsilon$ needs to be chosen much larger than that of Adam. How sensitive is this choice? Why is this the case? $\epsilon$ really should be for numerical stability in the original Adam. Does your analysis suggest such a large value for $\epsilon$?
limitations: limitations: 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. 
Therefore, I recommend "borderline accept" for this work.
paper_summary: paper_summary: The paper titled "ADOPT: Modified Adam Can Converge with Any β2 with the Optimal Rate" introduces a new adaptive gradient method named ADOPT. This method aims to address the non-convergence issues of the Adam optimization algorithm. Adam, despite its popularity in deep learning, does not theoretically converge unless the hyperparameter β2 is chosen in a problem-dependent manner. The paper proposes ADOPT, which achieves the optimal convergence rate with any choice of β2, without relying on the bounded noise assumption. ADOPT modifies Adam by removing the current gradient from the second moment estimate and changing the order of the momentum update and normalization. The paper also presents extensive numerical experiments showing ADOPT's superior performance across various tasks, including image classification, generative modeling, natural language processing, and deep reinforcement learning.
strength_weakness: strength_weakness: 1. 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.
2. 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).).
3. The paper does not thoroughly discuss the potential computational overhead introduced by ADOPT compared to Adam and other optimizers.
questions: questions: 1. 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?
2. 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?
3. 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?
limitations: limitations: see the weaknesses.
paper_summary: paper_summary: The submitted work analyzes the divergence of Adam and RMSprop in smooth nonconvex settings. The authors propose a new optimizer ADOPT, whose convergence does not depend on the second-moment coefficient $\beta_2$.
The proposed optimizer is evaluated in toy settings, image classification, language modeling (finetuning settings), and reinforcement learning.
strength_weakness: strength_weakness: * 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.
* 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.
questions: questions: * 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.
* 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?
limitations: limitations: See Weaknesses above.

Comments: title: A Gentle Reminder to Reviewer AhBk
A Gentle Reminder to Reviewer f5HN
Response to the Additional Comment by Reviewer qABQ
Response to the Additional Comment to Reviewer bGLP
Response to the Comment to Reviewer bGLP
Thanks
comment: 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 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 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 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**

| Batch size \ LR | 6e-5  |1e-4 | 6e-4        |
| ----------------- | ------- |------ | --------   |
| 96                     |  3.21 | 3.17 | 7.54        |
| 480                   |     -    | 3.31 | **3.02**   |

**ADOPT**

| Batch size \ LR | 6e-5      |1e-4      | 6e-4 |
| ----------------- | --------- |--------- | ----- |
| 96                     | 3.11       |  3.09      | 3.17 |
| 480                   | -            | **2.98** | 3.00 |

> 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!
Thanks to the authors for the additional results. I suggested looking at the test loss vs. learning rate curves because Adam can achieve the same performance as ADOPT but with a smaller learning rate. This would just imply that ADOPT achieves optimal performance at a high learning rate. I am still unclear about the following:

1. For the optimal hyperparameters, such as learning rate, does ADOPT always perform better than Adam?
2. Does ADOPT require a smaller warmup (or no warmup)? Is warmup used in the GPT experiments?
3. 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?
Thank you for your comment.

As you pointed out, Adam's loss spikes in the GPT-2 experiment may simply be due to the learning rate being too high, so we also ran the experiment with the learning rate lowered from the default of 6e-4 to 1e-4 for the case of a small batch size. The results are shown in the table below; for Adam, lowering the learning rate only slows down the timing of loss spikes and eventually causes divergence in training. ADOPT, on the other hand, is able to perform stable learning at all settings. 
Further reductions in the learning rate resulted in significantly slower training, so we did not include them in the table.

| Optimizer   |  LR  | 50K iters   | 100K iters     | 150K iters | 200K iters |
| ------------- | ---- | :----------: | :------------:  | :-----------: | :----------: |
| Adam         | 6e-4 | 7.64          | 7.54             |  -                |                 |
| Adam         | 1e-4 | 3.26          | 3.17             | 7.09           | 7.56         |
| ADOPT       | 6e-4 | 3.22          | 3.17             | 3.13           | 3.10         |
| ADOPT       | 1e-4 | 3.16          | 3.09             | 3.04           | 3.02         |

This result is also consistent with theory: the constant term for Adam's convergence bounds in Theorem 3.1 is independent of the learning rate, so a failure in training cannot be prevented by reducing the learning rate. As you stated, Theorem 3.1 is for RMSprop, which does not account for momentum, but there is a similar constant term in Adam's convergence bounds in prior studies (e.g., [1]) that is independent of the learning rate. Thus, this result is a practical example of how Adam can become unstable when training on a large-scale data set or with a small batch size, as the gradient noise increases in such cases. Such instability is not unique to this experimental setting, as it has often been reported to be observed in the training of larger LLMs [2, 3].

We hope that this will address your concerns about our submission.
We would be glad to respond to any further comments that you may have during the discussion period.

Thanks!


### References

[1] 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.

[2] Molybog, Igor, et al. "A theory on adam instability in large-scale machine learning." arXiv preprint arXiv:2304.09871 (2023).

[3] Takase, Sho, et al. "Spike No More: Stabilizing the Pre-training of Large Language Models." arXiv preprint arXiv:2312.16903 (2023).
Thanks for your responses. 
I read through them, and it partially addresses my concerns.
Hence, I'll increase the score to 6.

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.

[1] Ahn and Cutkosky, "Adam with model exponential moving average is effective for nonconvex optimization" (https://arxiv.org/pdf/2403.02648)
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.