Heavy-ball momentum and batch size in streaming least squares: matching bounds on the tuned speedup over SGD
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.
Authors
- Yitzchak Shmalo (ORCID: https://orcid.org/0000-0002-2886-3049)
Institutions
- Hebrew University of Jerusalem (IL)
Publication Details
- Journal
- Zenodo (CERN European Organization for Nuclear Research)
- Published
- 2026-11-08
- DOI
- https://doi.org/10.5281/zenodo.24069894
- Primary Topic
- Stochastic Gradient Optimization Techniques
- Type
- preprint