MoMo: Momentum Models for Adaptive Learning Rates

Fabian Schaipp, Ruben Ohana, Michael Eickenberg, Aaron Defazio, Robert M. Gower

Keywords: adaptive learning rates, adaptive step sizes, stochastic gradient, SGD, polyak step size, model based methods, momentum

TL;DR: How to combine momentum with model based stochastic gradient methods

Abstract:

Training a modern machine learning architecture on a new task requires extensive
learning-rate tuning, which comes at a high computational cost. In this work, we develop new adaptive learning rates that can be used with any momentum method, and require less tuning to perform well. We first develop MoMo, a Momentum Model
based adaptive learning rate for SGD-M (Stochastic gradient descent with momentum). MoMo uses momentum estimates of the batch losses and gradients sampled at each iteration to build a model of the loss function. Our model also makes use of
any known lower bound of the loss function by using truncation, e.g. most losses
are lower-bounded by zero. We then approximately minimize this model at each
iteration to compute the next step. We show how MoMo can be used in combination
with any momentum-based method, and showcase this by developing MoMo-Adam - which is Adam with our new model-based adaptive learning rate. Additionally, for losses with unknown lower bounds, we develop on-the-fly estimates of a lower
bound, that are incorporated in our model. Through extensive numerical experiments, we demonstrate that MoMo and MoMo-Adam improve over SGD-M and Adam in terms of accuracy and robustness to hyperparameter tuning for training image
classifiers on MNIST, CIFAR10, CIFAR100, Imagenet, recommender systems
on the Criteo dataset, and a transformer model on the translation task IWSLT14.

Primary Area: optimization

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 4134