LEAP: Optimization Hierarchical Federated Learning on Non-IID Data with Coalition Formation Game

1 Introduction

As a novel distributed machine learning paradigm, FL [McMahan et al., 2017] has gained the attention of many fields, such as the Internet of Things (IoT) [Rahman et al., 2023], smart transportation [Pandya et al., 2023], and healthcare [Zhang et al., 2023], to break down the information silos while enabling privacy preservation. With the power of FL, Artificial Intelligence (AI) can effectively handle machine learning tasks involving decentralized data, which draws upon the advantages of distribution machine learning, while also significantly eliminating the privacy risks. During the training process, only model parameters are transmitted without involving local data, breaking information silos and greatly improving training efficiency [Wang et al., 2023]. However, FL performance is affected by various factors, both in the training and transmission phases [Xu et al., 2018]. During the training phase, FL involves hundreds and thousands of clients, and the data distribution of each client is severely different due to diverse user behavior patterns and data collection methods. Consequently, the local data of an individual client fails to represent the overall data distribution of the environment, leading to significant reductions in model performance and compromising the model's generalization capability. In the transmission phase, there is high communication latency and instability between clients and central server (CS) by mass and frequent data transfers, reducing model training efficiency and even increasing the risk of data leakage.

∗The full version is at https://arxiv.org/abs/2405.00579.

†Corresponding Author.

In recent years, numerous solutions have been proposed to address data heterogeneity and communication bottlenecks. For a FL task, the client's data distribution is very important to the performance of the FL model. Significant deviations in the data distribution among clients can severely impair FL model learning and performance. For example, non-IID data causes a number of FL model performance issues, including decreased accuracy, sluggish model convergence, and model communication delays. Reinforcement learning [Xia et al., 2020] and data augmentation [Chen et al., 2023] have been proposed to address the non-IID challenges. These studies overemphasize the importance of individual clients, ignoring the performance improvement of the model benefited from aggregating local updates. In addition, despite their significance, most of these approaches require auxiliary models or extra data transmissions in FL, potentially introducing additional complexities. To alleviate communication pressures, various techniques such as model compression [Zhu et al., 2023], gradient sparsity [Lin et al., 2023], and over-the-air computation [An et al., 2023] were proposed. Although these methods can effectively reduce communication overhead, they may still result in bottlenecks in communication at the CS each communication round. This is due to the fact that during some training epochs, the CS still receives large model weight updates.

Inspired by [Liu et al., 2020], our aim is reduce the impact of non-IID distribution on HFL training by increasing the degree of IID of data distribution between ESs. Additionally, we seek to further optimize the resource allocation scheme to reduce the communication latency. The combination of these two objectives gives rise to an extremely complex and difficult problem, which presents the following three challenges: First, different edge association relationships represent different edge data distributions. Once edge association changes, the data distribution will evolve in an unpredictable direction, potentially reducing or increasing the degree of edge IID. Therefore, the impact of changes in association relationships on changes in data distribution is vague and uncertain. Second, straggler problem caused by worst-performing client. Edge aggregation latency is susceptible to the communication performance of the worst-performing client in synchronized FL. Dynamic edge association relationships make it difficult to capture communication performance information for each edge coalition. As a result, targeted resource optimization is impossible. Third, contradiction between task execution latency requirements and clients' energy consumption. Sufficient resource investment can meet the task requirements but may cause excessive overhead on clients, which is impractical. Consequently, striking a balance that satisfies the needs of both parties simultaneously poses a formidable challenge.

To tackle the abovementioned challenges, we propose a novel optimization method for HFL based on coaLition formation gamE and grAdient Projection method, named LEAP, which not only effectively reduces the impact of cross-edge non-IID, but also improves the communication efficiency. The main contributions of our work are as follows:

• Theoretically, we frame a complex optimization problem that focuses on the effect of multi-dimensional properties (i.e., time delay, energy consumption, and data distribution) on the performance of HFL. We strategically decouple this problem by transforming data distribution optimization into edge correlation analysis and further optimizing heterogeneous resource allocation.
• Methodologically, we construct a coalition formation game by analyzing the relationship between edge association and edge data distribution similarity, and prove the existence of stable coalitions. Moreover, we utilize the gradient projection method to calculate the optimal bandwidth allocation for each coalition, and further determine the transmission performance of heterogeneous clients to ensure that the latency requirements of the tasks are met.
• Experimentally, we validate the effectiveness of LEAP on four real datasets and baselines, it is able to achieve 20.62% improvement in accuracy compared to Meanshift algorithm. Moreover, experiments demonstrate that our optimization method is able to reduce the transmission energy consumption by at least 2.4 times while ensuring that the maximum latency requirement is met.

2 RelatedWork

In this section, we briefly discuss related work on non-IID data and communication bottleneck in FL.

Non-IID Data. In FL, the non-IID data, caused by the heterogeneity among clients, poses a challenge for training robust FL models due to its impact on slowing down the convergence of the global model [Yao et al., 2019]. Many efforts have been made to address this issue. For instance, Arisdakessian et al. [Arisdakessian et al., 2023] proposed a trust-based coalitional FL approach, which mitigates non-IID problems by sharing data of coalition masters. Lu et al. [Lu et al., 2023] utilized the Mean-shift algorithm to cluster clients according to data distribution and then selected clients from different clusters to participate in training. Shin et al. [Shin et al., 2020] proposed a novel approach that uses XorMixup hybrid data enhancement technology. This approach generates synthetic yet realistic sample data on the server, aiming to solve the issue of unbalanced training datasets in one-shot FL. However, it will bring a large computational burden.

Data-sharing operation raises privacy concerns for clients, thus limiting its application scenarios and it is difficult to operate under privacy-preserving FL. The method of selecting clients through clustering does mitigate the non-IID problem, but it does not guarantee that the final selection result is optimal. In contrast, our work can find optimal edge association relationship in no additional model training without considering the raw data leakage problem.

Communication Bottleneck. In FL, communication cost is a significant factor that affects overall efficiency and effectiveness, while the uplink transmission rate of the underlying client is a major bottleneck in the training process. Many researchers proposed related solutions. For example, Mills et al. [Mills et al., 2019] focused on enhancing communication efficiency in FL by combining a distributed Adam optimizer with a compression technique. They emphasized reducing the uploaded data size during training rounds to mitigate communication costs. Building upon of model compression, Liu et al. [Liu et al., 2021] applied it to wireless FL to alleviate local computation and communication bottlenecks.

The abovementioned studies were conducted on cloudbased FL systems, whereby the CS receives the local model from the clients. However, in cloud-based FL, the transmission distance can often be considerable, resulting in unstable and undependable communication among the clients and CS. In our study, we make full use of abundant bandwidth resources of ESs and design communication optimization method for heterogeneous clients to solve the resource allocation problem in HFL, as well as to improve the effectiveness of HFL in heterogeneous client environments.

3 SystemModel and ProblemFormulation

In this section, we introduce the workflow of HFL, refine its multidimensional properties, and give an explicit definition of the optimization problem.

3.1 HFL Framework

We consider a HFL framework that consists of a set N = {1, · · · ,N} of clients, a setM= {1, · · · ,M} of ESs, and a CS. The data set of the client n is denoted as Dn = {Xn, Yn}, where Xn = {xn,1, · · · , xn,Dn } are the training dataset, Yn = {yn,1, · · · , yn,Dn } are the corresponding label set, and |Dn| is the number of training data owned by the client n. The CS aims to train a model, with parameters denoted by a vector ω, over K iterations to minimize the global loss LK (ω). The ESs are employed to facilitate the uplink transmission of parameter updates by distributing orthogonal resource blocks to their clients. Then, each client can only associate with one ES to perform the model training. We define Gm as the set of clients that associate with ES m, i.e., Gm = {n ∈ N : am,n = 1}, where A = [a]M×N is the edge association matrix and Gm ∩ Gm′ = ∅ for m ̸= m′ . The HFL iteration i consists of four main steps as follows [Ng et al., 2022]:

• Local Training: Each client receives the intermediate model from ES m denoted by ωi , to train a local model using its dataset.
• Local Model Parameter Transmission: After every τc rounds of local updates, clients transmit the updated local model ωi,τc n,m to the associated ES m.
• Edge Aggregation: ES m aggregates the local model parameters from its associated clients to derive the intermediate model ωi,1 m , which is transmitted back to the clients for the next edge iteration.
• Global Aggregation: At the end of predefined intervals τe, each ES transmits the intermediate model ωi,τe m to CS for aggregation to derive the updated global model ωi+1 and transmits a new global model back to clients for the next global iteration.

The entire process described above will continue until a predetermined number of global training rounds τg is reached.

3.2 Multi-Dimensional Properties in HFL

The efficiency and sustainability of HFL are affected by execution time, energy cost and data quality, and these three comprehensive properties consider the impacts of different aspects on FL systems. We hence give a formal definition of each property as follows:

Definition 1. The muti-dimensional properties of FL V are represented as a 3-tuple (T , E, J ), i.e., execution time T , energy consumption E, and data distribution similarity J .

• T is the execution time of a task, which includes computation latency T C and communication latency T U , i.e.,



where t is an edge iteration, cn is the number of CPU cycles for training unit data, fn is the CPU cycle frequency that determines the computational power, and Z is the model size. Clients upload local models to ESs via frequency domain multiple access (FDMA). The bandwidth allocation matrix is defined as B = [B]1×M . Bm is the bandwidth allocated for ES m and BUm,n is the bandwidth allocated for client n to upload local model. Rn,m represents the uplink transmission rate of client n:

where pn,m denotes the transmission power of the client n, hn,m is the channel gain between client n and ES m, and N0 is the power of additive white Gaussian noise.

• E defines the energy consumption, which contains the training energy E C and transmission energy E U .

where φ is the effective capacitance parameter of the computing chipset.

• J denotes the data distribution similarity cross-edge that measured by Jensen Shannon Divergence (JSD) [Men´endez et al., 1997] . JSD and data distribution similarity are negatively correlated. A lower JSD implies that the two sets of data are more likely to fulfill the assumption of being IID. We only compute the JSD value once between the two ESs because of the symmetry of JSD, i.e., JS(Qa,Qb) = JS(Qb,Qa) and JSD ∈ [0, 1].

where Qa, Qb are the probability distributions of the data under ES a and ES b, KL(·) denotes the KLD (Kullback-Leibler Divergence) [van Erven and Harremos, 2014]. Mab denotes the mean distribution of Qa and Qb:

Remark: We mainly consider the case of equal number of local data in this paper. And we are based on a synchronized FL scenario, so the execution latency in a round of global iteration under a ES depends on the last completed client, as shown in Eq. (4). Due to the high transmission power of ES, the aggregation time and downlink transmission time are ignored compared to the local training and upload time. The clients under the same ES share bandwidth resources equally, which means that BU 1,m = · · · = BU n,m, n ∈ Gm.

3.3 ProblemFormulation

For a FL task, high-quality data can enhance model performance, while longer time delays and higher energy consumption bring negative impacts. For ease of representation, we define the utility of the network as a function of J and E.

where λ1 and λ2 are weighting parameters. Based on the multi-dimensional properties, we model the problem as fol-

where Eq. (16b) indicates that each client can only associate with one ES at a time. Eq. (16c) and Eq. (16d) are the bandwidth constraints of uplink channels. The transmission power constraint is given by Eq. (16e) and Eq. (16f). I in Eq. (16f) is the maximum execution latency of the currently executed FL task.

LEAP decomposes the original problem P1 into several subproblems, as shown in Figure.1, which can be solved one by one by combining coalition formation game and gradient projection method. In the coalition formation game, we design a coalition-friendly preference rule F to determine the optimal edge association relationship A∗ . Based on a stable coalition structure, to reduce the communication delay and energy consumption, LEAP optimizes the bandwidth allocation B∗ using the gradient projection method. In addition to this, LEAP captures the heterogeneous client communication resource conditions and determines the optimal transmission power p∗ of the clients based on satisfying the maximum delay of the task.

4 Optimal Solution Based on LEAP

In this section, we address the previously formulated problem P1. We start by identifying a stable coalition partition, and then optimize bandwidth allocation and transmission power based on this result.

4.1 Optimal Data Distribution

To achieve the goal of optimal network utility, we need to enhance the similarity of data distribution between ESs. Coalition game possesses an excellent tool for revealing the coalition formulation process. We model the problem of minimizing the JSD value of data distribution among ESs as a coalition formation game.

Definition 2. A coalition formation game C is represented as a 4-tuple (N, O, F, A), i.e., a set N of clients, a coalition partition O, a preference relation F, and a game strategy profile A.

• N: A set N = {1, · · · ,N} of players.
• O: A coalition partition O = {G}M1 , where Gm ⊆ O, ∪Mm=1Gm = N and m denotes the index of the coalition or ES.
• F: A preference relation ≻n is a complete, reflexive, and transitive binary relation over the set of all coalitions that client n may join in, i.e., G1≻nG2 indicates that client n strictly prefers joining coalition G1 over coalition G2.
• A: A strategy profile of edge association of clients, i.e., am,n = 1 means client n associates with ES m ∈ M.

When determining client preferences for multiple coalitions, the order of coalition preferences can be determined according to different rules. For example, the "Selfishness Rule", which only considers individual's choices, and the "Pareto Rule", which never harms the choices of any member in original coalition and new coalition [Zhang et al., 2018]. The former completely ignores the development of other clients in same coalition, posing a risk of harming coalition partition. The latter is too strict for the development of clients and coalitions. To minimize J S, we use a preference rule that places greater emphasis on the collective welfare of the entire coalition, called coalition-friendly preference rule, which the definition is as follows,

Definition 3. If there are two potential coalitions that client n can join, i.e., Ga, Gb ∈ O, then the preference relation is

where JS nGb→Ga means the JS value after client n leaves original coalition Gb to join the new coalition Ga, and JS nGb means the JS value before client n leaves.

Algorithm 1: Coalition Formation Game for Data Distributions Adjustment

• 11 until coalition partition converges or l = Lmax ;

Under the coalition-friendly preference rule, clients aim towards a globally optimal solution by considering the reduction of JS before and after the switch. Based on the preference relations given in Eq. (17), we define the switch rule:

Definition 4. Given a partition O = {G}M1 , the client n ∈ Ga decides to leave the original coalition Ga and move to another coalition Gb, b ̸= a , if and only if Gb ∪ {n}≻nGa. The new coalition partition can be described as ˜ O → {(O\ {Ga, Gb}) ∪ (Ga\ {n}) ∪ (Gb ∪ {n})}.

The coalition-friendly preference rule is considered from a coalition standpoint, which can be viewed as a partially collaboration. It is critical to investigate the stability under it.

Definition 5. If there exists a potential function ϕ such that the difference between potential function and utility function remains constant when the client's association relationship changes, the game is an exact potential game.

Theorem 1. The coalition formation game C is an exact potential game.

According to Theorem 1, the coalition formation game C has at least a stable coalition partition. To obtain the solution of the game, we will focus on the algorithm for forming an effective coalition partition, which shown in Algorithm 1. In the coalition formation algorithm, a client n is selected to undergo a comparative update based on the switch rule defined in Definition 4. This rule determines whether the client should leave its current coalition or join another coalition (lines 3-7). We assume that client n leaves current coalition and compute JS of each situation that client n joins in other coalition respectively based on Eq. (13). Therefore, according to the result of assumption, the prioritization of each situation or coalition can be determined. We choose the case that yields the lowest J S, and then client n leaves the current coalition Gm, joins the new coalition Gm′ if the two coalition are not the same (lines 8-9). Coalition partition will be updated due to this switching (line 10). However, if JS increases after the switch operation, client will remain in the current coalition. The iterative process described above repeats until form a stable partition of coalitions O∗ = {G ∗ }M1 where no exchange exists that can bring down JS in current partition Ocr or reach the maximum iteration rounds. We need to perform (m-1)m 2 calculations for JS. However, in reality, the computation is equally distributed to each ES, so the final time complexity is O(M). This is hardly a burden for a high-performance ES that can respond quickly to clients.

4.2 Optimal Bandwidth Allocation

Based on the final coalition partition, the delay of local training is determined. According to Eqs. (7) to (11), energy consumption is proportional to the execution time of training. From Eq. (2) and Eq. (16c), we can observe that the transmission delay is minimized when p∗n = pmax n ,∀n ∈ N . Hence, the joint optimization problem Eq. (16a) is distilled into a single-object optimization problem:

It can be observed that T U n (Bm) is a convex function with respect to Bm from Eq. (2). We assume that the worst transmission case in each coalition is n0m and the clients in each coalition have the same status as the worst case. The solution of the communication minimization problem P2, i.e., optimal bandwidth allocation for coalitions, is denoted as B∗ . Then, the λ2E can reach the minimum value when B = B∗ . Because a strictly convex function has at most one minimum, by setting nm = n0m, ∀m ∈ M and p∗n,m = pmax n,m , ∀n ∈ N, the optimization problem is transformed as:

Lemma 1. If gi(x) is convex function, max(min)gi(x), ∑ gi(bix) and ∑ bigi(x) are also convex functions.

Since the objective function of the optimization problem is the sum of a convex function, according to Lemma 1, the objective function is also convex with respect to B. We can apply the gradient projection method (GP) [Chu et al., 2023] to allocate bandwidth. The GP method is summarized in Algorithm 2. Through several iterations (lines 3-6), we can obtain the optimal bandwidth allocation in P2 and P3.

Algorithm 2: Bandwidth Allocation

Input: B(0), step size η, accuracy tolerance ϵ, and iteration number J max 1 j = 0; 2 repeat 3 Gradient: ∇G(Bj); 4 Projection: PΩB ; 5 Update B: Bj+1 ←- PΩB (Bj - η∇G(Bj)); 6 j = j + 1; 7 until objective value converges or j = J max ;

4.3 Optimal Transmit Power

Once the bandwidth allocation matrix B∗ and stable coalition partition G ∗ are determined, the optimization of transmit power for each client can be formulated as follows:

Eq. (21b) gives the clients' transmission power range and Eq. (21c) emphasizes the constraint of maximum execution latency, so the solution needs to satisfy both of them.

Theorem 2. There exists an optimal solution p∗n,m for problem P4, i.e.,

where

5 Experiments

In this section, we conduct extensive experiments to assess the performance of LEAP. We first introduce the experimental setups, and then compare and analyze the effectiveness of our approach to SOTAs.

5.1 Experimental Setup

Datasets and Models. We evaluate the performances of LEAP on two commonly adopted learning models and four real datasets: a LR (Logistic Regression) model on MNIST dataset [Lecun et al., 1998] and a CNN (Convolutional Neural Network) with two convolution layers and three fully connected layers on CIFAR-10 [Krizhevsky, 2009], SVHN [Netzer et al., 2011] and CINIC-10 [Darlow et al., 2018].

Parameter Settings. We set two scale settings on each dataset, 5 ESs with 50 clients or 3 ESs with 10 clients. For each dataset, 5 rounds of local training and 12 rounds of edge iterations are conducted. Among these setting, 100 rounds of global iterations are performed for MNIST and CINIC-10 datasets, while 200 rounds of global iterations for CIFAR-10

and SVHN. We set learning rate to 0.01 for MNIST, CIFAR10 and CINIC-10, and 0.005 for SVHN. The momentum is set within the range of [0, 0.9] and weight decay is 0.005.

Baselines. Four baselines are considered for comparison with LEAP, consist of two clustering algorithms, a coalition formation method and a model aggregation method.

• Mean-Shift [Lu et al., 2023]: It is a density-based nonparametric clustering algorithm. One advantage is that it does not require specifying the number of clusters in advance, as it automatically determines this in the process.
• K-Means [Lim et al., 2022]: It is an iterative clustering algorithm that partitions data points into K clusters (K is pre-specified) and assigns each data point to the nearest cluster center based on distance.
• RH [Ng et al., 2022]: It is a reputation-aware hedonic coalition formation algorithm, in which clients form stable coalition partitions with selfish preferences based on the reputation of cluster heads and their own utility.
• MA [Zhang et al., 2021] [Shi et al., 2022]: It is a model aggregation method based on marginal losses. By setting marginal loss thresholds, it becomes possible to identify and reduce the impact of low-quality models on the aggregation process.

5.2 Experimental Results

Validating the effectiveness of mitigating the degree of cross-edge non-IID. Figures 2(a) to 2(c) show the distribution of data for each coalition during the coalition formation process, with the color of each cell indicating the percentage of such data under that coalition. The initial JS is 0.69 with two label categories of each coalition. As the client switching process progresses, the data distributions under each coalition become increasingly similar, with the final JS reaching 0, which means that the distribution in each coalition is same. Figure 2(d) shows the complete variation of JS during the client switching process. Each switching operation demonstrates a consistent decreasing trend of JS.

Comparing with K-means and Mean-shift algorithms. Figure. 3 shows the accuracy under different methods and different data distributions. Comparing the initial state, the average accuracy based on the final distribution is improved by 2.9%, 33.3%, 47.6%, and 26.2% in the four datasets, respectively. Based on the same initial conditions, Mean-shift algorithm divides the data into five clusters. While randomly assigning the clients to ESs, the optimal client combination cannot be ensured because of the duplicate labels within the clusters. A similar issue arises when using the K-means algorithm. In addition, the K-means algorithm requires specifying the number of clusters in advance, which further hampers its applicability. It is clear that the final result after optimization based on our method is significantly improved compared to the other two methods, because our method is always in the direction of better when adjusting the combination of data distribution.

Comparing with RH and MA. The the initial correlation relation of Table 1 is based on the experimental results in [Ng et al., 2022] that presented RH, with 10 clients and 3 ESs. Compared to RH and MA, our approach still performs well.

This is because RA performs association formation with selfish client preference rule without considering the impact on the coalition partition. MA discards some model parameters below the loss threshold when aggregating based on marginal losses, resulting in data wasting.

Verifing the effectiveness on resource allocation. We calculate average transmission energy consumed per round of edge aggregation with random bandwidth allocation (RB), random transmission power (RP), and a combination of RB and RP (RB RP). From Figure. 4(a), we can observe that LEAP achieves a significant reduction in transmission energy consumption. We notice that in some cases RP is lower, but it fails to satisfy the maximum execution delay. From Figure. 4(b), the randomly determined transmission power is below the optimal value several times, so it fails to satisfy the delay requirement despite producing lower energy consumption.

6 Conclusion

In this paper, a novel optimization method LEAP, which has a lightweight implementation, was proposed to address the impact of muti-dimentional properties on HFL, i.e., data contribution, consumption of time and energy. Due to the stochastic characteristics of data distribution under ESs, edge association was combined with LEAP and a coalition formation game was built to model data distribution under different associations. To reduce the degree of non-IID of cross-edge data, a coalition-friendly preference rule was employed, and the existence of stable coalition partitions was proved. Further, the gradient projection method (GP) was utilized to reduce task execution time in heterogeneous resources within stable coalitions, improving the communication efficiency. Finally, extensive experiments were conducted on various real datasets to validate the effectiveness of LEAP.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China (No. 62372343, 62072411), in part by the Zhejiang Provincial Natural Science Foundation of China (No. LR21F020001), and in part by the Key Research and Development Program of Hubei Province (No. 2023BEB024).