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

Introduction

Until recently, machine learning models were largely trained in the data center setting 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 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 cannot be directly sent over to the cloud. Instead, the nodes locally perform a few iterations of training using local-update stochastic gradient descent (SGD) , 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 and methods to overcome the adverse effects of data and computational heterogeneity are proposed in , 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 , 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 , 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 because 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 this paper. This idea has been explored in a couple of recent empirical studies . proposed grouping clients based on hardware and wireless resources in order to save communication resources. (which we include as a benchmark in our experiments) proposed client selection with local loss, and proposed utilizing the progression of clients’ weights. But these schemes are limited to empirical demonstration without a rigorous analysis of how selection bias affects convergence speed.

Another relevant line of work 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, 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\times$ 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\times$ less clients participating in each round as compared to random selection, Power-of-choice gives $2\times$ faster convergence and 5% higher test accuracy.

Problem Formulation

Consider a cross-device federated learning setup with total $K$ clients, where client $k$ has a local dataset $\mathcal{B}_{k}$ consisting $|\mathcal{B}_{k}|=D_{k}$ data samples. The clients are connected via a central aggregating server, and seek to collectively find the model parameter $\mathbf{w}$ that minimizes the empirical risk:

where $f(\mathbf{w},\xi)$ is the composite loss function for sample $\xi$ and parameter vector $\mathbf{w}$. The term $p_{k}=D_{k}/{\sum_{k=1}^{K}D_{k}}$ is the fraction of data at the $k$-th client, and $F_{k}(\mathbf{w})=\frac{1}{|\mathcal{B}_{k}|}\sum_{\xi\in\mathcal{B}_{k}}f(\mathbf{w},\xi)$ is the local objective function of client $k$. In federated learning, the vectors $\mathbf{w}^{*}$, and $\mathbf{w}_{k}^{*}$ for $k=1,\dots,K$ that minimize $F(\mathbf{w})$ and $F_{k}(\mathbf{w})$ respectively can be very different from each other. We define $F^{*}=\min_{\mathbf{w}}F(\mathbf{w})=F(\mathbf{w}^{*})$ and $F_{k}^{*}=\min_{\mathbf{w}}F_{k}(\mathbf{w})=F_{k}(\mathbf{w}_{k}^{*})$.

Federated Averaging with Partial Client Participation. The most common algorithm to solve (1) is federated averaging (FedAvg) proposed in . 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 $\tau$ iterations of local SGD 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.

where $\mathbf{w}_{k}^{(t+1)}$ denotes the local model parameters of client $k$ at iteration $t$, $\eta_{t}$ is the learning rate, and $g_{k}(\mathbf{w}_{k}^{(t)},\xi_{k}^{(t)})=\frac{1}{b}\sum_{\xi\in\xi_{k}^{(t)}}\nabla f(\mathbf{w}_{k}^{(t)},\xi)$ is the stochastic gradient over mini-batch $\xi_{k}^{(t)}$ of size $b$ that is randomly sampled from client $k$’s local dataset $\mathcal{B}_{k}$. Moreover, $\overline{\mathbf{w}}^{(t+1)}$ denotes the global model at server. Although $\overline{\mathbf{w}}^{(t)}$ is only updated after every $\tau$ iterations, for the purpose of convergence analysis we consider a virtual sequence of $\overline{\mathbf{w}}^{(t)}$ that is updated at each iteration as follows:

with $\overline{\mathbf{g}}^{(t)}=\frac{1}{m}\sum_{k\in\mathcal{S}^{(t)}}g_{k}(\mathbf{w}_{k}^{(t)},\xi_{k}^{(t)})$. Note that in 2 and 3 we do not weight the client models by their dataset fractions $p_{k}$ because $p_{k}$ is considered in the client selection scheme used to decide the set $\mathcal{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 . The set $\mathcal{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 $\mathcal{S}^{(t)}$ behave as different clients, that is, they perform local updates independently.

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. For example, in the two-client example in Figure 1, we set $\mathcal{S}^{(t+1)}=\arg\max_{k\in[K]}F_{k}(\overline{\mathbf{w}}^{(t)})$, a single client with the highest local loss at the current global model. In this toy example, the selection strategy cannot 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 $\pi$ as a function that maps the current global model $\mathbf{w}$ to a selected set of clients $\mathcal{S}(\pi,\mathbf{w})$.

Convergence Analysis

In this section we analyze the convergence of federated averaging with partial device participation for any client selection strategy $\pi$ as defined above. This analysis reveals that biased client selection can give faster convergence, albeit at the risk of having a non-vanishing gap between the true optimum $\mathbf{w}^{*}=\arg\min F(\mathbf{w})$ and $\lim_{t\rightarrow\infty}\overline{\mathbf{w}}^{(t)}$. We use this insight in Section 4 to design client selection strategies that strike a balance between convergence speed and bias.

First we introduce the assumptions and definitions utilized for our convergence analysis.

$F_{1},~{}...,~{}F_{k}$ are all $L-$smooth, i.e., for all $\mathbf{v}$ and $\mathbf{w}$, $F_{k}(\mathbf{v})\leq F_{k}(\mathbf{w})+(\mathbf{v}-\mathbf{w})^{T}\nabla F_{k}(\mathbf{w})+\frac{L}{2}\|\mathbf{v}-\mathbf{w}\|_{2}^{2}$.

$F_{1},~{}...,~{}F_{k}$ are all $\mu-$strongly convex, i.e., for all $\mathbf{v}$ and $\mathbf{w}$, $F_{k}(\mathbf{v})\geq F_{k}(\mathbf{w})+(\mathbf{v}-\mathbf{w})^{T}\nabla F_{k}(\mathbf{w})+\frac{\mu}{2}\|\mathbf{v}-\mathbf{w}\|_{2}^{2}$.

Next, we introduce two metrics, the local-global objective gap and the selection skew, which feature prominently in the convergence analysis presented in Theorem 3.1.

For the global optimum $\mathbf{w}^{*}=\arg\min_{\mathbf{w}}F(\mathbf{w})$ and local optimum $\mathbf{w}_{k}^{*}=\arg\min_{\mathbf{w}}F_{k}(\mathbf{w})$ we define the local-global objective gap as

Note that $\Gamma$ is an inherent property of the local and global objective functions, and it is independent of the client selection strategy. A larger $\Gamma$ implies higher data heterogeneity. If $\Gamma=0$ then it implies that the local and global optimal values are consistent, and there is no solution bias due to the client selection strategy (see Theorem 3.1). Next, we define another metric called selection skew, which captures the effect of the client selection strategy on the local-global objective gap.

For any ${k}\in\mathcal{S}(\pi,\mathbf{w})$ we define,

Since $\rho(\mathcal{S}(\pi,\mathbf{w}),\mathbf{w}^{\prime})$ is a function of versions of the global model $\mathbf{w}$ and $\mathbf{w}^{\prime}$, which change during training, we define two related metrics that are independent of $\mathbf{w}$ and $\mathbf{w}^{\prime}$. These metrics enable us to obtain a conservative error bound in the convergence analysis.

where $\mathbf{w}^{*}=\arg\min_{\mathbf{w}}F(\mathbf{w})$. From (6), we have $\overline{\rho}\leq\widetilde{\rho}$ for any client selection strategy $\pi$.

Effect of the Client Selection Strategy on $\overline{\rho}$ and $\widetilde{\rho}$. For the unbiased client selection strategy $\pi_{\text{rand}}$ we have $\rho(\mathcal{S}(\pi_{\text{rand}},\mathbf{w}),\mathbf{w}^{\prime})=1$ for all $\mathbf{w}$ and $\mathbf{w}^{\prime}$ since the numerator and denominator of (5) become equal, and $\overline{\rho}=\widetilde{\rho}=1$. For a client selection strategy $\pi$ that chooses clients with higher $F_{k}(\mathbf{w})$ more often, $\overline{\rho}$ and $\widetilde{\rho}$ will be larger (and $\geq 1$). In the convergence analysis we show that a larger $\overline{\rho}$ implies faster convergence, albeit with a potential error gap, which is proportional to $(\widetilde{\rho}/\overline{\rho}-1)$. Motivated by this, in Section 4 we present an adaptive client selection strategy that prefers selecting clients with higher loss $F_{k}(\mathbf{w})$ and achieves faster convergence speed with low solution bias.

2 Main Convergence Result

Here, we present the convergence results for any client selection strategy $\pi$ for federated averaging with partial device participation in terms of local-global objective gap $\Gamma$, and selection skew $\overline{\rho},\widetilde{\rho}$.

Under Assumptions 3.1 to 3.4, for learning rate $\eta_{t}=\frac{1}{\mu(t+\gamma)}$ with $\gamma=\frac{4L}{\mu}$, and any client selection strategy $\pi$, the error after $T$ iterations of federated averaging with partial device participation satisfies

To the best of our knowledge, Theorem 3.1 provides the first convergence analysis of federated averaging with a biased client selection strategy $\pi$. We also show the results for fixed learning rate in Appendix A. The proof is presented in Appendix C. In the following paragraphs, we discuss the effects of the two terms in (7) in detail.

Large $\overline{\rho}$ and Faster Convergence. A key insight from Theorem 3.1 is that a larger selection skew $\overline{\rho}$ results in faster convergence at the rate $\mathcal{O}(\frac{1}{T\overline{\rho}})$. Note that since we obtain $\overline{\rho}$ (defined in (6)) by taking a minimum of the selection skew $\rho(\mathcal{S}(\pi,\mathbf{w}),\mathbf{w}^{\prime})$ over $\mathbf{w},\mathbf{w}^{\prime}$, this is a conservative bound on the true convergence rate. In practice, since the selection skew $\rho(\mathcal{S}(\pi,\mathbf{w}),\mathbf{w}^{\prime})$ changes during training depending on the current global model $\mathbf{w}$ and the local models $\mathbf{w}^{\prime}$, the true convergence rate can be improved by a factor larger than and at least equal to $\overline{\rho}$.

Non-vanishing Bias Term. The second term $Q(\overline{\rho},\widetilde{\rho})=\frac{8L\Gamma}{3\mu}\left(\frac{\widetilde{\rho}}{\overline{\rho}}-1\right)$ in (7) denotes the solution bias, which is dependent on the selection strategy. By the definitions of $\overline{\rho}$ and $\widetilde{\rho}$, it follows that $\widetilde{\rho}\geq\overline{\rho}$, which implies that $Q(\overline{\rho},\widetilde{\rho})\geq 0$. For an unbiased selection strategy, we have $\overline{\rho}=\widetilde{\rho}=1$, $Q(\overline{\rho},\widetilde{\rho})=0$, and hence (7) recovers previous bound for unbiased selection strategy as . For $\overline{\rho}>1$, while we gain faster convergence rate by a factor of $\overline{\rho}$, we cannot guarantee $Q(\overline{\rho},\widetilde{\rho})=0$. Thus, there is a trade-off between the convergence speed and the solution bias. Later in the experimental results, we show that even with biased selection strategies, the term $\frac{\widetilde{\rho}}{\overline{\rho}}-1$ in $Q(\overline{\rho},\widetilde{\rho})$ can be close to , and hence $Q(\overline{\rho},\widetilde{\rho})$ has a negligible effect on the final error floor.

Proposed Power-of-choice Client Selection Strategy

From (5) and (6) we discover that a selection strategy $\pi$ that prefers clients with larger $F_{k}(\mathbf{w})-F_{k}^{*}$ will result in a larger $\overline{\rho}$, yielding faster convergence. Using this insight, a naive client selection strategy can be choosing the clients with highest local loss $F_{k}(\mathbf{w})$. However, a larger selection skew $\overline{\rho}$ may result in a larger $\overline{\rho}/\widetilde{\rho}$, i.e., a larger non-vanishing error term. This naive selection strategy has another drawback – to find the current local loss $F_{k}(\mathbf{w})$, it requires sending the current global model to all $K$ clients and having them evaluate $F_{k}$ and sending it back. This additional communication and computation cost can be prohibitively high because the number of clients $K$ is typically very large, and these clients have limited communication and computation capabilities.

In this section, we use these insights regarding the trade-off between convergence speed, solution bias and communication/computation overhead to propose the Power-of-choice client selection strategy. Power-of-choice is based on the power of $d$ choices load balancing strategy , which is extensively used in queueing systems. In the Power-of-choice client selection strategy (denoted by $\pi_{\text{pow-d}}$), the central server chooses the active client set $\mathcal{S}^{(t)}$ as follows:

Sample the Candidate Client Set. The central server samples a candidate set $\mathcal{A}$ of $d~{}(m\leq d\leq K)$ clients without replacement such that client $k$ is chosen with probability $p_{k}$, the fraction of data at the $k$-th client for $k=1,\dots K$.

Estimate Local Losses. The server sends the current global model $\overline{\mathbf{w}}^{(t)}$ to the clients in set $\mathcal{A}$, and these clients compute and send back to the central server their local loss $F_{k}(\overline{\mathbf{w}}^{(t)})$.

Select Highest Loss Clients. From the candidate set $\mathcal{A}$, the central server constructs the active client set $\mathcal{S}^{(t)}$ by selecting $m=\max(CK,1)$ clients with the largest values $F_{k}(\overline{\mathbf{w}})$, with ties broken at random. These $\mathcal{S}^{(t)}$ clients participate in the training during the next round, consisting of iterations $t+1$, $t+2$, …$t+\tau$.

Variations of $\pi_{\text{pow-d}}$. The three steps of $\pi_{\text{pow-d}}$ can be flexibly modified to take into account practical considerations. For example, intermittent client availability can be accounted for in step 1 by constructing set $\mathcal{A}$ only from the set of available clients in that round. We demonstrate the performance of $\pi_{\text{pow-d}}$ with intermittent client availability in Section G.3. The local computation cost and server-client communication cost in step 2 can be reduced or eliminated by the following proposed variants of $\pi_{\text{pow-d}}$ (see Appendix F for their pseudo-codes).

Computation-efficient Variant $\pi_{\text{cpow-d}}$: To save local computation cost, instead of evaluating the $F_{k}(\mathbf{w})$ by going through the entire local dataset $\mathcal{B}_{k}$, we use an estimate $\sum_{\xi\in\widehat{\xi}_{k}}f(\mathbf{w},\xi)/{|\widehat{\xi}_{k}|}$, where $\widehat{\xi}_{k}$ is the mini-batch of $b$ samples sampled uniformly at random from $\mathcal{B}_{k}$.

Communication- and Computation-efficient Variant $\pi_{\text{rpow-d}}$: To save both local computation and communication cost, the selected clients for each round sends their accumulated averaged loss over local iterations, i.e., $\frac{1}{\tau|\xi_{k}^{(l)}|}\sum_{l=t-\tau+1}^{t}\sum_{\xi\in\xi_{k}^{(l)}}f(\mathbf{w}_{k}^{(l)},\xi)$ when they send their local models to the server. The server uses the latest received value from each client as a proxy for $F_{k}(\mathbf{w})$ to select the clients. For the clients that have not been selected yet, the latest value is set to $\infty$.

Selection Skew of Power-of-choice Strategy. The size $d$ of the candidate client set $\mathcal{A}$ is an important parameter which controls the trade-off between convergence speed and solution bias. With $d=m$ we have random sampling without replacement in proportion of $p_{k}$. As $d$ increases, the selection skew $\overline{\rho}$ increases, giving faster error convergence at the risk of a higher error floor. However, note that the convergence analysis replaces $\rho(\mathbf{w},\mathbf{w}^{\prime})$ with $\overline{\rho}$ to get a conservative error bound. In practice, the convergence speed and the solution bias is dictated by $\rho(\overline{\mathbf{w}}^{(\tau\lfloor t/\tau\rfloor)},\overline{\mathbf{w}}^{(t)})$ which changes during training. With $\pi_{\text{pow-d}}$ which is biased towards higher local losses, we expect the selection skew $\rho(\mathbf{w},\mathbf{w}^{\prime})$ to reduce through the course of training. We conjecture that this is why $\pi_{\text{pow-d}}$ gives faster convergence as well as little or no solution bias in our experiments presented in Section 5.

Experimental Results

We evaluate our proposed $\pi_{\text{pow-d}}$ and its practical variants $\pi_{\text{cpow-d}}$ and $\pi_{\text{rpow-d}}$, by three sets of experiments: (1) quadratic optimization, (2) logistic regression on a synthetic federated dataset, Synthetic(1,1) , and (3) DNN trained on a non-iid partitioned FMNIST dataset . We also benchmark the selection strategy proposed by , active federated learning, denoted as $\pi_{\text{afl}}$. Details of the experimental setup are provided in Appendix F, and the code for all experiments are shared in the supplementary material.

Quadratic and Synthetic Simulation Results. In Figure 2(a), even with few clients ($K=30$), $\pi_{\text{pow-d}}$ converges faster than $\pi_{\text{rand}}$ with nearly negligible solution bias for small $d$. The convergence speed increases with the increase in $d$, at the cost of higher error floor due to the solution bias. For $K=100$, $\pi_{\text{pow-d}}$ shows convergence speed-up as with $K=30$, but the bias is smaller. Figure 2(b) shows the theoretical values $\overline{\rho}$ and $\widetilde{\rho}/\overline{\rho}$ which represents the convergence speed and the solution bias respectively in our convergence analysis. Compared to $\pi_{\text{rand}}$, $\pi_{\text{pow-d}}$ has higher $\overline{\rho}$ for all $d$ implying higher convergence speed than $\pi_{\text{rand}}$. By varying $d$ we can span different points on the trade-off between the convergence speed and bias. For $d=15$ and $K=100$, $\widetilde{\rho}/\overline{\rho}$ of $\pi_{\text{pow-d}}$ and $\pi_{\text{rand}}$ are approximately identical, but $\pi_{\text{pow-d}}$ has higher $\overline{\rho}$, implying that $\pi_{\text{pow-d}}$ can yield higher convergence speed with negligible solution bias. In Section G.1, we present the clients’ selected frequency ratio for $\pi_{\text{pow-d}}$ and $\pi_{\text{rand}}$ which gives novel insights regarding the difference between the two strategies.

For the synthetic dataset simulations, we present the global losses in Figure 3 for $\pi_{\text{rand}}$ and $\pi_{\text{pow-d}}$ for different $d$ and $m$. We show that $\pi_{\text{pow-d}}$ converges approximately 3$\times$ faster to the global loss $\approx 0.5$ than $\pi_{\text{rand}}$ when $d=10m$, with a slightly higher error floor. Even with $d=2m$, we get 2$\times$ faster convergence to global loss $\approx 0.5$ than $\pi_{\text{rand}}$.

Experiments with Heterogeneously Distributed FMNIST. As elaborated in Appendix F, $\alpha$ determines the data heterogeneity across clients. Smaller $\alpha$ indicates larger data heterogeneity. In Figure 4, we present the test accuracy and training losses for the different sampling strategies from the FMNIST experiments with $\alpha=0.3$ and $\alpha=2$. Observe that $\pi_{\text{pow-d}}$ achieves approximately $10\%$ and $5\%$ higher test accuracy than $\pi_{\text{rand}}$ and $\pi_{\text{afl}}$ respectively for both $\alpha=2$ and $\alpha=0.3$. For higher $\alpha$ (less data heterogeneity) larger $d$ (more selection skew) performs better than smaller $d$.

Figure 4(a) shows that this performance improvement due to the increase of $d$ eventually converges. For smaller $\alpha$, as in Figure 4(b), smaller $d=6$ performs better than larger $d$ which shows that too much solution bias is adversarial to the performance in the presence of large data heterogeneity. The observations on training loss are consistent with the test accuracy results.

Performance of the Communication- and Computation-Efficient variants. Next, we evaluate $\pi_{\text{cpow-d}}$ and $\pi_{\text{rpow-d}}$ which were introduced in Section 4. In Figure 5, for $\alpha=2$, $\pi_{\text{rpow-d}}$ and $\pi_{\text{cpow-d}}$ each yields approximately $5\%$ and $6\%$ higher accuracy than $\pi_{\text{rand}}$, but both yield lower accuracy than $\pi_{\text{pow-d}}$ that utilizes the highest computation and communication resources. For $\alpha=0.3$, $\pi_{\text{cpow-d}}$ and $\pi_{\text{rpow-d}}$ perform as well as $\pi_{\text{pow-d}}$ and give a $10\%$ accuracy improvement over $\pi_{\text{rand}}$. Moreover, $\pi_{\text{pow-d}},~{}\pi_{\text{rpow-d}}$ and $\pi_{\text{cpow-d}}$ all have higher accuracy and faster convergence than $\pi_{\text{afl}}$.

We evaluate the communication and computation efficiency of Power-of-choice by comparing different strategies in terms of $R_{60}$, the number of communication rounds required to reach test accuracy 60%, and $t_{\text{comp}}$, the average computation time (in seconds) spent per round. The computation time includes the the time taken by the central server to select the clients (including the computation time for the $d$ clients to compute their local loss values) and the time taken by selected clients to perform local updates. In Table LABEL:tab:comp, with only $C=0.03$ fraction of clients, $\pi_{\text{pow-d}},~{}\pi_{\text{cpow-d}},$ and $\pi_{\text{rpow-d}}$ have about $5\%$ higher test accuracy than $(\pi_{\text{rand}},C=0.1)$. The $R_{60}$ for $\pi_{\text{pow-d}},~{}\pi_{\text{cpow-d}},~{}\pi_{\text{rpow-d}}$ is $0.52$, $0.47$, $0.57$ times that of $(\pi_{\text{rand}},C=0.1)$ respectively. This implies that even for $\pi_{\text{rpow-d}}$ which does not incur any additional communication cost for client selection, we can get a $2\times$ reduction in the number of communication rounds using $1/3$ of clients compared to $(\pi_{\text{rand}},C=0.1)$ and still get higher test accuracy performance. Note that the computation time $t_{\text{comp}}$ for $\pi_{\text{cpow-d}}$ and $\pi_{\text{rpow-d}}$ with $C=0.03$ is smaller than that of $\pi_{\text{rand}}$ with $C=0.1$. In Section G.2, we show that the results for $\alpha=2$ are consistent with the $\alpha=0.3$ case shown in Table LABEL:tab:comp. In Section G.4, we also show that for $C=0.1$, the results are consistent with the $C=0.03$ case.

Concluding Remarks

In this work, we present the convergence guarantees for federated learning with partial device participation with any biased client selection strategy. We discover that biasing client selection can speed up the convergence at the rate $\mathcal{O}(\frac{1}{T\overline{\rho}})$ where $\overline{\rho}$ is the selection skew towards clients with higher local losses. Motivated by this insight, we propose the adaptive client selection strategy Power-of-choice. Extensive experiments validate that Power-of-choice yields 3$\times$ faster convergence and $10\%$ higher test accuracy than the baseline federated averaging with random selection. Even with using fewer clients than random selection, Power-of-choice converges 2 $\times$ faster with high test performance. An interesting future direction is to improve the fairness and robustness of Power-of-choice by modifying step 3 of the algorithm to use a different metric such as the clipped loss or the $q$-fair loss proposed instead of $F_{k}(\mathbf{w})$.

This research was generously supported in part by the Doctoral Study Abroad Scholarship from the Korean Government (Yae Jee Cho) and the Qualcomm Innovation fellowship (Jianyu Wang), NSF grant CCF-1850029 and the 2018 IBM Faculty Research Award. Experiments were conducted on clusters provided by the Parallel Data Lab at CMU.

References

Appendix A Additional Theorem

Under Assumptions 3.1 to 3.4, a fixed learning rate $\eta\leq\min\{\frac{1}{2\mu B},\frac{1}{4L}\}$ where $B=1+\frac{3\overline{\rho}}{8}$, and any client selection strategy $\pi$ as defined above, the error after $T$ iterations of federated averaging with partial device participation satisfies

As $T\rightarrow\infty$ the first term in (8) goes to 0 and the second term becomes the bias term for the fixed learning rate case. For a small $\eta$, we have that the bias term for the fixed learning rate case in Theorem A.1 is upper bounded by $\frac{8L\Gamma}{3\mu}\left(\frac{\widetilde{\rho}}{\overline{\rho}}-1\right)$ which is identical to the decaying-learning rate case. The proof is presented in Appendix D.

Appendix B Preliminaries for Proof of Theorem 3.1 and Theorem A.1

Suppose $F_{k}$ is $L-$smooth with global minimum at $\mathbf{w}_{k}^{*}$, then for any $\mathbf{w}_{k}$ in the domain of $F_{k}$, we have that

Observe from the update rule that $k,~{}{k^{\prime}}$ are in the same set $\mathcal{S}^{(t)}$ and hence the terms where $k={k^{\prime}}$ in the summation in (14) will be zero resulting in (15). Moreover for any arbitrary $t$ there is a $t_{0}$ such that $0\leq t-t_{0}<\tau$ that $\mathbf{w}_{k^{\prime}}^{(t_{0})}=\mathbf{w}_{k}^{(t_{0})}$ since the selected clients are updated with the global model at every $\tau$. Hence even for an arbitrary $t$ we have that the difference between $\|\mathbf{w}_{k^{\prime}}^{(t)}-\mathbf{w}_{k}^{(t)}\|^{2}$ is upper bounded by $\tau$ updates. With non-increasing $\eta_{t}$ over $t$ and $\eta_{t_{0}}\leq 2\eta_{t}$, (15) can be further bounded as,

where (22) is because there can be at most $m(m-1)$ pairs such that $k\neq{k^{\prime}}$ in $\mathcal{S}^{(t)}$. ∎

Appendix C Proof of Theorem 3.1

With $\overline{\mathbf{g}}^{(t)}=\frac{1}{m}\sum_{k\in\mathcal{S}^{(t)}}g_{k}(\mathbf{w}_{k}^{(t)},\xi_{k}^{(t)})$ as defined in Section 2, we have that

Lastly we can bound $A_{4}$ using the bound of variance of stochastic gradients as,

Using the bounds of $A_{1},A_{2},A_{3},A_{4}$ above we have that the expectation of the LHS of (27) is bounded as

where (44) is due to Lemma B.3. Now we aim to bound $A_{5}$ in (44). First we can represent $A_{5}$ in a different form as:

Now with $\eta_{t}<1/(4L)$ and $\nu_{t}=2\eta_{t}(1-2L\eta_{t})$, we have that $A_{6}$ can be rewritten and bounded as

where (48) is due to $\mu-$convexity, (49) is due to Lemma B.1 and the AM-GM inequality and Cauchy–Schwarz inequality, and (51) is due to the fact that $\frac{\nu_{t}(1-\eta_{t}\mu)}{2\eta_{t}}\leq 1$. Hence using this bound of $A_{6}$ we can upper bound $A_{5}$ as

where (54) is due to the definition of $\rho(\mathcal{S}(\pi,\mathbf{w}),\mathbf{w}^{\prime})$ in Definition 3.2 and (55) is due to the definition of $\Gamma$ in Definition 3.1 and the definitions of $\overline{\rho},~{}\widetilde{\rho}$ in Definition 3.2. We can expand $A_{7}$ in (55) as

where (60) is due to the $\mu-$convexity, (61) is due to $-2\eta_{t}(1-2L\eta_{t})(1-\eta_{t}L)\leq-\frac{3}{4}\eta_{t}$, and (62) is due to $-(1-2L\eta_{t})(1-\eta_{t}L)\leq-(1-3L\eta_{t})$. Hence we can finally bound $A_{5}$ as

By setting $\Delta_{t}\leq\frac{\psi}{t+\gamma}$, $\eta_{t}=\frac{\beta}{t+\gamma}$ and $\beta>\frac{1}{\mu B},~{}\gamma>0$ by induction we have that

Then by the L-smoothness of $F(\cdot)$, we have that

Appendix D Proof of Theorem A.1

With fixed learning rate $\eta_{t}=\eta$, we can rewrite (65) as

and with $\eta\leq\min\{\frac{1}{2\mu B},\frac{1}{4L}\}$ using recursion of (68) we have that

Using $\Delta_{t}\leq\frac{2}{\mu}(F(\overline{\mathbf{w}}^{(t)})-F^{*})$ and $L$-smoothness, we have that

Appendix E Extension: Generalization to different averaging schemes

While we considered a simple averaging scheme where $\overline{\mathbf{w}}^{(t+1)}=\frac{1}{m}\sum_{k\in\mathcal{S}^{(t)}}\left(\mathbf{w}_{k}^{(t)}-\eta_{t}g_{k}(\mathbf{w}_{k}^{(t)})\right)$, we can extend the averaging scheme to any scheme $\mathbf{q}$ such that the averaging weights $q_{k}$ are invariant in time and satisfies $\sum_{k\in\mathcal{S}^{(t)}}q_{k}=1$ for any $t$. Note that $\mathbf{q}$ includes the random sampling without replacement scheme introduced by where the clients are sampled uniformly at random without replacement with the averaging coefficients $q_{k}=p_{k}K/m$. With such averaging scheme $\mathbf{q}$, we denote the global model for the averaging scheme $q_{k}$ as $\widehat{\mathbf{w}}^{(t)}$, where $\widehat{\mathbf{w}}^{(t+1)}\triangleq\sum_{k\in\mathcal{S}^{(t)}}q_{k}\left(\mathbf{w}_{k}^{(t)}-\eta_{t}g_{k}(\mathbf{w}_{k}^{(t)})\right)$, and the update rule changes to

where $\widehat{\mathbf{g}}^{(t)}=\sum_{k\in\mathcal{S}^{(t)}}q_{k}g_{k}(\mathbf{w}_{k}^{(t)},\xi_{k}^{(t)})$. We show that the convergence analysis for the averaging scheme $\mathbf{q}$ is consistent with Theorem 3.1. In the case of the averaging scheme $\mathbf{q}$, we have that Lemma B.2 and Lemma B.3 shown in Appendix B, each becomes

Then, using the same method we used for the proof of Theorem 3.1, we have that

By defining the selection skew for averaging scheme $\mathbf{q}$ similar to Definition 5 as

With $\eta_{t}<1/(2L(1+m))$, using the same methodology for proof of Theorem 3.1 we have that $M$ becomes upper bounded as

Again, by setting $\widehat{\Delta}_{t}\leq\frac{\psi}{t+\gamma}$, $\eta_{t}=\frac{\beta}{t+\gamma}$ and $\beta>\frac{1}{\mu\widehat{B}},~{}\gamma>0$ by induction we have that

Then by the L-smoothness of $F(\cdot)$, we have that

With $\beta=\frac{m}{\mu}$, $\gamma=\frac{4m(1+m)L}{\mu}$ and $\eta_{t}=\frac{\beta}{t+\gamma}$, we have that

Appendix F Experiment Details

Quadratic Model Optimization. For the quadratic model optimization, we set each local objective function as strongly convex as follows:

where the global model is defined as $\overline{\mathbf{w}}^{(t+1)}=\frac{1}{m}\sum_{k\in\mathcal{S}^{(t)}}\mathbf{w}_{k}^{(t+1)}$. We sample $m=KC$ clients for every round where for each round the clients perform $\tau$ gradient descent local iterations with fixed learning rate $\eta$ and then these local models are averaged to update the global model. For all simulations we set $\tau=2,~{}v=5,~{}\eta=2\times 10^{-5}$.

For the estimation of $\overline{\rho}$ and $\widetilde{\rho}$ for the quadratic model, we get the estimates of the theoretical $\overline{\rho},~{}\widetilde{\rho}$ values by doing a grid search over a large range of possible $\mathbf{w},\mathbf{w}^{\prime}$ for $\rho(\mathcal{S}(\pi,\mathbf{w}),\mathbf{w}^{\prime})$ and $\rho(\mathcal{S}(\pi,\mathbf{w}),\mathbf{w}^{*})$ respectively. The distribution of $\mathcal{S}(\pi,\mathbf{w})$ is estimated by simulating 10000 iterations of client sampling for each $\pi$ and $\mathbf{w}$.

Logistic Regression on Synthetic Dataset. We conduct simulations on synthetic data which allows precise manipulation of heterogeneity. Using the methodology constructed in , we use the dataset with large data heterogeneity, Synthetic(1,1). We assume in total 30 devices where the local dataset sizes for each device follows the power law. We set the mini batch-size to 50 with $\tau=30$, and $\eta=0.05$, where $\eta$ is decayed to $\eta/2$ every 300 and 600 rounds.

DNN on FMNIST Dataset. We train a deep multi-layer perceptron network with two hidden layers on the FMNIST dataset . We construct the heterogeneous data partition amongst clients using the Dirichlet distribution $\text{Dir}_{K}(\alpha)$ , where $\alpha$ determines the degree of the data heterogeneity across clients (the data size imbalance and degree of label skew across clients). Smaller $alpha$ indicates larger data heterogeneity. For all experiments we use mini-batch size of 64, with $\tau=30$ and $\eta=0.005$, where $\eta$ is decayed by half for every 150, 300 rounds. We experiment with three different seeds for the randomness in the dataset partition across clients and present the averaged results.

All experiments are conducted with clusters equipped with one NVIDIA TitanX GPU. The number of clusters we use vary by $C$, the fraction of clients we select. The machines communicate amongst each other through Ethernet to transfer the model parameters and information necessary for client selection. Each machine is regarded as one client in the federated learning setting. The algorithms are implemented by PyTorch.

Pseudo-code of the variants of pow-d: cpow-d and rpow-d. We here present the pseudo-code for $\pi_{\text{cpow-d}}$ and $\pi_{\text{rpow-d}}$. Note that the pseudo-code for $\pi_{\text{cpow-d}}$ in Algorithm 1 can be generalized to the algorithm for $\pi_{\text{pow-d}}$, by changing $\frac{1}{|\widehat{\xi}_{k}|}\sum_{\xi\in\widehat{\xi}_{k}}f(\mathbf{w},\xi)$ to $F_{k}(\mathbf{w})$.

Appendix G Additional Experiment Results

We further visualize the difference between our proposed sampling strategy $\pi_{\text{pow-d}}$ and the baseline scheme $\pi_{\text{rand}}$ by showing the selected frequency ratio of the clients for $K=30,~{}C=0.1$ for the quadratic simulations in Figure 6. Note that the selected ratio for $\pi_{\text{rand}}$ reflects each client’s dataset size. We show that the selected frequencies of clients for $\pi_{\text{pow-d}}$ are not proportional to the data size of the clients, and we are selecting clients frequently even when they have relatively low data size like client 6 or 22. We are also not necessarily frequently selecting the clients that have the highest data size such as client 26. This aligns well with our main motivation of Power-of-choice that weighting the clients’ importance based on their data size does not achieve the best performance, and rather considering their local loss values along with the data size better represents their importance. Note that the selected frequency for $\pi_{\text{rand}}$ is less biased than $\pi_{\text{pow-d}}$.

G.2 Communication and Computation Efficiency with larger data heterogeneity

In Table LABEL:tab:comp2, we show the communication and computation efficiency of Power-of-choice for $\alpha=2$, as we showed for $\alpha=0.3$ in Table LABEL:tab:comp in Section 5. With $C=0.03$ fraction of clients, $\pi_{\text{pow-d}},~{}\pi_{\text{cpow-d}},$ and $\pi_{\text{rpow-d}}$ have better test accuracy of at least approximately 10% higher test accuracy performance than $(\pi_{\text{rand}},~{}C=0.1)$. $R_{60}$ for $\pi_{\text{pow-d}},~{}\pi_{\text{cpow-d}},~{}\pi_{\text{rpow-d}}$ is 0.61, 0.66, 0.73 times that of $(\pi_{\text{rand}},~{}C=0.1)$ respectively. This indicates that we can reduce the number of communication rounds by at least 0.6 using 1/3 of clients compared to $(\pi_{\text{rand}},~{}C=0.1)$ and still get higher test accuracy performance. The computation time $t_{\text{comp}}$ for $\pi_{\text{cpow-d}}$ and $\pi_{\text{rpow-d}}$ with $C=0.03$ is smaller than that of $(\pi_{\text{rand}},~{}C=0.1)$.

G.3 Intermittent Client Availability

In real world scenarios, certain clients may not be available due to varying availability of resources such as battery power or wireless connectivity. Hence we experiment with a virtual scenario, where amongst $K$ clients, for each communication round, we select clients alternately from one group out of two fixed groups, where each group has $0.5K$ clients. This altering selection reflects a more realistic client selection scenario where, for example, we have different time zones across clients. For each communication round, we select 0.1 portion of clients from the corresponding group uniformly at random and exclude them from the client selection process. This random exclusion of certain clients represents the randomness in the client availability within that group for cases such as low battery power or wireless connectivity. In Figure 7 we show that $\pi_{\text{pow-d}}$ and $\pi_{\text{rpow-d}}$ achieves 10% and 5% test accuracy improvement respectively compared to $\pi_{\text{rand}}$ for $\alpha=2$. For $\alpha=3$, both $\pi_{\text{pow-d}}$ and $\pi_{\text{rpow-d}}$ shows 10% improvement. Therefore, we demonstrate that Power-of-choice also performs well in a realistic scenario where clients are available intermittently.

G.4 Effect of the fraction of selected clients

In Figure 8, for larger $C=0.1$ with $\alpha=2$, the test accuracy improvement for $\pi_{\text{pow-d}}$ is even higher than the case of $C=0.03$ with approximately 15% improvement. $\pi_{\text{cpow-d}}$ performs slightly lower in test accuracy than $\pi_{\text{pow-d}}$ but still performs better than $\pi_{\text{rand}}$ and $\pi_{\text{afl}}$. $\pi_{\text{rpow-d}}$ performs as well as $\pi_{\text{afl}}$. For $\alpha=0.3$, $\pi_{\text{pow-d}}$, $\pi_{\text{cpow-d}}$, and $\pi_{\text{rpow-d}}$ have approximately equal test accuracy performance, higher than $\pi_{\text{rand}}$ by 5%. The Power-of-choice strategies all perform slightly better than $\pi_{\text{afl}}$. Therefore we show that Power-of-choice performs well for selecting a larger fraction of clients, i.e., when we have larger $C=0.1>0.03$.