Promoting Exploration in Memory-Augmented Adam using Critical Momenta
This addresses generalization issues in adaptive optimizers for deep learning practitioners, though it is incremental as it builds on existing Adam methods.
The paper tackles the problem of Adam optimizers converging to sharp minima, which hurts generalization, by proposing a memory-augmented version that uses critical momenta to encourage exploration toward flatter minima, improving performance on tasks like ImageNet classification and language modeling.
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 \url{https://github.com/chandar-lab/CMOptimizer}.