Promoting Exploration in Memory-Augmented Adam using Critical Momenta
Abstract
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}.
Cite
@article{arxiv.2307.09638,
title = {Promoting Exploration in Memory-Augmented Adam using Critical Momenta},
author = {Pranshu Malviya and Gonçalo Mordido and Aristide Baratin and Reza Babanezhad Harikandeh and Jerry Huang and Simon Lacoste-Julien and Razvan Pascanu and Sarath Chandar},
journal= {arXiv preprint arXiv:2307.09638},
year = {2024}
}
Comments
Published in Transactions on Machine Learning Research