Minibatch and Local SGD: Algorithmic Stability and Linear Speedup in Generalization
Yunwen Lei, Tao Sun, Mingrui Liu
TL;DR
The paper addresses how minibatch SGD and Local SGD behave in terms of generalization when trained across multiple passes over data. It introduces an on-average stability framework and an expectation-variance decomposition that incorporate training error, enabling generalization bounds that reflect the interpolation regime and do not rely on Lipschitz assumptions. The main contributions show that both methods attain linear speedup in generalization: minibatch SGD achieves a speedup with batch size $b$ and yields optimal excess risk rates in convex and strongly convex settings, while Local SGD achieves linear speedup with the number of machines $M$ under convex and strongly convex regimes. The results rely on novel analytical techniques, including a binomial reformulation of minibatch sampling and self-bounding properties, and provide a pathway to fast, scalable generalization guarantees for parallel SGD methods.
Abstract
The increasing scale of data propels the popularity of leveraging parallelism to speed up the optimization. Minibatch stochastic gradient descent (minibatch SGD) and local SGD are two popular methods for parallel optimization. The existing theoretical studies show a linear speedup of these methods with respect to the number of machines, which, however, is measured by optimization errors in a multi-pass setting. As a comparison, the stability and generalization of these methods are much less studied. In this paper, we study the stability and generalization analysis of minibatch and local SGD to understand their learnability by introducing an expectation-variance decomposition. We incorporate training errors into the stability analysis, which shows how small training errors help generalization for overparameterized models. We show minibatch and local SGD achieve a linear speedup to attain the optimal risk bounds.