Randomized Momentum Transport with Cheap Predictors

El Mahdi Chayti; Martin Jaggi

cs.LGOct 9, 2026

Randomized Momentum Transport with Cheap Predictors

Authors: El Mahdi Chayti, Martin Jaggi

Organizations: Machine Learning and Optimization Laboratory (MLO), EPFL

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.

Figures & tables

Unit ball of	Direction	$S_{\mathcal{K}}(\nabla f)$	$R$
$\ell_{2}$	$m/\lVert m\rVert_{2}$	$\lVert\nabla f\rVert_{2}$	$1$
$\ell_{\infty}$	$\operatorname{sign}(m)$	$\lVert\nabla f\rVert_{1}$	$\sqrt{d}$
Spectral	Polar factor	$\lVert\nabla f\rVert_{\rm nuc}$	$\sqrt{\min\{d_{1},d_{2}\}}$

Table 1: Direction sets and stationarity measures. Each unit ball has inradius

r_{\mathcal{K}}=1

. The spectral row concerns matrix parameters in

\mathbb{R}^{d_{1}\times d_{2}}

Figure 1: Leading gradient-budget ratio ( 30 ). Lower is better, and the dashed line is full correction. For

\omega<1

the minimum lies below one; the

\omega=1

curve is minimized at

q=1

Figure 2: Controlled quadratic. Left: run average of

\lVert\nabla f(x_{t})\rVert

relative to

q=1

, zero (solid) and diagonal (dashed) predictor, with paired

95\%

intervals over

10

seeds; the band marks

\hat{\omega}\in[0.9,1.7]

on the

q=1

runs. Right: residual ratio of the diagonal predictor (Appendix G ) and the conditional benchmark

1/(1+\hat{\omega})

of an exact predictor.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Step	Training loss	$b=1$	$b=4$	$b=32$	$b=256$
$0$	$3.11$	$0.088$	$0.33$	$2.6$	$21$
$500$	$1.58$	$0.037$	$0.15$	$1.3$	$11$
$2000$	$1.03$	$0.017$	$0.060$	$0.46$	$3.6$
$8000$	$0.38$	$0.004$	$0.020$	$0.17$	$1.4$

Appendix

Table 2: Estimated

\hat{\omega}

along the reference run, by batch size.

Batch	Oracle	$q=1$	$q=1/2$	$q=1/4$
$4$	STORM	$0.425$	$0.80\,[0.75,0.86]$	$0.70\,[0.67,0.73]$
$4$	HVP	$0.470$	$0.78\,[0.70,0.87]$	$0.66\,[0.60,0.71]$
$32$	STORM	$0.077$	$0.99\,[0.69,1.29]$	$0.82\,[0.74,0.90]$
$32$	HVP	$0.089$	$0.73\,[0.66,0.81]$	$0.57\,[0.53,0.61]$

Appendix

Table 3: MNIST1D: final training loss at

q=1

and ratio of the final loss at

q<1

to it, with paired

95\%

confidence intervals over

10

seeds, at equal oracle cost (

16000

gradient-equivalents at batch size

4

8000

at batch size

32

Batch	Oracle	$q=1$	$q=1/2$	$q=1/4$
$4$	STORM	$61.9$	$+1.55\,[+0.43,+2.67]$	$+1.31\,[+0.08,+2.54]$
$4$	HVP	$62.7$	$+0.40\,[-1.24,+2.04]$	$+2.28\,[+1.27,+3.29]$
$32$	STORM	$62.9$	$+0.18\,[-1.11,+1.47]$	$-1.75\,[-2.73,-0.77]$
$32$	HVP	$62.2$	$+0.90\,[-0.28,+2.08]$	$+1.29\,[+0.58,+2.00]$

Appendix

Table 4: MNIST1D: test accuracy at

q=1

(in

\%

, mean over

10

seeds) and paired difference at

q<1

from it (percentage points,

95\%

interval), at equal oracle cost; settings selected by training loss.

Batch	Method	$q$	Test accuracy	Difference from zero predictor	Residual ratio
$4$	STORM	$1/2$	$64.01\pm 1.61$	$+0.58\,[-0.68,+1.84]$	$0.995$
$4$	STORM	$1/4$	$63.18\pm 1.21$	$-0.01\,[-0.18,+0.16]$	$0.989$
$4$	HVP	$1/2$	$63.61\pm 1.32$	$+0.54\,[-0.31,+1.39]$	$0.995$
$4$	HVP	$1/4$	$65.06\pm 1.56$	$+0.11\,[-1.64,+1.86]$	$0.996$
$32$	STORM	$1/2$	$63.06\pm 1.85$	$-0.01\,[-1.35,+1.33]$	$0.983$
$32$	STORM	$1/4$	$61.37\pm 0.79$	$+0.23\,[-0.41,+0.87]$	$0.954$

Appendix

Table 5: MNIST1D, diagonal predictor with the settings of the zero-predictor run: test accuracy (mean and standard deviation over

10

seeds, in

\%

), paired difference from the zero predictor at the same

q

(percentage points,

95\%

interval), and residual ratio

\sum_{t}J_{t}\lVert C_{t}-B_{t}\rVert^{2}/\sum_{t}J_{t}\lVert C_{t}\rVert^{2}

(mean over seeds).

Batch	Method	Training loss	Test acc.	Iterations	$(\eta,\underline{\alpha})$
$4$	STORM, $q=1$	$0.425\pm 0.025$	$0.619$	$8000$	$(0.1,0.3)$
$4$	STORM, $q=1/2$	$0.341\pm 0.029$	$0.634$	$10710$	$(0.1,0.3)$
$4$	STORM, $q=1/4$	$0.297\pm 0.016$	$0.632$	$12840$	$(0.05,0.3)$
$4$	HVP, $q=1$	$0.470\pm 0.032$	$0.627$	$6700$	$(0.1,0.3)$
$4$	HVP, $q=1/2$	$0.365\pm 0.031$	$0.631$	$9450$	$(0.1,0.3)$
$4$	HVP, $q=1/4$	$0.308\pm 0.035$	$0.650$	$11910$	$(0.1,0.3)$

Appendix

Table 6: All MNIST1D runs: final training loss (mean and standard deviation over

10

seeds), test accuracy, number of iterations (mean over seeds), and selected

(\eta,\underline{\alpha})

Figure 4: MNIST1D, batch size

4

. Test accuracy against gradient-equivalent evaluations for gradient-difference (left) and HVP (right) transport at

q\in\{1,1/2,1/4\}

, with the zero predictor (solid; mean and standard deviation over

10

seeds) and the diagonal predictor (dashed; mean). Figure 5 shows batch size

32

Figure 5: MNIST1D, batch size

32

; same setting as Figure 4 .

$\delta$	$\hat{\omega}$ ( $q=1$ )	$q=1/2$	$q=1/4$	$q=1/2$	Residual ratio
		Zero predictor		Diagonal
$10$	$0.0011$	$0.92\,[0.91,0.94]$	$0.86\,[0.84,0.88]$	$0.92\,[0.91,0.94]$	$1.000$
$3$	$0.013$	$0.90\,[0.89,0.91]$	$0.83\,[0.82,0.84]$	$0.90\,[0.89,0.91]$	$0.993$
$1$	$0.11$	$0.91\,[0.90,0.92]$	$0.83\,[0.82,0.85]$	$0.91\,[0.90,0.92]$	$0.933$
$0.3$	$1.1$	$0.93\,[0.93,0.93]$	$0.89\,[0.88,0.89]$	$0.93\,[0.92,0.93]$	$0.489$
$0.2$	$0.91$	$1.00\,[0.99,1.00]$	$0.95\,[0.95,0.96]$	$0.99\,[0.98,1.00]$	$0.521$

Appendix

Table 7: Controlled quadratic: ratio of the run-average

\lVert\nabla f(x_{t})\rVert

q

to its value at

q=1

, with paired

95\%

confidence intervals over

10

seeds, and the residual ratio of the diagonal predictor at

q=1/2

		$q=1/2$				$q=1/4$
$\delta$	$\omega$	$\mathcal{Q}^{2/3}$	$2000$	$8000$	$32000$	$\mathcal{Q}^{2/3}$	$2000$	$8000$	$32000$
$0.25$	$16$	$1.569$	$1.137$	$1.160$	$1.211$	$1.874$	$1.236$	$1.273$	$1.356$
$0.5$	$4$	$1.187$	$1.051$	$1.052$	$1.070$	$1.347$	$1.121$	$1.128$	$1.163$
$1$	$1$	$1.016$	$0.981$	$0.977$	$0.982$	$1.091$	$1.022$	$1.020$	$1.034$
$2$	$0.25$	$0.958$	$0.949$	$0.948$	$0.951$	$0.996$	$0.975$	$0.976$	$0.982$
$4$	$0.0625$	$0.942$	$0.949$	$0.947$	$0.944$	$0.968$	$0.971$	$0.969$	$0.965$

Appendix

Table 8: Unnormalized STORM with the tuning of Theorem 4 and the variance-bounded oracle: ratio of the run-average

\lVert\nabla f(x_{t})\rVert^{2}

q

to its value at

q=1

(mean over

10

paired seeds) at budgets

2000

8000

, and

32000

, next to the fixed-budget benchmark

\mathcal{Q}(q;\omega)^{2/3}

$\delta$	$\hat{\omega}$ ( $q=1$ )	$q=1/2$	$q=1/4$	$q=1/2$	Residual ratio
		Zero predictor		Diagonal
$10$	$0.00023$	$0.79\,[0.60,0.98]$	$0.78\,[0.57,0.99]$	$0.77\,[0.60,0.94]$	$1.012$
$3$	$0.0017$	$0.95\,[0.83,1.08]$	$0.93\,[0.78,1.07]$	$0.96\,[0.84,1.07]$	$1.029$
$1$	$0.02$	$0.92\,[0.89,0.96]$	$0.87\,[0.84,0.89]$	$0.93\,[0.90,0.97]$	$1.349$
$0.3$	$0.21$	$1.04\,[0.99,1.09]$	$0.99\,[0.96,1.02]$	$1.01\,[0.95,1.08]$	$1.157$
$0.15$	$14$	$1.10\,[1.06,1.14]$	$1.08\,[1.05,1.12]$	$0.97\,[0.95,0.99]$	$0.638$

Appendix

Table 9: Symmetric curvature noise: ratio of the run-average

\lVert\nabla f(x_{t})\rVert

q

to its value at

q=1

, with paired

95\%

confidence intervals over

10

seeds, and the residual ratio of the diagonal predictor at

q=1/2

Method	$q$	$\eta$	$\underline{\alpha}$	Updates	Test accuracy	Difference from $q=1$	Time
STORM	$1$	$0.27$	$0.3$	$17550$	$87.34\pm 0.70$	baseline	$1053$
STORM	$1/2$	$0.27$	$0.3$	$23376$	$88.05\pm 1.13$	$+0.72\,[-0.01,+1.44]$	$1063$
STORM	$1/4$	$0.03$	$0.3$	$28091$	$81.64\pm 1.12$	$-5.69\,[-7.12,-4.26]$	$1105$
HVP	$1$	$0.09$	$0.1$	$10570$	$85.99\pm 0.38$	baseline	$934$
HVP	$1/2$	$0.09$	$0.3$	$16218$	$86.95\pm 0.22$	$+0.96\,[+0.64,+1.28]$	$966$
HVP	$1/4$	$0.09$	$0.3$	$22258$	$87.50\pm 0.22$	$+1.52\,[+0.96,+2.07]$	$1046$

Appendix

Table 10: CIFAR-10: selected

(\eta,\underline{\alpha})

, mean number of updates within the budget, test accuracy (mean and standard deviation over

5

seeds, in

\%

), paired difference from

q=1

(percentage points,

95\%

interval), and training time in seconds.

Method	$q$	Test accuracy	Difference from zero predictor	Residual ratio
STORM	$1/2$	$88.29\pm 0.33$	$+0.24\,[-1.10,+1.58]$	$0.992$
STORM	$1/4$	$80.99\pm 0.88$	$-0.65\,[-1.33,+0.02]$	$0.991$
HVP	$1/2$	$86.93\pm 0.32$	$-0.01\,[-0.50,+0.48]$	$0.996$
HVP	$1/4$	$87.51\pm 0.47$	$+0.00\,[-0.47,+0.47]$	$0.995$

Appendix

Table 11: CIFAR-10, diagonal predictor: test accuracy (mean and standard deviation over

5

seeds, in

\%

), paired difference from the zero predictor at the same

q

(percentage points,

95\%

interval), and residual ratio

\sum_{t}J_{t}\lVert C_{t}-B_{t}\rVert^{2}/\sum_{t}J_{t}\lVert C_{t}\rVert^{2}

(mean over seeds).

Explore similar work

May 7, 2026cs.LG

Accelerating LMO-Based Optimization via Implicit Gradient Transport

Recent optimizers such as Lion and Muon have demonstrated strong empirical performance by normalizing gradient momentum via linear minimization oracles (LMOs). While variance reduction has been explored to accelerate LMO-based methods, it typically incurs substantial computational overhead due to additional gradient evaluations. At the same time, the theoretical understanding of LMO-based methods remains fragmented across unconstrained and constrained formulations. Motivated by these limitations, we propose \emph{LMO-IGT}, a new class of stochastic LMO-based methods leveraging implicit gradient transport (IGT). We further introduce a unified framework for stochastic LMO-based optimization together with a new stationarity measure, the \emph{regularized support function} (RSF), which bridges gradient-norm and Frank--Wolfe-gap notions within a common framework. By evaluating stochastic gradients at transported points, LMO-IGT accelerates convergence while retaining the single-gradient-per-iteration structure of standard stochastic LMO. Our analysis establishes that stochastic LMO achieves an iteration complexity of

\mathcal{O}(\varepsilon^{-4})

, variance-reduced LMO achieves

\mathcal{O}(\varepsilon^{-3})

at the cost of additional gradient evaluations, and LMO-IGT achieves

\mathcal{O}(\varepsilon^{-3.5})

using only a single stochastic gradient per iteration. Empirically, LMO-IGT consistently improves over stochastic LMO counterparts with negligible overhead. Among its instantiations, Muon-IGT achieves the strongest overall performance across evaluated settings, demonstrating that IGT provides an effective and practical acceleration mechanism for modern LMO-based optimization.

Won-Jun Jang, Si-Hyeon Lee

School of Electrical Engineering · Korea Advanced Institute of Science and Technology (KAIST) · Daejeon, South Korea

Muon Optimizer Stochastic Optimization Convergence

Ask PDF

Feb 6, 2026math.OC

RanSOM: Second-Order Momentum with Randomized Scaling for Constrained and Unconstrained Optimization

Second-order momentum corrects the lag of momentum behind the current gradient with a Hessian-vector product (HVP), but existing corrections need either a Lipschitz Hessian or extra forward and backward passes at an auxiliary point. We show that randomizing the step length removes this trade-off: if the step is exponentially distributed with mean

η

, a Stein-type identity makes

η

times one HVP at the new iterate an estimate of the gradient change that is unbiased in expectation over the step, and this HVP shares its forward pass with the next gradient. The resulting method, RanSOM-E, works with normalized, sign, and Muon-type steps, and a Beta-distributed variant, RanSOM-B, keeps Frank-Wolfe iterates feasible. For objectives with Lipschitz gradient and heavy-tailed gradient and Hessian noise with tail indices

p, q \in (1,2]

, RanSOM-E attains the optimal oracle-complexity exponents for every pair

(p,q)

, including

p \neq q

, without a Lipschitz Hessian or clipping (

\mathcal{O}(\varepsilon^{-3})

for

p = q = 2

). On CIFAR-10 with ResNet-20, we detect no accuracy difference from existing second-order momentum methods, and its mean update time is lower than that of the auxiliary-point method.

El Mahdi Chayti, Mohamad Fakhouri

Machine Learning and Optimization Laboratory (MLO) EPFL, Switzerland

Second-Order Optimization Stochastic Optimization Convergence Heavy-Tailed Stochastic Optimization

Ask PDF

Jun 18, 2024cs.LG

Accelerated Stochastic Min-Max Optimization Based on Bias-corrected Momentum

Lower-bound analyses for nonconvex strongly-concave minimax optimization problems have shown that stochastic first-order algorithms require at least

\mathcal{O}(\varepsilon^{-4})

sample complexity to find an

\varepsilon

-stationary point. Some works indicate that this complexity can be improved to

\mathcal{O}(\varepsilon^{-3})

when the stochastic loss gradient is Lipschitz continuous. The question of achieving enhanced convergence rates under distinct conditions, remains open. In this work, we address this question for optimization problems that are nonconvex in the minimization variable and strongly concave or Polyak-Lojasiewicz (PL) in the maximization variable. We introduce novel bias-corrected momentum algorithms utilizing efficient Hessian-vector products. We establish convergence conditions and demonstrate a lower iteration complexity of

\mathcal{O}(\varepsilon^{-3})

for the proposed algorithms. The effectiveness of the proposed method is validated through applications to robust logistic regression and robust adaptive cruise control.

Haoyuan Cai, Sulaiman A. Alghunaim, Ali H. Sayed

Ecole Polytechnique Fédérale de Lausanne, Switzerland · Kuwait University, Kuwait

Momentum Methods Nonconvex Stochastic Optimization Min-Max Optimization Robust Optimization Nonconvex Optimization

Ask PDF