Random Reshuffling with Momentum for Nonconvex Problems: Iteration Complexity and Last Iterate Convergence

Random reshuffling with momentum (RRM) corresponds to the SGD optimizer with momentum option enabled, as found in many machine learning libraries like PyTorch and TensorFlow. Despite its widespread use, the convergence properties of RRM do not seem to be well understood. This work establishes new complexity bounds and asymptotic convergence guarantees for popular versions of RRM using stochastic heavy-ball momentum, Nesterov acceleration, and mini-batches in a general nonconvex setting. In particular, we prove that the base variant of RRM achieves the complexity ${\cal O}(n^{-1/3}((1-\beta^n)T)^{-2/3})$, where $n$ denotes the number of component functions, $\beta \in [0,1)$ is a momentum parameter, and $T$ is the total number of iterations. Furthermore, every accumulation point of the iterates $\{x^k\}_k$ generated by RRM is shown to be a stationary point of the problem. When the objective function is definable, the sequence of iterates $\{x^k\}_k$ is proven to converge to a single stationary point $x^*$ and we derive improved asymptotic complexity bounds.