Heavy-ball momentum and batch size in streaming least squares: matching bounds on the tuned speedup over SGD
Year: 2026
Type: preprint
Abstract: We consider realizable streaming least squares with Gaussian covariates $x\sim N(0,H)$ and mini-batches of size $B$, and compare stochastic heavy-ball momentum with plain stochastic gradient descent, each with its step size and momentum chosen to minimise the spectral radius $\rho$ of the exact second-moment recursion. Writing $\kappa=h_{\max}/h_{\min}$ for the condition number of $H$, we prove that the ratio of the optimal spectral gaps satisfies \[ \max\Bigl(1,\tfrac1{16}\,m\Bigr)\;\le\;\frac{1-\rho_{\mathrm{HB}}^*}{1-\rho_{\mathrm{SGD}}^*}\;\le\;\frac{128}{3}\,m, \qquad m=\min\Bigl(\max\Bigl(1,\frac{Bh_{\max}}{\operatorname{tr} H}\Bigr),\sqrt\kappa\Bigr). \] Momentum therefore gains at most a constant factor while $B\lesssim\operatorname{tr} H/h_{\max}$, gains a factor proportional to $B$ in the intermediate range, and gains the full factor $\sqrt\kappa$ once $B\gtrsim\sqrt\kappa\,\operatorname{tr} H/h_{\max}$. The upper bound $1-\rho_{\mathrm{HB}}^*\le8\min(Bh_{\min}/\operatorname{tr} H,1/\sqrt\kappa)$ is due to Morwani, Meterez, Nair and Kakade; the new ingredient is an explicit choice of parameters, critical damping at the smallest eigenvalue, that attains it up to a constant. Both directions rest on an exact renewal equation for the diagonal second moments, which also gives a closed-form necessary and sufficient condition for mean-square stability.
Author: Shmalo, Yitzchak
Institution: Hebrew University of Jerusalem
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
OpenAlex ID: W7225655361
Cites:
Cited by:
Related to: