Client Selection in Federated Learning: Convergence Analysis and Power-of-Choice Selection Strategies

1 INTRODUCTION . Until recently , machine learning models were largely trained in the data center setting ( Dean et al. , 2012 ) using powerful computing nodes , fast inter-node communication links , and large centrally available training datasets . The future of machine learning lies in moving both data collection as well as model training to the edge . The emerging paradigm of federated learning ( McMahan et al. , 2017 ; Kairouz et al. , 2019 ; Bonawitz et al. , 2019 ) considers a large number of resource-constrained mobile devices that collect training data from their environment . Due to limited communication capabilities and privacy concerns , these data can not be directly sent over to the cloud . Instead , the nodes locally perform a few iterations of training using local-update stochastic gradient descent ( SGD ) ( Yu et al. , 2018 ; Stich , 2018 ; Wang & Joshi , 2018 ; 2019 ) , and only send model updates periodically to the aggregating cloud server . Besides communication limitations , the key scalability challenge faced by the federated learning framework is that the client nodes can have highly heterogeneous local datasets and computation speeds . The effect of data heterogeneity on the convergence of local-update SGD is analyzed in several recent works ( Reddi et al. , 2020 ; Haddadpour & Mahdavi , 2019 ; Khaled et al. , 2020 ; Stich & Karimireddy , 2019 ; Woodworth et al. , 2020 ; Koloskova et al. , 2020 ; Huo et al. , 2020 ; Zhang et al. , 2020 ; Pathak & Wainwright , 2020 ; Malinovsky et al. , 2020 ; Sahu et al. , 2019 ) and methods to overcome the adverse effects of data and computational heterogeneity are proposed in ( Sahu et al. , 2019 ; Wang et al. , 2020 ; Karimireddy et al. , 2019 ) , among others . Partial Client Participation . Most of the recent works described above assume full client participation , that is , all nodes participate in every training round . In practice , only a small fraction of client nodes participate in each training round , which can exacerbate the adverse effects of data heterogeneity . While some existing convergence guarantees for full client participation and methods to tackle heterogeneity can be generalized to partial client participation ( Li et al. , 2020 ) , these generalizations are limited to unbiased client participation , where each client ’ s contribution to the expected global objective optimized in each round is proportional to its dataset size . In Ruan et al . ( 2020 ) , the authors analyze the convergence with flexible device participation , where devices can freely join or leave the training process or send incomplete updates to the server . However , adaptive client selection that is cognizant of the training progress at each client has not been understood yet . It is important to analyze and understand biased client selection strategies since they can sharply accelerate error convergence , and hence boost communication efficiency in heterogeneous environments by preferentially selecting clients with higher local loss values , as we show in our paper . This idea has been explored in recent empirical studies ( Goetz et al. , 2019 ; Laguel et al. , 2020 ; Ribero & Vikalo , 2020 ) . Nishio & Yonetani ( 2019 ) proposed grouping clients based on hardware and wireless resources in order to save communication resources . Goetz et al . ( 2019 ) ( which we include as a benchmark in our experiments ) proposed client selection with local loss , and Ribero & Vikalo ( 2020 ) proposed utilizing the progression of clients ’ weights . But these schemes are limited to empirical demonstration without a rigorous analysis of how selection skew affects convergence speed . Another relevant line of work ( Jiang et al. , 2019 ; Katharopoulos & Fleuret , 2018 ; Shah et al. , 2020 ; Salehi et al. , 2018 ) employs biased selection or importance sampling of data to speed-up convergence of classic centralized SGD – they propose preferentially selecting samples with highest loss or highest gradient norm to perform the next SGD iteration . In contrast , Shah et al . ( 2020 ) proposes biased selection of lower loss samples to improve robustness to outliers . Generalizing such strategies to the federated learning setting is a non-trivial and open problem because of the large-scale distributed and heterogeneous nature of the training data . Our Contributions . In this paper , we present the first ( to the best of our knowledge ) convergence analysis of federated learning with biased client selection that is cognizant of the training progress at each client . We discover that biasing the client selection towards clients with higher local losses increases the rate of convergence compared to unbiased client selection . Using this insight , we propose the POWER-OF-CHOICE client selection strategy and show by extensive experiments that POWER-OF-CHOICE yields up to 3× faster convergence with 10 % higher test performance than the standard federated averaging with random selection . POWER-OF-CHOICE is designed to incur minimal communication and computation overhead , enhancing resource efficiency in federated learning . In fact , we show that even with 3× less clients participating in each round as compared to random selection , POWER-OF-CHOICE gives 2× faster convergence and 5 % higher test accuracy . 2 PROBLEM FORMULATION . Consider a cross-device federated learning setup with total K clients , where client k has a local dataset Bk consisting |Bk| = Dk data samples . The clients are connected via a central aggregating server , and seek to collectively find the model parameter w that minimizes the empirical risk : F ( w ) = 1∑K k=1Dk K∑ k=1 ∑ ξ∈Bk f ( w , ξ ) = K∑ k=1 pkFk ( w ) ( 1 ) where f ( w , ξ ) is the composite loss function for sample ξ and parameter vector w. The term pk = Dk/ ∑K k=1Dk is the fraction of data at the k-th client , and Fk ( w ) = 1 |Bk| ∑ ξ∈Bk f ( w , ξ ) is the local objective function of client k. In federated learning , the vectors w∗ , and w∗k for k = 1 , . . . , K that minimize F ( w ) and Fk ( w ) respectively can be very different from each other . We define F ∗ = minw F ( w ) = F ( w∗ ) and F ∗k = minw Fk ( w ) = Fk ( w ∗ k ) . Federated Averaging with Partial Client Participation . The most common algorithm to solve ( 1 ) is federated averaging ( FedAvg ) proposed in McMahan et al . ( 2017 ) . The algorithm divides the training into communication rounds . At each round , to save communication cost at the central server , the global server only selects a fraction C of m = CK clients to participate in the training . Each selected/active client performs τ iterations of local SGD ( Stich , 2018 ; Wang & Joshi , 2018 ; Yu et al. , 2018 ) and sends its locally updated model back to the server . Then , the server updates the global model using the local models and broadcasts the global model to a new set of active clients . Formally , we index the local SGD iterations with t ≥ 0 . The set of active clients at iteration t is denoted by S ( t ) . Since active clients performs τ steps of local update , the active set S ( t ) also remains constant for every τ iterations . That is , if ( t+ 1 ) mod τ = 0 , then S ( t+1 ) = S ( t+2 ) = · · · = S ( t+τ ) . Accordingly , the update rule of FedAvg can be written as follows : w ( t+1 ) k = { w ( t ) k − ηtgk ( w ( t ) k , ξ ( t ) k ) for ( t+ 1 ) mod τ 6= 0 1 m ∑ j∈S ( t ) ( w ( t ) j − ηtgj ( w ( t ) j , ξ ( t ) j ) ) , w ( t+1 ) for ( t+ 1 ) mod τ = 0 ( 2 ) where w ( t+1 ) k denotes the local model parameters of client k at iteration t , ηt is the learning rate , and gk ( w ( t ) k , ξ ( t ) k ) = 1 b ∑ ξ∈ξ ( t ) k ∇f ( w ( t ) k , ξ ) is the stochastic gradient over mini-batch ξ ( t ) k of size b that is randomly sampled from client k ’ s local dataset Bk . Moreover , w ( t+1 ) denotes the global model at server . Although w ( t ) is only updated after every τ iterations , for the purpose of convergence analysis we consider a virtual sequence of w ( t ) that is updated at each iteration as follows w ( t+1 ) = w ( t ) − ηtg ( t ) = w ( t ) − ηt 1 m ∑ k∈S ( t ) gk ( w ( t ) k , ξ ( t ) k ) ( 3 ) with g ( t ) = 1m ∑ k∈S ( t ) gk ( w ( t ) k , ξ ( t ) k ) . Note that in ( 2 ) and ( 3 ) we do not weight the client models by their dataset fractions pk because pk is considered in the client selection scheme used to decide the set S ( t ) . Our convergence analysis can be generalized to when the global model is a weighted average instead of a simple average of client models , and we show in Appendix E that our convergence analysis also covers the sampling uniformly at random without replacement scheme proposed by Li et al . ( 2020 ) . The set S ( t ) can be sampled either with or without replacement . For sampling with replacement , we assume that multiple copies of the same client in the set S ( t ) behave as different clients , that is , they perform local updates independently . Client Selection Strategy . To guarantee FedAvg converges to the stationary points of the objective function ( 1 ) , most current analysis frameworks ( Li et al. , 2020 ; Karimireddy et al. , 2019 ; Wang et al. , 2020 ) consider a strategy that selects the set S ( t ) by sampling m clients at random ( with replacement ) such that client k is selected with probability pk , the fraction of data at that client . This sampling scheme is unbiased since it ensures that in expectation , the update rule ( 3 ) is the same as full client participation . Hence , it enjoys the same convergence properties as local-update SGD methods ( Stich , 2018 ; Wang & Joshi , 2018 ) . We denote this unbiased random client selection strategy as πrand . In this paper , we consider a class of biased client selection strategies that is cognizant of the global training progress which ( to the best of our knowledge ) has not been worked on before . Note that for any aggregation scheme and sampling scheme with partial client participation , if the expectation over the sampling scheme for the update rule of the global model is equal to the case of the update rule for full client participation , we distinguish this as an unbiased client participation scheme . For example in Horváth & Richtárik ( 2020 ) , even with a biased sampling scheme , with the normalizing aggregation the update rule is unbiased . Henceforth , we state that our paper encompasses both biased and unbiased update rules . In the two-client example in Figure 1 , we set S ( t+1 ) = arg maxk∈ [ K ] Fk ( w ( t ) ) , a single client with the highest local loss at the current global model . In this toy example , the selection strategy can not guarantee the updates ( 3 ) equals to the full client participation case in expectation . Nevertheless , it gives faster convergence to the global minimum than the random one . Motivated by this observation , we define a client selection strategy π as a function that maps the current global model w to a selected set of clients S ( π , w ) .

Summary: This work investigates federated optimization considering data heterogeneity, communication and computation limitations, and partial client participation. In contrast to past works, this paper focuses on deeper understanding of the effect of partial client participation on the convergence rate by considering biased client participation. The paper provides convergence analysis for any biased selection strategy, showing that the rate is composed of vanishing error term and non-vanishing bias term. The obtained rates explicitly show the effect of client selection strategy and the trade-off between convergence speed and the solution bias. Then it proposes a parametric family of biased selection strategy, called power-of-choice, which aims to speed up the convergence of the error term at the cost of possibly bigger bias term. Experiments are provided to highlight the benefits of the proposed pow-d strategy over the standard unbiased selection strategies.