Paper: Promoting Exploration in Memory-Augmented Adam using Critical Momenta (ICLR-2024)

Reviewer comment (Experiments):
My main concerns are regarding the memory cost of the approach and whether these costs are justified given the plethora of other alternatives to Adam. The paper does not sufficiently discuss the computational overhead introduced by the memory component, which is critical for practical deployments, especially in resource-constrained environments.

Author response:
We agree that our method introduces memory cost. However, it's important to recognize that there are (still) many deep learning applications that involve fewer parameters. In these cases, our method can deliver exploration advantages and result in better generalization. 
We also emphasize that our findings hold broader implications beyond buffer utilization. Our integration of memory-augmented optimizer concepts demonstrates how Adam can effectively navigate the loss surface, leading to the discovery of flatter regions. In essence, our work introduces a novel perspective on optimization, focusing on efficient exploration of the loss surface. Therefore tackling the memory issue can be a good potential future direction to explore. 
We would also like to highlight that our method does not introduce more hyper-parameters than the existing memory-augmented method CG. Furthermore, our grid search experiments, as detailed in Appendix A.2, indicate that, unlike CG, our method consistently performs optimally with C=5 and decay=0.7 across all cases. This observation suggests that there may be no necessity for conducting a grid search over these hyperparameters in real-world scenarios.

Reviewer comment (Theory):
The experimental results indicate some improvements, but these could potentially be attributed to the additional use of memory and hyperparameters. It would be beneficial for the paper to discuss whether Adam+CM offers any regret guarantees or optimality conditions given the added computational complexity, i.e., could there be another algorithm that does better?

Author response:
We agree that our method introduces memory cost. However, it's important to recognize that there are (still) many deep learning applications that involve fewer parameters. In these cases, our method can deliver exploration advantages and result in better generalization. 
We also emphasize that our findings hold broader implications beyond buffer utilization. Our integration of memory-augmented optimizer concepts demonstrates how Adam can effectively navigate the loss surface, leading to the discovery of flatter regions. In essence, our work introduces a novel perspective on optimization, focusing on efficient exploration of the loss surface. Therefore tackling the memory issue can be a good potential future direction to explore. 
We would also like to highlight that our method does not introduce more hyper-parameters than the existing memory-augmented method CG. Furthermore, our grid search experiments, as detailed in Appendix A.2, indicate that, unlike CG, our method consistently performs optimally with C=5 and decay=0.7 across all cases. This observation suggests that there may be no necessity for conducting a grid search over these hyperparameters in real-world scenarios.

Reviewer comment (Experiments):
One potential alternative baseline that comes to mind is to consider averaging the weights via approaches such as Polyak averaging or Stochastic Weight Averaging. These methods have been shown to improve generalization performance and only require one additional copy of the parameters if one employs the exponential moving average trick. The paper could be strengthened by comparing Adam+CM to these alternatives—perhaps the two approaches are orthogonal.

Author response:
We appreciate the suggestions made by the reviewer about comparing Adam+CM with alternatives like Polyak averaging or Stochastic Weight Averaging. However, as also pointed out by the reviewer, these two approaches are orthogonal and could seamlessly be integrated to improve performance further.

Reviewer comment (Evaluation):
I would consider additional measures of sharpness as well. For instance, it seems common to consider $|\lambda_{max}/\lambda_5|$ (Jastrzebski et al. (2020)) or measures of how quickly loss changes in a local neighborhood. The interaction between sharpness, distance travelled, and generalization performance seems quite complicated generally as it is also a function of the learning rate and other hyperparameters. This should be made clear in Figure 7 and Figure 8.

Author response:
As suggested by the reviewer, we ran our CIFAR10/100 experiments with another definition of sharpness defined in [1] that determines how quickly loss changes in a local neighborhood. We track m-sharpness for training ResNet34 with different optimizers on the same hyper-parameter setup. We draw similar conclusions as presented in Figure 8 (top-left) i.e., Adam+CM has lower m-sharpness than other baselines. We have added these results in Appendix A.4.3. 
We agree that the interaction between sharpness, distance travelled, and generalization performance is a function of learning rate and other hyperparameters. However since we show the plots in Figure 8 belonging to the best-performing setup for a given optimizer, these metrics are the indicator of why exactly Adam+CM achieves better overall performance. We wanted to highlight that it’s the property of Adam+CM that helps efficient exploration and finding flatter surfaces to obtain the best performance rather than a particular hyperparameter setup. We have made it clearer in Fig 7 and Fig 8 in the revision.

[1] Foret, P., Kleiner, A., Mobahi, H., & Neyshabur, B. (2020). Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412.

Reviewer comment (Experiments):
In general, as mentioned in weaknesses, I would be interested in a more thorough study on how additional memory can best be used. It seems like C is generally set to 5 in the grid search—can we get away with smaller $C$? What if we used a different value for CG and CM?

Author response:
As mentioned earlier, CM consistently performs optimally with C=5 and decay=0.7 across all cases, unlike CG which may require a larger C. However, upon the suggestion made by the reviewer, we reduced C further in CM, and we performed the same grid search. We observed that with a smaller C, the validation performance in CIFAR10/100 experiments slightly reduced. We have included the results in the Appendix where we show a sensitivity analysis using buffer size on CIFAR10/100 experiments.

Reviewer comment (Evaluation):
Although the experiments using artificial loss functions are promising, the results on actual deep neural networks is not as good. The fact that both CM methods overfit after convergence for LSTM on PTB in Figure 6 raises a major concern, and contradicts the main claim of the paper that CM methods can find a flatter minima.

Author response:
Thank you for your attentive review of our submission. Yes, for LSTM, the optimizers seem to suffer from over-fitting. A similar phenomenon was previously observed in the critical gradients paper [1].  Therefore, as one popular way to control aggressive exploration, we reduce the learning rate in this experiment. We apply a similar technique on other optimizers too which results in further improvement in their performances. However as observed in Figure 6, Adam+CM still outperforms by converging faster to a lower validation perplexity.

[1] McRae, P. A., Parthasarathi, P., Assran, M., & Chandar, S. (2022). Memory Augmented Optimizers for Deep Learning. Proceedings of International Conference on Learning Representations.

Reviewer comment (Evaluation):
The scores that are reported in this paper seem to be quite lower than what is reported in "papers with code". For example, EfficientNet-B0 on ImageNet should achieve a top-1 validation accuracy of 76.3% [https://paperswithcode.com/paper/efficientnet-rethinking-model-scaling-for], but the results reported in the paper range between 65.4-71.7%.

Author response:
We appreciate your attention to the results presented in our submission. Please note that EfficientNet-B0 on ImageNet achieves a top-1 validation accuracy of 76.3% using SGD with momentum. While we use the same setup that reproduces these results, we replace the SGD with Adam and other baselines for comparison with Adam+CM. This is because our objective is to specifically improve Adam by integrating exploration capabilities into it. This integration aims to result in the discovery of flatter solutions and improved generalization.

Reviewer comment (Experiments):
It would seem like the current experiments would be strongly affected by the batch size, but the batch size is not included in their hyperparameter search.

Author response:
We appreciate your recommendation for including batch size in the hyper-parameter search. We have added a section in Appendix (A.4.1) focused on comparing different batch sizes, particularly on CIFAR10/100 experiments. In summary, the results indicate that the reported batch size (=64) used in our experiments indeed results in the best performance when compared with other batch sizes.

Reviewer comment (Experiments):
The escape ratio in Figure 5 is monotonically increasing with the sharpness coefficient. What happens if it is increased further?

Author response:
As we increase the sharpness coefficient further, all the optimizers tend to escape that sharp minima more often and hence they all eventually achieve the escape ratio of 1. But through Fig 5, we wanted to highlight that the sharpness of that minima doesn’t need to be high for Adam+CM since it will be able to explore the loss landscape better by escaping such minima.

Reviewer comment (Experiments):
(minor) The paper would be more convincing if larger-scale empirical experiments on NLP were conducted.

Author response:
We would like to thank the reviewer for this suggestion. Experiments involving larger-scale NLP often come with the caveat of a sophisticated LR scheduler that is designed particularly for AdamW. This LR scheduler has a huge impact on the optimization trajectory and any modification in the optimization process may result in sub-optimal performance. Experiments with Adam+CM on transformers may, therefore, require a careful exploration of LR scheduling, a consideration that falls beyond the scope of our paper. In contrast, our primary objective in this work is to present a new perspective on optimization. We focus on integrating exploration capabilities into adaptive optimizers, that could result in finding flatter solutions and improved generalization.

Reviewer comment (Theory):
Prioritizing momenta with large gradient norms is not well motivated. It is unclear if and how this choice affects the performance of the method.

Author response:
**Response:**  Thank you for pointing out the need to explicitly provide a motivation for using gradient norms as priorities. The intuitive motivation using large gradient norms comes from the data-pruning approach, which is also discussed in recent works. For example, [1] mentions that by using a large gradient norm, their approach finds a smaller set of gradients to store that will aid in the optimization process. Furthermore, [2] identified that samples with larger gradient norms can be used to identify a smaller set of training data that is important for generalization, and hence prune the data as a pre-processing step. 

[1] McRae, P. A., Parthasarathi, P., Assran, M., & Chandar, S. (2022). Memory Augmented Optimizers for Deep Learning. Proceedings of International Conference on Learning Representations.

[2] Paul, M., Ganguli, S., & Dziugaite, G. K. (2021). Deep learning on a data diet: Finding important examples early in training. Advances in Neural Information Processing Systems, 34, 20596-20607

Reviewer comment (Theory):
The paper seems to associate flat minima with global minima or lower loss values in some of the toy examples, which is not always the case in practice. Indeed, the training loss of a flat minimum is usually higher than that of a sharp minimum on real-world datasets.

Author response:
**Response:** While we agree that a flat minimum doesn't always result in low loss values, several studies show that flat loss surfaces generalize better [3], [4]. This phenomenon is also seen in our experiments where we show a decrease in sharpness as well as improved validation performance on CIFAR10/100 datasets simultaneously. Moreover, we also show how changing hyper-parameter C affects the sharpness of the solution. 

[3] Foret, P., Kleiner, A., Mobahi, H., & Neyshabur, B. (2020). Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412.

[4] Andriushchenko, M., & Flammarion, N. (2022). Towards understanding sharpness-aware minimization. International Conference on Machine Learning (pp. 639-668). PMLR.

Reviewer comment (Experiments):
The proposed method is partly motivated by the performance gap between SGD and Adam, but is not compared to SGD in the experiments. Moreover, it would be interesting to see if critical momenta improve the performance of SGD as well.

Author response:
**Response:** The focus of our paper is on improving adaptive optimizers (particularly Adam).  In line with previous work on the topic, our analysis of existing optimizers has revealed that the poor generalization of adaptive optimizers in settings like image classification stems from a limited exploration of the loss surface. Our proposed method Adam+CM is a way to promote exploration while preserving the fast adaptive nature of such optimizers. Thus, we evaluate Adam+CM by comparing it with existing variants of Adam. 

Although we think this comparison extends beyond the paper’s intended scope,  we ran an experiment with a non-adaptive optimizer in response to the reviewer’s suggestion. We perform a hyperparameter grid search on CIFAR10/100 experiments and report the results in Appendix A.5. We observed that SGD, SGD+CG, and SGD+CM perform similarly but there is a speed up by both memory-augmented optimizers.

Reviewer comment (Experiments):
The baselines are not particularly strong, and it is a bit strange that Adam+SAM performs almost the same or even worse than vanilla Adam on image classification tasks.

Author response:
**Response:** This is indeed what we consistently observed in our experiments. Note that SAM was initially introduced to enhance the flatness of solutions in SGD. To the best of our knowledge, the combination of Adam and SAM, i.e. Adam+SAM, has not been widely adopted in practice by the community for image classification in ResNet-based models. Given that our approach aims to promote exploration to identify flatter solutions, it was nevertheless a logical choice to use the addition of SAM to Adam as a baseline for comparison.

Reviewer comment (Experiments):
The results shown in Fig. 5 could be sensitive to hyperparameters such as learning rate, which is not discussed in the paper.

Author response:
**Response:** The experiments associated with the toy examples in Fig 5 correspond to the same learning rate (=0.1) for comparison of our proposed optimizer with the baselines. This learning rate was obtained by performing a grid search with set [0.1, 0.01, 0.001, 0.0001]. We have included these hyper-parameters details in the appendix of the revision.