Momentum method saves work with fewer gradient corrections
Randomized Momentum Transport with Cheap Predictors
Machine Learning
Summary
Optimization methods that help computers learn faster often need extra calculations to adjust their steps, which can be slow. The authors found a way to randomly skip some of these costly corrections and use cheaper guesses instead, while still keeping the method accurate enough. This approach saves computing effort and can improve training results on real tasks like image recognition. They also showed their method mathematically and tested it on standard datasets.
What this means in practice
- •For machine learning engineers: Cut down the cost of gradient adjustment steps during training without losing accuracy by selectively using cheaper approximations.
- •For autonomous system developers: Improve efficiency of real-time learning algorithms by reducing expensive computations in momentum updates while maintaining convergence quality.
Authors
El Mahdi Chayti, Martin Jaggi
Abstract
Momentum methods such as STORM and second-order momentum track the gradient by adding a transport correction to momentum, at the cost of an extra oracle call per iteration. We introduce randomized momentum transport, which evaluates the correction with probability $q$ and otherwise uses a cheap prediction, with importance weighting that preserves the conditional mean and adds $(q^{-1}-1)$ times the prediction loss to the variance. Under transport-bias, second-moment, and residual-control assumptions, linear minimization oracle (LMO) updates, including normalized, sign, and exact spectral (Muon) directions, retain an $O(T^{-1/3})$ stationarity bound for each fixed $q>0$ with a single-gradient initialization. For unnormalized STORM, some $q<1$ lowers the leading coefficient of our sufficient gradient budget if and only if the gradient-difference noise modulus exceeds the smoothness constant, by at most about $10\%$. An online diagonal predictor admits a regret-to-convergence guarantee. At approximately equal oracle cost, querying a quarter of the corrections lowers the final training loss by $18$ to $43\%$ on MNIST1D and raises the test accuracy of Hessian-vector transport on CIFAR-10 by $1.5$ points.