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 proposes two variants of stochastic gradient algorithms without replacement. For smooth functions satisfying the PŁ condition, the authors showed that the proposed shuffling-based variants converge faster than their with-replacement counterparts (for the case with large number of epochs). Moreover, the authors showed that their provided convergence analysis is tight for the case with not large number of epochs.