Client Selection for Generalization in Accelerated Federated Learning:
A Multi-Armed Bandit Approach

Previous studies have shown that client selection has a great impact on the model’s performance and the training latency. The standard client selection method is to choose the clients randomly with probability proportional to the amount of data each client contains [2, 22]. The clients can differ from each other in several factors such as the size of their database, the amount of available resources, their channel state, the energy they can invest, and the value of their local loss function. Each iteration in FL systems ends when the last scheduled client uploads its updated model. Thus, variability between clients might lead to poor performance of the standard random selection method. As a result, more recent methods used observations of one or more of the aforementioned factors to infer the iteration latency of each client and subsequently select the appropriate group that leads to efficient training latency [23, 24, 25, 26, 27, 28]. In [23, 24, 25, 26] the focus was on the effect of client sampling on the learning performance. In [27, 28] the problem is investigated under the assumption that parameters related to the communication link and latency are known (e.g., channel state, transmission latency). As a result, the optimization problem in those studies is deterministic with respect to the client selection strategy. However, in many scenarios, client latencies are affected by other random factors as discussed earlier (e.g., client energy state, available computational resources, random transmission rate), and are unknown at the time of the client selection. Therefore, recent studies (including this paper) tackle the problem using a fundamentally different approach by treating clients’ latencies as unknown random variables drawn from unknown distributions, which leads to a stochastic optimization problem that raises the well-known exploration versus exploitation dilemma. On one hand, it is necessary to explore different actions (i.e., client selection sets) to explore the system state. On the other hand, it is imperative to exploit the information gathered so far to converge to the optimal selection strategy. This approach involves modeling and analyzing the problem as a MAB problem by treating the iteration latency or the local loss function of each client as a reward taken from an unknown distribution and given to the server (modeled as the gambler in MAB) [29, 30, 31, 32, 33]. In [29], the authors developed a naive MAB method as in the classic i.i.d. MAB that converges to a strategy that selects a small subset of the quickest clients (i.e., with the smallest expected latency). This method, however, suffers from over-fitting to the quickest clients’ data when applying to scheduling in FL systems. In [30, 31, 32], the focus was on reducing the loss, but the generalization issue remained open. Therefore, in [33] the authors suggested to use a fairness constraint to ensure that each client is selected in a certain proportion of the communication rounds. However, this method still converges to the group of the quickest clients except for the times of selecting the slower clients due to the fairness constraint. Therefore, we suggest a novel MAB formulation for the client selection problem consisting of a reward function that captures the generalization of the ML task. This allows us to solve the problem rigorously in this paper.

To solve the client selection problem in FL, it is needed to solve an online learning problem with the well-known exploration versus exploitation dilemma. On the one hand, the player (i.e., the server) should explore all arms (i.e., clients) in order to learn their state (i.e., latency distribution, generalization ability) which affects the overall training time of the FL task. On the other hand, it should exploit the information gathered so far to select the most rewarding subset of arms at each given time. In this paper we develop a rigorous MAB formulation to model this problem, and a novel learning algorithm to solve it. We provide rigorous theoretical analysis of the algorithm, as well as extensive numerical analysis. Below, we summarize the main contributions.

• •

  A novel MAB formulation for the client selection problem: We present a new MAB formulation that trades-off between the training latency and the generalization of the model in FL by client selection. This trade-off raises new challenges in the learning design. On the one hand, it is desired to reduce the training process time of the model, and on the other hand, it is desired to increase the ability of the model to generalize by avoiding over-fitting. The novel MAB formulation tackles this trade-off by selecting clients based on a time-varying reward influenced by the history of previous selections.

• •

  Algorithm development: We develop a novel upper confidence bound (UCB)-based algorithm to solve the problem, dubbed Bandit Scheduling for FL (BSFL). BSFL presents a fundamentally different approach of selecting the clients, as the reward function updates the contributions of each client to the FL task based on the history of previous selections. As a result, BSFL does not aim to converge to selecting a fixed subset of clients, but rather aims to increase the ability of the model to generalize. Furthermore, we provide concrete examples of the use of BSFL in both cases where the data is i.i.d. and balanced across clients, and where the data is non-i.i.d. and imbalanced. Since the UCB optimization in BSFL requires to solve a combinatorial problem to select the most rewarding subset of clients at each given time, the time complexity increases exponentially with the number of clients, which might be infeasible for a large number of clients. Hence, we utilize the optimization problem’s structure to devise a low-complexity algorithm, named accelerated lightweight simulated annealing (ALSA), which relies on a newly proposed accelerated lightweight simulated annealing technique that we develop for BSFL. Specifically, we show analytically that ALSA reduces the degree-order in the simulated annealing graph from $O(K^{2})$ to $O(K)$ (where $K$ is the number of clients), and still keeps the convergence property. This results in a significant accelerated convergence time requires to solve the UCB optimization in BSFL.

• •

  Performance analysis: We analyze the performance of BSFL theoretically and numerically. In terms of theoretical analysis, as commonly done in the MAB literature, we evaluate the performance by regret, defined as the loss of BSFL as compared to a genie that has complete knowledge about the latency means of all clients. We analyze the performance of BSFL rigorously, and show that it achieves a logarithmic regret with time. In terms of numerical analysis, we present extensive simulations using synthetic and real datasets. All simulations demonstrate that BSFL is superior to existing methods in both the regret and the prediction accuracy.

MAB problems are often illustrated with the example of a player or gambler facing a row of slot machines, selecting arms to pull at each time and obtaining rewards accordingly. In our context, the server acts as the player, and the clients act as the arms. At each iteration, the server selects a subset of clients, which contribute to the learning rate accordingly. Solving the MAB problem, and specifically the client selection problem in this paper, requires addressing the well-known exploration versus exploitation dilemma in online learning. On one hand, the server should explore all clients to learn their state, such as their latency distribution and generalization ability, which affect the overall training time of the FL task. On the other hand, the server should exploit the information gathered so far to select the most rewarding subset of clients at each given time.

The exploration versus exploitation dilemma has been rigorously addressed in MAB optimization [34]. As a result, recent studies have modeled and analyzed the client selection problem in FL as a MAB problem, treating the iteration latency or local loss function of each client as a reward obtained from an unknown distribution and given to the server [29, 30, 31, 32, 33], as discussed in Section I-A. However, all these methods converge to selecting the fastest clients (except for the times of selecting slower clients due to the fairness constraint in [33]). Consequently, they suffer from overfitting to the data of the quickest clients when applied to scheduling in FL systems, resulting in decreased ability to generalize and decreased prediction accuracy.

In this paper, we overcome this limitation by developing a novel MAB formulation for the client selection problem consisting of a time-varying reward function that captures both the client latency as well as the client ability to contribute to the generalization of the ML task. It should be noted that previous work on rotting bandits [35] considered a reward function that decreases with the activation rate of selected arms. However, the model here is fundamentally different, since the time-varying reward function depends on the overall subset selections of correlated clients in the systems. This leads to a novel combinatorial MAB (CMAB) formulation [36], in which the design and analysis is fundamentally different from rotting bandits. Specifically, our formulation trades-off rigorously between the generalization and latency, as each client in FL systems might have different characteristics such as computing power, task operations, channel states, etc. This results in different latency distributions in training and transmitting the local models across clients in the FL system [4]. On the one hand, to reduce the total iteration time at each iteration, it is desired to select clients with small latencies at each iteration. On the other hand, it is desired to select as many clients as possible during the training process (even slower clients) for the model to be robust and generalized for both i.i.d. balanced data and non-i.i.d. imbalanced data [37, 38, 33, 39]. In the next susbection we describe the generalization function used in the CMAB formulation.

For simplicity and maintaining generality, we denote by $\Theta$ all system constant parameters that contribute to the reward (which will be described later with examples). The level of contribution of client $k$ to the generalization ability and robustness of the model is evaluated by a generalization function $g_{k,\Theta}:\mathcal{H}_{t}\to[-1,1]$. It gives at each iteration $t$ a value that represents the contribution of client $k$ to the generalization of the global model, depending on the model constant parameters $\Theta$ and the selection history before time $t$, $\mathcal{H}_{t}$. We next present concrete examples of the function $g_{k,\Theta}$ in both cases of i.i.d. and balanced data, and non-i.i.d. and imbalanced data across clients.

We start by considering the case of i.i.d. balanced data. We use the selection history to extract a counter $c_{k,t}$ for each client that counts the number of iterations that client $k$ was selected until the beginning of iteration $t$, i.e., $c_{k,t}=\sum_{i=1}^{t-1}\mathbbm{1}_{\{k\in\mathcal{A}_{i}\}}$. In order to achieve high generalization in this case, it is important that each client is selected an equal number of times [37, 38, 22, 33, 39]. In this case, we let $\Theta$ store the number of clients and the number of channels, i.e., $\Theta=\{K,m\}$. Then, an effective design of the function $g_{k,\Theta}$ is given by:

where $\beta$ is a tuning parameter (natural number), and $sgn(\cdot)$ is the sign function that returns $\pm 1$. An illustration is presented in Fig. 2.

The key idea behind the design is that the more the ML model is trained with new data points or alternatively with data that was used too little, the more its ability to generalize increases [37, 39]. It is thus desired to select clients uniformly with rate $m/K$ over time. Therefore, the generalization function $g_{k,\Theta}$ is designed to be monotonically decreasing so that it provides incentive (i.e., positive value in terms of contributing to generalization) to select clients that were selected too little (i.e., $\frac{c_{k,t}}{t}<m/K$), and returns zero when the client’s counter reaches the desired selection rate (i.e., $\frac{c_{k,t}}{t}=m/K$). Clients that were selected too frequently are given negative values in terms of contributing to generalization.

In the second scenario, consider non-i.i.d. and imbalanced data, where the data quality varies between clients. In this case, it is desired to take into account the participation rate of each client, as well as the quantity and quality of the data, when selecting clients for transmission. The quantity of each client’s data (i.e., its database size $|X_{k}|$) is useful for the generalization, as it more likely to represent the true distribution of data [40]. The second aspect stems from the desire to train the model on reliable, high-quality, and meaningful data [41, 42, 43, 44]. Generally, in ML, and particularly FL, the data for each client might have different quality due to several factors such as device heterogeneity, environment heterogeneity, etc. We define the client’s data quality by $q_{k}\in[0,1]$, where $1$ stands for the best quality, and thus the generalization function $g_{k,\Theta}$ is designed to prioritize clients with higher data quality. Since the quality for each client (say client $k$) is aggregated over all its data points $X_{k}$, we denote the overall significance of client $k$ by $d_{k}=q_{k}\cdot|X_{k}|$. Therefore. the model parameters in this case is given by $\Theta=\bigl{\{}K,m,\{(|X_{k}|,q_{k}):k\in\mathbb{K}\}\bigl{\}}$, and the generalization function is given by:

In the next subsection we present the CMAB formulation for the client selection problem. In the empirical study in this paper, we used the forms of $g_{k,\Theta}$ as detailed above, which demonstrated very good performance. It should be noted, however, that the CMAB formulation is general. It does not depend on a specific form of $g_{k,\Theta}$, and other forms can be used. For example, it is possible to define $g_{k,\Theta}$ to be dependent on the local loss value of each client, which is advantageous in improving the convergence in the long term [45].

We now formulate the client selection problem as a CMAB problem with history-dependent reward. At each iteration $t$, the server (i.e., the player) selects a subset of clients $\mathcal{A}_{t}\in H(\mathbb{A}_{t})$ (i.e., arms), and at the end of the iteration observes the clients’ iteration latencies. The reward given at the end of iteration $t$ is defined by:

where $\alpha$ is a hyper-parameter of the generalization, balancing between the amount of the generalization and the iteration time, $\tau_{min}$ is the minimum iteration latency possible, which is the time needed to send and receive the model, as if the client required zero training time, and $\tau_{max}$ is the fixed latency bound as defined before (2). For purposes of analysis, we further assume that the reward given at the end of each iteration is quantized and the smallest difference between two different rewards is denoted by $\Delta_{min}$, as commonly assumed in analyzing MAB-based problems. Since the iteration time is determined by the maximal latency under the selected subset of clients $\mathcal{A}_{t}$, the first term on the RHS of (5) represents the reward due to the iteration time, and the second term represents the reward due to the generalization. The smaller the value of $\alpha$, the higher priority to select faster clients. On the one hand, this reduces the iteration time. On the other hand, it tends to select the same faster clients at each iteration, which reduces the generalization, and increases the total number of iterations needed to achieve reliable predictions.

An important insight from the CMAB formulation is that if a slow client is given high priority based on the reward function and is selected for transmission at the current iteration, then since the iteration time is determined by the slowest client, there is no motivation to choose fast clients at the same iteration. Instead, clients with high rewards in terms of generalization, even if they are slow, should be prioritized.

Note that the formulation of our CMAB problem is fundamentally different from the classical CMAB problem. Specifically, here the reward does not depend on the new observations solely (as in classic CMAB), but also on the selection history up to the current iteration. Thus, optimizing the reward in classic MAB leads to a fixed selection of arms, while in our problem the optimal selection is time-varying depending on the selection history, as will be detailed in the next subsection.

We aim to find the best selection policy $\Pi=\{\mathcal{A}_{1},\mathcal{A}_{2},...\}$ (where $\mathcal{A}_{t}$ stands for the selection at iteration $t$) which minimizes the training process time while preserving the ability of the model to perform the generalization. The performance of online learning algorithms are commonly evaluated by the regret, defined as the loss of an algorithm as compared to genie with side information on the system. Here, to measure the performance of our proposed policy we define the regret as follows. We denote the maximal mean reward at iteration $t$ that can be obtained by genie that has complete knowledge about the latency means of all clients by $r^{*}(t)$, where genie’s selection at time $t$ is given by:

where $\mu_{k}$ stands for the speed mean of client $k$ (i.e., $\mu_{k}=\mathbb{E}\bigl{[}\frac{\tau_{min}}{\tau_{k}}\bigl{]}$).

The regret of the policy (say $\Pi$) at time $n$ is defined by the accumulated loss by time $n$ between the reward obtained by genie and the expected reward obtained by $\Pi$:

The objective is thus to find a policy $\Pi$ that minimizes the regret order with time.

The BSFL algorithm is based on a novel UCB-type design for client selection in FL. The pseudocode for the BSFL algorithm is provided in Algorithm 1. We now discuss the steps of the BSFL algorithm in detail.

BSFL observes realizations of the random speed of the selected clients (i.e., $\frac{\tau_{min}}{\tau_{k,t}}$) and estimates its mean accordingly. Let

be the sample-mean speed of client $k$ after $t$ iterations. We design the UCB function of client $k$’s speed after $t$ iterations by:

To maintain an updated $\overline{\mu_{k,t}}$ and $\mbox{ucb}(k,t)$ for each client, each time a client is selected to participate in the FL iteration, the algorithm observes its speed $\frac{\tau_{min}}{\tau_{k,t}}$ and updates its counter and the speed’s sample-mean as follows:

Then, for each client $k\in\mathbb{K}$, the UCB function is updated as follows:

At the initialization step, for each client $k\in\mathbb{K}$, $c_{k,0}$ and $\overline{\mu_{k,0}}$ are set to 0, and $\mbox{ucb}(k,0)$ is set to infinity. Later, in the main loop, the algorithm selects clients at each iteration according to:

Afterwards, for each client $k\in\mathcal{A}_{t}$, the algorithm observes $\tau_{k,t}$ and updates $c_{k,t}$, $\overline{\mu_{k,t}}$ using (10), (11). At the end of every iteration, $g_{k,\Theta}(\mathcal{H}_{t+1})$ and $\mbox{ucb}(k,t)$ for all client $k\in\mathbb{K}$ are updated using (12). Note that two different trade-off mechanisms are manifested during the algorithms. The first is between exploration and exploitation of the client latencies, while the second is between the iteration latency and the generalization ability.

In this subsection, we analyze the regret (7) achieved by BSFL analytically, and show that it has a logarithmic order with time. Let $\Delta_{max}$ be the maximum possible difference between the expected reward obtained by genie and by any selection $\mathcal{S}$.

To evaluate the regret of BSFL, we define $r(\mathcal{S},\mathcal{H}_{t})$ as the reward that could have been obtained at the end of iteration $t$, given the selection history $\mathcal{H}_{t}$, if selection $\mathcal{S}$ had been made. Let $\Delta_{max}$ be the maximum difference between the expected reward obtained by $\mathcal{G}_{t}$ and by any selection $\mathcal{S}$, given any selection history, i.e.,

Note that $\Delta_{max}$ is bounded by:

where $\mu_{max}$ and $\mu_{min}$ are the expected speeds of the fastest and slowest clients, respectively.

The proof can be found in Appendix A.

Theorem 1 implies that BSFL achieves a logarithmic regret order with time $O(\ln n)$.

It is worth noting that the selection policy of Algorithm 1, i.e., finding the maximum value in the update rule stated in (13), can be computationally infeasible due to its exponential complexity. Therefore, in the following section, we will develop a practical solution to solve the maximization problem.

To implement SA-type algorithm to solve (13), we need to define a multi-state environment, in which each state $s$ has neighboring states denoted as $N_{s}$. Each state also has an associated energy, represented by $E(s)$. The objective of the algorithm is to find the state with the highest energy (or, alternatively, the lowest energy, depending on the problem formulation). To facilitate movement between states, a temperature parameter $T_{i}$ for time step $i$ is introduced. This temperature affects the probability of transitioning from the current state to a state with lower energy. The algorithm begins by randomly selecting an initial state $s_{0}$. It then proceeds through a series of time steps, during which the temperature is updated and a neighboring state $u\in N_{s_{i}}$ is chosen at random. The next state $s_{i+1}$ is then updated according to:

In our MAB-based FL scheduling setting, each client selection determines a state, which means that in each FL round $t$, the number of states of the stochastic optimization problem would be ${|\mathbb{A}_{t}|\choose m}$. The energy of each state would be:

As previously mentioned, the SA dynamics requires that each state has neighboring states. A direct implementation of the classic SA method results in a structure where any two client selections (which represent two states), $\mathcal{S}$ and $\mathcal{U}$, will be considered neighboring selections only if they differ in only one client. Formally,

where $m$ is the number of channels, same as before. It can be seen that through this construction, the neighborhood is symmetrical between the selections (states) and that for each selection there are $m(K-m)$ neighbors. From [46], setting the temperature at each time step to $T_{i}=\frac{\Delta_{max}}{\log{(i+1)}}$ guarantees convergence to the selection with the maximum energy as the number of time steps in the SA algorithm approaches infinity, where $\Delta_{max}$ is defined in (14). Despite the strong theoretical convergence guarantee, a recognized disadvantage of the SA method is its rate of convergence which can be quite slow. Thus, in the subsequent section, we present a novel SA-type algorithm that addresses this issue while still maintaining the theoretical convergence guarantee.

Building on the concept of SA-type dynamics, we present a novel accelerated SA-type method that capitalizes on the specific structure of the MAB-based FL scheduling problem. Our proposed method accelerates the optimization process by designing a lightweight search space, resulting in faster convergence compared to the classic SA technique. This is achieved by taking into account the unique features of the MAB-based FL scheduling problem.

When applying the classic SA method to our problem, the number of neighbors for each state is $m(K-m)$, which becomes $O(K^{2})$ when $m=O(K)$. In our proposed ALSA method, we exploit the characteristics of the multi-armed bandit optimization to reduce the number of connections between states and create a lightweight state graph for the search space. This modification allows for faster convergence while still maintaining the theoretical convergence guarantee of the SA method (which will be shown later). Specifically, we define a new neighborhood rule that only allows for client selections that differ by a single client from the current selection, and requires that this client has the lowest ucb value or the lowest $g_{k,\Theta}$ value in one of the selections. Formally,

It is important to note that each selection $\mathcal{S}$ has at most $2(K-m)$ other selections that differ by only one client $k$ which is in the set

$\Bigl{\{}\operatorname*{arg\,min}_{k\in\mathcal{S}}{\mbox{ucb}(k,t)},\operatorname*{arg\,min}_{k\in\mathcal{S}}{g_{k,\Theta}(\mathcal{H}_{t})}\Bigr{\}}$.

The neighborhoods in this formulation are symmetrical, and the average number of neighbors for each selection is no more than $4(K-m)$, resulting in a total of $O(K)$ neighbors regardless of the value of $m$.

Denote the state with the global maximum energy as $s^{*}$. The theoretical convergence of ALSA is shown next.

The proof can be found in Appendix B.

Theorem 2 implies that as the number of time steps for running ALSA increases, the result of executing ALSA will converge to the optimal solution of (13). By using ALSA, the execution of line 3 in the BSFL algorithm (Algorithm 1) becomes efficient.

The classic SA method, when applied to our problem, results in a graph structure in which each state has $O(K^{2})$ neighbors. However, by exploiting the characteristics of the MAB-based optimization, our proposed ALSA method is able to reduce the number of connections between states to create a lightweight graph for the search space, while still maintaining the theoretical convergence guarantee as previously demonstrated. This construction results in a graph structure in which each state has $O(K)$ neighbors.

To demonstrate the effectiveness of our proposed method, we conducted thousands of runs with varying numbers of clients and selection sizes. In the vast majority of runs (98.3%), ALSA reached a state with a higher energy within the fixed time period compared to the classic SA method. Furthermore, we observed that the majority of the additional edges in the classic SA method are between low-energy states or states with similar energy, which do not contribute to the convergence to the global maximum state and instead cause the algorithm to wander for a longer time between these low-energy states. This is illustrated in Figure 3, which shows a comparison between the graph structures created by ALSA and SA for a simulation of selecting 4 clients out of 8 with drawn $\mbox{ucb}(k,t)$ and $g_{k}$ values. As can be observed in Fig. 4, the judicious removal of excess edges in the graph created by ALSA leads to a significantly faster rate of convergence to high-energy states compared to SA.