English

Error dynamics of mini-batch gradient descent with random reshuffling for least squares regression

Machine Learning 2025-02-05 v2 Machine Learning Optimization and Control

Abstract

We study the discrete dynamics of mini-batch gradient descent with random reshuffling for least squares regression. We show that the training and generalization errors depend on a sample cross-covariance matrix ZZ between the original features XX and a set of new features X~\widetilde{X} in which each feature is modified by the mini-batches that appear before it during the learning process in an averaged way. Using this representation, we establish that the dynamics of mini-batch and full-batch gradient descent agree up to leading order with respect to the step size using the linear scaling rule. However, mini-batch gradient descent with random reshuffling exhibits a subtle dependence on the step size that a gradient flow analysis cannot detect, such as converging to a limit that depends on the step size. By comparing ZZ, a non-commutative polynomial of random matrices, with the sample covariance matrix of XX asymptotically, we demonstrate that batching affects the dynamics by resulting in a form of shrinkage on the spectrum.

Keywords

Cite

@article{arxiv.2406.03696,
  title  = {Error dynamics of mini-batch gradient descent with random reshuffling for least squares regression},
  author = {Jackie Lok and Rishi Sonthalia and Elizaveta Rebrova},
  journal= {arXiv preprint arXiv:2406.03696},
  year   = {2025}
}

Comments

33 pages. Accepted at ALT 2025

R2 v1 2026-06-28T16:55:16.133Z