Parallelizing Stochastic Gradient Descent for Least Squares Regression: mini-batching, averaging, and model misspecification

TL;DR

The paper analyzes parallelizing SGD for strongly convex stochastic least squares regression using mini-batching and tail-averaging, providing non-asymptotic excess risk bounds under mis-specification. It identifies a problem-dependent batch-size threshold $b_{thresh}$ up to which mini-batching yields near-linear speedups and demonstrates how doubling-batch-size epochs plus tail-averaging achieves near batch-gradient-descent generalization with substantially reduced serial depth. It also develops a non-asymptotic model averaging analysis, showing linear speedups on the variance and conditions under which parallelization is most effective, even in the mis-specified setting where optimal stepsizes depend on noise properties. The results are grounded in an operator-theoretic framework that bounds the bias-variance components and provide practical guidance for deploying highly parallel SGD in streaming contexts. Overall, the work offers a principled, non-asymptotic blueprint for achieving minimax rates with parallelized SGD in least squares regression, including robust behavior under model misspecification.

Abstract

This work characterizes the benefits of averaging schemes widely used in conjunction with stochastic gradient descent (SGD). In particular, this work provides a sharp analysis of: (1) mini-batching, a method of averaging many samples of a stochastic gradient to both reduce the variance of the stochastic gradient estimate and for parallelizing SGD and (2) tail-averaging, a method involving averaging the final few iterates of SGD to decrease the variance in SGD's final iterate. This work presents non-asymptotic excess risk bounds for these schemes for the stochastic approximation problem of least squares regression. Furthermore, this work establishes a precise problem-dependent extent to which mini-batch SGD yields provable near-linear parallelization speedups over SGD with batch size one. This allows for understanding learning rate versus batch size tradeoffs for the final iterate of an SGD method. These results are then utilized in providing a highly parallelizable SGD method that obtains the minimax risk with nearly the same number of serial updates as batch gradient descent, improving significantly over existing SGD methods. A non-asymptotic analysis of communication efficient parallelization schemes such as model-averaging/parameter mixing methods is then provided. Finally, this work sheds light on some fundamental differences in SGD's behavior when dealing with agnostic noise in the (non-realizable) least squares regression problem. In particular, the work shows that the stepsizes that ensure minimax risk for the agnostic case must be a function of the noise properties. This paper builds on the operator view of analyzing SGD methods, introduced by Defossez and Bach (2015), followed by developing a novel analysis in bounding these operators to characterize the excess risk. These techniques are of broader interest in analyzing computational aspects of stochastic approximation.

Figures (2)

• Figure 1: Effect of increased batch sizes on the Algorithm's generalization error. The variance decreases monotonically with increasing batch size. The bias indicates that the rate of decay increases till the optimal $b_{thresh}$. With $b=b_{\textrm{thresh}}$, mini-batch SGD obtains the same generalization error as batchsize $1$ using smaller number of iterations (i.e. smaller depth) compared to larger batch sizes.
• Figure 2: [Zoom in to see detail] Effect of tail-averaging with mini-batch size of $b_{thresh}=11$.

Theorems & Definitions (32)

• Theorem 1
• Corollary 2
• Lemma 3
• Lemma 4
• Theorem 5
• Theorem 6
• Lemma 7
• Lemma 8
• Lemma 9
• proof