PaperlayerView exact revision on arXiv

Source-linked AI summary

Randomized Momentum Transport with Cheap Predictors

El Mahdi Chayti, Martin Jaggi

arXiv:2610.13789v1cs.LGmath.OC

TL;DR

Momentum methods use transport corrections to track changing gradients, but evaluating each correction costs an extra oracle call. This paper randomizes correction queries and uses cheap predictions, preserving the conditional mean while controlling the added variance. It proves fixed-q stationarity guarantees and reports lower training loss on MNIST1D and higher HVP test accuracy on CIFAR-10 at equal oracle cost.

  • Problem

    Momentum methods track changing gradients with transport corrections, but each correction requires an extra oracle call per iteration.

    Source: §1 · p. 1

  • Method

    The method queries a transport correction with probability q and otherwise uses an importance-weighted cheap prediction, preserving the conditional mean while adding variance proportional to prediction loss.

    Source: §1 · p. 1

  • Results

    At equal oracle cost, randomized queries lower final training loss on MNIST1D for both transports and raise CIFAR-10 test accuracy for HVP transport; fixed-q LMO updates retain an O(T^-1/3) stationarity bound.

    Source: §1 · pp. 1–2

  • Takeaways & Limitations

    Querying corrections less often can trade oracle calls for iterations while retaining the stated stationarity rate for fixed q.

    Source: §7 · p. 8

    The Euclidean-norm analysis can have dimension-dependent factors for sign directions; measuring smoothness in the geometry’s own norm can yield sharper dimension dependence.

    Source: The same second-moment increment holds… · p. 4

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper randomizes transport-correction queries while preserving their conditional mean, then analyzes the resulting variance cost and convergence. It also studies when this saves gradient budget, how online prediction can help, and how the approach differs from related estimators.

1 INTRODUCTION

  • The randomized estimator preserves the correction’s conditional mean and covariance with the fresh gradient, adding exactly (q^-1−1) times the prediction loss to variance.The correction is queried with probability q; otherwise the cheap prediction is used.Source: §1 · p. 1
  • LMO updates with gradient-difference or HVP transport retain an O(T^-1/3) stationarity bound for fixed q > 0 with decreasing momentum weights and single-gradient initialization.The bound has explicit constants; q enters through (q^-1−1) times average prediction loss, while initial error contributes O(T^-1).Source: §1 · p. 1
  • For unnormalized STORM with a zero predictor, some q < 1 lowers the leading sufficient gradient budget if and only if δ > L, by at most about 10%.The expected gradient cost is 1 + q per iteration, and the leading budget term is order ε^-3.Source: §1 · p. 1
  • A diagonal predictor with O(d) state is trained only on queried corrections and has a regret-to-convergence guarantee against fixed comparators.When a comparator has uniformly smaller expected loss than the zero predictor, the learned bound improves once q(T−1) exceeds a sufficient threshold linear in d.Source: §1 · pp. 1–2
  • At equal oracle cost, randomized queries lower final MNIST1D training loss for both transports and raise CIFAR-10 test accuracy for HVP transport.On a controlled-curvature-noise quadratic, the benefit reverses as curvature noise decreases, and the diagonal predictor removes most of the reducible residual.Source: §1 · p. 2

1.1 Related Work

  • Unlike PAGE and MARINA, which switch between gradient refreshes and recursive corrections, this estimator takes a fresh gradient every iteration and randomizes only the transport correction.Its predictor is a shift for a transport correction within one momentum recursion, with accuracy controlled through online regret.Source: §1.1 · p. 2

2 RANDOMIZED TRANSPORT

The method initializes momentum from one stochastic gradient, then updates it using fresh gradients and randomized transport corrections while choosing directions through an LMO. It covers gradient-difference and HVP corrections, with oracle work accounted for separately from predictor and direction overhead.

2 RANDOMIZED TRANSPORT

  • Momentum is initialized with m0 = g0, then updated using fresh stochastic gradients and the randomized correction; q = 1 recovers the original transport recurrence.The iterates use measurable LMO directions and satisfy a step-length bound determined by the Euclidean radius of the direction set.Source: §2 · pp. 2–3
  • The query probability q is fixed, and each query decision is conditionally independent of the correction, predictor, and fresh gradient given the past.The predictor may use past data and cheap information from the current sample, but not the correction whose query it controls.Source: §2 · p. 3

2.1 Two Transport Oracles

  • For gradient-difference transport, the correction uses the same fresh sample at the current and previous iterates and costs one extra gradient at the previous iterate.Under the stated centered similarity condition, its transport bias parameter is zero and its second-moment parameter is δ.Source: §2.1 · p. 3
  • For HVP transport, Ct = Htst; a Lipschitz Hessian yields a Taylor-drift bound, while an unbiased HVP at the current iterate can reuse the current gradient’s derivative graph.The stated parameters are b = ρ/2 and v = τ; both HVP endpoints and gradient differences satisfy ∥μt∥ ≤ L∥st∥.Source: §2.1 · p. 3

2.2 Oracle Accounting

  • For gradient-difference transport, expected oracle work is (1 + q)T − q; for HVP transport it is T + cHq(T−1).Predictor, direction, and systems overhead are counted separately in runtime comparisons.Source: §2.2 · p. 3

3 RESIDUAL-DEPENDENT CONVERGENCE

The analysis decomposes tracking and descent errors to bound LMO stationarity under randomized transport. With a uniform residual control, the rate remains O(T^-1/3) for every fixed q > 0, while prediction loss contributes an explicit query-dependent variance term.

The same second-moment increment holds for centering

3 RESIDUAL-DEPENDENT CONVERGENCE

  • Theorem 2 bounds the LMO stationarity measure by terms arising from transport variance, descent, and momentum tracking under Assumptions 1–3.The proof treats tracking error as a weighted initial error, drift sum, and martingale sum; decreasing momentum weights control the initial-error contribution.Source: The same second-moment increment holds… · p. 4

The same second-moment increment holds for centering

  • With η and α scaled as T^-2/3 and a uniform residual bound, the complete stationarity bound is O(T^-1/3) for fixed q > 0.For a zero predictor, both transport oracles satisfy the residual control with a bound based on v² + M²; in general prediction accuracy enters through v² + (q^-1−1)κ̄².Source: The same second-moment increment holds… · p. 4
  • The update geometry enters through the Euclidean radius R and inradius rK, converting tracking error into the set’s stationarity measure.For sign directions the alignment multiplier is dimension-dependent; measuring smoothness in the norm associated with K can yield sharper dimension dependence.Source: The same second-moment increment holds… · p. 4
  • The bound also covers two inexact-LMO models, with additive penalties for suboptimal alignment or for Euclidean distance from an exact direction.The distance-based model requires ϵL < rK.Source: The same second-moment increment holds… · p. 4

4 GRADIENT COMPLEXITY

The section derives a finite-horizon sufficient gradient budget for unnormalized randomized STORM and compares its leading coefficient with full correction. The comparison identifies when query skipping can improve that certificate and how large the gain can be.

4 GRADIENT COMPLEXITY

  • Theorem 4 gives an expected gradient budget to reach expected gradient norm ε for unnormalized STORM with zero predictor and gradient-difference correction.Its guarantee uses a decreasing momentum schedule and accounts for the query probability in the correction cost.Source: §4 · p. 5
  • The horizon-dependent guarantee is evaluated for T ≥ Tcap using η = (Δ/(KqT))^1/3, with a separate derivation for smaller horizons.The resulting ε-iteration guarantee has expected gradient budget proportional to (1 + q)T.Source: §4 · p. 5
  • For ω < 1, the leading sufficient-budget coefficient can fall below the full-correction benchmark, with a maximum reduction of about 9.7%.The ratio comparison optimizes c separately at each q; the figure marks full correction with a dashed line and shows the ω = 1 curve minimized at q = 1.Source: §4 · p. 5, Figure 1 · p. 5
  • Skipping a query randomizes both the centered correction and its mean; the mean contribution includes a term proportional to (1−q)L²/q.A predictable surrogate Bt can remove part of that mean contribution, and if a correction costs cC gradient-equivalent units, the query-cost factor becomes 1 + cCq.Source: §4 · p. 5

5 LEARNING CHEAP PREDICTORS

This section studies how cheap predictors can reduce the variance cost of randomized transport, then gives an online diagonal predictor and its regret-to-convergence guarantee. The resulting convergence and oracle bounds quantify when learning improves on the zero predictor.

5.1 What a Predictor Can Remove

  • A predictor cannot reduce conditional oracle variance; it can reduce only the mean mismatch, and linear predictions of the parameter displacement have residuals that vanish when parameters do not move.For a known quadratic component, the corresponding exact predictor removes that component's transport from the residual.Source: §5.1 · p. 6

5.2 An Online Diagonal Predictor

  • The online diagonal predictor learns only from queried corrections and uses O(d) memory and arithmetic, with no additional oracle queries.Its update uses a convex prediction loss and may assign negative coefficients.Source: §5.2 · p. 6
  • The fixed-comparator regret guarantee requires no approximation assumption, and the learned bound is never worse than the zero-predictor bound by more than the regret term.The comparator is fixed across iterations, with its loss evaluated along the learner's trajectory.Source: §5.2 · p. 6

5.3 Convergence and Oracle Complexity

  • The complete learned-predictor bound retains O(T^-1/3) stationarity and gives an expected gradient count for reaching ε-stationarity.The section also accounts for HVP work separately and the predictor's O(d) arithmetic and O(d) state.Source: §5.3 · p. 7
  • For q < 1, the leading learned-predictor gradient-budget coefficient is strictly below the zero-predictor coefficient when the comparator loss κ is below VC.The comparisons include the 1 + q gradient-work factor; the guarantee uses a horizon-wide bound on κT.Source: §5.3 · p. 7

6 EXPERIMENTS

The experiments compare randomized transport with full correction at equal oracle cost across MNIST1D, CIFAR-10, and a controlled quadratic. Lower correction-query probabilities often improve outcomes, but the effects vary by method and noise regime.

MNIST1D

  • At q = 1/4, MNIST1D final loss relative to q = 1 is 0.70 for STORM and 0.66 for HVP at batch size 4, and 0.82 and 0.57 at batch size 32.Every q < 1 reaches lower final loss at equal oracle cost except STORM at q = 1/2 and batch size 32, where the difference is not significant.Source: §6 · pp. 7–8
  • On MNIST1D, the diagonal predictor removes at most 4.6% of the residual and shows no statistically resolved test-accuracy change.The residual reduction is larger at batch size 32, where ω̂ is larger.Source: §6 · p. 8

Controlled curvature noise

  • At δ = 3, q = 1/4 lowers the controlled-quadratic run-average gradient norm by 17%; at δ = 0.03, it raises it by 37%.The sign changes where the measured ω̂ on q = 1 runs lies between 1.15 and 1.7; these normalized runs are qualitatively consistent with a mechanism stated for unnormalized STORM.Source: §6 · p. 8, Figure 2 · p. 8

CIFAR-10

  • +1.52 points ([0.96, 2.07]) in CIFAR-10 test accuracy at q = 1/4 for HVP transport; STORM changes by -5.69 points ([-7.12, -4.26]) at q = 1/4.The HVP comparison is at an equal budget of 35100 gradient-equivalents; the diagonal predictor has no statistically resolved test-accuracy change.Source: §6 · p. 8

7 CONCLUSION

The conclusion emphasizes that transport corrections need not be queried every iteration: randomized querying preserves the conditional mean and retains the fixed-q stationarity rate. The paper also reports the authors' use of a language model during manuscript preparation.

7 CONCLUSION

  • Querying transport corrections with probability q and otherwise using a cheap prediction preserves the conditional mean, adds prediction-loss-dependent variance, and retains O(T^-1/3) stationarity for fixed q.The conclusion notes that saved queries buy iterations, which helps most with expensive or noisy corrections.Source: §7 · p. 8

AI Use Statement

  • A language model assisted with manuscript restructuring and editing, proof checking, experiment-code drafting, and related-work suggestions; the authors verified its output and take responsibility for the work.Source: AI Use Statement · p. 9

SUPPLEMENTARY MATERIAL

The supplementary material establishes transport identities and bounds for gradient-difference and second-order corrections, then develops tracking bounds used for LMO convergence. Deterministic schedules preserve martingale orthogonality in the tracking argument.

A TRANSPORT IDENTITIES AND SPECIAL CASES

  • The importance-weighted correction has conditional mean μt, and for predictable Bt its residual separates oracle variance from the mean mismatch.The identity also permits predictors using cheap current-sample information.Source: A TRANSPORT IDENTITIES AND SPECIAL CASES · p. 11

A.1 Gradient-Difference Transport

  • For gradient-difference transport, smoothness bounds the correction norm by L times the displacement norm, while the centered variance bound uses δ.Source: §A.1 · p. 11

A.2 Second-Order Transport at Either Endpoint

  • Second-order transport at either endpoint satisfies the transport-bias and mean-norm bounds, with zero-predictor centered variance at most [τ^2 + (1 - q)L^2]q^-1||st||^2.Source: §A.2 · p. 11

B.1 General Deterministic Momentum Schedule

  • The weighted tracking bound yields an LMO stationarity bound, and deterministic schedules preserve martingale orthogonality because the noise weights are deterministic.The stationarity bound adds the initial objective gap, smoothness, and LMO-radius terms.Source: §B.1 · p. 12

B.2 Bounds for the Decreasing Schedule

  • The decreasing-schedule bounds use telescoping products and geometric-series estimates, including the case α = 1.Source: §B.2 · p. 13

B.3 Proof of Theorem 2

These proof sections establish descent and tracking bounds, extend the analysis to inexact directions, and derive a finite-horizon gradient budget for unnormalized STORM. The optimized leading-constant formulas have a stated positive-parameter scope.

B.3 Proof of Theorem 2

  • The descent argument yields (16) using only m0 through mT−1; when T = 1, all correction sums are empty.The proof invokes the descent statement of Lemma 8.Source: §B.3 · p. 13

B.4 Inexact Directions

  • Inexact directions add an error term to the descent bound, while the tracking bound replaces R with R+ϵL.The descent inequality gains ηϵt, and the resulting stationarity bound is scaled by (1−ϵL/rK).Source: §B.4 · p. 14

C PROOF OF THE UNNORMALIZED STORM BOUND

  • The unnormalized STORM proof telescopes a potential combining objective gap and tracking error, then substitutes the resulting bound to prove Theorem 4.The potential is Φt = E[f(xt) − f∗] + (η/αt)E∥et∥².Source: C PROOF OF THE UNNORMALIZED STORM BOUND · pp. 14–15

C.1 Leading Constants and Feasible Parameters

  • The finite-accuracy guarantee gives E∥∇f(xU)∥ ≤ ε and an expected gradient budget; the optimized formulas assume ∆, σ, and δ are positive.Theorem 4 itself also holds when any of those quantities vanishes.Source: §C.1 · p. 16

C.2 Proof of Proposition 5

The leading-budget comparison shows when reducing the correction-query probability can help, and quantifies its limiting best ratio as the noise parameter tends to zero.

C.2 Proof of Proposition 5

  • When ω < 1, decreasing q slightly from one lowers the ratio; for every q < 1, the ratio is nondecreasing in ω.The converse comparison reduces to showing Q(q; 1) ≥ 1.Source: §C.2 · p. 16
  • 0.9026797 is the limiting budget ratio at q ≈ 0.4320408 as ω → 0.The interior minimizer satisfies 2q^3/2 + q − 1 = 0.Source: §C.2 · p. 16

D PROOF OF THE ONLINE-PREDICTOR GUARANTEE

The online-predictor analysis turns regret against fixed comparators into a convergence certificate for transport residuals. It also identifies a condition for learning to improve certificates and gives finite oracle-complexity bounds.

D PROOF OF THE ONLINE-PREDICTOR GUARANTEE

D.1 Explicit Prediction Loss

  • The fixed-comparator regret bound yields the explicit prediction-loss certificate (37), which contributes to the bound in (40).The proof takes the infimum over fixed comparators outside the expectations.Source: §D.1 · p. 17

D.2 When Learning Pays Off

  • When condition (42) holds, the residual certificate improves strictly, and the stationarity certificate also improves for q < 1 and α < 1.Both certificates use the same parameters and constants, while bounding their respective trajectories.Source: §D.2 · p. 17

D.4 Finite Oracle Complexity

  • The finite oracle-complexity result gives an asymptotic leading coefficient of (1+q)A^3 and counts expected q(Tε−1) Hessian-vector products for the second-order variant.That variant uses exactly Tε fresh gradients; the predictor makes no oracle calls.Source: §D.4 · p. 17

E MARS SCALING AND RESIDUAL SCALING

  • The appendix analyzes two correction-scaling variants: one scales the entire correction, while the other scales only the prediction residual.Both use fixed q within a run.Source: E MARS SCALING AND RESIDUAL SCALING · p. 18

E.1 Scaling the Entire Correction

The appendix analyzes scaling the whole correction versus scaling only its prediction residual, deriving convergence bounds with different conditions on the scale and prediction error. Experiments report lower MNIST1D training-loss ratios at q = 1/4 and a controlled-quadratic crossover whose location is only bounded within an interval.

E.1 Scaling the Entire Correction

  • For standard scaling, the second tuning gives a convergent bound for fixed γ ≠ 1, while the first applies when the scale approaches one at the stated rate.At q = 1 the predictor cancels; γ = 1 and γ = 0 recover randomized transport and ordinary momentum, respectively.Source: §E.1 · pp. 18–19

E.2 Scaling Only the Prediction Residual

  • Residual scaling multiplies the correction residual by γ and adds bias proportional to the mean-prediction error.The sufficient mean-prediction condition bounds that error by a multiple of the squared update norm.Source: §E.2 · p. 20

E.3 Complete Rates for Both Scaling Variants

  • The complete-rate bounds can retain a non-decaying term, which vanishes when ℓ* = 0.Accurate mean prediction sets ℓ* = 0 and recovers the earlier bound.Source: §E.3 · pp. 21–22

E.4 Oracle Counts and Implementation

  • Both scaling variants use T fresh gradients and T − 1 Bernoulli decisions; gradient-difference transport costs (1 + q)T − q gradients, while HVP transport uses q(T − 1) HVPs.At γ = 0, standard scaling needs no correction queries, whereas residual scaling still queries corrections to train future predictors.Source: §E.4 · p. 22

F ADDITIONAL SCOPE

  • Predictable query probabilities preserve one-step identities, but adaptive sampling, cost, and online-learning interactions require separate treatment.The stated rates and binomial counts are for fixed q.Source: F ADDITIONAL SCOPE · p. 22

G.1 MNIST1D

  • Table 3 reports MNIST1D final-training-loss ratios for q < 1 versus q = 1 at equal oracle cost: 16,000 gradient-equivalents at batch size 4 and 8,000 at batch size 32.Source: Table 3 · p. 23
  • MNIST1D test-accuracy differences generally agree in sign with the loss ratios, except STORM at q = 1/4 and batch size 32, where lower training loss accompanies lower test accuracy.The experiments use final iterates at realized budgets, unlike the finite-horizon guarantees for a uniformly sampled iterate.Source: §G.1 · p. 23, Table 4 · p. 23

G.2 Controlled Quadratic

  • On the controlled quadratic, the global sign change occurs when L²/δ² is between 11 and 44, as normalized steps mostly follow low-curvature directions.The estimated δ scales as b^−1/2 with batch size; the experiment does not verify a horizon-uniform variance bound.Source: §G.1 · p. 23, §G.2 · p. 24

G.3 Unnormalized STORM with the Tuning of Theorem 4

Across the controlled-quadratic and CIFAR-10 experiments, randomized correction queries sometimes improve outcomes at comparable oracle work, but the direction depends on the noise regime and transport. The experiments also qualify these comparisons with state-dependent-noise and wall-time caveats.

G.3 Unnormalized STORM with the Tuning of Theorem 4

  • With Theorem 4 parameters and no tuning, all 45 controlled settings had mean run-average ∥∇f(x_t)∥^2 below bound (24).The sweep used L = 1, δ ∈ {0.25, 0.5, 1, 2, 4}, and budgets N ∈ {2000, 8000, 32000}.Source: §G.3 · pp. 24–25
  • In 9 of 10 (ω, q) settings, measured fixed-budget changes matched the direction of Q(q; ω)^(2/3) − 1; for ω < 1, ratios were within 2.1% of the benchmark.The exception was ω = 1 and q = 1/2: benchmark 1.016 versus measured ratio 0.98; these sufficient-bound comparisons do not bound actual error ratios.Source: §G.3 · p. 25, Table 8 · p. 26

G.4 Symmetric Curvature Noise

  • The symmetric-curvature runs have state-dependent gradient variance, so they illustrate the mechanism rather than verify a rate with uniform σ.Source: §G.4 · p. 26
  • Under symmetric curvature noise, q < 1 lowers the metric at large noise (0.87 at δ = 1, q = 1/4) but raises it at small noise, up to 1.49.At small noise, the diagonal predictor restores parity with q = 1.Source: §G.4 · p. 26, Table 9 · p. 26

G.5 CIFAR-10

  • CIFAR-10 used a randomly initialized ResNet-18, and within-epoch batches were not conditionally unbiased for the full training objective.The paper notes that this means part of Assumption 1 does not hold exactly.Source: §G.5 · p. 27
  • For CIFAR-10, both randomized HVP settings improve test accuracy over q = 1, with the larger gain at q = 1/4; STORM’s q = 1/2 interval includes zero and q = 1/4 is worse.At equal oracle work, randomized runs take longer in wall time because they perform more updates.Source: §G.5 · p. 27, Table 10 · p. 27
  • On CIFAR-10, the diagonal predictor removes 0.4% to 0.9% of the importance-weighted residual, while all test-accuracy intervals contain zero.The predictor used frozen zero-predictor settings and was not further tuned.Source: §G.5 · p. 27, Table 11 · p. 28
Loading 2610.13789v1…