Minibatch vs Local SGD with Shuffling: Tight Convergence Bounds and Beyond

1 INTRODUCTION . Distributed learning within the framework of federated learning ( Konečnỳ et al. , 2016 ; McMahan et al. , 2017 ) has witnessed increasing interest recently . A key property of this framework is that models are trained locally using only private data on devices/machines distributed across a network , while parameter updates are aggregated and synchronized at a server.1 Communication is often the key bottleneck for federated learning , which drives the search for algorithms that can train fast while requiring less communication—see Li et al . ( 2020a ) ; Kairouz et al . ( 2021 ) for recent surveys . A basic algorithm for federated learning is local stochastic gradient descent ( SGD ) , also known as federated averaging . The goal is to minimize the global objective that is an average of the local objectives . In local SGD , we have M machines and a server . After each round of communication , each of the M machines locally runs B steps of SGD on its local objective . Every B iterations , the server aggregates the updated local iterates from the machines , averages them , and then synchronizes the machines with the average . Convergence analysis of local SGD and its variants has drawn great interest recently ( Dieuleveut & Patel , 2019 ; Haddadpour et al. , 2019 ; Haddadpour & Mahdavi , 2019 ; Stich , 2019 ; Yu et al. , 2019 ; Li et al. , 2020b ; c ; Koloskova et al. , 2020 ; Khaled et al. , 2020 ; Spiridonoff et al. , 2020 ; Karimireddy et al. , 2020 ; Stich & Karimireddy , 2020 ; Qu et al. , 2020 ) . Of the many , the biggest motivation for our paper comes from the line of work by Woodworth et al . ( 2020a ; b ; 2021 ) . In ( Woodworth et al. , 2020a ; b ) , minibatch SGD is studied as a simple yet powerful baseline for this intermittent communication setting . Instead of locally updating the iterates B times , minibatch SGD aggregates B gradients ( evaluated at the last synced iterate ) from each of the M machines , forms a minibatch of size MB , and then updates the shared iterate . Given the same M and B , local SGD and minibatch SGD have the same number of gradient computations per round of communication , so it is worthwhile to understand which converges faster . Woodworth et al . ( 2020a ; b ) point out that many existing analyses on local SGD show inferior convergence rate compared to minibatch SGD . Through their new upper and lower bounds , they identify regimes where local SGD can be faster than minibatch SGD . While the theory of local and minibatch SGD has seen recent progress , there is still a gap between what is analyzed versus what is actually used . Most theoretical results assume independent and 1A distinctive feature of federated learning is that not all devices necessarily participate in the updates ; however , we focus on the full participation setting in this paper . unbiased gradient estimates obtained via with-replacement sampling of stochastic gradients ( i.e. , choosing training data indices uniformly at random ) . In contrast , most practitioners use withoutreplacement sampling , where they shuffle indices randomly and access them sequentially . Convergence analysis of without-replacement methods is challenging because gradients sampled within an epoch lack independence . As a result , the standard theory based on independent gradient estimates does not apply to shuffling-based methods . While shuffling-based methods are believed to be faster in practice ( Bottou , 2009 ) , broad theoretical understanding of such methods remains elusive , except for noteworthy recent progress mainly focusing on the analysis of SGD ( Gürbüzbalaban et al. , 2019 ; Haochen & Sra , 2019 ; Nagaraj et al. , 2019 ; Nguyen et al. , 2020 ; Safran & Shamir , 2020 ; 2021 ; Rajput et al. , 2020 ; 2021 ; Ahn et al. , 2020 ; Mishchenko et al. , 2020 ; 2021 ; Tran et al. , 2021 ) . These results indicate that in the large-epoch regime ( where the number of epochs is greater than some threshold ) , without-replacement SGD converges faster than with-replacement SGD . 1.1 OUR CONTRIBUTIONS . We analyze convergence rates of without-replacement versions of local and minibatch SGD , where local component functions are reshuffled at every epoch . We call the respective algorithms local RR ( Algorithm 1 ) and minibatch RR ( Algorithm 2 ) , and their with-replacement counterparts local SGD and minibatch SGD . Our key contributions are as follows : • In Section 3 , we present convergence bounds on minibatch and local RR for L-smooth functions satisfying the µ-Polyak-Łojasiewicz condition ( Theorems 1 & 2 ) . Our theorems give highprobability bounds , a departure from the common in-expectation bounds in the literature . We show that minibatch and local RR converge faster than minibatch and local SGD when the number of epochs is sufficiently large . We also identify a regime where local RR converges as fast as minibatch RR : when synchronization happens frequently enough and local objectives are not too heterogeneous . See also Appendix A for a detailed comparison with existing upper bounds . • In Section 4 , we prove that the upper bounds obtained in Section 3 are tight , in all factors except L and µ . We present Theorems 3 & 4 and Proposition 5 which show lower bounds that match the upper bound up to a factor of L2/µ2 . Our lower bound on local RR indicates that if the synchronization interval B is too large , then local RR has no gain from parallel computation . • In Section 5 , we propose a simple modification called synchronized shuffling that allows us to bypass the lower bounds in Section 4 , at the cost of a slight increase in communication . By having the server broadcast random permutations to local machines , we show that in near-homogeneous settings , the modified algorithms converge faster than the lower bounds ( Theorems 6 & 7 ) . • In Appendix C , we present numerical experiments that corroborate our theoretical findings . 2 PROBLEM SETUP . Notation . For a natural number a ∈ N , let [ a ] : = { 1 , 2 , . . . , a } . Let Sa be the set of all permutations of [ a ] . Since our indices start from 1 , we redefine the modulo operation between a ∈ Z and b ∈ N as a mod b : = a− ba−1b cb , to make a mod b ∈ [ b ] . Optimization task . Consider M machines , each with its objective Fm ( x ) : = 1N ∑N i=1 f m i ( x ) , for m ∈ [ M ] . The m-th machine has access only to the gradients of its own N local components fm1 ( x ) , . . . , f m N ( x ) . In this setting , we wish to minimize the global objective function which is an average of the local objectives : F ( x ) : = 1M ∑M m=1 F m ( x ) = 1MN ∑M m=1 ∑N i=1 f m i ( x ) . Further , we assume that each individual component function fmi is L-smooth , so that fmi ( y ) ≤ fmi ( x ) + 〈∇fmi ( x ) , y − x〉+ L2 ‖y − x‖ 2 , for all x , y ∈ Rd , ( 1 ) and that the global objective F satisfies the µ-Polyak-Łojasiewicz ( PŁ ) condition.2 1 2 ‖∇F ( x ) ‖ 2 ≥ µ ( F ( x ) − F ∗ ) for all x ∈ Rd , where µ > 0 . ( 2 ) Algorithms . Under the above setting , we analyze local RR ( Algorithm 1 ) and minibatch RR ( Algorithm 2 ) and characterize their worst-case convergence rates.3 The algorithms are run overK epochs , 2PŁ functions can be thought as a nonconvex generalization of strongly convex functions . 3In Algorithms 1 and 2 , consider SYNCSHUF as FALSE for now . We will discuss SYNCSHUF in Section 5 . Algorithm 1 Local RR ( with and without SYNCSHUF ) Input : Initialization y0 , step-size η , # machines M , # components N , # epochs K , sync interval B . 1 : Initialize xm1,0 : = y0 for all m ∈ [ M ] . 2 : for k ∈ [ K ] do 3 : if SYNCSHUF = TRUE then . Local RR with SYNCSHUF 4 : Sample σ ∼ Unif ( SN ) , π ∼ Unif ( SM ) . 5 : Set σmk ( i ) : = σ ( ( i+ N M π ( m ) ) mod N ) for all m ∈ [ M ] , i ∈ [ N ] . 6 : else . Local RR 7 : Sample σmk ∼ Unif ( SN ) independently and locally , for all m ∈ [ M ] . 8 : end if 9 : for i ∈ [ N ] do 10 : for m ∈ [ M ] do locally 11 : Update xmk , i : = x m k , i−1 − η∇fmσm k ( i ) ( x m k , i−1 ) . 12 : end for 13 : if B divides i then 14 : Aggregate and average yk , i B : = 1 M ∑M m=1 x m k , i . 15 : Synchronize xmk , i : = yk , i B , for all m ∈ [ M ] . 16 : end if 17 : end for 18 : xmk+1,0 : = yk , N B , for all m ∈ [ M ] . 19 : end for 20 : return the last iterate yK , N B . Algorithm 2 Minibatch RR ( with and without SYNCSHUF ) Input : Initialization x0 , step-size η , # machines M , # components N , # epochs K , sync interval B . 1 : Initialize x1,0 : = x0 . 2 : for k ∈ [ K ] do 3 : if SYNCSHUF = TRUE then . Minibatch RR with SYNCSHUF 4 : Sample σ ∼ Unif ( SN ) , π ∼ Unif ( SM ) . 5 : Set σmk ( i ) : = σ ( ( i+ N M π ( m ) ) mod N ) for all m ∈ [ M ] , i ∈ [ N ] . 6 : else . Minibatch RR 7 : Sample σmk ∼ Unif ( SN ) independently and locally , for all m ∈ [ M ] . 8 : end if 9 : for i ∈ [ N B ] do 10 : Update xk , i : = xk , i−1 − ηM ∑M m=1 1 B ∑iB j= ( i−1 ) B+1 ∇fmσm k ( j ) ( xk , i−1 ) ︸ ︷︷ ︸ averaging done locally . 11 : end for 12 : xk+1,0 : = xk , N B . 13 : end for 14 : return the last iterate xK , N B . i.e. , K passes over the entire component functions . At the beginning of epoch k , each machine m shuffles its local component functions { fmi } Ni=1 using a random permutation σmk ∼ Unif ( SN ) . In local RR , each machine makes B local RR updates to its iterate by sequentially accessing its shuffled component functions , before the server aggregates iterates from all the machines and then synchronizes the machines with the average iterate . In minibatch RR , instead of making B local updates , each machine collects B gradients evaluated at the last iterate , and the server aggregates them to make an update using theseMB gradients . Since these two algorithms use the same amount of communication and local gradients , minibatch RR is a simple yet powerful baseline for local RR . Below , we collect our assumptions on the algorithm parameters used throughout the paper . Assumption 1 ( Algorithm parameters ) . We assume M ≥ 1 , N ≥ 2 , and K ≥ 1 . Also , assume that B divides N . We restrict 1 ≤ B ≤ N2 for minibatch RR because B = N makes the algorithm equal to GD . We also assume 2 ≤ B ≤ N for local RR because B = 1 makes the two algorithms the same . We choose a constant step-size scheme , i.e. , η > 0 is kept constant over all updates . We next state assumptions on intra- and inter-machine deviations used in this paper.4 4Assumptions 2 , 3 & 4 require that they hold for the whole Rd . We discuss ways to avoid it in Appendix D.7 . Assumption 2 ( Intra-machine deviation ) . There exists ν ≥ 0 such that for allm ∈ [ M ] and i ∈ [ N ] , ‖∇fmi ( x ) −∇Fm ( x ) ‖ ≤ ν , for all x ∈ Rd . Assumption 2 requires that the difference between the gradient of each local component function fmi ( x ) and its corresponding local objective function F m ( x ) is uniformly bounded . It models the variance of local components fmi within each machine . While the uniform boundedness requirement may look strong , we use this assumption to prove high-probability upper bounds , which are stronger than the common in-expectation bounds . See Appendix A for comparisons with other assumptions , and also Appendix D.7 for ways to avoid uniform boundedness over the entire Rd . The next two assumptions capture the deviation across different machines , i.e. , the degree of heterogeneity , in two different levels of granularity : objective-wise and component-wise . Assumption 3 ( Objective-wise inter-machine deviation ) . There exist τ ≥ 0 and ρ ≥ 1 such that 1 M ∑M m=1 ‖∇Fm ( x ) ‖ ≤ τ + ρ ‖∇F ( x ) ‖ , for all x ∈ Rd . Assumption 3 models the heterogeneity by bounding the mean of ‖∇Fm‖ by a constant plus a multiplicative factor times ‖∇F‖ . The assumption includes the homogeneous case ( i.e. , F 1 = · · · = FM = F ) by τ = 0 and ρ = 1 . Assumption 3 is weaker than many other heterogeneity assumptions in the literature ( e.g. , Karimireddy et al . ( 2020 ) ) ; see Appendix A for detailed comparisons . Assumption 3 measures heterogeneity by only considering the local objectives Fm , not the local components fmi . We consider a more fine-grained notion of heterogeneity in Assumption 4 : Assumption 4 ( Component-wise inter-machine deviation ) . For all i ∈ [ N ] , let f̄i : = 1M ∑M m=1 f m i . There exist λ ≥ 0 such that for all m ∈ [ M ] and i ∈ [ N ] , ∥∥∇fmi ( x ) −∇f̄i ( x ) ∥∥ ≤ λ , for all x ∈ Rd . Assumption 4 states that the gradients of the i-th components of local machines are “ close ” to each other . The assumption subsumes the component-wise homogeneous setting , i.e. , f1i = f 2 i = · · · = fMi , by λ = 0 . In distributed learning , this choice corresponds to the setting where each machine has the same training dataset . Assumption 4 with λ > 0 is also relevant to the case where each device has a slightly perturbed ( e.g. , by data augmentation techniques ) version of a certain dataset . It is straightforward to check that Assumption 4 implies Assumption 3 with τ = λ and ρ = 1 . We conclude this section by defining the function classes we study in this paper . Definition 1 ( Function classes ) . We consider two classes of global objective functions F , also taking into account their local objectives Fm and local components fmi . We assume throughout that f m i are differentiable and F is bounded from below . Fobj ( L , µ , ν , τ , ρ ) : = { F | F is µ-PŁ ; fmi are L-smooth ; F , Fm , fmi satisfy Assumptions 2 & 3 } , Fcmp ( L , µ , ν , λ ) : = { F | F is µ-PŁ ; fmi are L-smooth ; F , Fm , fmi satisfy Assumptions 2 & 4 } . Notice that Fobj ( L , µ , ν , τ , ρ ) ⊃ Fcmp ( L , µ , ν , τ ) for any ρ ≥ 1 . We only make the PŁ assumption on the global objective F , not on the local objectives Fm nor on the local components fmi . Using L and µ , we define the condition number κ : = L/µ ≥ 1 .

Summary: This paper studies two popular algorithms: local SGD and minibatch SGD using the data shuffling technique, which means sampling without replacement. Authors provide analysis under PL condition and show that in some cases these methods with shuffling outperform classical local SGD and minibatch SGD. Additionally, this paper provides lower bounds for minibatch RR and local RR in homogeneous and heterogeneous settings. At the end of the paper, a new technique of using synchronized permutations is proposed. In an almost homogeneous setting, this approach can improve rates.