Optimal Client Sampling for Federated Learning

It is well understood that client-master communication can be a primary bottleneck1 in Federated Learning . In this work , we address this issue with a novel client sub-2 sampling scheme , where we restrict the number of clients allowed to communicate3 their updates back to the master node . In each communication round , all partici-4 pated clients compute their updates , but only the ones with “ important ” updates5 communicate back to the master . We show that importance can be measured using6 only the norm of the update and give a formula for optimal client participation.7 This formula minimizes the distance between the full update , where all clients8 participate , and our limited update , where the number of participating clients is9 restricted . In addition , we provide a simple algorithm that approximates the optimal10 formula for client participation which only requires secure aggregation and thus11 does not compromise client privacy . We show both theoretically and empirically12 that our approach leads to superior performance for Distributed SGD ( DSGD ) and13 Federated Averaging ( FedAvg ) compared to the baseline where participating clients14 are sampled uniformly . Our approach is orthogonal to and compatible with ex-15 isting methods for reducing communication overhead , such as local methods and16 communication compression methods.17 1 Introduction18 We consider the standard cross-device Federated Learning ( FL ) setting [ 13 ] , where the objective is of19 the form20 min x∈Rd [ f ( x ) : = n∑ i=1 wifi ( x ) ] , ( 1 ) where x ∈ Rd represents the parameters of a statistical model we aim to find , n is the total number of21 clients , fi : Rd → R is a continuously differentiable local loss function which depends on the data Di22 owned by client i via fi = Eξ∼Di [ f ( x , ξ ) ] , and wi ≥ 0 are client weights such that ∑n i=1 wi = 1.23 We assume the classical FL setup in which a central master ( server ) orchestrates the training by24 securely aggregating updates from clients without seeing the raw data.25 1.1 Communication as the Bottleneck26 It is well understood that cost of communication can be the primary bottleneck in Federated Learning.27 Indeed , wireless links and other end-user internet connections typically operate at lower rates than28 intra-datacenter or inter-datacenter links and can be potentially expensive and unreliable . Moreover,29 the capacity of the aggregating master and other FL system considerations impose direct or indirect30 constrains on the number of clients that are allowed to participate in each communication round.31 These considerations have led to significant interest in reducing the communication bandwidth of FL32 systems.33 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . 1.1.1 Local Methods34 One of the most popular strategies is to reduce the frequency of communication and put more35 emphasis on computation . This is usually achieved by asking the devices to perform multiple local36 steps before communicating their updates . A prototype method in this category is the Federated37 Averaging ( FedAvg ) algorithm [ 23 ] . The original work was a heuristic , offering no theoretical38 guarantees , which motivated the community to try to understand the method and various existing and39 new variants theoretically [ 35 , 21 , 15 , 37 , 17 , 9 ] .40 1.1.2 Communication Compression41 Another popular approach is to reduce the size of the object ( typically gradients ) communicated from42 clients to the master . These techniques are usually referred to as gradient/communication compression.43 In this approach , instead of transmitting the full-dimensional gradient/update vector g ∈ Rd , one44 transmits a compressed vector C ( g ) , where C : Rd → Rd is a ( possibly random ) operator chosen45 such that C ( g ) can be represented using fewer bits , for instance by using limited bit representation46 ( quantization ) or by enforcing sparsity ( sparsification ) . A particularly popular class of quantization47 operators is based on random dithering [ 7 , 30 ] ; see [ 1 , 41 , 42 , 28 ] . A new variant of random dithering48 developed in [ 10 ] offers an exponential improvement on standard dithering . Sparse vectors can be49 obtained by random sparsification techniques that randomly mask the input vectors and preserve a50 constant number of coordinates only [ 40 , 18 , 36 , 24 , 39 ] . There is also a line of work [ 10 , 3 ] where a51 combination of sparsification and quantization was proposed to obtain a more aggressive combined52 effect.53 1.2 Related Work54 Importance sampling methods for optimization have been studied extensively in the last few years in55 several contexts , including convex optimization and deep learning . LASVM developed in [ 5 ] , which is56 an online algorithm that uses importance sampling to train kernelized support vector machines . The57 first importance sampling for randomized coordinate descent methods was proposed in a seminal58 paper in [ 26 ] . It was showed in [ 29 ] that the proposed sampling is optimal . Later , several extensions59 and improvements followed [ 33 , 20 , 6 , 27 , 2 , 38 ] . Another branch of work studies sample complexity.60 In [ 25 , 43 ] , the authors make a connection with the variance of the gradient estimates of SGD and61 show that the optimal sampling distribution is proportional to the per-sample gradient norm . In terms62 of computation , obtaining this distribution is as hard as the computation of the full gradient , thus63 it is not practical . For simpler problems , one can sample proportionally to the norms of the inputs,64 which can be linked to the Lipschitz constants of the per-sample loss function for linear and logistic65 regression . For instance , it was shown in [ 11 ] that static optimal sampling can be constructed even66 for mini-batches and the probability is proportional to these Lipschitz constants under the assumption67 that these constants of the per-sample loss function are known . Unfortunately , importance measures68 such as smoothness of the gradient are often hard to compute/estimate for more complicated models69 such as those arising in deep learning , where most of the importance sampling schemes are based70 on heuristics . A manually designed sampling scheme was proposed in [ 4 ] . It was inspired by the71 perceived way that human children learn ; in practice , they provide the network with examples of72 increasing difficulty in an arbitrary manner . In a diametrically opposite approach , it is common for73 deep embedding learning to sample hard examples because of the plethora of easy non-informative74 ones [ 32 , 34 ] . Other approaches use a history of losses for previously seen samples to create the75 sampling distribution and sample either proportionally to the loss or based on the loss ranking [ 31 , 22 ] .76 In [ 16 ] , the authors propose to sample based on the gradient norm of a small uniformly sampled77 subset of samples.78 In our work , we avoid all the aforementioned problems as our motivation is not to reduce computation,79 which is not the main bottleneck of Federated Learning , but to use importance sampling to decrease80 the number of bits communicated . This , as we show in Section 2 , allows us to construct optimal81 adaptive sampling ; that is , we do not need to rely on any heuristics , historical losses , or partial82 information.83 1.3 Contributions84 In this work , we propose a new approach to addressing the communication bandwidth issues appearing85 in FL . Our approach is based on the observation that in the situation where partial participation86 is desired and a budget on the number of participating clients is applied , careful selection of the87 participating clients can lead to better communication complexity , and hence faster training . In other88 words , we claim that in any given communication round , some clients will have “ more informative ” 89 updates than others and that the training procedure will benefit from capitalizing on this fact by90 ignoring some of the worthless updates.91 In particular , we propose a principled optimal client sampling scheme capable of identifying the most92 informative clients in any given communication round . Our scheme works by minimizing the variance93 of the stochastic gradient produced by the partial participation procedure , which then translates to94 a probable reduction in the number of communication rounds . To the best of our knowledge , this95 approach was not considered before . Moreover , our proposal is orthogonal to and hence combinable96 with existing approaches to communication reduction such as communication compression or local97 updates ( Section 3.2 ) .98 Our contributions can be summarized as follows:99 • we propose a novel adaptive partial participation strategy for reducing communication in100 FL that works by a careful selection of the clients that are allowed to communicate their101 updates to the master node in any given communication round ; 102 • our adaptive client sampling procedure is optimal in the sense that it minimizes the variance103 of the master update ; 104 • we propose an approximation to our optimal adaptive sampling strategy which only requires105 aggregation , thus allows for secure aggregation and stateless clients ; 106 • we show theoretically that our approach allows for larger learning rates for Distributed SGD107 and FedAvg algorithms than the baseline which performs uniform client sampling , and as a108 result leads to better communication complexity.109 • we show empirically that the performance of our approach is superior to uniform sampling110 and is close to full participation.111 2 Smart Client Sampling for Reducing Communication112 We now describe our client sampling strategy for reducing the communication bottleneck in Federated113 Learning . Each client i participating in round k computes an update vector Uki ∈ Rd . For simplicity114 and ease of exposition , we assume that all clients i ∈ [ n ] : = { 1 , 2 , . . . , n } are available in each round.115 However , we would like to point out that this is not a limiting factor , and all presented theory can be116 easily extended to the case of partial participation with an arbitrary distribution . In our framework,117 only a subset of clients communicates their updates to the master node in each communication round118 in order to reduce the number of transmitted bits.119 In order to provide analysis in this framework , we consider a general partial participation frame-120 work [ 12 ] , where we assume that the subset of participating clients is determined by an arbitrary121 random set-valued mapping S ( a “ sampling ” ) with values in 2 [ n ] . A sampling S is uniquely defined122 by assigning probabilities to all 2n subsets of [ n ] . With each sampling S we associate a probability123 matrix P ∈ Rn×n defined by Pij : = Prob ( { i , j } ⊆ S ) . The probability vector associated with S is124 the vector composed of the diagonal entries of P : p = ( p1 , . . . , pn ) ∈ Rn , where pi : = Prob ( i ∈ S ) .125 We say that S is proper if pi > 0 for all i . It is easy to show that b : = E [ |S| ] = Trace ( P ) = ∑n i=1 pi,126 and hence b can be seen as the expected number of clients participating in each communication round.127 Given parameters p1 , . . . , pn ∈ [ 0 , 1 ] , consider a random set S ⊆ [ n ] generated as follows : for each128 i ∈ [ n ] , we include i in S with probability pi . This is called independent sampling , since the event129 i ∈ S is independent of j ∈ S for any i 6= j.130 While our client sampling strategy can be adapted to essentially any underlying learning method , we131 give details here for DSGD:132 xk+1 = xk − ηkGk , Gk : = ∑ i∈Sk wi pki Uki , ( 2 ) where Sk ∼ Sk and Uki = gki is an unbiased estimator of∇fi ( xk ) . The scaling factor 1pki is necessary133 in order to obtain an unbiased estimator of the true update , i.e. , ESk [ Gk ] = ∑n i=1 wiU k i .134 2.1 Optimal Client Sampling135 We start with a simple observation that the variance of our gradient estimator Gk can be decomposed136 as137 E [ ∥∥Gk −∇f ( xk ) ∥∥2 ] = E ∥∥∥∥∥Gk − n∑ i=1 wiU k i ∥∥∥∥∥ 2 + E ∥∥∥∥∥ n∑ i=1 wiU k i −∇f ( xk ) ∥∥∥∥∥ 2 . Note that the second term on the right-hand side is independent of the sampling procedure and138 the first term is zero if every client sends its update ( i.e. , if pki = 1 for all i ) . In order to provide139 meaningful results , we restrict the expected number of clients to communicate in each round by140 bounding bk : = ∑n i=1 p k i by some positive integer m ≤ n. This raises the following question : What141 is the sampling procedure that minimizes ( 3 ) for any given m ? We answer this question using the142 following technical lemma:143 Lemma 1 . Let ζ1 , ζ2 , . . . , ζn be vectors in Rd and w1 , w2 , . . . , wn be non-negative real numbers144 such that ∑n i=1 wi = 1 . Define ζ̃ : = ∑n i=1 wiζi . Let S be a proper sampling . If v ∈ Rn is such that145 P− pp > Diag ( p1v1 , p2v2 , . . . , pnvn ) , ( 3 ) then146 E ∥∥∥∥∥∑ i∈S wiζi pi − ζ̃ ∥∥∥∥∥ 2 ≤ n∑ i=1 w2i vi pi ‖ζi‖2 , ( 4 ) where the expectation is taken over S. Whenever ( 3 ) holds , it must be the case that vi ≥ 1− pi.147 It turns out that given probabilities { pi } , among all samplings S satisfying pi = Prob ( i ∈ S ) , the148 independent sampling minimizes the left-hand side of ( 4 ) . This is due to two nice properties : a ) any149 independent sampling admits optimal choice of v , i.e. , vi = 1− pi for all i , and b ) for independent150 sampling ( 4 ) holds as equality . In the context of our method , these properties can be written as151 E ∥∥∥∥∥Gk − n∑ i=1 wiU k i ∥∥∥∥∥ 2 = E [ n∑ i=1 w2i 1− pki pki ∥∥Uki ∥∥2 ] . ( 5 ) It now only remains to find the parameters { pki } defining the optimal independent sampling , i.e. , one152 that minimizes ( 5 ) subject to the constraints 0 ≤ pki ≤ 1 and bk : = ∑n i=1 p k i ≤ m. It turns out that153 this problem has the following closed-form solution:154 pki = ( m+ l − n ) ‖Ũki ‖∑l j=1 ∥∥∥Ũk ( j ) ∥∥∥ , if i /∈ Ak , 1 , if i ∈ Ak , ( 6 ) where Ũki : = wiU k i , and ∥∥∥Ũk ( j ) ∥∥∥ is the j-th largest value in { ∥∥∥Ũki ∥∥∥ } n i=1 , l is the largest integer for155 which 0 < m+ l − n ≤ ∑l i=1‖Ũk ( i ) ‖∥∥∥Ũk ( l ) ∥∥∥ ( note that this inequality at least holds for l = n−m+ 1 ) , and156 Ak contains indices i such that ∥∥∥Ũki ∥∥∥ ≥ ∥∥∥Ũk ( l+1 ) ∥∥∥ . We summarize this procedure in Algorithm 1.157 2.2 Secure Aggregation158 Note that in the case l = n , the optimal probabilities pki = m ‖Ũki ‖∑n j=1‖Ũkj ‖ can be computed easily : the159 master aggregates the norm of each update and then sends the sum back to the clients . However , if160 l < n , in order to compute optimal probabilities , the master would need to identify the norm of every161 Algorithm 1 Optimal Client Sampling ( OCS ) . 1 : Input : expected batch size m 2 : each client i computes a local update Uki ( in parallel ) 3 : each client i sends the norm of its update uki = wi ∥∥Uki ∥∥ to the master ( in parallel ) 4 : master computes optimal probabilities pki using equation ( 6 ) 5 : master broadcasts pki to all clients 6 : each client i sends its update wi pki Uki to the master with probability p k i ( in parallel ) update and perform partial sorting , which can be computationally expensive and also slightly violates162 the privacy requirements of clients in FL.163 Therefore , we develop an algorithm for approximately solving the problem , which only requires to164 perform aggregation at the master node without compromising privacy of any client . The construction165 of this algorithm is similar to [ 40 ] . We first set p̃ki = m‖Ũki ‖∑n j=1‖Ũkj ‖ and pki = min { p̃ki , 1 } . In an ideal166 situation , this would be sufficient . However , due to the truncation operation , the expected minibatch167 size bk = ∑n i=1 p k i ≤ ∑n i=1 m‖gki ‖∑n j=1‖gkj ‖ = m can be strictly less than m if p̃ki > 1 holds true for at168 least one i . Hence , we employ an iterative procedure to fix this gap by rescaling the probabilities169 which are smaller than 1 , as summarized in Algorithm 2 . This algorithm is much easier to implement170 and computationally more efficient on parallel computing architectures . In addition , it only requires a171 secure aggregation procedure on the master , which is essential in privacy preserving FL , and thus it is172 compatible with existing FL software and hardware . We realize that Algorithm 2 brings some extra173 communication costs , but this is not an issue as it only requires to communicate O ( jmax ) extra floats174 for each client . We pick jmax = O ( 1 ) , and thus it is negligible for large models of size d.175 Remark 1 . We realize that our algorithm requires two communication rounds per optimization round,176 but the first round is negligible due to the minimal number of communicated bits as argued above.177 3 Convergence Guarantees178 In this section , we provide convergence analysis of DSGD and FedAvg with our optimal client sampling179 technique and compare it with full participation and independent uniform sampling of m clients.180 We use standard assumptions [ 14 ] and assume throughout that f has a unique minimizer x ? with181 f ? = f ( x ? ) > −∞ . We further assume that f is µ-strongly convex and fi ’ s are L-smooth and182 convex . Detailed definitions of convexity and smoothness can be found in the Appendix . Note that183 nothing prevents us from extending the results in this section to convex and non-convex cases with a184 similar standard analysis , since our proposed method only affects the aggregation step as described in185 Section 2 , which is independent of the strong convexity assumption.186 Assumption 1 ( Gradient oracle for DSGD ) . The stochastic gradient estimator gki = ∇fi ( xk ) + ξki of187 the local gradient∇fi ( xk ) , for each round k and all i = 1 , . . . , n , satisfies188 E [ ξki ] = 0 ( 7 ) and189 E [ ∥∥ξki ∥∥2 |xki ] ≤M ∥∥∇fi ( xk ) ∥∥2 + σ2 , for some M ≥ 0 . ( 8 ) This further implies that E [ 1 n ∑n i=1 g k i | xk ] = ∇f ( xk ) .190 Assumption 2 ( Gradient oracle for FedAvg ) . The stochastic gradient estimator gi ( yki , r ) =191 ∇fi ( yki , r ) + ξki , r of the local gradient∇fi ( yki , r ) , for each round k , each local step r = 0 , . . . , R and192 all i = 1 , . . . , n , satisfies193 E [ ξki , r ] = 0 ( 9 ) and194 E [ ∥∥ξki , r∥∥2 |yki , r ] ≤M ∥∥∇fi ( yki , r ) ∥∥2 + σ2 , for some M ≥ 0 , ( 10 ) where yki,0 = x k and yki , r = y k i , r−1 − ηlgi ( yki , r ) , r = 1 , · · · , R.195 Algorithm 2 Approximate Optimal Client Sampling ( AOCS ) . 1 : Input : expected batch size m , maximum number of iteration jmax 2 : each client i computes an update Uki ( in parallel ) 3 : each client i sends the norm of its update uki = wi ∥∥Uki ∥∥ to the master ( in parallel ) 4 : master aggregates uk = ∑n i=1 u k i 5 : master broadcasts uk to all clients 6 : each client i computes pki = min { muki uk , 1 } ( in parallel ) 7 : for j = 1 , · · · , jmax do 8 : each client i sends tki = ( 1 , p k i ) to the master if p k i < 1 ; else sends t k i = ( 0 , 0 ) ( in parallel ) 9 : master aggregates ( Ik , P k ) = ∑n i=1 t k i 10 : master computes Ck = ( m−n+I k ) Pk 11 : master broadcasts Ck to all clients 12 : each client i recalibrates pki = min { Ckpki , 1 } if pki < 1 ( in parallel ) 13 : if Ck ≤ 1 then 14 : break 15 : end if 16 : end for 17 : each clients i sends its update wi pki Uki to master with probability p k i ( in parallel ) We also define two quantities , which appear in our convergence guarantees:196 Ri : = fi ( x ? ) − f ? i , rk : = xk − x ? , ( 11 ) where f ? i is the functional value of fi at its optimum . Ri represents the mismatch between the local197 and global minimizer , and rk captures the distance of the current point to the minimizer of f .198 Equipped with these assumptions , we are ready to proceed with our convergence guarantees . We start199 with the definition of the improvement factor200 αk : = E [ ∥∥∥∑i∈Sk wipki Uki −∑ni=1 wiUki ∥∥∥2 ] E [ ∥∥∥∑i∈Uk wipUi Uki −∑ni=1 wiUki ∥∥∥2 ] , ( 12 ) where Sk ∼ Sk with pki defined in ( 6 ) and Uk ∼ U is an independent uniform sampling with201 pUi = m/n . By construction , α k is less than or equal to one , as Sk minimizes the variance term . In202 addition , αk can reach zero in the case where there are at most m non-zero updates . If αk = 0,203 our method performs as if all updates were communicated . In the worst-case αk = 1 , our method204 performs as if we picked m updates uniformly at random , and one can not do better due to the205 structure of the updates Uki . In the following subsections , we analyze specific methods for solving the206 optimization problem ( 1 ) under the aforementioned assumptions . The proofs and detailed description207 are deferred to the Appendix.208 Fairness . Based on our sampling strategy , it might be tempting to assume that the obtained solution209 could exhibit fairness issues . In our convergence analysis , we show that this is not the case , as our210 proposed methods converge to the optimal solution . Hence , as long as the original objective has no211 inherent issue with fairness , our methods do not exhibit any fairness issues . Besides , our algorithm212 can be used in conjunction with other “ more fair ” objectives , e.g. , tilted ERM [ 19 ] .213 3.1 Distributed SGD with Optimal Client Sampling214 We begin with the convergence analysis for DSGD ( see ( 2 ) ) with optimal client sampling.215 Algorithm 3 FedAvg with Optimal Client Sampling . 1 : Input : initial global model x1 , global and local step-sizes ηkg , ηkl 2 : for each round k = 1 , . . . , K do 3 : master broadcasts xk to all clients i ∈ [ n ] 4 : for each client i ∈ [ n ] ( in parallel ) do 5 : initialize local model yki,0 ← xk 6 : for r = 1 , . . . , R do 7 : compute mini-batch gradient gi ( yki , r−1 ) 8 : update yki , r ← yki , r−1 − ηkl gi ( yki , r−1 ) 9 : end for 10 : compute Uki : = ∆y k i = x k − yki , R 11 : compute pki using Algorithm 1 or 2 12 : send wi pki ∆yki to master with probability p k i 13 : end for 14 : master computes ∆xk = ∑ i∈Sk wi pki ∆yki 15 : master updates global model xk+1 ← xk − ηkg∆xk 16 : end for Theorem 2 . Let fi be L-smooth and convex for all i = 1 , . . . , n. Let f be µ-strongly convex . Suppose216 that Assumption 1 holds . Choose ηk ∈ ( 0 , γ k ( 1+maxi∈ [ n ] { wi } M ) L ) , where217 γk : = m αk ( n−m ) +m ∈ [ m n , 1 ] , k = 0 , . . . , K − 1 . Define218 β1 : = n∑ i=1 w2i ( 2L ( 1 +M ) Ri + σ 2 ) and β2 : = 2L n∑ i=1 w2iRi . Then , the iterates of DSGD with optimal client sampling ( 6 ) satisfy219 E [ ∥∥rk+1∥∥2 ] ≤ ( 1− µηk ) E [ ∥∥rk∥∥2 ] + ( ηk ) 2 ( β1 γk − β2 ) . ( 13 ) Interpretation . In order to understand the results of Theorem 2 , we first look at the best and worst220 case scenarios . In the best case scenario , we have γk = 1 for all k. This implies that there is no221 loss of speed comparing to the method with full participation . It is indeed confirmed by our theory222 as our obtained recursion recovers the best-known rate of DSGD in the full participation regime [ 8 ] .223 Similarly , in the worst case , we have γk = m/n for all k ’ s , which corresponds to uniform sampling224 with sample size m and our recursion recovers the best-know rate for DSGD in this regime . This is225 expected as ( 12 ) implies that each update Uki is equivalent , thus we can not hope for better rate than226 the uniform sampling . In the general scenario , our obtain recursion sits somewhere between full227 and uniform partial participation , where the actual position is determined by γk which capture the228 distribution of updates ( here gradients ) on clients . For instance , with a larger number of γk ’ s tending229 to 1 , we are closer to full participation regime . Similarly , with more γk ’ s tending to m/n , we are230 closer to the rate of partial participation.231 3.2 FedAvg with Optimal Client Sampling232 One of the most common approaches to optimization for Federated Learning is Federated Averaging233 ( FedAvg ) [ 23 ] , an adaption of local-update to parallel SGD . In FedAvg , each client runs some234 number of SGD steps locally , and then local updates are averaged to form the global update which is235 then used for the global model on the master . Pseudo-code that adapts the standard FedAvg algorithm236 to our framework is given in Algorithm 3.237 Theorem 3 . Assume that fi is L-smooth and µ-strongly convex for all i = 1 , . . . , n and Assumption 2238 holds . Let ηk : = Rηkl η k g be the effective step-size and η k g ≥ √ γk∑ i w 2 i , where239 γk : = m αk ( n−m ) +m ∈ [ m n , 1 ] . If ηk ≤ 18 min { 1 L ( 2+M/R ) , γk ( 1+maxi∈ [ n ] { wi } ( 1+M/R ) ) L } , then the iterates of FedAvg ( R ≥ 2 ) with240 optimal client sampling ( 6 ) satisfy241 3 8 E [ ( f ( xk ) − f ? ) ] ≤ 1 ηk ( 1− µη k 2 ) E [ ∥∥rk∥∥2 ] − 1 ηk E [ ∥∥rk+1∥∥2 ] + ηkβk1 + ( ηk ) 2β2 , where242 βk1 : = 2σ2 γkR n∑ i=1 w2i + 4L ( M R + 1− γk ) n∑ i=1 w2iRi and β2 : = 72L 2 ( 1 + M R ) n∑ i=1 wiRi . Interpretation . Similar to DSGD , the convergence guarantees of FedAvg with optimal client sam-243 pling ( Algorithm 3 ) sits somewhere between the performances of those with full and uniform partial244 participations , where the actual position is again determined by the distribution of updates which245 directly impact αk ’ s that are linked to γk ’ s . In the edge cases , i.e . γk = 1 ( best case ) or γk = m/n246 ( worst case ) , we recover the state-of-the-art complexity guarantees provided in [ 15 ] in both regimes.247 Note that our results are slightly more general , as [ 15 ] assumes M = 0 and wi = 1/n.248 4 Experiments249 In this section , we empirically evaluate our optimal client sampling method , comparing it with 1 ) the250 baseline where participating clients are sampled uniformly from available clients in each round and251 2 ) full participation where all available clients participate . We simulate the cross-device FL setting252 and train our models using TensorFlow Federated ( TFF ) 1 . For all three methods , we report validation253 accuracy and ( local ) training loss ( vertical axis ) as a function of the number of communication254 rounds and the number of bits communicated from clients to the master ( horizontal axis ) . Each figure255 displays the mean performance with standard error over 5 independent runs . For a fair comparison,256 we use the same random seed for the three compared methods in a single run and vary random seeds257 across different runs.258 Setup . We conclude an evaluation on FedAvg where we extend the TFF implementation of FedAvg2259 to fit our framework . For the model , we use the two-layer Convolutional Neural Network ( CNN ) 260 1https : //github.com/tensorflow/federated 2https : //github.com/tensorflow/federated/tree/master/tensorflow_federated/python/ examples/simple_fedavg provided in the implementation . The default dataset is Federated EMNIST with only digits , but as this261 is a well-balanced dataset with mostly the same quality data on each client , we modify it by removing262 some clients or some of their training images , in order to better simulate conditions in which our263 proposed methods bring significant theoretical improvements . As a result , we produce 3 unbalanced264 datasets as summarized in Figure 1 , on which we train the CNN model . For validation , we use the265 unchanged validation set in the Federated EMNIST dataset , which consists of 40 , 832 validation266 images . In each communication round of FedAvg , n = 32 clients are sampled uniformly from the267 client pool , each of which then performs several SGD steps on its local training images for 1 epoch268 with batch size 20 . For partial participation , the expected number of clients allowed to communicate269 their updates back to the master is set to m = 3 for all the experiments . We use constant step sizes,270 where we set ηg = 1 and tune ηl from the set of values { 2−1 , 2−2 , 2−3 , 2−4 , 2−5 } using a holdout271 set . We implement our sampling procedure using Algorithm 2 , as this supports stateless clients and272 secure aggregation . We include extra communication costs in our results , where we set jmax = 4.273 More details of the hyper-parameters that we use can be found in the Appendix.274 Results and Discussions . As predicted by our theory , the performance of FedAvg with our proposed275 optimal client sampling strategy is in between the performances of that with full and uniform partial276 participation . Figures 2 , 3 and 4 ( red curves : optimal sampling ; blue curves : uniform sampling ; green277 curves : full participation ) show that , for all three datasets , the optimal sampling strategy performs278 slightly worse than but is still competitive with the full participation strategy in terms of the number279 of communication rounds – it almost reached the performance of full participation while only less280 than 10 % of the available clients communicate their updates back to the master . Note that the uniform281 sampling strategy performs significantly worse , which indicates that a careful choice of sampling282 probabilities can go a long way towards closing the gap between the performance of naive uniform283 sampling and full participation.284 More importantly , and this was the main motivation of our work , our optimal sampling strategy is285 significantly better than both the uniform sampling and full participation strategies when we compare286 validation accuracy as a function of the number of bits communicated from clients to the master.287 For instance , in case of Dataset 1 ( Figure 2 ) , while our optimal sampling approach reached around288 85 % validation accuracy after 26 × 108 communicated bits , neither the full nor the uniform sampling289 strategies are able to exceed 40 % validation accuracy within the same communication budget . Indeed,290 to reach the same 85 % validation accuracy , full participation approach needs to communicate more291 than 29 × 108 bits , i.e. , 8× more , and uniform sampling approach needs to communicate about the292 same number of bits as full participation or even more . The results for Datasets 2 and 3 are of a293 similar qualitative nature , showing that these conclusions are robust across the datasets considered.294 In the Appendix , we include additional figures which show the current best validation accuracy as a295 function of the number of communication rounds and the number of bits communicated from clients296 to the master.297 References298 [ 1 ] Dan Alistarh , Jerry Li , Ryota Tomioka , and Milan Vojnovic . QSGD : Randomized quantization299 for communication-optimal stochastic gradient descent . arXiv preprint arXiv:1610.02132 , 2016.300 [ 2 ] Zeyuan Allen-Zhu , Zheng Qu , Peter Richtárik , and Yang Yuan . Even faster accelerated301 coordinate descent using non-uniform sampling . In International Conference on Machine302 Learning , pages 1110–1119 , 2016.303 [ 3 ] Debraj Basu , Deepesh Data , Can Karakus , and Suhas Diggavi . Qsparse-local-SGD : Dis-304 tributed SGD with quantization , sparsification and local computations . In Advances in Neural305 Information Processing Systems , pages 14668–14679 , 2019.306 [ 4 ] Yoshua Bengio , Jérôme Louradour , Ronan Collobert , and Jason Weston . Curriculum learning.307 In Proceedings of the 26th annual international conference on machine learning , pages 41–48,308 2009.309 [ 5 ] Antoine Bordes , Seyda Ertekin , Jason Weston , and Léon Bottou . Fast kernel classifiers with310 online and active learning . Journal of Machine Learning Research , 6 ( Sep ) :1579–1619 , 2005.311 [ 6 ] Olivier Fercoq and Peter Richtárik . Accelerated , parallel , and proximal coordinate descent.312 SIAM Journal on Optimization , 25 ( 4 ) :1997–2023 , 2015.313 [ 7 ] WM Goodall . Television by pulse code modulation . Bell System Technical Journal , 30 ( 1 ) :33–49,314 1951.315 [ 8 ] Robert Mansel Gower , Nicolas Loizou , Xun Qian , Alibek Sailanbayev , Egor Shulgin , and Peter316 Richtárik . SGD : General analysis and improved rates . Proceedings of the 36th International317 Conference on Machine Learning , Long Beach , California , 2019.318 [ 9 ] Filip Hanzely and Peter Richtárik . Federated learning of a mixture of global and local models.319 arXiv:2002.05516 , 2020.320 [ 10 ] Samuel Horváth , Chen-Yu Ho , L ’ udovit Horváth , Atal Narayan Sahu , Marco Canini , and Peter321 Richtárik . Natural compression for distributed deep learning . arXiv preprint arXiv:1905.10988,322 2019.323 [ 11 ] Samuel Horváth and Peter Richtárik . Nonconvex variance reduced optimization with arbitrary324 sampling . Proceedings of the 36th International Conference on Machine Learning , 2019.325 [ 12 ] Samuel Horváth and Peter Richtárik . A better alternative to error feedback for communication-326 efficient distributed learning . arXiv preprint arXiv:2006.11077 , 2020.327 [ 13 ] Peter Kairouz , H Brendan McMahan , Brendan Avent , Aurélien Bellet , Mehdi Bennis , Ar-328 jun Nitin Bhagoji , Keith Bonawitz , Zachary Charles , Graham Cormode , Rachel Cummings,329 et al . Advances and open problems in federated learning . arXiv preprint arXiv:1912.04977,330 2019.331 [ 14 ] Hamed Karimi , Julie Nutini , and Mark Schmidt . Linear convergence of gradient and proximal-332 gradient methods under the polyak-łojasiewicz condition . In Joint European Conference on333 Machine Learning and Knowledge Discovery in Databases , pages 795–811 . Springer , 2016.334 [ 15 ] Sai Praneeth Karimireddy , Satyen Kale , Mehryar Mohri , Sashank J Reddi , Sebastian U Stich,335 and Ananda Theertha Suresh . Scaffold : Stochastic controlled averaging for on-device federated336 learning . arXiv preprint arXiv:1910.06378 , 2019.337 [ 16 ] Angelos Katharopoulos and François Fleuret . Not all samples are created equal : Deep learning338 with importance sampling . arXiv preprint arXiv:1803.00942 , 2018.339 [ 17 ] Ahmed Khaled , Konstantin Mishchenko , and Peter Richtárik . Tighter theory for local SGD on340 identical and heterogeneous data . In The 23rd International Conference on Artificial Intelligence341 and Statistics ( AISTATS 2020 ) , 2020.342 [ 18 ] Jakub Konečný and Peter Richtárik . Randomized distributed mean estimation : Accuracy vs.343 communication . Frontiers in Applied Mathematics and Statistics , 4:62 , 2018.344 [ 19 ] Tian Li , Ahmad Beirami , Maziar Sanjabi , and Virginia Smith . Tilted empirical risk minimization.345 In International Conference on Learning Representations , 2021.346 [ 20 ] Qihang Lin , Zhaosong Lu , and Lin Xiao . An accelerated proximal coordinate gradient method.347 In Advances in Neural Information Processing Systems , pages 3059–3067 , 2014.348 [ 21 ] Tao Lin , Sebastian U Stich , Kumar Kshitij Patel , and Martin Jaggi . Don ’ t use large mini-batches,349 use local SGD . arXiv preprint arXiv:1808.07217 , 2018.350 [ 22 ] Ilya Loshchilov and Frank Hutter . Online batch selection for faster training of neural networks.351 arXiv preprint arXiv:1511.06343 , 2015.352 [ 23 ] Brendan McMahan , Eider Moore , Daniel Ramage , Seth Hampson , and Blaise Aguera y Arcas.353 Communication-efficient learning of deep networks from decentralized data . In Artificial354 Intelligence and Statistics , pages 1273–1282 , 2017.355 [ 24 ] Konstantin Mishchenko , Filip Hanzely , and Peter Richtárik . 99 % of parallel optimization is356 inevitably a waste of time . arXiv preprint arXiv:1901.09437 , 2019.357 [ 25 ] Deanna Needell , Rachel Ward , and Nati Srebro . Stochastic gradient descent , weighted sampling,358 and the randomized kaczmarz algorithm . In Advances in neural information processing systems,359 pages 1017–1025 , 2014.360 [ 26 ] Yu Nesterov . Efficiency of coordinate descent methods on huge-scale optimization problems.361 SIAM Journal on Optimization , 22 ( 2 ) :341–362 , 2012.362 [ 27 ] Zheng Qu , Peter Richtárik , and Tong Zhang . Quartz : Randomized dual coordinate ascent with363 arbitrary sampling . In Advances in Neural Information Processing Systems 28 , pages 865–873,364 2015.365 [ 28 ] Ali Ramezani-Kebrya , Fartash Faghri , and Daniel M Roy . NUQSGD : Improved communication366 efficiency for data-parallel SGD via nonuniform quantization . arXiv preprint arXiv:1908.06077,367 2019.368 [ 29 ] Peter Richtárik and Martin Takáč . Iteration complexity of randomized block-coordinate descent369 methods for minimizing a composite function . Mathematical Programming , 144 ( 1-2 ) :1–38,370 2014.371 [ 30 ] Lawrence Roberts . Picture coding using pseudo-random noise . IRE Transactions on Information372 Theory , 8 ( 2 ) :145–154 , 1962.373 [ 31 ] Tom Schaul , John Quan , Ioannis Antonoglou , and David Silver . Prioritized experience replay.374 arXiv preprint arXiv:1511.05952 , 2015.375 [ 32 ] Florian Schroff , Dmitry Kalenichenko , and James Philbin . Facenet : A unified embedding for376 face recognition and clustering . In Proceedings of the IEEE conference on computer vision and377 pattern recognition , pages 815–823 , 2015.378 [ 33 ] Shai Shalev-Shwartz and Tong Zhang . Accelerated proximal stochastic dual coordinate ascent379 for regularized loss minimization . In International conference on machine learning , pages380 64–72 , 2014.381 [ 34 ] Edgar Simo-Serra , Eduard Trulls , Luis Ferraz , Iasonas Kokkinos , Pascal Fua , and Francesc382 Moreno-Noguer . Discriminative learning of deep convolutional feature point descriptors . In383 Proceedings of the IEEE International Conference on Computer Vision , pages 118–126 , 2015.384 [ 35 ] Sebastian U Stich . Local SGD converges fast and communicates little . ICLR 2019 - International385 Conference on Learning Representations , 2019.386 [ 36 ] Sebastian U Stich , Jean-Baptiste Cordonnier , and Martin Jaggi . Sparsified SGD with memory.387 In Advances in Neural Information Processing Systems , pages 4447–4458 , 2018.388 [ 37 ] Sebastian U Stich and Sai Praneeth Karimireddy . The error-feedback framework : Better rates389 for SGD with delayed gradients and compressed communication . ICLR 2020 - International390 Conference on Learning Representations , 2020.391 [ 38 ] Sebastian U Stich , Anant Raj , and Martin Jaggi . Safe adaptive importance sampling . In392 Advances in Neural Information Processing Systems , pages 4381–4391 , 2017.393 [ 39 ] Thijs Vogels , Sai Praneeth Karimireddy , and Martin Jaggi . PowerSGD : Practical low-rank gra-394 dient compression for distributed optimization . In Advances in Neural Information Processing395 Systems , pages 14236–14245 , 2019.396 [ 40 ] Jianqiao Wangni , Jialei Wang , Ji Liu , and Tong Zhang . Gradient sparsification for397 communication-efficient distributed optimization . In Advances in Neural Information Pro-398 cessing Systems , pages 1299–1309 , 2018.399 [ 41 ] Wei Wen , Cong Xu , Feng Yan , Chunpeng Wu , Yandan Wang , Yiran Chen , and Hai Li . Terngrad:400 Ternary gradients to reduce communication in distributed deep learning . In Advances in Neural401 Information Processing Systems , pages 1509–1519 , 2017.402 [ 42 ] Hantian Zhang , Jerry Li , Kaan Kara , Dan Alistarh , Ji Liu , and Ce Zhang . Zipml : Training linear403 models with end-to-end low precision , and a little bit of deep learning . In Proceedings of the404 34th International Conference on Machine Learning-Volume 70 , pages 4035–4043 . JMLR . org,405 2017.406 [ 43 ] Peilin Zhao and Tong Zhang . Stochastic optimization with importance sampling for regularized407 loss minimization . In international conference on machine learning , pages 1–9 , 2015.408 Checklist409 1 . For all authors ... 410 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s411 contributions and scope ? [ Yes ] 412 ( b ) Did you describe the limitations of your work ? [ Yes ] See Remark 1 , where we413 acknowledge that our algorithm requires two ( although the first one is negligible ) 414 communication rounds per iteration.415 ( c ) Did you discuss any potential negative societal impacts of your work ? [ Yes ] We416 discussed potential fairness issues in Section 3.417 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to418 them ? [ Yes ] 419 2 . If you are including theoretical results ... 420 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] See Section 3421 and Appendix A.422 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] See Appendix C-F for423 complete proofs . We also provided interpretations of our theorems in the main paper424 and discussed the relationship between our results and related results in the literature.425 3 . If you ran experiments ... 426 ( a ) Did you include the code , data , and instructions needed to reproduce the main experi-427 mental results ( either in the supplemental material or as a URL ) ? [ Yes ] See supplemen-428 tal material.429 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they430 were chosen ) ? [ Yes ] See Section 4 and Appendix B431 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-432 ments multiple times ) ? [ Yes ] We ran every experiment 5 times with different random433 seeds and reported results with error bars.434 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type435 of GPUs , internal cluster , or cloud provider ) ? [ N/A ] Since we only run simulations,436 this is not applicable.437 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 438 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] See Section 4439 ( b ) Did you mention the license of the assets ? [ Yes ] All the data and assets we used in this440 manuscript are open-source.441 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 442 We included an anonymized URL in the supplemental material for the datasets used.443 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re444 using/curating ? [ N/A ] 445 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable446 information or offensive content ? [ N/A ] 447 5 . If you used crowdsourcing or conducted research with human subjects ... 448 ( a ) Did you include the full text of instructions given to participants and screenshots , if449 applicable ? [ N/A ] 450 ( b ) Did you describe any potential participant risks , with links to Institutional Review451 Board ( IRB ) approvals , if applicable ? [ N/A ] 452 ( c ) Did you include the estimated hourly wage paid to participants and the total amount453 spent on participant compensation ? [ N/A ] 454

Summary: This paper proposes a strategy to _optimally_ sample clients on a given FL global round. The sampling strategy selects more informative updates by their magnitude. This criterion minimizes the variance of the sampling scheme.   Experiments show this method outperforms uniform sampling.