A view of mini-batch SGD via generating functions: conditions of convergence, phase transitions, benefit from negative momenta
This provides theoretical insights into SGD dynamics for machine learning practitioners, though it is incremental as it builds on existing analysis methods.
The paper tackled the analysis of mini-batch SGD with momentum for linear models, developing a new analytic framework to derive explicit conditions for convergence and loss asymptotics, and found that negative momenta can achieve optimal convergence rates, with experimental verification on datasets like MNIST and CIFAR10 showing good agreement.
Mini-batch SGD with momentum is a fundamental algorithm for learning large predictive models. In this paper we develop a new analytic framework to analyze noise-averaged properties of mini-batch SGD for linear models at constant learning rates, momenta and sizes of batches. Our key idea is to consider the dynamics of the second moments of model parameters for a special family of "Spectrally Expressible" approximations. This allows to obtain an explicit expression for the generating function of the sequence of loss values. By analyzing this generating function, we find, in particular, that 1) the SGD dynamics exhibits several convergent and divergent regimes depending on the spectral distributions of the problem; 2) the convergent regimes admit explicit stability conditions, and explicit loss asymptotics in the case of power-law spectral distributions; 3) the optimal convergence rate can be achieved at negative momenta. We verify our theoretical predictions by extensive experiments with MNIST, CIFAR10 and synthetic problems, and find a good quantitative agreement.