An Even More Optimal Stochastic Optimization Algorithm: Minibatching and Interpolation Learning

We present and analyze an algorithm for optimizing smooth and convex or strongly convex objectives using minibatch stochastic gradient estimates . The algorithm is optimal with respect to its dependence on both the minibatch size and minimum expected loss simultaneously . This improves over the optimal method of Lan [ 17 ] , which is insensitive to the minimum expected loss ; over the optimistic acceleration of Cotter et al . [ 10 ] , which has suboptimal dependence on the minibatch size ; and over the algorithm of Liu and Belkin [ 19 ] , which is limited to least squares problems and is also similarly suboptimal with respect to the minibatch size . Applied to interpolation learning , the improvement over Cotter et al . and Liu and Belkin translates to a linear , rather than square-root , parallelization speedup . 1 Introduction . The massive scale of many modern machine learning models and datasets give rise to complex , highdimensional training objectives that can be very computationally expensive to optimize . To reduce the computational cost , it is therefore important to devise optimization algorithms that can leverage parallelism to reduce the amount of time needed to train . Stochastic first-order methods , which use stochastic estimates of the gradient of the training objective , are by far the most common approach and these methods can be directly improved by using minibatch stochastic gradient estimates , which are easy to parallelize across multiple computing cores or devices . Accordingly , we propose and analyze an optimal accelerated minibatch stochastic gradient descent algorithm . Our analysis exploits that while training machine learning models , it is typically possible to drive the loss either all the way to zero , or at least very close to zero , and the performance of our algorithm improves as the minimum value of the loss approaches zero . This property has multiple names in different contexts . In learning theory , it is common to show fast rates under the assumption of “ realizability ” —meaning the data can be fit perfectly by the model , i.e . the loss can be driven to zero . Even when the problem is not exactly realizable , it is sometimes possible to derive “ optimistic rates ” that interpolate between fast rates for realizable learning and the slower agnostic rates [ 33 ] . In the context of “ interpolation learning , ” there has recently been great interest in understanding training “ overparametrized ” models—which have many more parameters than there are training examples , generally meaning that many settings of the parameters would attain zero training loss—both in terms of optimization [ 2–4 , 9 , 15 ] and generalization [ 6 , 35 , 38 ] . Finally , in the optimization literature , there have been efforts to prove optimistic rates depending on the minimum value of the objective or on the variance of the stochastic gradients a the minimizer [ 10 , 19–21 , 23 , 31 ] . Regardless of the name , these ideas are all based on the same fundamental concept of exploiting the fact that the minimum value of the objective is nearly zero . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Our contributions . In Section 3 , we present and analyze an accelerated minibatch SGD algorithm for optimizing smooth and convex objectives using minibatch stochastic gradient estimates . Our method closely resembles the methods of Lan [ 17 ] and Cotter et al . [ 10 ] , but with different stepsizes and momentum parameters , and a tighter analysis . Importantly , our algorithm enjoys a linear speedup in the minibatch size all the way up to a critical threshold beyond which larger minibatches do not help . In contrast , Lan and Cotter et al. ’ s bounds have a worse , sublinear speedup and a correspondingly higher critical threshold . In Section 4 , we show that a modified version of our algorithm can attain substantially faster convergence when the objective satisfies a certain quadratic growth condition , which is a relaxation of strong convexity . As part of our analysis , we simplify and generalize a restarting technique [ 13 , 18 , 24 , 26 , 29 , 30 ] , which we show amounts to a reduction from strongly convex optimization to convex optimization . The reduction in the other direction is a well-known tool in optimization analysis , but our result shows that it goes both ways and it may be of more general interest . In Sections 5 and 6 , we prove that our methods are optimal with respect to both the minibatch size and also the minimum value of the loss in the settings we consider . We then explain how our guarantees demonstrate a linear speedup in the minibatch size , which improves over a sublinear speedup in previous work [ 10 , 19 ] . Finally , in Section 7 , we extend our results to a related setting where the bound on the minimum value of the loss is replaced by a bound on the variance of the stochastic gradients at the optimum [ as in , e.g . 8 , 14 , 21 , 23 , 31 ] . Under this condition , we establish the optimal error achievable by any learning rule , including non-first-order methods , and we show that the optimal convergence rate is nevertheless achieved by SGD—a first-order method without acceleration . Further , we show that the accelerated optimization rate , T−2 , is unattainable using minibatches of size 1 , but with larger minibatches , our accelerated minibatch SGD method can match the optimal error of SGD using a substantially smaller parallel runtime but the same number of samples . 2 Setting and Background . We consider a generic stochastic optimization problem min w∈Rd { L ( w ) : = Ez∼D [ ` ( w ; z ) ] } ( 1 ) This problem captures , for instance , supervised machine learning where z = ( x , y ) is feature vector and label pair ; ` ( w ; ( x , y ) ) is the loss of a model parametrized by w on that example ; and L ( w ) is the expected loss . We note that there are two possible interpretations of L depending on what D corresponds to . We can take D to be the uniform distribution over a set of training examples , in which case L is the training loss . Alternatively , we can take D to be the population distribution , in which case L is the population risk . An advantage of the latter view is that optimization guarantees directly imply good performance on the population , however , computing an independent , unbiased stochastic gradient estimate requires a fresh sample from the distribution , so only “ one-pass ” methods are possible . Nevertheless , our algorithm and analysis apply equally well in either viewpoint . We consider optimizing objectives on Rd , but we are most interested in dimension-free analysis , i.e . one that does not explicitly depend on the dimension , and our algorithm ’ s guarantees hold even in infinite dimension . Indeed , in many applications including machine learning , the dimension can be very large , so an explicit reliance on the dimension would often lead to impractical bounds . We study the family of optimization algorithms that attempt to minimize ( 1 ) using T sequential stochastic gradient estimates with minibatches of size b of the form g ( wt ) = 1 b b∑ i=1 ∇ ` ( wt ; zi ) for i.i.d . z1 , . . . , zb ∼ D ( 2 ) at points w1 , . . . , wT of the algorithm ’ s choice . Later , we will argue that our proposed method is optimal with respect to all of the algorithms in this family . Because it is easy to parallelize the computation of the b stochastic gradients in the minibatch , T roughly captures the runtime of the algorithm , while n = bT , the total number of samples used , captures its sample complexity . In our analysis , we will rely on several assumptions about the losses ` and the expected loss L. Assumption 1 . For almost every z ∼ D , ` ( w ; z ) is non-negative , convex , and H-smooth w.r.t . w , i.e . ∀w , u , z ` ( u ; z ) + 〈∇ ` ( u ; z ) , w − u〉 ≤ ` ( w ; z ) ≤ ` ( u ; z ) + 〈∇ ` ( u ; z ) , w − u〉+ H 2 ‖w − u‖2 This assumption holds , for instance , for training linear models with a smooth , convex , and nonnegative loss function , as in least squares problems or logistic regression . Assumption 2 . The expected loss L is convex and has minimum value L∗ = minw L ( w ) , which is attained at a point w∗ with ‖w∗‖ ≤ B . The minimum value of the loss , L∗ , is the key quantity in our analysis , and we will prove an “ optimistic rate ” for our algorithm , meaning it performs increasingly well as L∗ → 0 . The key idea is to use L∗ to bound the variance of the stochastic gradient estimates at the point w∗ . Following prior work [ 10 , 32 , 33 ] , we observe that under Assumptions 1 and 2 , we can upper bound ( see Lemma 3 ) E‖∇ ` ( w∗ ; z ) −∇L ( w∗ ) ‖2 ≤ 2HL∗ ( 3 ) The H-Lipschitzness of ∇ ` allows us to upper bound the variance of the stochastic gradients at other , non-minimizing points w too . In contrast , a common practice in the optimization literature is to simply assert that the variance of the stochastic gradient estimates is uniformly upper bounded , i.e . sup w E‖ ` ( w ; z ) −∇L ( w ) ‖2 ≤ σ2 ( 4 ) This bound is very convenient , but it can fail to hold even in very simple cases like least squares regression , where the variance grows with ‖w‖2 . Therefore , while faster rates may be achieved under the condition ( 4 ) , this is quite strong and may not correspond well with problems of interest . For many optimization problems , particularly those arising from training machine learning models , the minimum of the loss , L∗ , can be expected to be small . For example , machine learning models are often trained in the interpolation regime [ 2–4 , 6 , 9 , 15 , 35 , 38 ] , where the training loss is underdetermined , so there are many settings of the parameters which achieve exactly zero training loss . When D is the empirical distribution ( see above ) , the interpolation regime therefore corresponds to L∗ = 0 . However , even if D is the population distribution , machine learning problems are often realizable , or nearly realizable , meaning the population risk can also be driven to zero or close to it . In the optimization literature , it is often possible to show enhanced guarantees under the favorable condition of λ-strong convexity . However , the assumption of strong convexity is somewhat at odds with our goal of studying the interpolation or near-interpolation setting . For example , the simplest example of interpolation learning is underdetermined least squares , where the number of observations is less than the dimension ; in this case , the empirical covariance is degenerate so the problem is not strongly convex . Furthermore , assuming strong convexity and L∗ ≈ 0 and ` is non-negative is very strong , and puts very strong constraints on the objective function which may often fail to hold . We consider instead a relaxation of strong convexity , which only requires the objective to grow faster than the squared distance to the set of minimizers , which need not be a singleton [ 7 , 12 , 22 ] . This condition does hold for underdetermined least squares—with λ equal to the smallest non-zero eigenvalue of the covariance—and for many other problems with non-unique minimizers , and it turns out to be sufficient to achieve the faster rates that typically arise from strong convexity . Assumption 3 . The expected loss L is convex ; it has minimum value L∗ = minw L ( w ) ; L ( 0 ) −L∗ ≤ ∆ ; and L satisfies the following growth condition for all w : L ( w ) − L∗ ≥ λ 2 min w∗∈arg minw L ( w ) ‖w − w∗‖2

Summary: This paper presents and analyzes an accelerated algorithm for convex and strongly convex minimization when the objective function is an expectation of a convex function. It claims that the algorithm is optimal w.r.t. Mini-batch size and the convergence rate. The presented algorithm reaches the best-known rate for the accelerated stochastic algorithm and also shows the dependence on the mini-batch size is also optimal.