Parallelizing Stochastic Gradient Descent for Least Squares Regression: mini-batching, averaging, and 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.

39 pages. Published in the Journal of Machine Learning Research (JMLR)

Cited by in corpus (40)

• Robust Aggregation for Federated Learning
• Don't Use Large Mini-Batches, Use Local SGD
• Local SGD Converges Fast and Communicates Little
• On the Convergence of Local Descent Methods in Federated Learning
• Measuring the Effects of Data Parallelism on Neural Network Training
• On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift
• The Error-Feedback Framework: Better Rates for SGD with Delayed Gradients and Compressed Communication
• The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Procedure For Least Squares
• Ensemble of Averages: Improving Model Selection and Boosting Performance in Domain Generalization
• Is Local SGD Better than Minibatch SGD?
• A Survey of Optimization Methods from a Machine Learning Perspective
• The Implicit Regularization of Stochastic Gradient Flow for Least Squares
• HiGrad: Uncertainty Quantification for Online Learning and Stochastic Approximation
• Understanding and Improving Model Averaging in Federated Learning on Heterogeneous Data
• Federated Accelerated Stochastic Gradient Descent
• On Linear Stochastic Approximation: Fine-grained Polyak-Ruppert and Non-Asymptotic Concentration
• Robust, Accurate Stochastic Optimization for Variational Inference
• Communication-Efficient Distributed Stochastic AUC Maximization with Deep Neural Networks
• Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model
• Reducing the variance in online optimization by transporting past gradients
• Benign Overfitting of Constant-Stepsize SGD for Linear Regression
• Learning Rates as a Function of Batch Size: A Random Matrix Theory Approach to Neural Network Training
• On the Effectiveness of Richardson Extrapolation in Machine Learning
• Least Squares Regression with Markovian Data: Fundamental Limits and Algorithms
• Stochastic Training is Not Necessary for Generalization
• Beating SGD Saturation with Tail-Averaging and Minibatching
• The Benefits of Implicit Regularization from SGD in Least Squares Problems
• Do optimization methods in deep learning applications matter?
• Optimal Mini-Batch Size Selection for Fast Gradient Descent
• On the Double Descent of Random Features Models Trained with SGD
• Making Coherence Out of Nothing At All: Measuring the Evolution of Gradient Alignment
• The Minimax Complexity of Distributed Optimization
• Sample Efficient Linear Meta-Learning by Alternating Minimization
• Critical Parameters for Scalable Distributed Learning with Large Batches and Asynchronous Updates
• Convergence Analysis of Accelerated Stochastic Gradient Descent under the Growth Condition
• Last Iterate Risk Bounds of SGD with Decaying Stepsize for Overparameterized Linear Regression
• An elementary analysis of ridge regression with random design
• Improved SVRG for quadratic functions
• Towards Understanding Generalization via Decomposing Excess Risk Dynamics
• Some Limit Properties of Markov Chains Induced by Stochastic Recursive Algorithms