On the Impact of Client Sampling on Federated Learning Convergence

1 INTRODUCTION . Federated Learning ( FL ) has gained popularity in the last years as it enables different clients to jointly learn a global model without sharing their respective data . Among the different FL approaches , federated averaging ( FedAvg ) has emerged as the most popular optimization scheme ( McMahan et al. , 2017 ) . An optimization round of FedAvg requires data owners , also called clients , to receive from the server the current global model which they update on a fixed amount of Stochastic Gradient Descent ( SGD ) steps before sending it back to the server . The new global model is then created as the weighted average of the client updates , according to their data ratio . FedAvg was first proven to converge experimentally ( McMahan et al. , 2017 ) , before theoretical guarantees were provided for any non-iid federated dataset ( Wang et al. , 2020a ; Karimireddy et al. , 2020 ; Haddadpour & Mahdavi , 2019 ; Khaled et al. , 2020 ) . A drawback of naive implementations of FedAvg consists in requiring the participation of all the clients at every optimization round . As a consequence , the efficiency of the optimization is limited by the communication speed of the slowest client , as well as by the server communication capabilities . To mitigate this issue , the original FedAvg algorithm already contemplated the possibility of considering a random subset of m clients at each FL round . It has been subsequently shown that , to ensure the convergence of FL to its optimum , clients must be sampled such that in expectation the resulting global model is identical to the one obtained when considering all the clients ( Wang et al. , 2020a ; Cho et al. , 2020 ) . Clients sampling schemes compliant with this requirement are thus called unbiased . Due to its simplicity and flexibility , the current default unbiased sampling scheme consists in sampling m clients according to a Multinomial Distribution ( MD ) , where the sampling probability depends on the respective data ratio ( Li et al. , 2020a ; Wang et al. , 2020a ; Li et al. , 2020c ; Haddadpour & Mahdavi , 2019 ; Li et al. , 2020b ; Wang & Joshi , 2018 ; Fraboni et al. , 2021 ) . Nevertheless , when clients have identical amount of data , clients can also be sampled uniformly without replacement ( Li et al. , 2020c ; Karimireddy et al. , 2020 ; Reddi et al. , 2021 ; Rizk et al. , 2020 ) . In this case , Uniform sampling has been experimentally shown to yield better results than MD sampling ( Li et al. , 2020c ) . In spite of the practical usage of sampling in FL , the impact of different unbiased client sampling schemes on the convergence and speed of FL remains to date under-investigated . In ( Fraboni et al. , 2021 ) , MD sampling was extended to account for collections of sampling distributions with varying client sampling probability , to define clustered sampling . From a theoretical perspective , this approach was proven to have identical convergence guarantees of MD sampling , with albeit experimental improvement justified by lower variance of the clients ’ aggregation weights . Another study investigated the convergence guarantees of FedAvg under different sampling schemes , through a penalized optimization framework ( Li et al. , 2020c ) . These studies reflect the ongoing interest and need for a theoretical framework to elucidate the impact of client sampling on FL convergence . The main contribution of this work consists in deriving a decomposition theorem for the convergence of FL , allowing to clearly quantify the impact of client sampling on the global model update at any FL round . This contribution has important theoretical and practical implications . First , we demonstrate the dependence of FL convergence on the variance of the aggregation weights . Second , we prove for the first time that the convergence speed is also impacted through sampling by the resulting covariance between aggregation weights . Contrarily to the convergence bound illustrated in ( Li et al. , 2020c ) , our theorem provides the exact decomposition of a given convergence step , thus enabling accurate quantification of the role of client sampling and heterogeneity in FL . From a practical point of view , we establish both theoretically and experimentally that client sampling schemes based on aggregation weights with sum different than 1 are less efficient . We also prove that MD sampling is outperformed by Uniform sampling only when clients have identical data ratio . Finally , we show that the comparison between different client sampling schemes is appropriate only when considering a small number of clients . Our theory ultimately shows that MD sampling should be used as default sampling scheme , due to the favorable statistical properties and to the resilience to FL applications with varying data ratio and heterogeneity . Our work is structured as follows . In Section 2 , we provide formal definitions for FL , unbiased client sampling , and for the server aggregation scheme . In Section 3 , we introduce our decomposition theorem ( Theorem 1 ) relating the convergence of FL to the aggregation weight variance of the client sampling scheme . Consistently with our theory , in Section 4 , we experimentally demonstrate the importance of the clients aggregation weights variance and covariance on the convergence speed , and conclude by recommending Uniform sampling for FL applications with identical client ratio , and MD sampling otherwise . 2 BACKGROUND . Before investigating in Section 3 the impact of client sampling on FL convergence , we recapitulate in Section 2 the current theory behind FL aggregation schemes for clients local updates . We then introduce a formalization for unbiased client sampling . 2.1 AGGREGATING CLIENTS LOCAL UPDATES . In FL , we consider a set I of n clients each respectively owning a dataset Di composed of ni samples . FL aims at optimizing the average of each clients local loss function weighted by pi such that ∑n i=1 pi = 1 , i.e . L ( θ ) = n∑ i=1 piLi ( θ ) , ( 1 ) where θ represents the model parameters . The weight pi can be interpreted as the importance given by the server to client i in the federated optimization problem . While any combination of { pi } is possible , we note that in practice , either ( a ) every device has equal importance , i.e . pi = 1/n , or ( b ) every data point is equally important , i.e . pi = ni/M with M = ∑n i=1 ni . Unless stated otherwise , in the rest of this work , we consider to be in case ( b ) , i.e . ∃i , pi ̸= 1/n . In this setting , to estimate a global model across clients , FedAvg ( McMahan et al. , 2017 ) is an iterative training strategy based on the aggregation of local model parameters . At each iteration step t , the server sends the current global model parameters θt to the clients . Each client updates the respective model by minimizing the local cost function Li ( θ ) through a fixed amount K of SGD steps initialized with θt . Subsequently each client returns the updated local parameters θt+1i to the server . The global model parameters θt+1 at the iteration step t+1 are then estimated as a weighted average : θt+1 = n∑ i=1 piθ t+1 i . ( 2 ) To alleviate the clients workload and reduce the amount of overall communications , the server often considers m ≤ n clients at every iteration . In heterogeneous datasets containing many workers , the percentage of sampled clients m/n can be small , and thus induce important variability in the new global model , as each FL optimization step necessarily leads to an improvement on the m sampled clients to the detriment of the non-sampled ones . To solve this issue , Reddi et al . ( 2021 ) ; Karimireddy et al . ( 2020 ) ; Wang et al . ( 2020b ) propose considering an additional learning rate ηg to better account for the clients update at a given iteration . We denote by ωi ( St ) the stochastic aggregation weight of client i given the subset of sampled clients St at iteration t . The server aggregation scheme can be written as : θt+1 = θt + ηg n∑ i=1 ωi ( St ) ( θ t+1 i − θ t ) . ( 3 ) 2.2 UNBIASED CLIENTS SAMPLING . While FedAvg was originally based on the uniform sampling of clients ( McMahan et al. , 2017 ) , this scheme has been proven to be biased and converge to a suboptimal minima of problem ( 1 ) ( Wang et al. , 2020a ; Cho et al. , 2020 ; Li et al. , 2020c ) . This was the motivation for Li et al . ( 2020c ) to introduce the notion of unbiasedness , where clients are considered in expectation subject to their importance pi , according to Definition 1 below . Unbiased sampling guarantees the optimization of the original FL cost function , while minimizing the number of active clients per FL round . We note that unbiased sampling is not necessarily related to the clients distribution , as this would require to know beforehand the specificity of the clients ’ datasets . Unbiased sampling methods ( Li et al. , 2020a ; c ; Fraboni et al. , 2021 ) are currently among the standard approaches to FL , as opposed to biased approaches , known to over- or under-represent clients and lead to suboptimal convergence properties ( McMahan et al. , 2017 ; Nishio & Yonetani , 2019 ; Jeon et al. , 2020 ; Cho et al. , 2020 ) , or to methods requiring additional computation work from clients ( Chen et al. , 2020a ) . Definition 1 ( Unbiased Sampling ) . A client sampling scheme is said unbiased if the expected value of the client aggregation is equal to the global deterministic aggregation obtained when considering all the clients , i.e . ESt [ n∑ i=1 wi ( St ) θ t i ] : = n∑ i=1 piθ t i , ( 4 ) where wj ( St ) is the aggregation weight of client j for subset of clients St . The sampling distribution uniquely defines the statistical properties of stochastic weights . In this setting , unbiased sampling guarantees the equivalence between deterministic and stochastic weights in expectation . Unbiased schemes of primary importance in FL are MD and Uniform sampling , for which we can derive a close form formula for the aggregation weights : MD sampling . This scheme considers l1 , ... , lm to be the m iid sampled clients from a Multinomial Distribution with support on { 1 , ... , m } satisfying P ( lk = i ) = pi ( Wang et al. , 2020a ; Li et al. , 2020a ; c ; Haddadpour & Mahdavi , 2019 ; Li et al. , 2020b ; Wang & Joshi , 2018 ; Fraboni et al. , 2021 ) . By definition , we have ∑n i=1 pi = 1 , and the clients aggregation weights take the form : ωi ( St ) = 1 m m∑ k=1 I ( lk = i ) . ( 5 ) Uniform sampling . This scheme samples m clients uniformly without replacement . Since in this case a client is sampled with probability p ( { i ∈ St } ) = m/n , the requirement of Definition 1 implies : ωi ( St ) = I ( i ∈ St ) n m pi . ( 6 ) Table 1 : Synthesis of statistical properties of different sampling schemes . Sampling Var [ ωi ( St ) ] α Var [ ∑n i=1 ωi ( St ) ] Full participation = 0 = 0 = 0 MD = − 1mp 2 i + 1 mpi = 1/m = 0 Uniform = ( n m − 1 ) p2i = n−m m ( n−1 ) = n−m m ( n−1 ) [ n ∑n i=1 p 2 i − 1 ] We note that this formulation for Uniform sampling is a generalization of the scheme previously used for FL applications with identical client importance , i.e . pi = 1/n ( Karimireddy et al. , 2020 ; Li et al. , 2020c ; Reddi et al. , 2021 ; Rizk et al. , 2020 ) . We note that Var [ ∑n i=1 ωi ( St ) ] = 0 if and only if pi = 1/n for all the clients as , indeed , ∑n i=1 ωi ( St ) = m n m 1 n = 1 With reference to equation ( 3 ) , we note that by setting ηg = 1 , and by imposing the condition ∀St , ∑n i=1 ωi ( St ) = 1 , we retrieve equation ( 2 ) . This condition is satisfied for example by MD sampling and Uniform sampling for identical clients importance . We finally note that the covariance of the aggregation weights for both MD and Uniform sampling is such that there exists an α such that ∀i ̸= j , Cov [ ωi ( St ) , ωj ( St ) ] = −αpipj . We provide in Table 1 the derivation of α and the resulting covariance for these two schemes with calculus details in Appendix A . Furthermore , this property is common to a variety of sampling schemes , for example based on Binomial or Poisson Binomial distributions ( detailed derivations can be found in Appendix A ) . Following this consideration , in addition to Definition 1 , in the rest of this work we assume the additional requirement for a client sampling scheme to satisfy ∃α ≥ 0 , ∀i ̸= j , Cov [ ωi ( St ) , ωj ( St ) ] = −αpipj .

Summary: The paper studies a framework for studying the impact of client selection strategies for unbiased aggregation at the server for federated learning. The framework separates the impact of client selection, client model drift, global model drift thereby providing new insights regarding the impact of client selection. In particular, the paper proposed a decomposition theorem which quantifies  the dependence on client sampling. Finally, experimental results on Shakespeare back the theory and demonstrate the effectiveness of one sampling choice over another based on the loss function setup in a FL framework.