Batched Stochastic Gradient Descent with Weighted Sampling

We analyze a batched variant of Stochastic Gradient Descent (SGD) with weighted sampling distribution for smooth and non-smooth objective functions. We show that by distributing the batches computationally, a significant speedup in the convergence rate is provably possible compared to either batched sampling or weighted sampling alone. We propose several computationally efficient schemes to approximate the optimal weights, and compute proposed sampling distributions explicitly for the least squares and hinge loss problems. We show both analytically and experimentally that substantial gains can be obtained. 1. MATHEMATICAL FORMULATION We consider minimizing an objective function of the form F (x) = 1 n n ∑ i=1 fi (x) = E fi (x). (1.1) One important such objective function is the least squares objective for linear systems. Given an n ×m matrix A with rows a1, . , an and a vector b ∈Rn , one searches for the least squares solution xLS given by xLS def = argmin x∈Rm 1 2 ‖Ax −b‖2 = argmin x∈Rm 1 n n ∑ i=1 n 2 (bi −〈ai , x〉)2 = argmin x∈Rm E fi (x), (1.2) where the functionals are defined by fi (x) = 2 (bi −〈ai , x〉)2. Another important example is the setting of support vector machines where one wishes to minimize the hinge loss objective given by xHL def = argmin w∈Rm 1 n n ∑ i=1 [1− yi 〈w , xi 〉]++ λ 2 ‖w‖2. (1.3) Here, the data is given by the matrix X with rows x1, . , xn and the labels yi ∈ {−1,1}. The function [z]+ def = max(0, z) denotes the positive part. We view the problem (1.3) in the form (1.1) with fi (w ) = [1− yi 〈w , xi 〉]+ and regularizer λ2 ‖w‖2. The stochastic gradient descent (SGD) method solves problems of the form (1.1) by iteratively moving in the gradient direction of a randomly selected functional. SGD can be described succinctly by the update rule: xk+1 ← xk −γ∇ fik (xk ), where index ik is selected randomly in the kth iteration, and an initial estimation x0 is chosen arbitrarily. Typical implementations of SGD select the functionals uniformly at random, although if the problem at hand allows a one-pass preprocessing of the functionals, certain weighted sampling distributions preferring functionals with larger variation can provide better convergence (see e.g. [NSW16, ZZ15] and references therein). In particular, Needell et al. show that selecting a functional with probability proportional to the Lipschitz constant of its gradient yields a convergence rate depending on the average of all such Lipschitz constants, rather than the supremum [NSW16]. An analogous result in the same work shows that for non-smooth functionals, the probabilities should be chosen proportional to the Lipschitz constant of the functional itself. Date: August 30, 2016. 1 ar X iv :1 60 8. 07 64 1v 1 [ m at h. N A ] 2 7 A ug 2 01 6 2 DEANNA NEEDELL AND RACHEL WARD Another variant of SGD utilizes so-called mini-batches; in this variant, a batch of functionals is selected in each iteration rather than a single one [CSSS11, AD11, DGBSX12, TBRS13]. The computations over the batches can then be run in parallel and speedups in the convergence are often quite significant. Contribution. Our main contribution is to propose a weighted sampling scheme to be used in minibatch SGD. We show that when the batches can be implemented in parallel, significant speedup in convergence is possible. In particular, we analyze the convergence using computed distributions for the least squares and hinge loss objectives, the latter being especially challenging since it is non-smooth. We demonstrate theoretically and empirically that weighting the distribution and utilizing batches of functionals per iteration together form a complementary approach to accelerating convergence. Organization. We next briefly discuss some related work on SGD, weighted distributions, and batching methods. We then combine these ideas into one cohesive framework and discuss the benefits in various settings. Section 2 focuses on the impact of weighting the distribution. In Section 3 we analyze SGD with weighting and batches for smooth objective functions, considering the least squares objective as a motivating example. We analyze the non-smooth case along with the hinge loss objective function in Section 4. We display experimental results for the least squares problem in Section 5 that serve to highlight the relative tradeoffs of using both batches and weighting, along with different computational approaches. We conclude in Section 6. Related work. Stochastic gradient descent, stemming from the work [RM51], has recently received renewed attention for its effectiveness in treating large-scale problems arising in machine learning [BB11, Bot10, NJLS09, SSS08]. Importance sampling in stochastic gradient descent, as in the case of minibatching (which we also refer to simply as batching here), also leads to variance reduction in stochastic gradient methods and, in terms of theory, leads to improvement of the leading constant in the complexity estimate, typically via replacing the maximum of certain data-dependent quantities by their average. Such theoretical guarantees were shown for the case of solving least squares problems where stochastic gradient descent coincides with the randomized Kaczmarz method in [SV09]. This method was extended to handle noisy linear systems in [Nee10]. Later, this strategy was extended to the more general setting of smooth and strongly convex objectives in [NSW16], building on an analysis of stochastic gradient descent in [BM11]. Later, [ZZ15] considered a similar importance sampling strategy for convex but not necessarily smooth objective functions. Importance sampling has also been considered in the related setting of stochastic coordinate descent/ascent methods [Nes12, RT15, QRZ15, CQR15]. Other papers exploring advantages of importance sampling in various adaptations of stochastic gradient descent include but are not limited to [LS13, SRB13, XZ14, DB15]. Mini-batching in stochastic gradient methods refers to pooling together several random examples in the estimate of the gradient, as opposed to just a single random example at a time, effectively reducing the variance of each iteration [SSSSC11]. On the other hand, each iteration also increases in complexity as the size of the batch grows. However, if parallel processing is available, the computation can be done concurrently at each step, so that the “per-iteration cost" with batching is not higher than without batching. Ideally, one would like the consequence of using batch size b to result in a convergence rate speed-up by factor of b, but this is not always the case [BCNW12]. Still, [TBRS13] showed that by incorporating parallelization or multiple cores, this strategy can only improve on the convergence rate over standard stochastic gradient, and can improve the convergence rate by a factor of the batch size in certain situations, such as when the matrix has nearly orthonormal rows. Other recent papers exploring the advantages of mini-batching in different settings of stochastic optimization include [CSSS11, DGBSX12, NW13, KLRT16, LZCS14]. The recent paper [CR16] also considered the combination of importance sampling and mini-batching for a stochastic dual coordinate ascent algorithm in the general setting of empirical risk minimization, wherein the function to minimize is smooth and convex. There the authors provide a theoretical optimal sampling strategy that is not practical to implement but can be approximated via alternating minimization. They also provide a computationally efficient formula that yields better sample complexity than