# Promoting Exploration in Memory-Augmented Adam using Critical Momenta

arXiv: https://arxiv.org/abs/2307.09638 · HTML: https://mlsanity.com/paper/2307.09638/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 38 of its 65 references that are in the map.
- First posted on arXiv 18 Jul 2023.

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Exploring Generalization in Deep Learning](https://mlsanity.com/paper/1706.08947/) (2017)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectives](https://mlsanity.com/paper/1407.0202/) (2014)
- [Entropy-SGD: Biasing Gradient Descent Into Wide Valleys](https://mlsanity.com/paper/1611.01838/) (2016)
- [Lookahead Optimizer: k steps forward, 1 step back](https://mlsanity.com/paper/1907.08610/) (2019)
- [ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks](https://mlsanity.com/paper/2102.11600/) (2021)
- [AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients](https://mlsanity.com/paper/2010.07468/) (2020)
- [When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations](https://mlsanity.com/paper/2106.01548/) (2021)
- [The Marginal Value of Adaptive Gradient Methods in Machine Learning](https://mlsanity.com/paper/1705.08292/) (2017)
- [Surrogate Gap Minimization Improves Sharpness-Aware Training](https://mlsanity.com/paper/2203.08065/) (2022)
- [Efficient Sharpness-aware Minimization for Improved Training of Neural Networks](https://mlsanity.com/paper/2110.03141/) (2021)
- [An Empirical Investigation of the Role of Pre-training in Lifelong Learning](https://mlsanity.com/paper/2112.09153/) (2021)
- [Towards Efficient and Scalable Sharpness-Aware Minimization](https://mlsanity.com/paper/2203.02714/) (2022)
- [Towards Theoretically Understanding Why SGD Generalizes Better Than ADAM in Deep Learning](https://mlsanity.com/paper/2010.05627/) (2020)
- [Analysis and Design of Optimization Algorithms via Integral Quadratic Constraints](https://mlsanity.com/paper/1408.3595/) (2014)
- [On Empirical Comparisons of Optimizers for Deep Learning](https://mlsanity.com/paper/1910.05446/) (2019)
- [Adaptive Gradient Methods at the Edge of Stability](https://mlsanity.com/paper/2207.14484/) (2022)
- [Sharpness-Aware Training for Free](https://mlsanity.com/paper/2205.14083/) (2022)

## Abstract (arXiv)

Adaptive gradient-based optimizers, notably Adam, have left their mark in training large-scale deep learning models, offering fast convergence and robustness to hyperparameter settings. However, they often struggle with generalization, attributed to their tendency to converge to sharp minima in the loss landscape. To address this, we propose a new memory-augmented version of Adam that encourages exploration towards flatter minima by incorporating a buffer of critical momentum terms during training. This buffer prompts the optimizer to overshoot beyond narrow minima, promoting exploration. Through comprehensive analysis in simple settings, we illustrate the efficacy of our approach in increasing exploration and bias towards flatter minima. We empirically demonstrate that it can improve model performance for image classification on ImageNet and CIFAR10/100, language modelling on Penn Treebank, and online learning tasks on TinyImageNet and 5-dataset. Our code is available at urlhttps://github.com/chandar-lab/CMOptimizer.