Randomized Momentum Transport with Cheap Predictors
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.
Related
- Randomized average block coordinate descent method with heavy-ball momentum for tensor least squares problem under the t-product
- Annealed Sinkhorn with Momentum: Certified Unregularized Optimal Transport in Linear Memory
- Muon with Nesterov Momentum: Heavy-Tailed Noise and (Randomized) Inexact Polar Decomposition
- Fast, cheap and sustainable: armament portfolio arbitration under attrition, lead time and political discontinuity
- (Cheap) Stochastic Policy Gradient Converges with High Probability for Linear Quadratic Regulator
BibTeX
@misc{chayti2026randomized,
title = {{Randomized Momentum Transport with Cheap Predictors}},
author = {El Mahdi Chayti and Martin Jaggi},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.13789},
eprint = {2610.13789},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.13789},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Randomized Momentum Transport with Cheap Predictors
AU - El Mahdi Chayti
AU - Martin Jaggi
PY - 2026
DA - 2026/10/09
PB - arXiv
DO - 10.48550/arxiv.2610.13789
UR - https://arxiv.org/abs/2610.13789
AB - 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.
ER -