Heterogeneous Federated Learning via Grouped Sequential-to-Parallel Training

Under our grouped FL setting, FL clients are grouped such that clients in the same group have heterogeneous data but the overall data distributions among the groups are homogeneous. Due to the difference in data heterogeneity, the training modes within and among groups are designed separately. We refer to this jointly designed FL training mode as the “collaborative training mode”.

Intuitively, the homogeneous groups can be trained in a simple parallel mode PTM because the heterogeneity of their data has been eliminated by client grouping. Instead, for FL clients in the same group whose local data are still skewed, the sequential mode STM can be useful. In STM, FL clients train the model in a sequential manner. A FL client receives the model from its predecessor client and delivers the local trained model to its successor client to continue training. In the special case of training with only one local epoch (i.e., $e=1$), STM is equivalent to centralized SGD, which gives it robustness against data heterogeneity.

This naive collaborative training mode is static and has limitations. Therefore, we extend it to propose a more dynamic approach STP. As shown in Figure 1, STP reassigns FL clients into $f(r)$ groups and shuffles their order in each round $r$, where $f$ is a pre-specified group number growth function, with the goal to dynamically adjust the degree of parallelism. Then, STP can be smoothly transformed from (full) sequential mode to (full) parallel mode. This design can prevent catastrophic forgetting caused by the long “chain of clients” that causes the FL model to forget the data of previous clients and overfit the data of subsequent clients, and can also prevent the FL model from learning interfering information such as the order of clients. Moreover, the growing number of groups improves the parallelism efficiency, which promotes convergence and speeds up training when the global FL model is close to convergence.

The pseudo code of STP is given in Algorithm 1. In round $r$, STP divides all FL clients into $f(r)$ groups using the ICG grouping algorithm (Line 5), which will be described in Section 3.3. Due to the similarity of data among groups, each group can independently represent the global distribution, so only a small proportion of $\kappa$ groups are required to participate in each round of training (Line 7). The first FL client in each group pulls the global model from the FL server (Line 9), and trains its local model using mini-batch SGD for one epoch (Line 11). The trained local model is then delivered to the next FL client to continue training (Line 12), until the last FL client is reached. The last FL client in each group sends the trained model to the FL server (Line 14). Models from all groups are aggregated to update the global FL model (Line 16). The above steps repeat until the maximum training round $R$ is reached. Finally, the well-trained global FL model is obtained (Line 18).

The choice of the growth function for the number of groups, $f$, is critical for the performance of STP. We give three representative growth functions, including linear (smooth grow), logarithmic (fast first and slow later), and exponential (slow first and fast later) growth functions:

where the real number coefficient $\alpha$ controls the growth rate, and the integer coefficient $\beta$ controls the initial number of groups and the growth span. We recommend to initialize $\alpha$, $\beta$ to a moderate value and explore the best setting in an empirical manner.

In this section, we describe the FedGSP framework that enables dynamic STP. FedGSP is generally a grouped FL framework that supports dynamic group management, as shown in Figure 2. The basic components include a top server (which acts as an FL server and performs functions related to group assignment) and a large number of FL clients. FL clients can be smart devices with certain available computing and communication capabilities, such as smart phones, laptops, mobile robots and drones. They collect data from the surrounding environment and use the data to train local ML models.

In addition, FedGSP creates group managers to facilitate the management of the growing number of groups in STP. The group managers can be virtual function nodes deployed in the same machine as the top server. Whenever a new group is built, a new group manager is created to assist the top server to manage this group by performing the following tasks:

1. Collect distribution information. The group manager needs to collect class distributions of FL clients and report them to the top server. These meta information will be used to assign FL clients to $f(r)$ groups via ICG.

2. Coordinate model training. The group manager needs to coordinate the sequential training of FL clients in its group, as well as the parallel training with other groups, according to the rules of STP. Specifically, it needs to shuffle the order of clients and report resulting model to the top server for aggregation.

3. Schedule model transmission. In applications such as Industrial IoT systems, wireless devices can directly communicate with each other through wireless sensor networks (WSNs). However, this cannot be realized in most scenarios. Therefore, the group manager needs to act as a communication relay to schedule the transmission of ML models from one client to another.

As required by STP, the equally sized groups containing heterogeneous FL clients should have similar overall class distributions. To achieve this goal, in this section, we first formalize the FL client grouping problem which is NP-hard, and then explain how to simplify to propose the ICG approach.

Therefore, it is almost impossible to find the optimal solution within a polynomial time. To address this issue, we adopt the centroid equivalence theorem to simplify the original problem to a constrained clustering problem.

Consider a constrained clustering problem with $K$ points and $L$ clusters, where the size of all clusters is strictly the same $K/L$.

Proposition 2 indicates that there exists a grouping strategy $\tilde{\mathbf{x}}$ and $\mathcal{V}_{m_{1}}=\sum_{k=1}^{K}\tilde{\mathbf{x}}_{m_{1}}^{k}\mathcal{V}^{k}=LC^{m_{1}}$, $\mathcal{V}_{m_{2}}=\sum_{k=1}^{K}\tilde{\mathbf{x}}_{m_{2}}^{k}\mathcal{V}^{k}=LC^{m_{2}}$ ($\forall m_{1}\neq m_{2}$), so that the objective in Eq. (4) turns to $z=\sum_{m_{1}\neq m_{2}}L<C^{m_{1}},C^{m_{2}}>$ and the expectation value reaches 0. This motivates us to use the constrained clustering model to solve $\tilde{\mathbf{x}}$ in the objective Eq. (4). Therefore, we consider the constrained clustering problem below,

where $\mathbf{y}\in\mathbb{I}^{L\times K}$ is a selector variable, $\mathbf{y}_{l}^{k}=1$ means that client $k$ is assigned to cluster $l$ while 0 means not, $C_{l}$ represents the centroid of cluster $l$. Eq. (10) is the standard clustering objective, which aims to assign $K$ clients to $L$ clusters so that the sum of the squared $l_{2}$-norm distance between the class distribution vector $\mathcal{V}^{k}$ and its nearest cluster centroid $C_{l}$ is minimized. Constraint (11) ensures that each cluster has the same size $\frac{K}{L}$. Constraint (12) ensures that each client can only be assigned to one cluster at a time. In this simplified problem, Constraint (7) is relaxed to $\sum_{m=1}^{M}{\mathbf{x}_{m}^{k}}\leq 1$ to satisfy the assumption that $K/L$ is divisible.

The above constrained clustering problem can be modeled as a minimum cost flow (MCF) problem and solved by network simplex algorithms [19], such as SimpleMinCostFlow in Google OR-Tools. Then, we can alternately perform cluster assignment and cluster update to optimize $\mathbf{y}_{l}^{k}$ and $C_{l}(\forall k,l)$, respectively. Finally, we construct $M$ groups, each group consists of one client randomly sampled from each cluster without replacement, so that their group centroids are expected to coincide with the global centroid. The pseudo code is given in Algorithm 2. ICG has a complexity of $\mathcal{O}(\frac{K^{6}\mathcal{F}\tau}{M^{2}}\log{Kd})$, where $d=\max\{\sigma_{l}^{2}|\forall l\in[1,L]\}$, and $K,M,\mathcal{F},\tau$ are the number of clients, groups, categories, and iterations, respectively. In our experiment, ICG is quite fast, and it can complete group assignment within only 0.1 seconds, with $K=364,M=52,\mathcal{F}=62$ and $\tau=10$.

Environment and Hyperparameter Setup. The experiment platform contains $K=368$ FL clients. The most commonly used FEMNIST[11] is selected as the benchmark dataset, which is specially designed for non-i.i.d. FL environment and is constructed by dividing 805,263 digit and character samples into 3,550 FL clients in a non-uniform class distribution, with an average of $n=226$ samples per client. For the resource-limited mobile devices, a lightweight neural network composed of 2 convolutional layers and 2 fully connected layers with a total of 6.3 million parameters is adopted as the training model. The standard mini-batch SGD is used by FL clients to train their local models, with the learning rate $\eta=0.01$, the batch size $b=5$ and the local epoch $e=1$. We test FedGSP for $R=500$ rounds. By default, we set the group sampling rate $\kappa=0.3$, the group number growth function $f=\textsc{Log}$ and the corresponding coefficients $\alpha=2$, $\beta=10$. The values of $\kappa$, $f$, $\alpha$, $\beta$ will be further tuned in the experiment to observe their performance influence.

Benchmark Algorithms. In order to highlight the effect of the proposed STP and ICG separately, we remove them from FedGSP to obtain the naive version, NaiveGSP. Then, we compare the performance of the following versions of FedGSP through ablation studies:

1. 1.

   NaiveGSP: FL clients are randomly assigned to a fixed number of groups, the clients in the group are trained in sequence and the groups are trained in parallel (e.g., Astraea[9]).

2. 2.

   NaiveGSP+ICG: The ICG grouping algorithm is adopted in NaiveGSP to assign FL clients to a fixed number of groups strategically.

3. 3.

   NaiveGSP+ICG+STP (FedGSP): On the basis of NaiveGSP+ICG, FL clients are reassigned to a growing number of groups in each round as required by STP.

In addition, seven state-of-the-art baselines are experimentally compared with FedGSP. They are FedProx [12], FedMMD [13], FedFusion [14], IDA [15], and FedAdagrad, FedAdam, FedYogi from [16].

Evaluation Metrics. In addition to the fundamental test accuracy and test loss, we also define the following metrics to assist in performance evaluation.

Class Probability Distance (CPD). The maximum mean discrepancy (MMD) distance is a probability measure in the reproducing kernel Hilbert space. We define CPD as the kernel two-sample estimation with Gaussian radial basis kernel $\mathcal{K}$[20] to measure the difference in class probability (i.e., normalized class distribution) $\mathcal{P}=\mathrm{norm}(\mathcal{V}_{m_{1}}),\mathcal{Q}=\mathrm{norm}(\mathcal{V}_{m_{2}})$ between two groups $m_{1},m_{2}$. Generally, the smaller the CPD, the smaller the data heterogeneity between two groups, and therefore the better the grouping strategy.

Computational Time. We define $T_{\mathrm{comp}}$ in Eq. (15) to estimate the computational time cost, where the number of floating point operations (FLOPs) is $\mathcal{N}_{\mathrm{calc}}=96$M FLOPs per sample and $\mathcal{N}_{\mathrm{aggr}}=6.3$M FLOPs for global aggregation, and $\mathcal{T}_{\mathrm{FLOPS}}=567$G FLOPs per second is the computing throughput of the Qualcomm Snapdragon 835 smartphone chip equipped with Adreno 540 GPU.

Communication Time and Traffic. We define $T_{\mathrm{comm}}$ in Eq. (16) to estimate the communication time cost and $D_{\mathrm{comm}}$ in Eq. (17) to estimate the total traffic, where the FL model size is $\mathcal{M}=25.2$MB, the inbound and outbound transmission rates are $\mathcal{R}_{\mathrm{in}}=\mathcal{R}_{\mathrm{out}}=567\mathrm{Mbps}$ (tested in the Internet by AWS EC2 r4.large 2 vCPUs with disabled enhanced networking). Eq. (16) to Eq. (17) consider only the cross-WAN traffic between FL clients and group managers, but the traffic between the top server and group managers is ignored because they are deployed in the same physical machine.

Please note that Eq. (15) to Eq. (16) are theoretical metrics, which do not consider memory I/O cost, network congestion, and platform configurations such as different versions of CUDNN/MKLDNN libraries.

The effect of ICG and STP. We first compare the CPD of FedAvg[21], NaiveGSP and NaiveGSP+ICG in Figure 3a. These CPDs are calculated between every pair of FL clients. The results show that NaiveGSP+ICG reduces the median CPD of FedAvg by $82\%$ and NaiveGSP by $41\%$. We also show their accuracy performance in Figure 3b. The baseline NaiveGSP quickly converges but only achieves the accuracy similar to FedAvg. Instead, NaiveGSP+ICG improves the accuracy by $6\%$. This shows that reducing the data heterogeneity among groups can indeed effectively improve FL performance in the presence of non-i.i.d. data. Although NaiveGSP+ICG is already very effective, it still has defects. Figure 3c shows a rise in the loss value of NaiveGSP+ICG, which indicates that it has been overfitted. That is because the training mode of NaiveGSP+ICG is static, it may learn the client order and forget the previous data. Instead, the dynamic FedGSP overcomes overfitting and eventually converges to a higher accuracy $85.4\%$, which proves the effectiveness of combining STP and ICG.

The effect of the growth function $f$ and its coefficients $\alpha,\beta$. To explore the performance influence of different group number growth functions $f$, we conduct a grid search on $f=\{\textsc{Linear},\textsc{Log},\textsc{Exp}\}$ and $\alpha,\beta$. The test loss heatmap is shown in Figure 4. The results show that the logarithmic growth function achieves smaller loss 0.453 with $\alpha=2,\beta=10$ among 3 candidate functions. Besides, we found that both lower and higher $\alpha,\beta$ lead to higher loss values. The reasons may be that a slow increase in the number of groups leads to more STM and results in overfitting, while a rapid increase in the number of groups makes FedGSP degenerate into FedAvg prematurely and suffers the damage of data heterogeneity. Therefore, we recommend $\alpha\cdot\beta$ to be a moderate value, as shown in the green area.

The effect of the group sampling rate $\kappa$. $\kappa$ controls the participation rate of groups (also the participation rate of FL clients) in each round. We set $\kappa=\{0.1,0.2,0.3,0.5,1.0\}$ to observe its effect on accuracy and time cost. Figure 5a shows the robustness of accuracy to different values of $\kappa$. This is expected because ICG forces the data of each group to become homogeneous, which enables each group to individually represent the global data. In addition, Figures 5b and 5c show that $\kappa$ has a negligible effect on computational time $T_{\mathrm{comp}}$, but a proportional effect on communication time $T_{\mathrm{comm}}$ because a larger $\kappa$ means more model data are involved in data transmission. Therefore, we recommend that only $\kappa\in[0.1,0.3]$ of groups are sampled to participate in FL in each round to reduce the overall time cost. In our experiments, we set $\kappa=0.3$ by default.

The performance comparison of FedGSP. We compare FedGSP with seven state-of-the-art approaches and summarize their test accuracy, test loss and training rounds (required to reach the accuracy of $80\%$) in Table 1. The results show that FedGSP achieves $5.3\%$ higher accuracy than FedAvg and reaches the accuracy of $80\%$ within only $34$ rounds. Moreover, FedGSP outperforms all the comparison approaches, with an average of $3.7\%$ higher accuracy, $0.123$ lower loss and $84\%$ less rounds, which shows its effectiveness to improve FL performance in the presence of non-i.i.d. data.

The time and traffic cost of FedGSP. Figure 6 visualizes the time cost and total traffic of FedGSP and FedAvg when they reach the accuracy of $80\%$. The time cost consists of computational time $T_{\mathrm{comp}}(R)$ and communication time $T_{\mathrm{comm}}(R)$, of which $T_{\mathrm{comm}}(R)$ accounts for the majority due to the huge data traffic from hundreds of FL clients has exacerbated the bandwidth bottleneck of the cloud server. Figure 6 also shows that FedGSP spends $93\%$ less time and traffic than FedAvg, which benefits from a cliff-like reduction in the number of training rounds $R$ (only 34 rounds to reach the accuracy of $80\%$). Therefore, FedGSP is not only accurate, but also training- and communication-efficient.