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 improves the complexity bound w.r.t. minibatch speedup given by Cotter et al and Liu et al. for minibatch accelerated gradient descent for certain objectives (under sharpness condition). Generally the paper is well-organized with solid theoretical analysis and is easy to follow.