A General Analysis of Example-Selection for Stochastic Gradient Descent

1 INTRODUCTION . To minimize a differentiable function f : Rd → R , stochastic gradient descent ( SGD ) iteratively updates a parameter vector w ∈ Rd starting at some w0 by running wt+1 = wt − αt∇f ( wt ; xt ) , ( 1 ) where αt is the step size at iteration t , xt is a data example ( often a minibatch of data ) chosen by some process—typically by subsampling a training dataset—for SGD to use at iteration t , and ∇f ( wt ; xt ) is an example gradient , which we hope will be a good approximation for the gradient of the objective ∇f ( wt ) . In the standard setup , the x ’ s are drawn from a dataset D of size n , and f ( w ) = 1n ∑ x∈D f ( w ; x ) , which is often referred to as Empirical Risk Minimization ( ERM ) . The order in which the example sequence x0 , x1 , . . . are chosen is known to affect the convergence of SGD . For instance , compare so-called “ random reshuffling ” to with-replacement sampling : optimizing a strongly convex objective with T total iterations , random reshuffling samples the xt from D without replacement and achieves an accelerated convergence rate of 1/T 2 , while with-replacement sampling yields a convergence rate of 1/T ( Bottou , 2012 ; Recht & Ré , 2012 ; Gürbüzbalaban et al. , 2021 ) . Similar accelerated rates have been shown in other settings and for other example orders , such as shuffling the dataset once ( Nguyen et al. , 2020 ; Ahn et al. , 2020 ; Mishchenko et al. , 2020 ) . However , these analyses have mostly focused only on specific example-selection schemes , and study only the case of ERM-type finite-sum objectives . This does not help us understand how new example orders ( e.g. , the data echoing method of Choi et al . ( 2019 ) ) affect the convergence of SGD . This paper develops a general condition on the example gradients themselves that is sufficient to provide a convergence rate for SGD . Intuitively , our main result is : the convergence rate of SGD depends on how fast the averages of consecutive example gradients∇f ( w ; xt ) converge to the full objective gradient ∇f ( w ) . With-replacement sampling causes the average of m consecutive gradient examples starting at any timestep τ , 1m ∑τ+m−1 t=τ ∇f ( w ; xt ) , to converge to∇f ( w ) at a rate of O ( 1/m ) in terms of the norm squared : we show SGD with any example sequence that fulfills this condition will converge at the same asymptotic rate as with-replacement sampling . Alternatively , if that average converges at the faster Õ ( 1/m2 ) rate typical of Quasi-Monte-Carlo ( QMC ) ( Caflisch , 1998 ) , then we show SGD enjoys the accelerated rate that random reshuffling gets . Our contributions are as follows : • We propose a new condition on the example gradients—average gradient error—and provide convergence analysis using this general condition for both non-convex and strongly-convex problems . We justify the validity of this condition on synthetic experiments ( Section 3 ) . • We show that many commonly used example orderings—shuffle once , random reshuffling , data echoing , Markov Chain Gradient Descent—can be analyzed as special cases under our theoretical results , which match or improve upon their existing rates in the literature ( Section 4 ) . • We propose two new algorithms : ( 1 ) QMC-based data augmentation that transforms examples via a low-discrepancy sequence , improving generalization ; and ( 2 ) a greedy algorithm that sorts the examples before each epoch based on our new average gradient error metric ( Section 5 ) . • Empirically , we evaluate our two algorithms on several image classification benchmarks including MNIST , CIFAR10/100 and ImageNet . We show with QMC-based data augmentation , a higher validation accuracy can be achieved without hyperparameter tuning—this suggests that QMC may be a good default driver to use with data augmentation for deep learning in general . Meanwhile , the greedy algorithm converges faster both in terms of iteration and wall-clock time ( Section 6 ) . 2 RELATED WORK . Example ordering in stochastic optimization . Traditional example ordering in SGD is carried out in a with-replacement fashion , which is used to ensure unbiased estimation of the full gradient ( Robbins & Monro , 1951 ; Bach & Moulines , 2011 ; Zhang , 2004 ; Bottou et al. , 2018 ; Drori & Shamir , 2020 ) . Significant attention has been paid to importance sampling with respect to various measures , such as Lipschitz constants ( Schmidt et al. , 2017 ; Needell et al. , 2014 ) , example gradient norms and bounds ( Zhao & Zhang , 2015 ; Alain et al. , 2015 ; Papa et al. , 2015 ; Lee et al. , 2019 ) , individual losses ( Kawaguchi & Lu , 2020 ; Loshchilov & Hutter , 2015 ) , and data heterogeneity ( Lu et al. , 2021 ) . Without-replacement sampling , however , is more common in practice and empirically allows faster convergence ( Bottou , 2012 ) . Among the most popular without-replacement approaches are shuffle once ( SO ) ( Bertsekas , 2011 ; Gürbüzbalaban et al. , 2019 ) and random reshuffling ( RR ) ( Ying et al. , 2017 ) . In theory , Recht & Ré ( 2012 ) undertook the first investigation on convergence of RR via the noncommutative arithmetic-geometric mean conjecture , to which subsequent works provide counter examples ( Yun et al. , 2021 ; De Sa , 2020 ) . HaoChen & Sra ( 2019 ) performed an epoch-wise acceleration analysis on RR while Gürbüzbalaban et al . ( 2021 ) considered its convergence over infinite epochs . Safran & Shamir ( 2020 ) analyzed the lower bounds for both SO and RR methods , and their results are further polished by Mishchenko et al . ( 2020 ) via the Bregman divergence bound . Aside from manual ordering , another line of research focuses on perturbed example ordering from data echoing ( Choi et al. , 2019 ; Agarwal et al. , 2020 ) . Quasi-Monte Carlo . Quasi-Monte Carlo ( QMC ) is a variant of Monte Carlo ( MC ) methods that uses a low-discrepancy sequence instead of a pseudorandom sequence . QMC has been successfully applied in a wide variety of domains including computer graphics ( Keller , 1995 ) , finance ( Joy et al. , 1996 ) , and computational biology ( Cieslak et al. , 2008 ) . In machine learning , using QMC in place of MC can significantly improve many techniques including variational inference ( Buchholz et al. , 2018 ; Liu & Owen , 2021 ) , feature mapping ( Yang et al. , 2014 ; Avron et al. , 2016 ) , normalizing flows ( Wenzel et al. , 2018 ) , deep learning based PDE ( Chen et al. , 2019 ) , and time series analysis ( Philipson et al. , 2020 ) . For stochastic optimization , early works like Homem-de Mello ( 2008 ) and Pennanen ( 2005 ) established the asymptotic convergence of a QMC sequence in terms of the training set size , while Jank ( 2005 ) proposed replacing MC with QMC in computing the E-step of the EM algorithm . Similar to our motivation , Buchholz et al . ( 2018 ) analyzed the convergence of SGD when samples are drawn using QMC . Their approach differs significantly from ours in that their method draws an independent unbiased length-b QMC sequence for each minibatch , while our examples come from contiguous subsequences of a length-T QMC sequence that is used across all iterations . 3 EXAMPLE-GRADIENT AVERAGES AND SGD CONVERGENCE . In this section , we describe our setup , define the average gradient error condition we are proposing , and state our main result . Our objective is to minimize a continuously differentiable function f : Rd → R using examples xt from some set X retrieved at each iteration t. We make the usual assumption that both the loss gradient and the example gradients are L-Lipschitz continuous . Assumption 1 ( L-Smoothness ) . For some L < ∞ , for any u , v ∈ Rd and any example x ∈ X , ‖∇f ( u ; x ) −∇f ( v ; x ) ‖ ≤ L · ‖u− v‖ and ‖∇f ( u ) −∇f ( v ) ‖ ≤ L · ‖u− v‖ . We propose a new condition on the average gradient error . Informally , this condition bounds how averages of consecutive example gradients approximate the objective gradient . Formally , Assumption 2 . In the context of Equation ( 1 ) , we say that the example sequence x0 , x1 , . . . is ( γ , C , Φ ) -concentrating for γ ∈ [ 1 , 2 ] , C > 0 , Φ ≥ 0 , if for any timestep τ ≥ 0 and any m > 0 , ∥∥∥∥∥ 1m τ+m−1∑ t=τ ∇f ( wτ ; xt ) −∇f ( wτ ) ∥∥∥∥∥ 2 ≤ 1 mγ ( C2 + Φ2‖∇f ( wτ ) ‖2 ) , ( 2 ) where wτ is the weight parameter vector arrived at after τ SGD update steps in Equation ( 1 ) . The constants C and Φ here may depend on the total number of iterations T : specifically , they may absorb logarithmic factors like log ( τ +m ) 2s typical in a QMC error bound , where s is the dimension of the sample space . In addition , when the { xt } ’ s are random , we will show that this assumption holds with high probability . Furthermore , to show this assumption holds , it suffices to show that the inequality holds for any w : our requirement that it holds only for the specific wτ arrived at by SGD is weaker . We can develop intuition about Assumption 2 by considering familiar cases : • For general with-replacement sampling , the sum in ( 2 ) is a sum of independent random variables ; so , a concentration argument would yield Assumption 2 with γ = 1 with high probability . • For general without-replacement sampling from a dataset of size n , any whole-epoch subsequences from the sum in ( 2 ) will cancel to zero , so we expect that sum to have magnitudeO ( n ) independent of m , and thus Assumption 2 should hold with γ = 2 ( Propositions 1 to 4 ) . • For low-discrepancy sampling , we would expect Assumption 2 to hold with γ = 2 , along with a multiplicative log ( τ +m ) 2s term : this is the classic error rate for QMC ( Proposition 6 ) . In Section 4 , we will make this intuition rigorous with high probability under the assumption that the error of the example gradients is bounded . This intuition illustrates both how Assumption 2 can cover previously analyzed settings and how it can generalize to cases not previously studied ( e.g. , QMC data augmentation ) . Also observe that Assumption 2 is easily adapted to minibatch SGD : if it holds for the single-example case , it should also hold with modified constants for minibatches of size b ( consisting of averages of b consecutive example gradients ) by substituting m 7→ mb.1 Since this is straightforward , for simplicity of presentation we focus our theoretical analysis on the batch-size-1 case . Our analysis is based on the following key lemma , which bounds the evolution of SGD over an “ analysis phase ” of m steps . Lemma 1 . Suppose our setup satisfies Assumptions 1 and 2 and that we use a constant step size α . For all timesteps τ ≥ 0 , let m > 0 be some integer such that 3αm1−γ/2Φ ≤ αm ≤ 16L . Then the objective at timestep τ +m is bounded by f ( wτ+m ) ≤ f ( wτ ) − 14αm‖∇f ( wτ ) ‖ 2 + 2αm1−γC2 . By applying this lemma inductively and choosing the step size appropriately , we can derive convergence rates for SGD in a variety of settings . Note that while here for simplicity we state results for a fixed step size , Lemma 1 can also be used to get essentially the same rates for diminishing step size schemes , which we analyze in Appendices A.2 and A.3 . In what follows , we let f∗ be the global minimum of f , let ∆ : = f ( w0 ) − f∗ , and use Õ to hide logarithmic terms in the problem parameters such as C , Φ , L , ∆ , and , while treating γ as a constant . Theorem 1 ( Non-convex case ) . Suppose that our setup satisfies Assumptions 1 and 2 , and let > 0 be any target error . Using SGD ( 1 ) with a constant step size α = 16L ⌈ ( 4C/ + 3Φ ) 2/γ ⌉−1 , the number of steps T needed to achieve mint=0 , ··· , T−1‖∇f ( wt ) ‖2 ≤ 2 is at most T = ⌈ 48L∆ 2 ⌉ · ⌈ ( 4C + 3Φ ) 2/γ⌉ = Õ ( C 2/γL∆ 2+2/γ + Φ 2/γL∆ 2 + C 2/γ 2/γ + Φ2/γ ) . 1The only technical subtlety is that the iterates wτ of SGD would vary with minibatch sizes . Observe that when γ = 1 , this gives the standard −4 rate expected for non-convex SGD , and for γ = 2 , this gives us the accelerated −3 rate of shuffling methods . We can get an even faster rate if f satisfies the µ-Polyak-Łojasiewicz ( PL ) condition ‖∇f ( w ) ‖2 ≥ 2µ ( f ( w ) − f∗ ) , which generalizes strong convexity and has been applied before in the study of sample orders ( Mishchenko et al. , 2020 ; Ahn et al. , 2020 ) : in particular the following theorem holds for µ-strongly convex functions . Theorem 2 . Suppose that f satisfies the µ-PL condition and our setup satisfies Assumptions 1 and 2 . Let > 0 be any target error , and κ = L/µ be the condition number of the problem . Using SGD ( 1 ) with a constant step size α = 16L ⌈ ( 8C2/ ( µ 2 ) + 9Φ2 ) 1/γ ⌉−1 , the number of steps T needed to guarantee f ( wT ) − f∗ ≤ 2 is at most T = ⌈ 12κ log ( 2∆ 2 ) ⌉ · ⌈ ( 8C2 µ 2 + 9Φ 2 ) 1/γ⌉ = Õ ( C2/γκ µ1/γ 2/γ + κΦ2/γ + κ ) . Observe that when γ = 1 , this recovers the ordinary T = κ −2 rate we usually get for strongly convex SGD , and when γ = 2 we get a faster rate of κ −1 . These theorems together show that our Assumption 2 is sufficient to show the convergence of SGD , and the convergence-rate parameter γ of the assumption translates to affect the convergence rate of SGD . This validates our intuition that faster convergence of averages of consecutive example gradients to the full gradient∇f ( w ) leads to faster convergence of SGD . Synthetic experiments . We quickly validate these results on a synthetic 10-dimensional strongly-convex problem . The first setting we consider is the expected risk minimization of R ( w ) = E [ ( xTw − y ) 2 ] , where x ∼ N ( 0 , Id ) and y |x ∼ N ( xTw∗ , 1 ) for some optimal value w∗ . We run SGD in an online fashion—at iteration t we draw samples ( xt , yt ) from the underlying distribution to compute∇R ( wt ; xt ) used in the update . We compare drawing these samples independently at random against drawing using a QMC sequence ( Sobol in [ 0 , 1 ] d ) via an inverse transform , using for both cases the same diminishing step size scheme selected to minimize the expected risk for the random-sampling case—the optimal step size scheme for vanilla SGD . The online plot of Fig . 1 ( a ) shows that the convergence rate is strictly superior with QMC , which achieves a O ( 1/t2 ) rate compared to the O ( 1/t ) rate of random sampling , which is what Theorem 2 predicts . We also evaluate the offline setting , where we draw n independent examples from the same distribution to form a training set , and minimize the empirical risk Rn ( w ) = 1n ∑n i=1 ( x T iw − yi ) 2 . This corresponds to a least squares problem with optimal solution w∗n . We run SGD in an epoch-wise manner for K epochs , in which we compare sampling from the training set uniformly with replacement ( IID Uniform ) , random reshuffling ( RR ) , shuffle once ( SO ) , and sampling using one QMC sequence in [ 0 , 1 ] of length T = nK followed by a mapping to example indices ( Sobol ) . In the offline plot of Fig . 1 ( a ) , we see that the low-discrepancy methods all yield an accelerated rate compared to sampling with replacement , which again validates our theory . Additional details on the experiment setup , hyperparameters , and diminishing step size derivation can be found in Appendix A.1.3 . With the same synthetic setup , we also verified our bound in Assumption 2 by measuring the average gradient error over a sequence of examples . The fixed point wτ is arbitrarily set to the origin . For the online setting ( left of Fig . 1 ( b ) ) , we use the same set of examples as in the SGD experiments above starting at t = 0 . Similarly , in the offline setting , we go through the examples epoch-wise . As we can see , the sample orderings given by QMC , RR and SO ( offline only ) indeed give us an accelerated rate of decrease in the average gradient error as we increase m ( γ ≈ 2 ) , justifying our main assumption .

Summary: 1: This paper theoretically builds the following connection: if the averages of consecutive stochastic gradients converge faster to the full gradient, then the SGD with the corresponding sampling strategy will have a faster convergence rate.  2: Applying this theoretical finding on a few commonly used sampling strategies, the convergence rates established by prior works can be recovered.  3: Two algorithms are proposed based on this theory.