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

Accelerating LMO-Based Optimization via Implicit Gradient Transport
May 7, 2026Won-Jun Jang, Si-Hyeon LeeMuon Optimizer Stochastic Optimization Convergence
School of Electrical Engineering · Korea Advanced Institute of Science and Technology (KAIST) · Daejeon, South Korea
RanSOM: Second-Order Momentum with Randomized Scaling for Constrained and Unconstrained Optimization
Feb 6, 2026El Mahdi Chayti, Mohamad FakhouriSecond-Order Optimization Stochastic Optimization Convergence
Machine Learning and Optimization Laboratory (MLO) EPFL, Switzerland
Accelerated Stochastic Min-Max Optimization Based on Bias-corrected Momentum
Jun 18, 2024Haoyuan Cai, Sulaiman A. Alghunaim, Ali H. SayedMomentum Methods Nonconvex Stochastic Optimization
Ecole Polytechnique Fédérale de Lausanne, Switzerland · Kuwait University, Kuwait

May 7, 2026, cs.LGJ/K move · Enter open · S save

Accelerating LMO-Based Optimization via Implicit Gradient Transport

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Accelerating LMO-Based Optimization via Implicit Gradient Transport

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

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