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 first proposes a proximal variant of the gradient method with shuffling strategy for a class of composite finite-sum minimization problems, and then specifies it for an optimization problem in federated learning (FL). In fact, the proposed algorithm performs n gradient descent updates per epoch and then applies the proximal operator at the end of each epoch. Both single shuffling and randomized reshuffling strategies can be used. The authors analyze the convergence rates of this algorithm under the strong convexity of the losses and the strong convexity of the regularizer. Next, they apply the proposed algorithm to FL, with and without regularizer. This variant is similar to FedAvg, but it replaces SGD by a shuffling gradient scheme for local updates. They establish the convergence of this variant for both the strongly convex case (where the regularizer is strongly convex) and the non-convex case in both heterogeneity and homogeneity settings, respectively. Numerical experiments are conducted to illustrate the algorithm and compare it with two existing methods: local SGD and SCAFFOLD.