Tackling System and Statistical Heterogeneity for Federated Learning with Adaptive Client Sampling

I Introduction

Federated learning (FL) enables many clientsWe use “device” and “client” interchangeably in this paper. to collaboratively train a model under the coordination of a central server while keeping the training data decentralized and private (e.g., ). Compared to traditional distributed machine learning techniques, FL has two unique features (e.g., ), as shown in Fig. 1. First, clients are massively distributed and with diverse and low communication rates (known as system heterogeneity), where stragglers can slow down the physical training time.As suggested , we consider mainstream synchronized FL in this paper due to its composability with other techniques (such as secure aggregation protocols and differential privacy). Second, the training data are distributed in a non-i.i.d. and unbalanced fashion across the clients (known as statistical heterogeneity), which negatively affects the convergence behavior.

Due to limited communication bandwidth and across geographically dispersed devices, FL algorithms (e.g., the de facto FedAvg algorithm in ) usually perform multiple local iterations on a fraction of randomly sampled clients (known as partial participation) and then aggregates their resulting local model updates via the central server periodically . Recent works have provided theoretical convergence analysis that demonstrates the effectiveness of FL with partial participation in various non-i.i.d. settings .

However, these prior works have focused on sampling schemes that select clients uniformly at random or proportional to the clients’ data sizes, which often suffer from slow convergence with respect to wall-clock (physical) timeWe use wall-clock time to distinguish from the number of training rounds. due to high degrees of the system and statistical heterogeneity. This is because the total FL time depends on both the number of training rounds for reaching the target precision and the physical time in each round . Although uniform sampling guarantees that the aggregated model update in each round is unbiased towards that with full client participation, the aggregated model may have a high variance due to data heterogeneity, thus, requiring more training rounds to converge to a target precision. Moreover, considering clients’ heterogeneous communication delay, uniform sampling also suffers from the straggling effect, as the probability of sampling a straggler within the sampled subset in each round can be relatively high,For example, suppose there are 100 clients with only 5 stragglers, then the probability of sampling at least a straggler for uniformly sampling 10 clients in each round is more than 40%. thus yielding a long per-round time.

One effective way of speeding up the convergence with respect to the number of training rounds is to choose clients according to some sampling distribution where “important” clients have high probabilities . For example, recent works adopted importance sampling approaches based on clients’ statistical property . However, their sampling schemes did not account for the heterogeneous physical time in each round, especially under straggling circumstances. Another line of works aims to minimize the learning time via optimizing client selection and scheduling based on their heterogeneous system resources . However, their optimization schemes did not consider how client selection schemes influence the convergence behavior due to data heterogeneity and thus may negatively affect the total learning time.

In a nutshell, the fundamental limitation of existing works is the lack of joint consideration of the impact of the inherent system heterogeneity and statistical heterogeneity on client sampling. In other words, clients with valuable data may have poor communication capabilities, whereas those who communicate fast may have low-quality data. This motivates us to study the following key question.

Key Question: How to design an optimal client sampling scheme that tackles both system and statistical heterogeneity to achieve fast convergence with respect to wall-clock time?

The challenge of this question is threefold: (1) It is difficult to obtain an analytical FL convergence result for arbitrary client sampling probabilities. (2) The total learning time minimization problem can be complex and non-convex due to the straggling effect. (3) The optimal client sampling solution contains unknown parameters from the convergence result, which we can only estimate during the learning process (known as the chicken-and-egg problem).

In light of the above discussion, we state the main results and key contributions of this paper as follows:

Optimal Client Sampling for Heterogeneous FL: We study how to design the optimal client sampling strategy to minimize FL wall-clock time with convergence guarantees. To the best of our knowledge, this is the first work that aims to optimize client sampling probabilities to address both system and statistical heterogeneity.

Convergence Bound for Arbitrary Sampling: Using an adaptive client sampling and model aggregation design, we obtain a new tractable convergence upper bound for FL algorithms with arbitrary client sampling probabilities. This enables us to establish the analytical relationship between the total learning time and client sampling probabilities and formulate a non-convex training time minimization problem.

Optimization Algorithm and Sampling Principle: We propose a low-cost substitute sampling approach to learn the convergence-related unknown parameters and develop an efficient algorithm to approximately solve the non-convex problem with low computational complexity. Our solution characterizes the impact of communication time (system heterogeneity) and data quantity and quality (statistical heterogeneity) on the optimal client sampling design.

Simulation and Prototype Experimentation: We evaluate the performance of our proposed algorithms through both a simulated environment and a hardware prototype. Experimental results demonstrate that for both convex and non-convex learning models, our proposed sampling scheme significantly reduces the convergence time compared to several baseline sampling schemes. For example, with our hardware prototype and the EMNIST dataset, our sampling scheme spends $73$% less time than baseline uniform sampling for reaching the same target loss.

II Related Work

Active client sampling and selection play a crucial role in addressing the statistical and system heterogeneity challenges in cross-device FL. In the existing literature, the research efforts in speeding up the training process mainly focus on two aspects: importance sampling and resource-aware optimization-based approaches.

define The goal of importance sampling is to reduce the variance in traditional optimization algorithms based on stochastic gradient descent (SGD), where SGD draws data samples uniformly at random during the learning process (e.g., ). Recent works have adopted this idea in FL systems to improve communication efficiency via designing client sampling strategy. Specifically, clients with “important” data would have higher probabilities to be sampled in each round. For example, existing works use clients’ local gradient information (e.g., ) or local losses (e.g., ) to measure the importance of clients’ data. However, these schemes did not consider the speed of error convergence with respect to wall-clock time, especially the straggling effect due to heterogeneous transmission delays.

Another line of works aims to minimize wall-clock time via resource-aware optimization-based approaches, such as CPU frequency allocation (e.g., ), and communication bandwidth allocation (e.g., ), straggler-aware client scheduling (e.g., ), parameters control (e.g., ), and task offloading (e.g., ). While these papers provided some novel insights, their optimization approaches did not consider how client sampling affects the total wall-clock time and thus are orthogonal to our work.

Unlike all the above-mentioned works, our work focuses on how to design the optimal client sampling strategy that tackles both system and statistical heterogeneity to minimize the wall-clock time with convergence guarantees. In addition, most existing works on FL are based on computer simulations. In contrast, we implement our algorithm in an actual hardware prototype with resource-constrained devices, which allows us to capture real system operations.

The organization of the rest of the paper is as follows. Section III introduces the system model and problem formulation. Section IV presents our new error-convergence bound with arbitrary client sampling. Section V gives the optimal client sampling algorithm and solution insights. Section VI provides the simulation and prototype experimental results. We conclude this paper in Section VII.

III Preliminaries and System Model

We start by summarizing the basics of FL and its de facto algorithm FedAvg with unbiased client sampling. Then, we introduce the proposed adaptive client sampling for statistical and system heterogeneity based on FedAvg. Finally, we present our formulated optimization problem.

Consider a federated learning system involving a set of $\mathcal{N}={1,\ldots,N}$ clients, coordinated by a central server. Each client $i$ has $n_{i}$ local training data samples ($\mathbf{x}_{i,1},\ldots,\mathbf{x}_{i,n_{i}}$), and the total number of training data across $N$ devices is $n_{\textnormal{tot}}:=\sum\nolimits_{i=1}^{N}n_{i}$. Further, define $f(\cdot,\cdot)$ as a loss function where ${{f}\left(\mathbf{w};\mathbf{x}_{i,j}\right)}$ indicates how the machine learning model parameter $\mathbf{w}$ performs on the input data sample $\mathbf{x}_{i,j}$. Thus, the local loss function of client $i$ can be defined as

Denote $p_{i}=\frac{n_{i}}{n_{\textnormal{tot}}}$ as the weight of the $i$-th device such that $\sum\nolimits_{i=1}^{N}p_{i}=1$. Then, by denoting $F\left(\mathbf{w}\right)$ as the global loss function, the goal of FL is to solve the following optimization problem :

The most popular and de facto optimization algorithm to solve (2) is FedAvg . Here, denoting $r$ as the index of an FL round, we describe one round (e.g., the $r$-th) of the FedAvg algorithm as follows:

The server uniformly at random samples a subset of $K$ clients (i.e., $K\!:=\!\left|\mathcal{K}^{r}\right|$ with $\mathcal{K}^{r}\!\subseteq\!\mathcal{N}$) and broadcasts the latest model $\mathbf{w}^{r}$ to the selected clients.

Each sampled client $i$ chooses $\mathbf{w}_{i}^{r,0}\!=\!\mathbf{w}^{r}$, and runs $E$ steps$E$ is originally defined as epochs of SGD in . In this paper, we denote $E$ as the number of local iterations for theoretical analysis. of local SGD on (1) to compute an updated model $\mathbf{w}_{i}^{r,E}$. Then, the sampled client lets $\mathbf{w}_{i}^{r+1}\!=\!\mathbf{w}_{i}^{r,E}$ and send it back to the server.

The server aggregates (with weight $p_{i}$) the clients’ updated model and computes a new global model $\mathbf{w}^{\tau+1}$.

The above process repeats for many rounds until the global loss converges.

Recent works have demonstrated the effectiveness of FedAvg with theoretical convergence guarantees in various settings . However, these works assume that the server samples clients either uniformly at random or proportional to data size, which may slow down the wall-clock time for convergence due to the straggling effect and non-i.i.d. data . Thus, a careful client sampling design should tackle both system and statistical heterogeneity for fast convergence.

III-B System Model of FL with Client Sampling 𝐪𝐪\mathbf{q}

We aim to sample clients according to a probability distribution $\mathbf{q}\!=\!\{q_{i},\forall i\!\in\!\mathcal{N}\}$, where $0\!<\!q_{i}\!<\!1$ and $\sum_{i=1}^{N}q_{i}\!=\!1$. Through optimizing $\mathbf{q}$, we want to address system and statistical heterogeneity so as to minimize the wall-clock time for convergence. We describe the system model as follows.

Following recent works , we assume that the server establishes the sampled client set $\mathcal{K(\mathbf{q})}^{r}$ by sampling $K$ times with replacement from the total $N$ clients, where $\mathcal{K(\mathbf{q})}^{r}$ is a multiset in which a client may appear more than once. The aggregation weight of each client $i$ is multiplied by the number of times it appears in $\mathcal{K(\mathbf{q})}^{r}$.

) Statistical Heterogeneity Model

We consider the standard FL setting where the training data are distributed in an unbalanced and non-i.i.d. fashion among clients.

) System Heterogeneity Model

Following the same setup of and , we denote $t_{i}$ as the round time of client $i$, which includes both local model computation time and global communication time. For simplicity, we assume that $t_{i}$ remains the same across different rounds for each client $i$, while for different clients $i$ and $j$, $t_{i}$ and $t_{j}$ can be different. The extension to time-varying $t_{i}$ is left for future work. Without loss of generality, as illustrated in Fig. 1, we sort all $N$ clients in the ascending order $\{t_{i}\}$, such that

) Total Wall-clock Time Model

We consider the mainstream synchronized FL model where each sampled client performs multiple (e.g., $E$) steps of local SGD before sending back their model updates to the server (e.g., ). For synchronous FL, the per-round time is limited by the slowest client (known as straggler). Thus, the per-round time $T^{(r)}(\mathbf{q})$ of the entire FL process is

Therefore, the total learning time $T_{\textnormal{tot}}(\mathbf{q},R)$ after $R$ rounds is

III-C Problem Formulation

In Section IV and Section V, we address these two challenges, respectively, and propose approximate algorithms to find an approximate solution to Problem P1 efficiently.

IV Convergence Bound for Arbitrary Sampling

In this section, we address the first challenge by deriving a new tractable convergence bound for arbitrary client sampling probabilities.

To ensure a tractable convergence analysis, we first state several assumptions on the local objective functions ${F_{i}}\left(\mathbf{w}\right)$.

L-smooth: For each client $i\!\in\!\mathcal{N}$, $F_{i}$ is $L$-smooth, i.e., $\|\nabla\!f(\mathbf{v})\!-\!\nabla\!f(\mathbf{w})\|\!\leq\!L\|\mathbf{v}\!-\!\mathbf{w}\|$ for all $\mathbf{v}$ and $\mathbf{w}$.

Strongly-convex: For each client $i\in\mathcal{N}$, $F_{i}$ is $\mu$-strongly convex, i.e., $F_{i}(\mathbf{v})\geq F_{i}(\mathbf{w})+(\mathbf{v}-$ $\mathbf{w})^{T}\nabla F_{i}(\mathbf{w})+\frac{\mu}{2}\|\mathbf{v}-\mathbf{w}\|_{2}^{2}$ for all $\mathbf{v}$ and $\mathbf{w}$.

IV-B Aggregation with Arbitrary Client Sampling Probabilities

This section shows how to aggregate clients’ model updates under sampling probabilities $\mathbf{q}$, such that the aggregated global model is unbiased compared to that with full client participation, which leads to our convergence result.

We first define the virtual weighted aggregated model with full client participation in round $r$ as

With this, we can derive the following result.

(Adaptive Client Sampling and Model Aggregation) When clients $\mathcal{K}(\mathbf{q})^{r}$ are sampled with probability $\mathbf{q}=\{q_{1},\ldots q_{N}\}$ and their local updates are aggregated as $\mathbf{w}^{r+1}\leftarrow\mathbf{w}^{r}+\sum_{i\in\mathcal{K}(\mathbf{q})^{r}}\frac{p_{i}}{Kq_{i}}\left(\mathbf{w}_{i}^{r+1}-\mathbf{w}^{r}\right)$, we have

Proof Sketch. The basic idea is to take expectation over the aggregated global model of the sampled clients $\mathcal{K}(\mathbf{q})^{r}$, and with some mathematical derivations, we have (8). ∎

Remark: The key insight of our sampling and aggregation is that since we sample different clients with different probabilities (e.g., $q_{i}$ for client $i$), we need to inversely re-weight their updated model in the aggregation step (e.g., $\frac{1}{q_{i}}$ for client $i$), such that the aggregated model is still unbiased towards that with full client participation. We summarize how the server performs client sampling and model aggregation in Algorithm 1, where the main differences compared to the de facto FedAvg in are the Sampling (Line 1) and Aggregation (Line 1) procedures. Notably, Algorithm 1 recovers FedAvg algorithm with uniform sampling when letting $q_{i}=\frac{1}{N}$, or with weighted sampling when letting $q_{i}=p_{i}$ in .

IV-C Main Convergence Result for Arbitrary Client Sampling

Based on Lemma 1, we present the main convergence result for arbitrary client sampling in Theorem 1.

(Convergence Upper Bound) Let Assumptions 1 to 4 hold, $\gamma=\max\{\frac{8L}{\mu},E\}$, and decaying learning rate $\eta_{r}=\frac{2}{\mu(\gamma+r)}$. For given client sampling probabilities $\mathbf{q}=\{q_{1},\ldots,q_{N}\}$ and the corresponding aggregation described in Lemma 1, the optimality gap after $R$ rounds satisfies

where $\alpha\!=\!\frac{8LE}{\mu^{2}K}$ and $\beta\!=\!\frac{2L}{\mu^{2}E}B+\frac{12L^{2}}{\mu^{2}E}\Gamma+\frac{4L^{2}}{\mu E}\left\|\mathbf{w}_{0}\!-\!\mathbf{w}^{*}\right\|^{2}$, with $B\!=\!\sum\limits_{i=1}^{N}p_{i}^{2}\sigma_{i}^{2}\!+\!8\!\sum\limits_{i=1}^{N}p_{i}G_{i}^{2}E^{2}\!$ and $\Gamma\!=\!F^{*}\!-\!\sum_{i=1}^{N}p_{i}F_{i}^{*}$.

Proof Sketch. First, following the similar proof of convergence under full client participation in , we show that

For FL with homogeneous communication time i.e., $t_{i}=t_{0}$, for all $i\in\mathcal{N}$, the optimal client sampling probabilities $\mathbf{q}$ for Problem P1 is

Thus, for a target precision $\epsilon$, computing the optimal sampling $\mathbf{q}$ for minimizing the upper bound of $R$ is equivalent to solving

The problem in (13) can be easily solved with the Lagrange multiplier method in closed form as shown in (11).

In the following, we show how to leverage the derived convergence bound in Theorem 1 to design the optimal client sampling for the general heterogeneous system of Problem P1.

V Optimal Adaptive Client Sampling Algorithm

We first show that the probability of client $i$ being the slowest one (e.g., straggler) amongst the $K$ sampled clients in each round is $(\sum\nolimits_{j=1}^{i}q_{j})^{K}\!-\!(\sum\nolimits_{j=1}^{i-1}q_{j})^{K}$. Since we sample devices according to $\mathbf{q}$, taking the expectation of all $N$ clients over time $q_{i}$ gives (15), and for $R$ rounds we have (14). ∎

V-B Approximate Optimization Problem for Problem P1

Based on Theorem 2, and by letting the analytical convergence upper bound in (9) satisfy the convergence constraint,Optimization using upper bound as an approximation has also been adopted in . the original Problem P1 can be approximated as

Combining with (9), we can see that Problem P2 is more constrained than Problem P1, i.e., any feasible solution of Problem P2 is also feasible for Problem P1.

We further relax $R$ as a continuous variable to theoretically analyze Problem P2. For this relaxed problem, suppose ($\mathbf{q}^{*}$, $R^{*}$) is the optimal solution, then we must have

This is because if (17) holds with an inequality, we can always find an $R^{\prime}<R^{*}$ that satisfies (17) with equality, but the solution ($q^{*}$, $R^{\prime}$) can further reduce the objective function value. Therefore, for the optimal $R$, (17) always holds, and we can obtain $R$ from (17) and substitute it into the objective of Problem P2. Then, the objective of Problem P2 is

whichFor ease of analysis, we omit $\epsilon$ as it is a constant multiplied by the entire objective function. is only associated with client sampling probabilities $\mathbf{q}$.

Case 1: For homogeneous $t_{i}$ ($t_{i}=t_{0},\forall i\in\mathcal{N}$), we have

Case 2: For heterogeneous $t_{i}$ with $K=1$, we have

Remark: The objective function of Problem P3 is in a more straightforward form compared to Problem P2. However, to solve for the optimal sampling probabilities $\mathbf{q}$, we need to know the value of the parameters in (22), e.g., $G_{i}$, $\alpha$, and $\beta$.We assume that clients’ heterogeneous time $t_{i}$ and their dataset size $p_{i}$ can be measured offline.

In the following, we solve Problem P3 as an approximation of the original Problem P1. Our empirical results in Section VI demonstrate that the solution obtained from solving Problem P3 achieves superior total wall-clock time performances compared to baseline client sampling schemes.

V-C Solving Problem P3

Problem P3 is challenging to solve because we can only obtain the unknown parameters $G_{i}$, $\alpha$ and $\beta$ during the training process of FL. In this subsection, we first show how to estimate these unknown parameters. Then, we develop an efficient algorithm to solve Problem P3. We summarize the overall algorithm in Algorithm 2. Finally, we identify some insightful solution properties.

We first show how to estimate $\frac{\alpha}{\beta}$ via a substitute sampling scheme.We only need to estimate the value of $\frac{\alpha}{\beta}$ instead of $\alpha$ and $\beta$ each, because we can divide parameter $\alpha$ on the objective of Problem P3 without affecting the optimal sapling solution. Then, we show that we can indirectly acquire the knowledge of $G_{i}$ during the estimation process of $\frac{\alpha}{\beta}$.

The basic idea is to utilize the derived convergence upper bound in (9) to approximately solve $\frac{\alpha}{\beta}$ as a single variable, via performing Algorithm 1 with two baseline sampling schemes: uniform sampling $\mathbf{q_{1}}$ with ${q}_{i}=\frac{1}{N}$ and weighted sampling $\mathbf{q_{2}}$ with ${q}_{i}=p_{i}$, respectively.

Note that we only let sampling schemes $\mathbf{q_{1}}$ and $\mathbf{q_{2}}$ run until a pre-defined loss $F_{s}$ is reached (instead of running all the way until they converge to the precision $\epsilon$), because our goal is to find and run with the optimal sampling scheme $\mathbf{q^{*}}$ so that we can achieve the target precision with the minimum wall-clock time.

Specifically, suppose $R_{\mathbf{q_{1}},s}$ and $R_{\mathbf{q_{2}},s}$ are the number of rounds for reaching the pre-defined loss $F_{s}$ for schemes $\mathbf{q_{1}}$ and $\mathbf{q_{2}}$, respectively. Considering that the training loss decreases quickly at the beginning and slowly when approaching convergence, the values of $R_{\mathbf{q_{1}}}$ and $R_{\mathbf{q_{2}}}$ are normally very small compared to the required number of rounds for reaching the target loss (with precision $\epsilon$). According to (9), we have

Then, we can obtain $\frac{\alpha}{\beta}$ from (24) once we know the value of $G_{i}$. Notably, we can estimate $G_{i}$ during the procedure of estimating $\frac{\alpha}{\beta}$. The idea is to let the sampled clients send back the norm of their local SGD along with their returned local model updates, and then the server updates $G_{i}$ with the received norm values. This approach does not add much communication overhead, since we only need to additionally transmit the value of the gradient norm (e.g., only a few bits for quantization) instead of the full gradient information. In addition, instead of retraining the model using the calculated $\mathbf{q}^{*}$ from the initial parameter $\mathbf{w_{0}}$, we can continue to train the global model after the estimation process, to avoid repeated training and reduce the overall training time.

In practice, due to the sampling variance, we may set several different $F_{s}$ to obtain an averaged estimation of $\frac{\alpha}{\beta}$. The overall estimation process corresponds to Lines 2–2 of Algorithm 2.

We first identify the property of Problem P3 and then show how to compute $\mathbf{q}^{*}$.

To solve Problem P3, we define a new control variable

where $t_{1}\leq M\leq t_{N}$. Then, we rewrite Problem P3 as

For any fixed feasible value of $M\in[t_{1},t_{N}]$, Problem P4 is convex with $\mathbf{q}$, because the objective function is strictly convex and the constraints are linear.

We will solve Problem P4 in two steps. First, for any fixed $M$, we will solve for the optimal $\mathbf{q}^{*}(M)$ in Problem P4, via a convex optimization tool, e.g., CVX . This allows us to write the objective function of Problem P4 as $g(\mathbf{q}^{\ast}(M),M)$. Then we will further solve the problem by using a linear search method with a fixed step-size $\epsilon_{0}$ over the interval $[t_{1},t_{N}]$, where we use the optimal $M^{*}(\epsilon_{0})$ and the corresponding $\mathbf{q}^{*}(M^{*}(\epsilon_{0}))$ in the search domain to approximate the optimal $M^{*}$ and $\mathbf{q}^{*}$ in Problem P4. This optimization process corresponds to Lines 2–2 of Algorithm 2.

Remark: Our optimization algorithm is efficient in the sense that the linear search domain $[t_{1},t_{N}]$ is independent of the scale of the problem, e.g., number of $N$.

) ​Property of Optimal Client Sampling

​​ Next we show some interesting properties of the optimal sampling strategy.

Suppose $\mathbf{q}^{*}$ is the optimal solution of Problem P3. For two different clients $i$ and $j$, if $t_{i}\leq t_{j}$ and $p_{i}G_{i}\geq p_{j}G_{j}$, then $q_{i}^{*}\geq q_{j}^{*}$.

Theorem 4 shows that the optimal client sampling strategy should allocate higher probabilities to those who have smaller $t_{i}$ and larger product value of $p_{i}G_{i}$, which characterizes the impact and interplay between system heterogeneity and statistical heterogeneity. Although it may be infeasible to derive an analytical relationship regarding the exact impact of $t_{i}$ and $p_{i}G_{i}$ on $\mathbf{q}^{*}$ due to non-convexity of Problem P3, we show by Corollary 2 that we can obtain the closed-form solution of $\mathbf{q}^{*}$ with $t_{i}$ and $p_{i}G_{i}$ when $\frac{\beta}{\alpha}\rightarrow 0$.

When $\frac{\beta}{\alpha}\rightarrow 0$, the global optimal solution of Problem P3 is

If $\frac{\beta}{\alpha}\rightarrow 0$, because $\sum\nolimits_{i=1}^{N}q_{i}t_{i}$ is bounded between $\left[t_{1},t_{N}\right]$, we have $\left(\sum\nolimits_{i=1}^{N}q_{i}t_{i}\right)\frac{\beta}{\alpha}\rightarrow 0$. Then, the objective of Problem P3 can be written as

Hence, the minimum of (28) is $\left(\sum_{i=1}^{N}\sqrt{t_{i}}\cdot{p_{i}G_{i}}\right)^{\!2}$, which is independent of $\mathbf{q}$. The equality of (29) holds if and only if $\sqrt{q_{i}t_{i}}=c\cdot\frac{p_{i}G_{i}}{\sqrt{q_{i}}}$ (for an arbitrary scalar $c$). Noting $\sum_{i=1}^{N}q_{i}=1$ yields (27) and concludes the proof.

Though Corollary 2 is valid only for the special case of $\frac{\beta}{\alpha}\rightarrow 0$, the global optimal sampling solution in (27) characterizes an analytical interplay between the system heterogeneity (${t_{i}}$) and statistical heterogeneity ($p_{i}G_{i}$). Particularly, when $t_{i}\!=\!t_{0}$, for each $i\!\in\mathcal{N}$, $\mathbf{q}^{*}$ in (27) recovers the optimal sampling solution in Corollary 1 for homogeneous systems.

VI Experimental Evaluation

In this section, we empirically evaluate the performance of our proposed client sampling scheme (Algorithm 2) and compare it with four other benchmarks in each round: 1) full participation, 2) uniform sampling, 3) weighted sampling, and 4) statistical sampling where we sample clients according to Corollary 1. Benchmarks 1–3 are widely adopted for convergence guarantees in . The fourth baseline is an offline variant of the proposed schemes in .The client sampling in is weighted by the norm of the local stochastic gradient in each round, which frequently requires the knowledge of stochastic gradient from all clients to calculate the sampling probabilities.

In the following, we first present the evaluation setup and then show the experimental results.

We conduct experiments both on a networked hardware prototype system and in a simulated environment.The prototype implementation allows us to capture real system operation time, and the simulation system allows us to simulate large-scale FL environments with manipulative parameters. As illustrated in Fig. 2, our prototype system consists of $N=40$ Raspberry Pis serving as clients and a laptop computer acting as the central server. All devices are interconnected via an enterprise-grade Wi-Fi router. We develop a TCP-based socket interface for the communication between the server and clients with bandwidth control. In the simulated system, we simulate $N=100$ virtual devices and a virtual central server.

) Datasets and Models

We evaluate our results on two real datasets and a synthetic dataset. For the real dataset, we adopted the widely used MNIST dataset and EMNIST dataset . For the synthetic dataset, we follow a similar setup to that in , which generates $60$-dimensional random vectors as input data. We adopt both the convex multinomial logistic regression model and the non-convex convolutional neural network (CNN) model with LeNet-5 architecture .

) Implementation

Prototype Setup: We conduct the first experiment on the prototype system using logistic regression and the EMNIST dataset. To generate heterogeneous data partition, similar to , we randomly subsample $33,036$ lower case character samples from the EMNIST dataset and distribute among $N\!=\!40$ edge devices in an unbalanced (i.e., different devices have different numbers of data samples, following a power-law distribution) and non-i.i.d. fashion (i.e., each device has a randomly chosen number of classes, ranging from $1$ to $10$).The number of samples and the number of classes are randomly matched, such that clients with more data samples may not have more classes.

Simulation Setup 1: We conduct the second experiment in the simulated system using logistic regression and the Synthetic dataset. To simulate a heterogeneous setting, we use the non-i.i.d. $Synthetic\ (1,1)$ setting. We generate $20,509$ data samples and distribute them among $N\!=\!100$ clients in an unbalanced power-law distribution.

Simulation Setup 2: We conduct the third experiment in the simulated system using CNN and the MNIST dataset, where we randomly subsample $15,129$ data samples from MNIST and distribute them among $N\!=\!100$ clients in an unbalanced (following the power-law distribution) and non-i.i.d. (i.e., each device has 1–6 classes) fashion.

) Training Parameters

For all experiments, we initialize our model with $\mathbf{w}_{0}\!=\!\mathbf{0}$ and use an SGD batch size of $b=24$. We use an initial learning rate of $\eta_{0}=0.1$ with a decay rate of $\frac{{\eta}_{0}}{1+r}$, where $r$ is the communication round index. We adopt the similar FedAvg settings as in , where we sample $10$% of all clients in each round, i.e., $K\!=\!4$ for Prototype Setup and $K\!=\!10$ for Simulation Setups, with each client performing $E=50$ local iterations.We also conduct experiments both on Prototype and Simulation Setups with variant $E$ and $K$, which show a similar performance as the experiments in this paper, and due to page limitations, we do not illustrate them all.

) Heterogeneous System Parameters

For the Prototype Setup, to enable a heterogeneous communication time, we control clients’ communication bandwidth and generate a uniform distribution $t_{i}\sim\mathcal{U}(0.187,7.159)$ seconds, with a mean of $3.648$ seconds and the standard deviation of $2.071$ seconds. For the simulation system, we generate the client transmission delays with an exponential distribution, i.e., $t_{i}\!\sim\!\exp{(1)}$ seconds, with both mean and standard deviation as $1$ second.

VI-B Performance Results

We evaluate the wall-clock time performances of both the global training loss and test accuracy on the aggregated model in each round for all sampling schemes. We average each experiment over $50$ independent runs. For a fair comparison, we use the same random seed to compare sampling schemes in a single run and vary random seeds across different runs.

Fig. 3–5 show the results of Prototype Setup, Simulation Setup 1, and Simulation Setup 2, respectively. We summarize the key observations as follows.

As predicted by our theory, Figs. 3(a)–5(a) show that our proposed sampling scheme achieves the same target loss with significantly less time, compared to the baseline sampling schemes. Specifically, for Prototype Setup in Fig. 3(a), our proposed sampling scheme spends around $73$% less time than full sampling and uniform sampling and around $66$% less time than weighted sampling and statistical sampling for reaching the same target loss. Fig. 5(a) highlights the fact that our proposed sampling works well with the non-convex CNN model, under which the naive uniform sampling cannot reach the target loss within $900$ seconds, indicating the importance of a careful client sampling design. Table I summarizes the superior performances of our proposed sampling scheme in wall-clock time for reaching target loss in all three setups.

) Accuracy with Wall-clock Time

As shown in Fig. 3(b)–5(b), our proposed sampling scheme achieves the target test accuracyIn Fig. 3(b), Fig. 4(b), and Fig. 5(b), the target test accuracy corresponds to the test accuracy result when our proposed scheme reaches the target loss. much faster than the other benchmarks. Notably, for Simulation Setup 1 with the target test accuracy of $75.3$% in Fig. 4(b), our proposed sampling scheme takes around $70$% less time than full sampling and around $46$% less time than the other sampling schemes. We can also observe the superior test accuracy performance of our proposed sampling schemes in Prototype Setup and non-convex Simulation Setup $2$ in Fig. 3(b) and Fig. 5(b), respectively.

) Loss with Number of Rounds

Fig. 3(c)–5(c) show that our proposed sampling scheme requires more training rounds for reaching the target loss compared to baseline statistical sampling and full participation schemes. This observation is expected since our proposed sampling scheme aims to minimize the wall-clock time instead of the number of rounds. Nevertheless, we notice that statistical sampling performs better than the other sampling schemes, which verifies Corollary 1 since the performance of loss with respect to the number of rounds is equivalent to that with respect to wall-clock time for homogeneous systems.

VII Conclusion and Future Work

In this work, we studied the optimal client sampling strategy that addresses the system and statistical heterogeneity in FL to minimize the wall-clock convergence time. We obtained a new tractable convergence bound for FL algorithms with arbitrary client sampling probabilities. Based on the bound, we formulated a non-convex wall-clock time minimization problem. We developed an efficient algorithm to learn the unknown parameters in the convergence bound and designed a low-complexity algorithm to approximately solve the non-convex problem. Our solution characterizes the interplay between clients’ communication delays (system heterogeneity) and data importance (statistical heterogeneity), and their impact on the optimal client sampling design. Experimental results validated the superiority of our proposed scheme compared to several baselines in speeding up wall-clock convergence time.

References