Optimal Client Sampling for Federated Learning

Introduction

We consider the standard cross-device federated learning (FL) setting (Kairouz et al., 2019), where the objective is of the form

It is well understood that communication cost can be a primary bottleneck in cross-device FL, since typical clients are mobile phones or different IoT devices that have limited bandwidth and availability for connection (Van Berkel, 2009; Huang et al., 2013). Indeed, wireless links and other end-user internet connections typically operate at lower rates than intra-datacenter or inter-datacenter links and can be potentially expensive and unreliable. Moreover, the capacity of the aggregating master and other FL system considerations imposes direct or indirect constrains on the number of clients allowed to participate in each communication round. These considerations have led to significant interest in reducing the communication bandwidth of FL systems.

Local Methods. One of the most popular strategies is to reduce the frequency of communication and put more emphasis on computation. This is usually achieved by asking the devices to perform multiple local steps before communicating their updates. A prototype method in this category is the Federated Averaging (FedAvg) algorithm (McMahan et al., 2017), an adaption of local-update to parallel SGD, where each client runs some number of SGD steps locally before local updates are averaged to form the global update for the global model on the master. The original work was a heuristic, offering no theoretical guarantees, which motivated the community to try to understand the method and various existing and new variants theoretically (Stich, 2019; Lin et al., 2018; Karimireddy et al., 2019; Stich & Karimireddy, 2020; Khaled et al., 2020; Hanzely & Richtárik, 2020).

Client Sampling/Selection Methods. In the situation where partial participation is desired and a budget on the number of participating clients is applied, a careful selection of the participating clients can lead to better communication complexity, and hence faster training. In other words, some clients will have “more informative” updates than others in any given communication round, and thus the training procedure will benefit from capitalizing on this fact by ignoring some of the worthless updates (see Figure 1). We refer the readers to Section 4.1 for discussions on existing client sampling methods in FL and their limitations.

2 Contributions

We address the communication bandwidth issues appearing in FL by designing a principled optimal client sampling scheme with client privacy and system practicality in mind. We show that the ideas presented in the previous works on efficient sampling (Horváth & Richtárik, 2019) and sparsification (Wang et al., 2018; Wangni et al., 2018) can be adapted to be compatible with FL and can be used to construct a principled optimal client sampling scheme which is capable of identifying the most informative clients in any given communication round. Our contributions can be summarized as follows:

Inspired by Horváth & Richtárik (2019), we propose an adaptive partial participation strategy for reducing communication in FL. This strategy relies on a careful selection of clients that are allowed to communicate their updates back to the master in any given communication round, which then translates to a reduction in the number of communicated bits. We obtain this strategy by properly applying the sampling procedure from Horváth & Richtárik (2019) to the FL framework.

Specifically, building upon the importance sampling results in Horváth & Richtárik (2019, Lemma 1), we obtain an optimal adaptive client sampling procedure in the sense that it minimizes the variance of the master update for any budget $m$ on the number of participating clients, which generalizes the theoretical results in Zhao & Zhang (2015) that only applies to $m=1$.

Inspired by the greedy algorithm from Wangni et al. (2018, Algorithm 3) which was originally designed for gradient sparsification, we obtain an approximation to our optimal sampling strategy which only requires aggregation, fulfilling two core privacy requirements of FL: to our knowledge, our method is the first principled importance client sampling strategy that is compatible with both secure aggregation and stateless clients.

Our optimal sampling method is orthogonal to and hence compatible with existing approaches to communication reduction such as communication compression and/or local updates (cf. Section 3.2).

We provide convergence guarantees for our approach with Distributed SGD (DSGD) and Federated Averaging (FedAvg), relaxing a number of strong assumptions employed in prior works. We show both theoretically and empirically that the performance of our approach is superior to uniform sampling and can be close to full participation.

We show both theoretically and empirically that our approach allows for larger learning rates than the baseline which performs uniform client sampling, which results in better communication complexity and hence faster convergence.

3 Organization of the Paper

Section 2 describes the proposed optimal client sampling strategy for reducing the communication bottleneck in federated learning. Section 3 provides convergence analyses for DSGD and FedAvg with our optimal client sampling scheme in both convex and non-convex settings. Section 4 reviews prior works that are closely or broadly related to our proposed method. Section 5 empirically evaluates our optimal client sampling method on standard federated datasets. Section 6 summarizes the paper and lists some directions for future work.

Smart Client Sampling for Reducing Communication

This section describes the proposed optimal client sampling strategy for reducing the communication bottleneck in federated learning.

While our client sampling strategy can be adapted to essentially any underlying learning method, we give details here for DSGD as an illustrative example, where the master update in each communication round is of the form

A simple observation is that the variance of our gradient estimator ${\bf G}^{k}$ can be decomposed into

where the second term on the right-hand side is independent of the sampling procedure, and the first term is zero if every client sends its update (i.e., if $p_{i}^{k}=1$ for all $i$). In order to provide meaningful results, we restrict the expected number of clients to communicate in each round by bounding $b^{k}\coloneqq\sum_{i=1}^{n}p_{i}^{k}$ by some positive integer $m\leq n$. This raises the following question: What is the sampling procedure that minimizes (3) for any given $m$?

To answer this question, we connect Equation (3) to previous works on importance sampling (Horváth & Richtárik, 2019) and gradient sparsification (Wangni et al., 2018; Wang et al., 2018)Wangni et al. (2018) consider a slightly different problem, where they minimize the communication budget with constraints on the variance.. Despite difference in motivation, these works solve up to a scale the equivalent mathematical problem, based on which we answer the aforementioned question by the following technical lemma (see Appendix A for a proof):

where the expectation is taken over $S$. Whenever (4) holds, it must be the case that $v_{i}\geq 1-p_{i}$.

It turns out that given probabilities $\{p_{i}\}$, among all samplings $S$ satisfying $p_{i}=\operatorname{Prob}(i\in S)$, the independent sampling (i.e., $p_{ij}=\operatorname{Prob}(i,j\in S)=\operatorname{Prob}(i\in S)\operatorname{Prob}(j\in S)=p_{i}p_{j}$) minimizes the left-hand side of (5). This is due to two nice properties: a) any independent sampling admits the optimal choice of $v$, i.e., $v_{i}=1-p_{i}$ for all $i$, and b) (5) holds as equality for independent sampling. In the context of our method, these properties can be written as

It now only remains to find the parameters $\{p_{i}^{k}\}$ defining the optimal independent sampling, i.e., one that minimizes (6) subject to the constraints $0\leq p_{i}^{k}\leq 1$ and $b^{k}\coloneqq\sum_{i=1}^{n}p_{i}^{k}\leq m$. It turns out that this problem has the following closed-form solution (see Appendix B for a proof):

Optimizing the left-hand side of (5) does not guarantee the proposed sampling to be optimal with respect to the right-hand side of (5) in the general case. For this to hold, our sampling needs to be independent, which is not a very restrictive condition, especially considering that enforcing independent sampling across clients accommodates the privacy requirements of FL. In addition, since (5) is tight, our sampling is optimal if one is allowed to communicate only norms (i.e., one float per client) as extra information. We stress that requiring optimality with respect to the left-hand side of (5) in the full general case is not practical, as it cannot be obtained without revealing, i.e., communicating, all clients’ full updates to the master.

2 Ensuring Compatibility with Secure Aggregation and Stateless Clients

We acknowledge that Algorithm 2 brings extra communication costs, as it requires all clients to send the norms of their updates $u_{i}^{k}$’s and probabilities $p_{i}^{k}$’s in each round. However, since these are single floats, this only costs ${\cal O}(j_{\max})$ extra floats for each client. Picking $j_{\max}={\cal O}(1)$, this is negligible for large models of size $d$. We also acknowledge that engaging in multiple synchronous rounds of communication (as in Algorithm 2) can be a bottleneck (Huba et al., 2022). This is not an issue in our work, as we focus on reducing the total communication cost. However, Algorithm 2 may be less useful under other setups or metrics.

Based on our sampling strategy, it might be tempting to assume that the obtained solution could exhibit fairness issues. In our convergence analyses below, we show that this is not the case, as our proposed methods converge to the optimal solution of the original problem. Hence, as long as the original objective has no inherent issue with fairness, our method does not exhibit any fairness issues. Besides, our algorithm can be used in conjunction with other “more fair” objectives, e.g., Tilted ERM (Li et al., 2021), if needed.

Convergence Guarantees

This section provides convergence analyses for DSGD and FedAvg with our optimal client sampling scheme in both convex and non-convex settings. We compare the convergence results of our scheme with those of full participation and independent uniform sampling with sample size $m$. We match the forms of our convergence bounds to those of the existing bounds in the literature to make them directly comparable. We do not compare the sample complexities of these methods, as such comparisons would be difficult due to their dependence on the actual updates which are unknown in advance and do not follow a specific distribution in general.

We use standard assumptions (Karimi et al., 2016), assuming throughout that $f$ has a unique minimizer $x^{\star}$ with $f^{\star}=f(x^{\star})>-\infty$ and $f_{i}$’s are $L$-smooth, i.e., $f_{i}$’s have $L$-Lipschitz continuous gradients. We first define convex functions and $L$-smooth functions.

We now state standard assumptions of the gradient oracles for DSGD and FedAvg.

The stochastic gradient estimator $g_{i}^{k}=\nabla f_{i}(x^{k})+\xi_{i}^{k}$ of the local gradient $\nabla f_{i}(x^{k})$, for each round $k$ and all $i=1,\dots,n$, satisfies

This further implies that ${\rm E}\left[\frac{1}{n}\sum_{i=1}^{n}g_{i}^{k}\;|\;x^{k}\right]=\nabla f(x^{k})$.

The stochastic gradient estimator $g_{i}(y_{i,r}^{k})=\nabla f_{i}(y_{i,r}^{k})+\xi_{i,r}^{k}$ of the local gradient $\nabla f_{i}(y_{i,r}^{k})$, for each round $k$, each local step $r=0,\dots,R$ and all $i=1,\dots,n$, satisfies

where $y_{i,0}^{k}=x^{k}$ and $y_{i,r}^{k}=y_{i,r-1}^{k}-\eta_{l}g_{i}(y_{i,r}^{k})$, for $r=1,\cdots,R$.

For non-convex objectives, one can construct counter-examples that would diverge for both DSGD and FedAvg if the sampling variance is not bounded. Therefore, we need to employ the following standard assumption of local gradients for bounding the sampling varianceThis assumption is not required for convex objectives, as one can show that the sampling variance is bounded using smoothness and convexity..

The gradients of local loss functions $f_{i}$ satisfy

Some works employ a more restrictive assumption which requires $\left\lVert\nabla f_{i}(x)-\nabla f(x)\right\rVert\leq\rho,~{}\forall i$, from which Assumption 9 can be derived, since $\sum_{i=1}^{n}w_{i}=1$. Therefore, Assumption 9 can be seen as an assumption on similarity among local gradients. Furthermore, this assumption does not require $w_{i}$’s to be lower-bounded, as clients with $w_{i}=0$ will never be sampled and thus can be removed from the objective.

We now define some important quantities for our convergence analyses.

We define the improvement factor of optimal client sampling over uniform sampling:

For simplicity of notation, we define the following quantities which will be useful for our convergence analyses:

where $f_{i}^{\star}$ is the functional value of $f_{i}$ at its optimum, $Z_{i}$ represents the mismatch between the local and global minimizer, and $r^{k}$ captures the distance between the current point and the minimizer of $f$.

We are now ready to proceed with our convergence analyses. In the following subsections, we provide convergence analyses of specific methods for solving the optimization problem (1). The proofs of the theorems are deferred to Appendices C and D.

This subsection presents convergence analyses for DSGD (2) with optimal client sampling in both convex and non-convex settings.

Let $f_{i}$ be $L$-smooth and convex for $i=1,\dots,n$. Let $f$ be $\mu$-strongly convex. Suppose that Assumption 7 holds. Choose $\eta^{k}\in\left(0,\frac{\gamma^{k}}{(1+WM)L}\right]$. Define

The iterates of DSGD with optimal client sampling (7) satisfy

We first look at the best and worst case scenarios. In the best case scenario, we have $\gamma^{k}=1$ for all $k$’s. This implies that there is no loss of speed comparing to the method with full participation. It is indeed confirmed by our theory as our obtained recursion recovers the best-known rate of DSGD in the full participation regime (Gower et al., 2019, Theorem 3.1). To provide a better intuition, we include a full derivation in this case. To match their (stronger) assumptions, we let $M=0$ and $w_{i}=\nicefrac{{1}}{{n}}$. In full participation, we have $\gamma^{k}=1$ for all $k$’s. Then, taking the same step size $\eta$ for all $k$ leads to

Applying the above inequality recursively yields

which is equivalent to the result in Gower et al. (2019, Theorem 3.1). Similarly, in the worst case, we have $\gamma^{k}=\nicefrac{{m}}{{n}}$ for all $k$’s, which corresponds to uniform sampling with sample size $m$, and our recursion recovers the best-known rate for DSGD in this regime. This is expected as (15) implies that every update $\mathbf{U}_{i}^{k}$ is equivalent, and thus it is theoretically impossible to obtain a better rate than that of uniform sampling in the worst case scenario. In the general scenario, our obtained recursion sits somewhere between full and uniform partial participation, where the actual position is determined by $\gamma^{k}$’s which capture the distribution of updates (here gradients) on the clients. For instance, with a larger number of $\gamma^{k}$’s tending to $1$, we are closer to the full participation regime. Similarly, with more $\gamma^{k}$’s tending to $\nicefrac{{m}}{{n}}$, we are closer to the rate of uniform partial participation.

Let $f_{i}$ be $L$-smooth for $i=1,\dots,n$. Suppose that Assumptions 7 and 9 hold. Let $\eta^{k}$ be the step size and define

The iterates of DSGD with optimal client sampling (7) satisfy

The iterate (24) recovers the standard form of the convergence result of DSGD for one recursion step in the non-convex setting. Similar to the previous results, this convergence bound sits between the best-known rate of full participation and uniform sampling (Bottou et al., 2018, Theorem 4.8).

2 Federated Averaging (FedAvg) with Optimal Client Sampling

Pseudo-code that adapts the standard FedAvg algorithm to our framework is provided in Algorithm 3. This subsection presents convergence analyses for FedAvg with optimal client sampling in both convex and non-convex settings.

Let $f_{i}$ be $L$-smooth and $\mu$-strongly convex for $i=1,\dots,n$. Suppose that Assumption 8 holds. Let $\eta^{k}\coloneqq R\eta_{l}^{k}\eta_{g}^{k}$ be the effective step-size and $\eta_{g}^{k}\geq\sqrt{\frac{\gamma^{k}}{\sum_{i}w_{i}^{2}}}$. Choose $\eta^{k}\in\left(0,\frac{1}{8}\min\left\{\frac{1}{L(2+\nicefrac{{M}}{{R}})},\frac{\gamma^{k}}{(1+W(1+\nicefrac{{M}}{{R}}))L}\right\}\right]$,

The iterates of FedAvg ($R\geq 2$) with optimal client sampling (7) satisfy

Let $f_{i}$ be $L$-smooth for all $i=1,\dots,n$. Suppose that Assumptions 8 and 9 hold. Let $\eta^{k}\coloneqq R\eta_{l}^{k}\eta_{g}^{k}$ be the effective step-size and $\eta_{g}^{k}\geq\sqrt{\frac{5\gamma^{k}}{4\sum_{i}w_{i}^{2}}}$. Choose $\eta^{k}\in\left(0,\frac{1}{8L(2+\nicefrac{{M}}{{R}})}\right]$. Define

The iterates of FedAvg ($R\geq 2$) with optimal client sampling (7) satisfy

The convergence guarantees from Theorems 17 and 18 sit somewhere between those for full and uniform partial participation. The actual position is again determined by the distribution of the updates which are linked to $\gamma^{k}$’s. In the edge cases, i.e., $\gamma^{k}=1$ (best case) or $\gamma^{k}=\nicefrac{{m}}{{n}}$ (worst case), we recover the state-of-the-art complexity guarantees provided in (Karimireddy et al., 2019, Theorem I) in both regimes. Note that our results are slightly more general, as Karimireddy et al. (2019) assumes $M=0$ and $w_{i}=\nicefrac{{1}}{{n}}$.

Related Work

This section reviews prior works that are closely or broadly related to our proposed method.

Several recent works have studied efficient importance client sampling methods in FL (Cho et al., 2020; Nguyen et al., 2020; Ribero & Vikalo, 2020; Lai et al., 2021; Luo et al., 2022). Unfortunately, none of these methods is principled, as they rely on heuristics, historical losses, or partial information, which can be seen as proxies for our optimal client sampling. Furthermore, they violate at least one of the core privacy requirements of FL (secure aggregation and/or stateless clients). Specifically, the client selection strategy proposed by Lai et al. (2021) is based on the heuristic of system and statistical utility of clients, which reveals the identity of clients; Ribero & Vikalo (2020) propose to model the progression of the model parameters by an Ornstein-Uhlenbeck process based on partial information, where the master needs to process the raw update from each client. The work of Cho et al. (2020) biases client selection towards clients with higher local losses, which reveals the state of each individual client.

In contrast, our proposed method is the first principled optimal client sampling strategy in the sense that it minimizes the variance of the master update and is compatible with core privacy requirements of FL. We note that the client sampling/selection techniques mentioned in this section could be made compatible with our framework presented in Section 2, but they would not lead to the optimal method as they are only proxies for optimal sampling.

2 Importance Sampling in Stochastic Optimization

Importance sampling methods for optimization have been studied extensively in the last few years in several contexts, including convex optimization and deep learning. LASVM developed in Bordes et al. (2005) is an online algorithm that uses importance sampling to train kernelized support vector machines. The first importance sampling for randomized coordinate descent methods was proposed in the seminal paper of Nesterov (2012). It was showed by Richtárik & Takáč (2014) that the proposed sampling is optimal. Later, several extensions and improvements followed, e.g., Shalev-Shwartz & Zhang (2014); Lin et al. (2014); Fercoq & Richtárik (2015); Qu et al. (2015); Allen-Zhu et al. (2016); Stich et al. (2017). Another branch of work studies sample complexity. In Needell et al. (2014); Zhao & Zhang (2015), the authors make a connection with the variance of the gradient estimates of SGD and show that the optimal sampling distribution is proportional to the per-sample gradient norm. However, obtaining this distribution is as expensive as computing the full gradient in terms of computation, and thus it is not practical. For simpler problems, one can sample proportionally to the norms of the inputs, which can be linked to the Lipschitz constants of the per-sample loss function for linear and logistic regression. For instance, it was shown by Horváth & Richtárik (2019) that static optimal sampling can be constructed even for mini-batches and the probability is proportional to these Lipschitz constants under the assumption that these constants of the per-sample loss function are known. Unfortunately, importance measures such as smoothness of the gradient are often hard to compute/estimate for more complicated models such as those arising in deep learning, where most of the importance sampling schemes are based on heuristics. For instance, a manually designed sampling scheme was proposed in Bengio et al. (2009). It was inspired by the perceived way that human children learn; in practice, they provide the network with examples of increasing difficulty in an arbitrary manner. In a diametrically opposite approach, it is common for deep embedding learning to sample hard examples because of the plethora of easy non-informative ones (Schroff et al., 2015; Simo-Serra et al., 2015). Other approaches use history of losses for previously seen samples to create the sampling distribution and sample either proportionally to the loss or based on the loss ranking (Schaul et al., 2015; Loshchilov & Hutter, 2015). Katharopoulos & Fleuret (2018) propose to sample based on the gradient norm of a small uniformly sampled subset of samples.

Although our proposed optimal sampling method adapts and extends the importance sampling results from Horváth & Richtárik (2019) to the distributed setting of FL, it does not suffer from any of the limitations discussed above, since the motivation of our work is to reduce communication rather than reduce computation. In particular, our method allows for any budge $m<n$ on the number of participating clients, which generalizes the theoretical results from Zhao & Zhang (2015) which only applies to the case $m=1$.

Experiments

This section empirically evaluates our optimal client sampling method on standard federated datasets from LEAF (Caldas et al., 2018).

We compare our method with 1) full participation where all available clients participate in each round; and 2) the baseline where participating clients are sampled uniformly from available clients in each round. We chose not to compare with other client sampling methods, as such comparisons would be unfair. This is because they violate the privacy requirements of FL: our method is the only importance client sampling strategy that is deployable to real-world FL systems (cf. Section 4.1).

We simulate the cross-device FL distributed setting and train our models using TensorFlow Federated (TFF). We conclude our evaluations using FedAvg with Algorithm 2, as it supports stateless clients and secure aggregationWe compared the results of Algorithms 1 and 2 for all experiments as a subroutine. Their results are identical, so we only show results for Algorithm 2 and argue that the performance loss caused by its approximation is negligible.. We extend the TFF implementation of FedAvg to fit our framework. For all three methods, we report validation accuracy and (local) training loss as a function of the number of communication rounds and the number of bits communicated from clients to the masterThe communication from the master to clients is not considered as a bottleneck and thus not included in the results. This is a standard consideration for distributed systems, as one-to-many communication primitives (i.e., from the master to clients) are several orders of magnitude faster than many-to-one communication primitives (i.e., from clients to the master). This gap is further exacerbated in FL due to the large number of clients and slow client connections. . Each figure displays the mean performance with standard deviation over 5 independent runs for each of the three compared methods. For a fair comparison, we use the same random seed for all three methods in a single run and vary random seeds across different runs. Detailed experimental settings and extra results can be found in Appendices F.1 and F.2. Our code together with datasets can be found at https://github.com/SamuelHorvath/FL-optimal-client-sampling.

2 Federated EMNIST Dataset

3 Shakespeare Dataset

We also evaluate our method on the Shakespeare text dataset for next character prediction. Unlike in the FEMNIST experiments, we do not change the number of examples held by each client in this dataset. The vocabulary set for this task consists of $86$ unique characters. The dataset contains $715$ clients, each corresponding to a character in Shakespeare’s plays. We divide the text into batches such that each batch contains $8$ example sequences of length $5$. We use a two-hidden-layer GRU model with 256 units in each hidden layer. We set $n\in\{32,128\},~{}m\in\{2,4,6,12\},~{}j_{max}=4$, and run several SGD steps for $1$ epoch on each client’s local dataset in every communication round. We use vanilla SGD optimizers with constant step sizes, with $\eta_{g}=1$ and $\eta_{l}$ tuned on a holdout set. For full participation and optimal sampling, it turns out that the optimal is $\eta_{l}=2^{-2}$. For uniform sampling, the optimal is $\eta_{l}=2^{-3}$. The main results are shown in Figures 6 and 7.

4 Discussions

As predicted by our theory, the performance of FedAvg with our proposed optimal client sampling strategy is in between that with full and uniform partial participation. For all datasets, the optimal sampling strategy performs slightly worse than but is still competitive with the full participation strategy in terms of the number of communication rounds: it almost reached the performance of full participation while only less than $10\%$ of the available clients communicate their updates back to the master (in the cases $m=2,3$). As we increase the expected number $m$ of sampled clients, the performance of optimal sampling increases accordingly, which is consistent with our theory (e.g., Theorem 18) and with the observations from Yang et al. (2021), and quickly becomes almost identical to that of full participation. Note that the uniform sampling strategy performs significantly worse, which indicates that a careful choice of sampling probabilities can go a long way towards closing the gap between the performance of naive uniform sampling and full participation.

Also, it can be seen that the performances of our optimal client sampling strategy with $m=6$ and $m=12$ match the performances of full participation in the cases $n=32$ and $n=128$, respectively, in terms of the number of communication rounds. We therefore conjecture that $m=\mathcal{O}(\sqrt{n})$ is sufficient for our optimal client sampling strategy to obtain identical validation accuracy to that of full participation in terms of the number of communication rounds.

More importantly, and this was the main motivation of our work, our optimal sampling strategy is significantly better than both the uniform sampling and full participation strategies when we compare validation accuracy as a function of the number of bits communicated from clients to the master. For instance, on FEMNIST Dataset 1 (Figure 5), while our optimal sampling approach with $m=3$ reached around 85% validation accuracy after $2^{6}\times 10^{8}$ communicated bits, neither the full sampling strategy nor the uniform sampling strategy with $m=3$ is able to exceed 40% validation accuracy within the same communication budget. Indeed, to reach the same 85% validation accuracy, full participation approach needs to communicate more than $2^{9}\times 10^{8}$ bits, i.e., $8\times$ more, and uniform sampling approach needs to communicate about the same number of bits as full participation or even more. The results for FEMNIST Datasets 2 and 3 and for the Shakespeare dataset are of a similar qualitative nature, showing that these conclusions are robust across the datasets considered.

Finally, it is also worth noting that the empirical results from Sections 5.2 and 5.3 confirm that our optimal sampling strategy allows for larger step sizes than uniform sampling, as the hyperparameter search returns larger step sizes $\eta_{l}$ for optimal sampling than for uniform sampling.

In Appendix G, we present an additional experiment on the Federated CIFAR100 dataset from LEAF. The Federated CIFAR100 dataset is a balanced dataset, where every client holds the same number of training images. In this setting, letting all clients perform $1$ epoch of local training means that all clients have the same number of local steps in each round. We show that our optimal client sampling scheme still achieves better performance than uniform sampling on this balanced dataset.

Conclusion and Future Work

In this work, we have proposed a principled optimal client sampling strategy to address the communication bottleneck issue of federated learning. Our optimal client sampling can be computed using a closed-form formula by aggregating only the norms of the updates. Furthermore, our method is the first principled importance client sampling strategy that is compatible with stateless clients and secure aggregation. We have obtained convergence guarantees for our method with DSGD and FedAvg with relaxed assumptions, and have performed empirical evaluations of our method on federated datasets from the LEAF database. The empirical results show that our method is superior to uniform sampling and close to full participation, which corroborates our theoretical analysis. We believe that our proposed optimal client sampling scheme will be useful in reducing communication costs in real-world FL systems.

Some directions for future work are as follows:

A straightforward extension would be to combine our proposed optimal sampling approach with communication compression methods to further reduce the sizes of communicated updates.

In the settings where the communication latency is high, our proposed method may not be effective in reducing the real communication time. It would be interesting to extend our optimal client sampling strategy to take into account the constraints of local clients (e.g., computational speed, network bandwidth, and communication latency).

We thank Jakub Konečný for helpful discussions and comments. Most of the work was done when WC was a research intern at KAUST and when SH was a PhD student at KAUST.

References

Appendix A Proof of Lemma 1

Our proof technique can be seen as an extended version of that in (Horváth & Richtárik, 2019). Let $1_{i\in S}=1$ if $i\in S$ and $1_{i\in S}=0$ otherwise. Likewise, let $1_{i,j\in S}=1$ if $i,j\in S$ and $1_{i,j\in S}=0$ otherwise. Note that ${\rm E}\left[1_{i\in S}\right]=p_{i}$ and ${\rm E}\left[1_{i,j\in S}\right]=p_{ij}$. Next, let us compute the mean of $X\coloneqq\sum_{i\in S}\frac{w_{i}\zeta_{i}}{p_{i}}$:

Since, by assumption, we have $\boldsymbol{P}-pp^{\top}\preceq\textbf{Diag}(p\circ v)$, we can further bound

To obtain (5), it remains to combine this with (30). The inequality $v_{i}\ \geq 1-p_{i}$ follows by comparing the diagonal elements of the two matrices in (4). Consider now the independent sampling. Clearly,

Appendix B The Improvement Factor for Optimal Client Sampling

The optimal probabilities are obtained by minimizing (31) w.r.t. $\{p_{i}^{k}\}_{i=1}^{n}$ subject to the constraints $0\leq p_{i}^{k}\leq 1$ and $m\geq b^{k}=\sum_{i=1}^{n}p_{i}^{k}$.

This proof uses an argument similar to that in the proof of Lemma 2 in Horváth & Richtárik (2019). We first show that (33) is the solution to the following optimization problem:

The Lagrangian of this optimization problem is given by

This indicates that (33) is also the solution to the original optimization problem (32). ∎

Plugging the optimal probabilities obtained in (33) into (31) gives

Putting them together gives the improvement factor:

Appendix C DSGD with Optimal Client Sampling

$L$-smoothness of $f_{i}$ and the assumption on the gradient imply that the inequality

holds for all $k\geq 0$. We first take expectations over $x^{k+1}$ conditioned on $x^{k}$ and over the sampling $S^{k}$:

Now choose any $0<\eta^{k}\leq\frac{m}{(\alpha^{k}(n-m)+m)(1+WM)L}$ and define

Taking full expectation yields the desired result:

C.2 Proof of Theorem 15

where the expectations are conditioned on $x^{k}$. In Appendix C.1, we already obtained the upper bound for the last term in equation (34):

Combining the inequalities above and taking full expectation yields equation (24). ∎

Appendix D FedAvg with Optimal Client Sampling

For any given $x$, $y$, and $z$, the two inequalities below follows by the smoothness and strong convexity of the function $h$:

Further, applying the relaxed triangle inequality gives

Combining all these inequalities together we have

Summations are always over $i\in[n]$ and $r\in[R]$ unless stated otherwise. Taking expectations over $x$ conditioned on the results prior to round $k$ and over the sampling $S$ gives

Applying Lemma 21 with $h=w_{i}f_{i}$, $x=y_{i,r-1}$, $y=x^{\star}$ and $z=x$ gives

where $\mathcal{E}$ is the drift caused by the local updates on the clients:

Using independence, zero mean and bounded second moment of the random variables $\xi_{i,r}$, we obtain

Combining the smoothness of $f_{i}$’s, the definition of $\mathcal{E}$, and Jensen’s inequality with definition $\gamma\coloneqq\frac{m}{\alpha(n-m)+m}$, we obtain

Putting these bounds on $\mathcal{A}_{1}$ and $\mathcal{A}_{2}$ together and using the fact that $1-W\leq\nicefrac{{1}}{{\gamma}}$ yields

Let $\eta\leq\frac{\gamma}{8(1+W(1+\nicefrac{{M}}{{R}}))L}$, then

Next, we need to bound the drift $\mathcal{E}$. For $R\geq 2$, we have

If we further restrict $\eta\leq\frac{1}{8L(2+\nicefrac{{M}}{{R}})}$, then for any $\eta_{g}\geq 1$, we have

Due to the upper bound on the step size $\eta\leq\frac{1}{8L(2+\nicefrac{{M}}{{R}})}$, we have the inequalities

Plugging the assumption $\eta_{g}^{k}\geq\sqrt{\frac{\gamma^{k}}{\sum_{i}w_{i}^{2}}}$ into the RHS of the above inequality completes the proof.

D.2 Proof of Theorem 18

Summations are always over $i\in[n]$ and $r\in[R]$ unless stated otherwise. Taking expectations conditioned on $x$ and using a similar argument as in the proof in Appendix C.2, we have

where $\mathcal{E}$ is the drift caused by the local updates on the clients as defined in (36).

In Appendix D.1, we also obtained the upper bound for the drift $\mathcal{E}$:

Since $8\eta L\leq 8\eta L(1+\nicefrac{{M}}{{R}})\leq 1$, we have

Applying the assumption that $\eta_{g}\geq\sqrt{\frac{5\gamma}{4\sum_{i}w_{i}^{2}}}$ and taking full expectations completes the proof:

Appendix E A Sketch of Results on Partial Participation

This section discusses how our analysis can be extended to the case where not all clients are available to participate in each round. As an illustrative example, we consider Distributed SGD (DSGD), i.e., $\mathbf{U}_{i}^{k}=g_{i}^{k}$.

If not all clients are available to participate in each communication round, we will assume that there is a known distribution of client availability ${\cal Q}$ such that in each step a subset ${\cal Q}^{k}\sim{\cal Q}$ of clients are available to participate in a given communication round $k$. We denote the probability that client $i$ is available in the current run by $q_{i}$, i.e., $q_{i}=\operatorname{Prob}(i\in{\cal Q}^{k})$. Under this setting, we can apply twice tower property of the expectation and obtain the following variance decomposition:

where we update the definition of ${\bf G}^{k}$

Note that $S^{k}\subseteq{\cal Q}^{k}$ as we can only sample from available clients. Furthermore, in the particular case where all clients are available, the above equations become identical to the ones that we present in the main paper.

Upper-bounding Equation (39) in an analogous way as we proceed in our analysis in Appendices C and D would complete the proof of convergence for these settings.

Appendix F Experimental Details

We detail the hyper-parameters used in the experiments on the FEMNIST datasets. For each experiment, we run $151$ communication rounds, reporting (local) training loss every round and validation accuracy every $5$ rounds. In each round, $n=32$ clients are sampled from the client pool, each of which then performs SGD for $1$ epoch on its local training images with batch size $20$. For partial participation, the expected number of clients allowed to communicate their updates back to the master is set to $m\in\{3,6\}$. We use vanilla SGD and constant step sizes for all experiments, where we set $\eta_{g}=1$ and tune $\eta_{l}$ from the set of value $\{2^{-1},2^{-2},2^{-3},2^{-4},2^{-5}\}$. If the optimal step size hits a boundary value, then we try one more step size by extending that boundary and repeat this until the optimal step size is not a boundary value. For full participation and optimal sampling, it turns out that $\eta_{l}=2^{-3}$ is the optimal local step size for all three datasets. For uniform sampling, the optimal is $\eta_{l}=2^{-5}$ for Dataset 1 and $\eta_{l}=2^{-4}$ for Datasets 2 and 3. For the extra communications in Algorithm 2, we set $j_{max}=4$.

We also present some additional figures of the experiment results. Figures 8, 9 and 10 show the current best validation accuracy as a function of the number of communication rounds and the number of bits communicated from clients to the master on Datasets 1, 2 and 3, respectively.

F.2 Shakespeare Dataset

We detail the hyper-parameters used in the experiments on the Shakespeare dataset. For each experiment, we run $151$ communication rounds, reporting (local) training loss every round and validation accuracy every $5$ rounds. In each round, $n\in\{32,128\}$ clients are sampled from the client pool, each of which then performs SGD for $1$ epoch on its local training data with batch size $8$ (each batch contains $8$ example sequences of length $5$). For partial participation, the expected number of clients allowed to communicate their updates back to the master is set to $m\in\{2,4,6,12\}$. We use vanilla SGD and constant step sizes for all experiments, where we set $\eta_{g}=1$ and tune $\eta_{l}$ from the set of value $\{2^{-1},2^{-2},2^{-3},2^{-4},2^{-5}\}$. If the optimal step size hits a boundary value, then we try one more step size by extending that boundary and repeat this until the optimal step size is not a boundary value. For full participation and optimal sampling, it turns out that $\eta_{l}=2^{-2}$ is the optimal local step size. For uniform sampling, the optimal is $\eta_{l}=2^{-3}$. For the extra communications in Algorithm 2, we set $j_{max}=4$.

We also present an additional figure of the experiment result. Figures 11 and 12 show the current best validation accuracy as a function of the number of communication rounds and the number of bits communicated from clients to the maste for the cases $n=32,128$, respectively.

Appendix G Additional Experiment on Federated CIFAR100 Dataset

We evaluate our method on the Federated CIFAR100 image dataset for image classification. The Federated CIFAR100 dataset is a balanced dataset, where every client holds the same number of training images. In each communication round, $n=32$ clients are sampled uniformly from the client pool, each of which then performs several SGD steps on its local training images for $1$ epoch with batch size $20$. This means that all clients have the same number of local steps in each round. For partial participation, the expected number of clients allowed to communicate their updates back to the master is set to $m=3$. We use vanilla SGD optimizers with constant step sizes for both clients and the master, with $\eta_{g}=1$ and $\eta_{l}$ tuned on a holdout set. For full participation and optimal sampling, it turns out that $\eta_{l}=1\times 10^{-3}$ is the optimal local step size. For uniform sampling, the optimal is $\eta_{l}=3\times 10^{-4}$. We set $j_{\max}=4$ and include the extra communication costs in our results. The main results are shown in Figure 13. It can be seen that our optimal client sampling scheme achieves better performance than uniform sampling on this balanced dataset. The performance gains of our method over uniform sampling come from the fact that the norms of the updates from some clients are larger than those from other clients even if all clients run the same number of local steps in each round.