Proximal and Federated Random Reshuffling

Random Reshuffling ( RR ) , also known as Stochastic Gradient Descent ( SGD ) 1 without replacement , is a popular and theoretically grounded method for finite-sum2 minimization . We propose two new algorithms : Proximal and Federated Random3 Reshuffling ( ProxRR and FedRR ) . The first algorithm , ProxRR , solves composite4 finite-sum minimization problems in which the objective is the sum of a ( potentially5 non-smooth ) convex regularizer and an average of n smooth objectives . ProxRR6 evaluates the proximal operator once per epoch only . When the proximal operator7 is expensive to compute , this small difference makes ProxRR up to n times faster8 than algorithms that evaluate the proximal operator in every iteration , such as9 proximal ( stochastic ) gradient descent . We give examples of practical optimization10 tasks where the proximal operator is difficult to compute and ProxRR has a clear11 advantage . One such task is federated or distributed optimization , where the evalu-12 ation of the proximal operator corresponds to communication across the network.13 We obtain our second algorithm , FedRR , as a special case of ProxRR applied to14 federated optimization , and prove it has a smaller communication footprint than15 either distributed gradient descent or Local SGD . Our theory covers both constant16 and decreasing stepsizes , and allows for importance resampling schemes that can17 improve conditioning , which may be of independent interest . Our theory covers18 both convex and nonconvex regimes . Finally , we corroborate our results with19 experiments on real data sets.20 1 Introduction21 Modern theory and practice of training supervised machine learning models is based on the paradigm22 of regularized empirical risk minimization ( ERM ) [ Shalev-Shwartz and Ben-David , 2014 ] . While the23 ultimate goal of supervised learning is to train models that generalize well to unseen data , in practice24 only a finite data set is available during training . Settling for a model merely minimizing the average25 loss on this training set—the empirical risk—is insufficient , as this often leads to over-fitting and poor26 generalization performance in practice . Due to this reason , empirical risk is virtually always amended27 with a suitably chosen regularizer whose role is to encode prior knowledge about the learning task at28 hand , thus biasing the training algorithm towards better performing models.29 The regularization framework is quite general and perhaps surprisingly it also allows us to consider30 methods for federated learning ( FL ) —a paradigm in which we aim at training model for a number of31 clients that do not want to reveal their data [ Konečný et al. , 2016 , McMahan et al. , 2017 , Kairouz32 et al. , 2019 ] . The training in FL usually happens on devices with only a small number of model33 updates being shared with a global host . To this end , Federated Averaging algorithm has emerged34 that performs Local SGD updates on the clients ’ devices and periodically aggregates their average.35 Its analysis usually requires special techniques and deliberately constructed sequences hindering the36 research in this direction . We shall see , however , that the convergence of our FedRR follows from37 merely applying our algorithm for regularized problems to a carefully chosen reformulation.38 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . Formally , regularized ERM problems are optimization problems of the form39 min x∈Rd [ P ( x ) : = 1n ∑n i=1 fi ( x ) + ψ ( x ) ] , ( 1 ) where fi : Rd → R is the loss of model parameterized by vector x ∈ Rd on the i-th training data40 point , and ψ : Rd → R ∪ { +∞ } is a regularizer . Let [ n ] : = { 1 , 2 , . . . , n } . We shall make the41 following assumption throughout the paper without explicitly mentioning it:42 Assumption 1 . The functions fi are Li-smooth , and the regularizer ψ is proper , closed and convex.43 Let Lmax : = maxi∈ [ n ] Li.44 In some results we will additionally assume that either the individual functions fi , or their average45 f : = 1n ∑ i fi , or the regularizer ψ are µ-strongly convex . Whenever we need such additional46 assumptions , we will make this explicitly clear . While all these concepts are standard , we review47 them briefly in Section A.48 Proximal SGD . When the number n of training data points is huge , as is increasingly common49 in practice , the most efficient algorithms for solving ( 1 ) are stochastic first-order methods , such50 as stochastic gradient descent ( SGD ) [ Bordes et al. , 2009 ] , in one or another of its many variants51 proposed in the last decade [ Shang et al. , 2018 , Pham et al. , 2020 ] . These method almost invariably52 rely on alternating stochastic gradient steps with the evaluation of the proximal operator53 proxγψ ( x ) : = argminz∈Rd { γψ ( z ) + 12‖z − x‖ 2 } . The simplest of these has the form54 xSGDk+1 = proxγkψ ( x SGD k − γk∇fik ( xSGDk ) ) , ( 2 ) where ik is an index from { 1 , 2 , . . . , n } chosen uniformly at random , and γk > 0 is a properly55 chosen learning rate . Our understanding of ( 2 ) is quite mature ; see [ Gorbunov et al. , 2020 ] for a56 general treatment which considers methods of this form in conjunction with more advanced stochastic57 gradient estimators in place of∇fik .58 Applications such as training sparse linear models [ Tibshirani , 1996 ] , nonnegative matrix factoriza-59 tion [ Lee and Seung , 1999 ] , image deblurring [ Rudin et al. , 1992 , Bredies et al. , 2010 ] , and training60 with group selection [ Yuan and Lin , 2006 ] all rely on the use of hand-crafted regularizes . For most of61 them , the proximal operator can be evaluated efficiently , and SGD is near or at the top of the list of62 efficient training algorithms.63 Random reshuffling . A particularly successful variant of SGD is based on the idea of random64 shuffling ( permutation ) of the training data followed by n iterations of the form ( 2 ) , with the index65 ik following the pre-selected permutation [ Bottou , 2012 ] . This process is repeated several times,66 each time using a new freshly sampled random permutation of the data , and the resulting method is67 known under the name Random Reshuffling ( RR ) . When the same permutation is used throughout,68 the technique is known under the name Shuffle-Once ( SO ) .69 One of the main advantages of this approach is rooted in its intrinsic ability to avoid cache misses when70 reading the data from memory , which enables a significantly faster implementation . Furthermore,71 RR is often observed to converge in fewer iterations than SGD in practice . This can intuitively be72 ascribed to the fact that while due to its sampling-with-replacement approach SGD can miss to learn73 from some data points in any given epoch , RR will learn from each data point in each epoch.74 Understanding the random reshuffling trick , and why it works , has been a non-trivial open problem75 for a long time [ Bottou , 2009 , Recht and Ré , 2012 , Gürbüzbalaban et al. , 2019 , Haochen and Sra,76 2019 ] . Until recent development which lead to a significant simplification of the convergence77 analysis technique and proofs [ Mishchenko et al. , 2020 ] , prior state of the art relied on long and78 elaborate proofs requiring sophisticated arguments and tools , such as analysis via the Wasserstein79 distance [ Nagaraj et al. , 2019 ] , and relied on a significant number of strong assumptions about80 the objective [ Shamir , 2016 , Haochen and Sra , 2019 ] . In alternative recent development , Ahn et al.81 [ 2020 ] also develop new tools for analyzing the convergence of random reshuffling , in particular using82 decreasing stepsizes and for objectives satisfying the Polyak-Łojasiewicz condition , a generalization83 of strong convexity [ Polyak , 1963 , Lojasiewicz , 1963 ] .84 The difficulty of analyzing RR has been the main obstacle in the development of even some of the85 most seemingly benign extensions of the method . Indeed , while all these are well understood in86 Algorithm 1 Proximal Random Reshuffling ( ProxRR ) and Shuffle-Once ( ProxSO ) Require : Stepsizes γt > 0 , initial vector x0 ∈ Rd , number of epochs T 1 : Sample a permutation π = ( π0u , π1 , . . . , πn−1 ) of [ n ] ( Do step 1 only for ProxSO ) 2 : for epochs t = 0 , 1 , . . . , T − 1 do 3 : Sample a permutation π = ( π0 , π1 , . . . , πn−1 ) of [ n ] ( Do step 3 only for ProxRR ) 4 : x0t = xt 5 : for i = 0 , 1 , . . . , n− 1 do 6 : xi+1t = x i t − γt∇fπi ( xit ) 7 : xt+1 = proxγtnψ ( x n t ) combination with its much simpler-to-analyze cousin SGD , to the best of our knowledge , there exists87 no theoretical analysis of proximal , parallel , and importance sampling variants of RR with both88 constant and decreasing stepsizes , and in most cases it is not even clear how should such methods be89 constructed . Empowered by and building on the recent advances of Mishchenko et al . [ 2020 ] , in this90 paper we address all these challenges.91 2 Contributions92 In this section we outline the key contributions of our work , and also offer a few intuitive explanations93 motivating some of the development.94 • New algorithm : ProxRR . Despite rich literature on Proximal SGD [ Gorbunov et al. , 2020 ] , it is95 not obvious how one should extend RR to solve problem ( 1 ) when a regularizer ψ is present . Indeed,96 the standard practice for SGD is to apply the proximal operator after each stochastic step [ Duchi and97 Singer , 2009 ] , i.e. , in analogy with ( 2 ) . On the other hand , RR is motivated by the fact that a data98 pass better approximates the full gradient step . If we applied the proximal operator after each step of99 RR , we would no longer approximate the full gradient after an epoch , as we illustrate next.100 Example 1 . Let n = 2 , ψ ( x ) = 12‖x‖ 2 , f1 ( x ) = 〈c1 , x〉 , f2 ( x ) = 〈c2 , x〉 with some c1 , c2 ∈ Rd,101 c1 6= c2 . Let x0 ∈ Rd , γ > 0 and define x1 = x0 − γ∇f1 ( x0 ) , x2 = x1 − γ∇f2 ( x1 ) . Then , we102 have prox2γψ ( x2 ) = prox2γψ ( x0 − 2γ∇f ( x0 ) ) . However , if x̃1 = proxγψ ( x0 − γ∇f1 ( x0 ) ) and103 x̃2 = proxγψ ( x1 − γ∇f2 ( x̃1 ) ) , then x̃2 6= prox2γψ ( x0 − 2γ∇f ( x0 ) ) .104 Motivated by this observation , we propose ProxRR ( Algorithm 1 ) , in which the proximal operator is105 applied at the end of each epoch of RR , i.e. , after each pass through all randomly reshuffled data . A106 notable property of Algorithm 1 is that only a single proximal operator evaluation is needed during107 each data pass . This is in sharp contrast with the way Proximal SGD works , and offers significant108 advantages in regimes where the evaluation of the proximal mapping is expensive ( e.g. , comparable109 to the evaluation of n gradients∇f1 , . . . , ∇fn ) .110 • Convergence of ProxRR ( for strongly convex functions or regularizer ) . We establish several111 convergence results for ProxRR , of which we highlight two here . Both offer a linear convergence rate112 with a fixed stepsize to a neighborhood of the solution . In both we reply on Assumption 1 . Firstly , in113 the case when in addition , each fi is µ-strongly convex , we prove the rate ( see Theorem 2 ) 114 E [ ‖xT − x∗‖2 ] ≤ ( 1− γµ ) nT ‖x0 − x∗‖2 + 2γ 2σ2rad µ , where γt = γ ≤ 1/Lmax is the stepsize , and σ2rad is a shuffling radius constant ( for precise definition,115 see ( 4 ) ) . In Theorem 1 we bound the shuffling radius in terms of ‖∇f ( x∗ ) ‖2 , n , Lmax and the more116 common quantity σ2∗ : = 1 n ∑n i=1 ‖∇fi ( x∗ ) −∇f ( x∗ ) ‖ 2 . Secondly , if ψ is µ-strongly convex , and117 we choose the stepsize γt = γ ≤ 1/Lmax , we prove the rate ( see Theorem 3 ) 118 E [ ‖xT − x∗‖2 ] ≤ ( 1 + 2γµn ) −T ‖x0 − x∗‖2 + γ 2σ2rad µ . Both mentioned rates show exponential ( linear in logarithmic scale ) convergence to a neighborhood119 whose size is proportional to γ2σ2rad . Since we can choose γ to be arbitrarily small or periodically120 decrease it , this implies that the iterates converge to x∗ in the limit . Moreover , we show in Section 4121 that when γ = O ( 1T ) the error is O ( 1 T 2 ) , which is superior to the O ( 1 T ) error of SGD.122 • Results for SO . All of our results apply to the Shuffle-Once algorithm as well . For simplicity , we123 center the discussion around RR , whose current theoretical guarantees in the nonconvex case are124 better than that of SO . Nevertheless , the other results are the same for both methods , and ProxRR is125 identical to ProxSO in terms of our theory too . A study of the empirical differences between RR and126 SO can be found in [ Mishchenko et al. , 2020 ] .127 • Application to Federated Learning . In Section 6 we describe an application of our results to128 federated learning [ Konečný et al. , 2016 , McMahan et al. , 2017 , Kairouz et al. , 2019 ] . In this way we129 obtain the FedRR method , which is similar to Local SGD , except the local solver is a single pass130 of RR over the local data . Empirically , FedRR can be vastly superior to Local SGD ( see Figure 2 ) .131 Remarkably , we also show that the rate of FedRR beats the best known lower bound for Local SGD132 due to [ Woodworth et al. , 2020 ] ( we needed to adapt it from the original online to the finite-sum133 setting we consider in this paper ) for large enough n. See Section F for more details.134 • Nonconvex analysis . In the nonconvex regime , and under suitable assumptions , we establish ( see135 Theorems 5 and 8 ) anO ( 1γT ) rate up to a neighborhood of sizeO ( γ 2 ) . For a certain stepsize it yields136 an O ( 1ε3 ) convergence rate.137 Besides the above results , we describe several extensions in the appendix , which we now outline.138 • Extension 1 : Decreasing stepsizes . The convergence of RR is not always exact and depends on139 the parameters of the objective . Similarly , if the shuffling radius σ2rad is positive , and we wish to find140 an ε-approximate solution , the optimal choice of a fixed stepsize for ProxRR will depend on ε. This141 deficiency can be fixed by using decreasing stepsizes in both vanilla RR [ Ahn et al. , 2020 ] and in142 SGD [ Stich , 2019 ] . We adopt the same technique to our setting . However , we depart from [ Ahn et al.,143 2020 ] by only adjusting the stepsize once per epoch rather than at every iteration , similarly to the144 concurrent work of Tran et al . [ 2020 ] on RR with momentum . For details , see Section I.145 • Extension 2 : Importance resampling for Proximal RR . While importance sampling is a well146 established technique for speeding up the convergence of SGD [ Zhao and Zhang , 2015 , Khaled and147 Richtárik , 2020 ] , no importance sampling variant of RR has been proposed nor analyzed . This is not148 surprising since the key property of importance sampling in SGD—unbiasedness—does not hold for149 RR . Our approach to equip ProxRR with importance sampling is via a reformulation of problem ( 1 ) 150 into a similar problem with a larger number of summands . In particular , for each i ∈ [ n ] we include151 ni copies of the function 1ni fi , and then take average of all N = ∑ i ni functions constructed this152 way . The value of ni depends on the “ importance ” of fi , described below . We then apply ProxRR153 to this reformulation . If fi is Li-smooth for all i ∈ [ n ] and we let L̄ : = 1n ∑ i Li , then we choose154 ni = dLi/L̄e . It is easy to show that N ≤ 2n , and hence our reformulation leads to at most a doubling155 of the number of functions forming the finite sum . However , the overall complexity of ProxRR156 applied to this reformulation will depend on L̄ instead of maxi Li ( see Theorem 10 ) , which can lead157 to a significant improvement . For details of the construction and our complexity results , see Section J.158 3 Preliminaries159 In our analysis , we build upon the notions of limit points and shuffling variance introduced by160 Mishchenko et al . [ 2020 ] for vanilla ( i.e. , non-proximal ) RR . Given a stepsize γ > 0 ( held constant161 during each epoch ) and a permutation π of { 1 , 2 , . . . , n } , the inner loop iterates of RR/SO converge162 to a neighborhood of intermediate limit points x1∗ , x 2 ∗ , . . . , x n ∗ defined by163 xi∗ : = x∗ − γ ∑i−1 j=0∇fπj ( x∗ ) , i = 1 , . . . , n. ( 3 ) The intuition behind this definition is fairly simple : if we performed i steps starting at x∗ , we would164 end up close to xi∗ . To quantify the closeness , we define the shuffling radius.165 Definition 1 ( Shuffling radius ) . Given a stepsize γ > 0 and a random permutation π of { 1 , 2 , . . . , n } 166 used in Algorithm 1 , define xi∗ = x i ∗ ( γ , π ) as in ( 3 ) . Then , the shuffling radius is defined by167 σ2rad ( γ ) : = maxi=0 , ... , n−1 [ 1 γ2Eπ [ Dfπi ( x i ∗ , x∗ ) ] ] , ( 4 ) where the expectation is taken with respect to the randomness in the permutation π . If there are168 multiple stepsizes γ1 , γ2 , . . . used in Algorithm 1 , we take the maximum of all of them as the shuffling169 radius , i.e. , σ2rad : = maxt≥1 σ 2 rad ( γt ) .170 The shuffling radius is related by a multiplicative factor in the stepsize to the shuffling variance171 introduced by Mishchenko et al . [ 2020 ] . When the stepsize is held fixed , the difference between the172 two notions is minimal . When the stepsize is decreasing , however , the shuffling radius is easier to173 work with , since it can be upper bounded by problem constants independent of the stepsizes.174 Armed with a special lemma for sampling without replacement , we can upper bound the shuffling175 radius using the smoothness constant Lmax , size of the vector ∇f ( x∗ ) , and the variance σ2∗ of the176 gradient vectors∇f1 ( x∗ ) , . . . , ∇fn ( x∗ ) .177 Theorem 1 ( Bounding the shuffling radius ) . For any stepsize γ > 0 and any random permutation π178 of { 1 , 2 , . . . , n } we have σ2rad ≤ Lmax 2 n ( n‖∇f ( x∗ ) ‖2 + 12σ 2 ∗ ) , where x∗ is a solution of Problem ( 1 ) 179 and σ2∗ is the population variance at the optimum180 σ2∗ : = 1 n ∑n i=1‖∇fi ( x∗ ) −∇f ( x∗ ) ‖ 2 . ( 5 ) All proofs are relegated to the supplementary material . In order to better understand the bound181 given by Theorem 1 , note that if there is no proximal operator ( i.e. , ψ = 0 ) then ∇f ( x∗ ) = 0 and182 we get that σ2rad ≤ Lmaxnσ 2 ∗ 4 . This recovers the existing upper bound on the shuffling variance of183 Mishchenko et al . [ 2020 ] for vanilla RR . On the other hand , if∇f ( x∗ ) 6= 0 then we get an additive184 term of size proportional to the squared norm of∇f ( x∗ ) .185 4 Theory for strongly convex losses f1 , . . . , fn186 Our first theorem establishes a convergence rate for Algorithm 1 applied with a constant stepsize to187 Problem ( 1 ) when each objective fi is strongly convex . This assumption is commonly satisfied in188 machine learning applications where each fi represents a regularized loss on some data points , as in189 ` 2 regularized linear regression and ` 2 regularized logistic regression.190 Theorem 2 . Let Assumption 1 be satisfied . Further , assume that each fi is µ-strongly convex . If191 Algorithm 1 is run with constant stepsize γt = γ ≤ 1/Lmax , then its iterates satisfy192 E [ ‖xT − x∗‖2 ] ≤ ( 1− γµ ) nT ‖x0 − x∗‖2 + 2γ 2σ2rad µ . We can convert the guarantee of Theorem 2 to a convergence rate by properly tuning the stepsize193 and using the upper bound of Theorem 1 on the shuffling radius . In particular , if we choose the194 stepsize as γ = min { 1 Lmax , √ εµ√ 2σrad } , and let κ : = Lmax/µ and r0 : = ‖x0 − x∗‖2 , then we obtain195 E [ ‖xT − x∗‖2 ] = O ( ε ) provided that the total number of iterations KRR = nT is at least196 KRR ≥ [ ( κ+ √ κn√ εµ ( √ n ‖∇f ( x∗ ) ‖+ σ∗ ) ] log ( 2r0 ε ) . ( 6 ) Comparison with vanilla RR . If there is no proximal operator , then ‖∇f ( x∗ ) ‖ = 0 and we recover197 the earlier result of Mishchenko et al . [ 2020 ] on the convergence of RR without proximal , which is198 optimal in ε up to logarithmic factors . On the other hand , when the proximal operator is nonzero,199 we get an extra term in the complexity proportional to ‖∇f ( x∗ ) ‖ : thus , even when all the functions200 are the same ( i.e. , σ∗ = 0 ) , we do not recover the linear convergence of Proximal Gradient Descent201 [ Karimi et al. , 2016 , Beck , 2017 ] . This can be easily explained by the fact that Algorithm 1 performs202 n gradient steps per one proximal step . Hence , even if f1 = · · · = fn , Algorithm 1 does not reduce203 to Proximal Gradient Descent . We note that other algorithms for composite optimization which may204 not take a proximal step at every iteration ( for example , using stochastic projection steps ) also suffer205 from the same dependence [ Patrascu and Irofti , 2021 ] .206 Comparison with proximal SGD . In order to compare ( 6 ) against the complexity of Proximal SGD207 ( Algorithm 2 ) , we recall that Proximal SGD achieves E [ ‖xK − x∗‖2 ] = O ( ε ) if either f or ψ is208 µ-strongly convex and209 KSGD ≥ ( κ+ σ2∗ εµ2 ) log ( 2r0 ε ) . ( 7 ) Algorithm 2 Proximal SGD Require : Stepsizes γk > 0 , initial vector x0 ∈ Rd , number of steps K 1 : for steps k = 0 , 1 , . . . , K − 1 do 2 : Sample ik uniformly at random from [ n ] 3 : xk+1 = proxγkψ ( xk − γk∇fik ( xk ) ) This result is standard [ Needell et al. , 2016 , Gower et al. , 2019 ] , with the exception that we do not210 know any proof in the literature for the case when ψ is strongly convex . For completeness , we prove211 it in Appendix C , but since our proof is a minor modification of that in [ Gower et al. , 2019 ] , we do212 not provide it here.213 By comparing KSGD ( given by ( 7 ) ) and KRR ( given by ( 6 ) ) , we see that ProxRR has milder214 dependence on ε than Proximal SGD . In particular , ProxRR converges faster whenever the target215 accuracy ε is small enough to satisfy ε ≤ 1Lmaxnµ ( σ4∗ n‖∇f ( x∗ ) ‖2+σ2∗ ) . Furthermore , ProxRR is much216 better when we consider proximal iteration complexity ( # of proximal operator access ) , in which case217 the complexity of ProxRR ( 6 ) is reduced by a factor of n ( because we take one proximal step every n218 iterations ) , while the proximal iteration complexity of Proximal SGD remains the same as ( 7 ) . In this219 case , ProxRR is better whenever the accuracy ε satisfies220 ε ≥ nLmaxµ [ n‖∇f ( x∗ ) ‖2 + σ2∗ ] or ε ≤ nLmaxµ [ σ4∗ n‖∇f ( x∗ ) ‖2+σ2∗ ] . We can see that if the target accuracy is large enough or small enough , and if the cost of proximal221 operators dominates the computation , ProxRR is much quicker to converge than Proximal SGD.222

Summary: This paper studies a proximal extension of random reshuffling algorithm, namely ProxRR for finite-sum composite optimization. ProxRR requires only one proximal steps per round, which reduces the complexities incurred by costly proximal operators. This paper also studies a federated version of ProxRR, namely FedRR, and proved theoretical advantages over Local SGD.