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: The paper proposes a new condition (metric) on the sequence of samples seen by SGD (with or without replacement), which can be used to prove tight convergence rates for strongly convex and non-convex objectives. The metric quantifies how quickly the average gradient of a subsequence of length $m$ converges to the true gradient. This metric is computed for popular schemes like shuffle once, and random reshuffling, as well as for random reshuffling with echoing and Markov chain gradient descent. Then, the paper proposes to combine quasi-Monte Carlo based data augmentation with random reshuffling, proves the condition for it, and shows empirically that it performs well. The paper also proposes a technique which greedily tries to optimize the metric in the proposed condition.