Parallelizing Stochastic Approximation Through Mini-Batching and Tail-Averaging

By Prateek Jain, Sham M. Kakade, Rahul Kidambi, Praneeth Netrapalli, Aaron Sidford. 2016

39 citations recorded by the source corpus.

This work characterizes the benefits of averaging techniques widely used in conjunction with stochastic gradient descent (SGD). In particular, this work sharply analyzes: (1) mini-batching, a method of averaging many samples of the gradient to both reduce the variance of a stochastic gradient estimate and for parallelizing SGD and (2) tail-averaging, a method involving averaging the final few iterates of SGD in order to decrease the variance in SGD's final iterate. This work presents the first tight non-asymptotic generalization error 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-batching can be used to yield provable near-linear parallelization speedups over SGD with batch size one. These results are utilized in providing a highly parallelizable SGD algorithm that obtains the optimal statistical error rate with nearly the same number of serial updates as batch gradient descent, which improves significantly over existing SGD-style methods. 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 optimal statistical error rates for the agnostic case must be a function of the noise properties. The central analysis tools used by this paper are obtained through generalizing the operator view of averaged SGD, introduced by Defossez and Bach (2015) followed by developing a novel analysis in bounding these operators to characterize the generalization error. These techniques may be of broader interest in analyzing various computational aspects of stochastic approximation.

Subfields: Artificial Intelligence, Statistics and Probability, Optimization & Training, Mathematical Optimization

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Statistics and Probability

In Optimization & Training

In Mathematical Optimization

Citations in the corpus

Communication trade-offs for synchronized distributed SGD with large step size (2019, methodology): “In other words, for α = 2/3, we need P ≤ X1/3: the number of machines as to be smaller than the number of iterations to the power 1/3, in other words, for 1000 iterations, one could only use 10 machines to reach the asymptotic regime where OSA performs similarly to MBA.”
The Step Decay Schedule: A Near Optimal, Geometrically Decaying Learning Rate Procedure (2019, methodology): “[For e.g. Appendix A.2.2 from Jain et al. (2016)] Bias-Variance tradeoff: Running SGD for T − steps starting from w 0 and a stepsize sequence { η t } T t =1 presents a ﬁnal iterate w T whose excess risk is upper-bounded as: where, Proof.”
Communication trade-offs for Local-SGD with large step size (2019, methodology): “Further, Jain et al. [20] have provided non-asymptotic results for least-square regression using similar Polyak-Juditsky analysis of the stochastic process, while our results apply to more general problems.”
On the Insufficiency of Existing Momentum Schemes for Stochastic Optimization (2018, methodology): “Note that the ranges in which parameters such as learning rate need to be searched differ based on batch size (Jain et al., 2016).”
The Effect of Network Width on the Performance of Large-batch Training (2018, methodology): “Mini-batch The choice of an optimal batch size has been studied for non-strongly convex models [21], least square regression [22], and SVMs [29].”
Exponential convergence of testing error for stochastic gradient methods (2017, unknown): “We present here another approach and consider the tail-averaged recursion, ¯ η tail n = 1 n/ 2 = (cid:98) n/ 2 η i (as proposed by Jain et al., 2016; Shamir, 2011).”
Optimality and Approximation with Policy Gradient Methods in Markov Decision Processes (2019, result): “A variant of this result also appears in Theorem 2 of Scherrer and Geist [2014], which itself can be viewed as a generalization of the approach in Kakade and Langford [2002]. In contrast to CPI [Kakade and Langford, 2002] and the more general boosting-based approach in Scherrer and Geist [2014], we ”
Bridging the gap between constant step size stochastic gradient descent and Markov chains (2017, unknown): “Analysis has been extended to “tail averaging” [25], to improve the dependence on the initial conditions.”
AdaBatch: Efficient Gradient Aggregation Rules for Sequential and Parallel Stochastic Gradient Methods (2017, background): “The case of averaged stochastic gradient descent with constant step size for least-squares regression has been studied in much detail and in that case it is possible to get an explicit expression for the convergence of the algorithm [11].”
Accelerating Stochastic Gradient Descent (2017, methodology): “…t iterations, which presents the first improvement over the O (1/κ) rate offered by SGD (Robbins and Monro, 1951)/averaged SGD (Polyak and Juditsky, 1992; Jain et al., 2016) for the least squares stochastic approximation problem, (ii) the second term in the error bound indicates that ASGD obtains…”
Gradient Diversity: a Key Ingredient for Scalable Distributed Learning (2017, methodology): “In particular, for least square problems, Jain et al. (2016) propose a bound on batch-size similar to our measure of gradient diversity; however, as mentioned in Section 1, our result holds for a wider range of problems including nonconvex setups, and can be used to motivate heuristics that result…”
Parallelizing Stochastic Gradient Descent for Least Squares Regression: Mini-batching, Averaging, and Model Misspecification (2016, methodology): “The proofs turn out to be tedious for γ ∈ γ b, max 2 , γ b, max and can be found in the initial version of this paper Jain et al. (2016b) and these were obtained through generalizing the operator view of analyzing SGD methods introduced by D´efossez and Bach (2015).”
Gradient Diversity Empowers Distributed Learning (2017, methodology): “Data-dependent thresholds for batch-size have been developed for some speciﬁc problems such as least squares [21] and SVM [41].”