FBLG: A Local Graph Based Approach for Handling Dual Skewed Non-IID Data in Federated Learning

1 Introduction

Federated learning has been widely applied in many fields. For example, in the medical field, federated learning can help medical institutions of different levels share patient data to provide more accurate treatment plans. In the Internet of Things (IoT), federated learning is used for shared training among multiple devices to improve the performance of smart devices. Additionally, federated learning has been applied in various fields such as finance and transportation to address the data privacy and distributed learning [Li et al., 2020a]. With the federated learning framework, clients can collaboratively train models without exposing data. Firstly, clients train models on local devices. Subsequently, the clients send updates of local models to the central server, which aggregates these updated local models and sends the resulting global model to each client [McMahan et al., 2017].

A major challenge in federated learning is the non-IID data, which arises when the data distributions on different devices are not independent and identically distributed [Liao et al., 2023; Shang et al., 2022]. The non-IID data can cause drift in the optimization of both local and global models, resulting in slower convergence [Karimireddy et al., 2020; Li et al., 2020c]. Existing study indicates that causes of nonIID data can be subdivided into five categories: feature distribution skew, label distribution skew, concept drift with different features, concept drift with different labels, and sample size skew [Kairouz et al., 2021]. However, nearly all existing federated learning methods focus on researching non-IID data with only one specific type of skew. For instance, FedLC [Zhang et al., 2022] proposed a fine-grained calibrated crossentropy loss to reduce the bias in local updates to improve the performance of global models with label distribution skews; Tijani et al.[2021] proposed a data extension strategy aimed at generating placeholders for absent classes within a local dataset to address the label distribution skew.

However, in real-world situations, lots of non-IID data comprises dual or even multiple skews. For example, consider three hospitals: a tertiary hospital, a children's hospital, and a tumor hospital. The distribution of tumor labels among these hospitals tends to be skewed, and simultaneously, the sample sizes also tend to be skewed due to differing sizes of these hospitals [Wu et al., 2023]. Existing studies [Hsu et al., 2019; Hsieh et al., 2020] have demonstrated that the sole skew caused by heterogeneous label distributions among clients can reduce the performance of the global model by 40%. Furthermore, the presence of dual skews caused by heterogeneous label distributions and sample sizes among clients can lead to performance degradation by 56% for the global model, highlighting the substantial impact of dual skews on global model performance. As the complexity of addressing non-IID data consisting of multiple skews will be significantly increased, it remains an open problem that is far from being resolved [Li et al., 2022]. So far, there are very limited studies focusing on non-IID data characterized by dual or multiple skews [Zhu et al., 2021]. Therefore, as an initial effort, our primary focus is to address the dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients. We hope this work will provide a solid foundation for handling more intricate cases with multiple skews which require in-depth research in the future.

In this paper, we propose a Federated learning algorithm Based on Local Graph, named FBLG that can simultaneously address the dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients. The main components of the algorithm include:

• (i) Addressing the skew caused by heterogeneous label distributions among clients. After each communication round, the server firstly sorts clients according to their local losses. Then, the server selects clients with larger local losses to construct a local graph and calculates the JS divergence among clients as the weights of edges in the local graph. Based on the local graph, clients with larger local losses and higher similarities are selected for aggregation to make the global model highly consistent.

(ii) Addressing the skew caused by heterogeneous sample sizes among clients. Considering that the model trained with more samples tends to carry more useful information, the sample size is used to further select clients with more samples when designing the objective function.

In summary, the algorithm considers selecting clients with larger local losses, higher similarities, and more samples for aggregation. Experimental results demonstrate the FBLG algorithm can achieve higher accuracy than existing baseline methods on four datasets when both the label distribution and sample size are skewed among clients. Theoretical analysis and experiments show the FBLG algorithm can converge quickly.

Our contributions in this paper are as follows:

• We propose the FBLG algorithm to address dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients.
• We construct a local graph based on clients' local losses and JS divergence among clients. Subsequently, similar clients are selected for aggregation based on the local graph to address the label distribution skew. Additionally, considering that models trained with more samples tend to contain more useful information, we use the sample size to select clients with more samples when designing the objective function to address the sample size skew.
• Through comparison experiments with nine existing baseline methods and theoretical analysis, it is verified that the FBLG algorithm can achieve higher accuracy and quicker convergence under situations with dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients.

2 RelatedWork

Federated learning is a distributed machine learning method that involves sharing data across multiple clients for model training while protecting individual privacy. Existing federated learning frameworks are divided into vertical federated learning and horizontal federated learning [Zhang et al., 2021a]. This paper is based on the horizontal federated learning framework, where participants typically share identical features while possessing distinct sample sets.

However, due to the difference in distributions of data owned by various clients participating in federated learning, non-IID data has become an important issue. The existence of non-IID data can lead to performance degradation in models, thus affecting the overall effectiveness of federated learning [Zhao et al., 2018; Li et al., 2022].

In recent years, some progress has been made to address non-IID data. However, current work mainly focuses on solving non-IID data with only one certain kind of skew among clients, such as label distribution skew, feature distribution skew, or sample size skew. To address the label distribution skew, Ramakrishna et al. [2022] proposed approximate inference methods for category label distribution based on parameter updates of clients; FedOV [Diao et al., 2023] deleted the original features of a few classes and learned them as adversaries. To address the feature distribution skew, FedBN [Li et al., 2021b] added a batch normalization layer to the local model to mitigate feature bias before model aggregation; FedRDN [Yan and Zhu, 2023] randomly injected statistical information from the entire federation's dataset into clients' data. To address the sample size skew, Wang et al. [2021] monitored and designed a new loss to make weight updates proportional to the number of samples in different categories; Zhang et al. [2021b] proposed a client selection system, enabling clients to decide participation in each training round based on their individual and global data distribution probabilities.

However, existing studies [Hsu et al., 2019; Hsieh et al., 2020] have shown that dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients degrade the performance more than the sole skew caused by heterogeneous label distributions among clients. Therefore, based on existing studies, this paper addresses the dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients to improve the effectiveness of federated learning.

3 Our Method

3.1 Preliminaries

In our work, we suppose there are N clients. For the k-th client, its sample set is Dk, the number of samples is nk, and the global model delivered by the server at t-th round is θt . Firstly, for ∀d ∈ Dk, let fd(θt ) be the local loss of each sample on the client, then the local loss for the k-th client is denoted as

For the k-th client, the optimization objective is to find the local model θ∗ k that minimizes the loss, denoted as

We use the stochastic gradient descent (SGD) to optimize the local loss Fk(θt ) on the k-th client to obtain the local model θt+1 k after a new round of local training. Then, we upload Fk(θt ) and θt+1 k at the same time to the server, which constructs the local graph. Based on the local graph, we select clients to achieve aggregation.

3.2 Framework of Our Proposed FBLG

This section introduces the framework of our proposed FBLG, which is shown in Fig. 1 and Algorithm 1. Assuming that the global model θt has been obtained after completing the t-th round of iteration, the (t + 1)-th round in FBLG includes the following steps. Each step is labeled with their corresponding line numbers in Algorithm 1.

• Step 1: The server delivers the global model θt (Line 3).
• Step 2: Each client performs local training based on θt . It obtains the local loss Fk(θt ) and the trained local model θt+1 k , then uploads them to the server (Lines 4 to 7).
• Step 3: The server constructs a local graph based on the local loss Fk(θt ) and the JS divergence among clients (Lines 9 to 10).
• Step 4: The server uses the local graph to select clients with larger local losses, higher similarities, and more samples (Lines 11 to 12).
• Step 5: The server aggregates selected clients and generates the global model θt+1 (Lines 13 to 16).

The framework of FBLG is consistent with most federated learning methods in stages of global model delivery, client local training, and model aggregation, while the key difference lies in the selection of clients (i.e. Steps 3 and 4).

3.3 Local Graph Construction

Next, we will introduce the construction of the local graph (i.e. Step 3). Firstly, we need to select the top M = C ×N

Algorithm 1: FBLG algorithm

clients with larger local losses from N clients, where C is the proportion of clients with larger local losses selected, and we denote these M clients as c1, c2, · · · , cM. Then, we use clients c1, c2, · · · , cM to construct a local graph and denote the local graph as G(γ, V ), where γ = {c1, c2, · · · , cM} is the node set of the local graph G(γ, V ), V ∈ RM×M is the adjacency matrix of the local graph G(γ, V ), and the element Vij is the weight between any two clients i and j, which is mainly used to characterize the similarity between any two clients i and j.

Existing similarity methods usually employ the cosine distance, which only focuses on the direction of vectors, disregarding their specific magnitudes. This can lead to poor results when computing the similarity of non-sparse data. Therefore, we use JS divergence to measure the similarity between clients. In addition, once the feature vectors of clients are given, we can easily get the adjacency matrix among clients. However, it is crucial to note that feature vectors may leak clients' sensitive information. Taking inspiration from FedGS [Wang et al., 2023], we introduce Gaussian noise in the computation of JS divergence. The computation of JS divergence is denoted as

where ei = θi(ε)[r] is the average of the r-th layer's network embedding on a batch for each client after we feed the batch of random Gaussian noise ε ∼ N(µ,Σ) to all the locally trained models, while µ and Σ are respectively the mean and covariance of a small validation dataset owned by the server.

Based on the JS divergence Sij between any two clients i and j, the weight Vij between any two clients i and j is denoted as

Based on the local graph G(γ, V ), the following explains how to select clients.

3.4 Client Selection

Next, we will introduce how to select clients with larger local losses, higher similarities, and more samples based on the local graph (i.e. Step 4). To achieve quicker error convergence of the model [Cho et al., 2020], we naturally prioritize the selection of clients with larger local losses, as we consider clients with larger local losses as one of our construction indicators when constructing the local graph.

Then, to address the label distribution skew among clients, we select clients with higher similarities for aggregation. Firstly, we calculate the shortest path distance matrix H = [hij ]M×M between each node pair in the local graph G(γ, V ) by the Floyd algorithm, where hij is the shortest path distance between any two clients i and j. Subsequently, based on the shortest path distance matrix H, we use an optimization objective to select a larger shortest path. The optimization objective (O1) is denoted as

where st+1 = {s1t+1 , s2t+1 , · · · , sMt+1} is the selection result at (t + 1)-th round, skt+1 ∈ (0, 1), 1 ≤ k ≤ M, when skt+1 = 1, it means that the k-th client is selected to participate in (t + 1)-th round of aggregation, conversely, when skt+1 = 0, it means that the k-th client is not selected, M is the number of nodes in the local graph G(γ, V ), and M(M - 1) is the number of node pairs in the local graph G(γ, V ).

According to the optimization objective of Eq. (5) and the definition of Eq. (4), we can observe that when the shortest path distance matrix H is relatively large, it implies that the JS divergence between clients needs to be as small as possible and hence we need to select clients with higher similarities.

Finally, to address the sample size skew among clients, we prioritize selecting clients with more samples, as models trained with more samples tend to carry more useful information. To this end, the optimization objective (O2) is denoted as

where ∑k∈st+1 nk is the total sample size of all clients.

In summary, to simultaneously address the dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients, we define the total optimization objective (O1 + O2) when selecting clients as

3.5 Global Aggregation

After selecting clients based on the local graph through Eq. (7), local models are aggregated using the traditional weighted aggregation method to produce the global model (i.e. Step 5). During the aggregation process, considering that clients with more samples should have a larger weight in the aggregation, we define the weights based on sample sizes within clients as

This approach allows us to further alleviate the sample size skew among clients by assigning larger weights to clients with larger sample sizes and smaller weights to clients with smaller sample sizes. Finally, the global model at (t + 1)-th round is denoted as

3.6 Convergence Analysis

We analyze the convergence of our proposed FBLG from a global perspective under the assumption that the loss function is non-convex. First, according to Bubeck et al. [2015], if the loss function ∇F (θ) is β-smooth, then for any θ, θ′ ∈ Rd , there exists ∥∇F (θ) -∇F (θ′ )∥ ≤ β ∥θ - θ′ ∥. Second, according to Tian et al. [2022], if F(θ) is locally convex, then for any θ, θ′ ∈ Rd , φ ∈ [0, 1], the distance between θ

and the locally optimum θ′ within a radius r(>0), there exists F (φθ + (1 - φ)θ′ ) ≤ φF (θ) + (1 - φ)F (θ′ ).

Next, we first prove that the model θ based on our proposed FBLG converges within the range selected by clients each time during the training process.

Theorem 1. If η < 1/β, β is a constant, then there exists ∥∥F (θt+1 ) - F (θ∗ )∥∥ ≤ ∥F (θt ) - F (θ∗ )∥, where F (θt ) represents the loss function of the global model at t-th round, θt and θ∗ represents the model and the optimized model defined in Eq. (2) on the server, respectively.

We then discuss the reasons why convergence improves when similar clients are aggregated. In federated learning, there are N clients with sample sets D1,D2, ...Dn, each of which belongs to one of the pl (l = 1, 2, ..., (≤ n)) distributions. Assuming that the stochastic gradient gl (·) obtained from the distribution pl at t-th round is unbiased, i.e. E[gl (θt )] = ∇F l (θt ). Since clients selected based on the FBLG are highly similar, it is natural to assume that data of selected clients come from almost the same distribution.

Theorem 2. Suppose that FBLG selects a set of local models trained with the same distribution of datasets. Compared with the FedAvg, we get E∥θt l - θ∗ l ∥2 ≤ E∥∥¯ θt - θ∗ l ∥∥2 , where θ∗ l is the optimized model for the dataset fitting the distribution pl , θt l denotes the global model aggregated by FBLG for the distribution pl , while ¯ θt denotes the uniform global model of FedAvg at t-th round.

We use the loss function to measure convergence and theoretically justify our proposed FBLG. If the loss function value converges stably to 0, it equivalently reflects that the trained model can converge to the optimal.

4 Evaluation

This section aims to answer the following research questions:

RQ1. How do the dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients affect the accuracy of federated learning?

RQ2. How is the performance of our proposed FBLG method compared with baseline methods?

RQ3. How do different components (including similarity metrics and objective functions) affect our proposed FBLG?

4.1 Experimental Settings

Datasets. According to most papers in federated learning [McMahan et al., 2017; Wang et al., 2020a; Huang et al., 2022], we validate our proposed FBLG algorithm on four commonly used datasets: MNIST [LeCun et al., 1998], Fashion-MNIST (FMNIST) [Xiao et al., 2017], CIFAR10 [Krizhevsky et al., 2009] and SVHN [Netzer et al., 2011]. To create data with dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients, we assign label distribution and sample size to 20 clients through Dirichlet distribution i.e. Yk ∼ Dir(αΥ), where α denotes the degree of skews among clients, Υ denotes the global label distribution [Hsu et al., 2019]. The smaller α indicates the data is more skewed. The sole label skew is set by shards [Wang et al., 2023]. Our codes, some supplementary experiments and appendices mentioned in the paper are available at https://github.com/YingLi-Y/FBLG.git.

Baselines. We compare our FBLG method with: (i) FedAvg [McMahan et al., 2017] based on weight aggregation, (ii) MDSample [Li et al., 2020b] based on the client's local data size, (iii) Power-Of-Choice [Cho et al., 2020] based on the client's local loss, (iv) FedProx [Li et al., 2020b] based on the proximal term, (v) Moon [Li et al., 2021a] based on the model comparison loss, (vi) Scaffold [Karimireddy et al., 2020] based on control variables, (vii) FedAvgM [Hsu et al., 2019] based on regularization, (viii) FedNova [Wang et al., 2020b] based on the speed of local training on the client, (ix) FedGS [Wang et al., 2023] based on the data distribution dependency graph and the sampling frequency of the client.

Parameter Settings. For each dataset, the number of samples selected by the client for one training session is B = 64, the number of local iterations is E = 5, the learning rate is η = 0.05, the learning rate is fully attenuated by a factor of 0.998, and the optimization algorithm used for local training of clients is SGD.

Implementation. All our experiments are run on theAIStation server with 1 NVIDIA A100-SXM4-40GB and 4 CPUs. All codes are implemented in Pytorch 1.12.1.

4.2 Impact of Dual Skewed Non-IIDData (RQ1)

We take the MNIST dataset as an example to more intuitively illustrate the label distribution skew among clients, the sample size skew among clients, and the dual skews caused by heterogeneous label distributions and sample sizes among clients. The visualization of the three types of data mentioned above is shown in Fig. 2, where the x-axis denotes the number of samples, the y-axis denotes the client ID, and the color of the block denotes the label type of samples.

• Fig. (a) visualizes the case where only the label distribution is skewed among clients when the number of shards is 12. Here, each client has a consistent sample size, but the label distribution of each client is inconsistent.
• Fig. (b) visualizes the case where only the sample size is skewed among clients when α = 17. Here, each client has a consistent label distribution but the sample size of each client is inconsistent.
• Fig. (c) visualizes the case where the label distribution and sample size are both skewed among clients when

α = 0.8. Here, some clients have 3 labels, while others have only 1 label, and the sample size per client is inconsistent.

To demonstrate the impact of data with only the label distribution skew among clients, only the sample size skew among clients, and the dual skews caused by heterogeneous label distributions and sample sizes among clients in federated learning, we use the classic FedAvg algorithm as an example. We plot the classification accuracies of the three skews with 20 clients over 20 communication rounds with the line graph, as shown in Fig. 3. We observe that only the label distribution skew among clients and only the sample size skew among clients can achieve relatively high accuracy within the first 4 communication rounds. Specifically, only the sample size skew among clients can even achieve 99.02% accuracy within 20 communication rounds. However, the dual skews caused by heterogeneous label distributions and sample sizes among clients can only reach 90.20% accuracy within 20 communication rounds and the convergence of the global model is unstable. It is clear that dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients has a greater impact on the performance of the global model.

4.3 Results of Performance Comparison (RQ2)

We conducted experiments on four datasets with 20 clients over 300 communication rounds when α = {0.8, 0.05, 0.01}, and the accuracy results are shown in Table 1. We observe that: (i) the accuracy of most baseline methods decreases as α decreases, which indicates that highly skewed data can seriously affect the performance of baseline methods. (ii) our proposed FBLG (C = 0.5) can achieve high accuracy on almost all four datasets, where C is the proportion of clients with larger local losses selected. When α = 0.05, it can improve at least 6.02% on the CIFAR10 dataset, and even achieve relatively high accuracy on extremely skewed data (i.e. α = 0.01), while existing nine baseline methods cannot. In addition, our proposed FBLG can achieve higher classification accuracy on MNIST and CIFAR10 datasets at C = 0.4 when α = 0.8. (iii) Although Power-Of-Choice and FedProx (ϕ = 0.1) can achieve the same high accuracy as FBLG (C = 0.5) proposed in this paper on two of the experimental results, Power-Of-Choice outperforms ours in an experimental result, we plot the test loss curves of each algorithm under extremely skewed data (i.e. α = 0.01) respectively on the FMNIST and CIFAR10 datasets, as shown in Fig. 4. We observe that the test loss of FBLG is small and converges quickly, consistently outperforming other baseline methods. Thus, our proposed FBLG can effectively address the impact of dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients. More details about the data visualization and test loss are in Appendix C.

4.4 Ablation Study (RQ3)

Similarity Metrics. We first verify the impact of the similarity metric based on JS divergence adopted in this paper on our FBLG's performance when α = 0.8. Here, we respectively replace the similarity metric with cosine distance and euclidean distance, then plot the impact of the three similarity metrics on the classification accuracy of our proposed FBLG into a table, as shown in Table 2. We observe that our proposed FBLG performs best with the similarity metric based on JS divergence, achieving up to 0.46% higher accuracy on the SVHN dataset. This preference is attributed to the robust nature of JS divergence in handling extreme values or outliers, making it more capable of addressing skewed data.

Objective Functions. We then compare the objective function that considers both the similarity and the number of samples among clients, the objective function that only considers the similarity among clients, and the objective function that only considers the number of samples among clients, and verify the impact of these three situations on the classification accuracy of our proposed FBLG algorithm when α = 0.8. We denote Eq. (5) that only considers the similarity among clients to address the label distribution skew among clients as FBLG (O1) and denote Eq. (6) that only considers the number of samples among clients to address the sample size skew among clients as FBLG (O2). Simultaneously, the Eq. (7) that considers both the similarity and the number of samples among clients to address dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients is denoted as FBLG (O1+O2). The impact of the above three situations on the classification accuracy of our proposed FBLG is drawn into a table, as shown in Table 3. We observe that considering both the similarity and the number of samples among clients in Eq. (7) can bring better performance than only considering the similarity among clients or only considering the number of samples among clients.

5 Conclusion

In this paper, we proposed a new federated learning algorithm based on local graph (FBLG) to address dual skewed non-IID data caused by heterogeneous label distributions and sample sizes among clients. Specifically, (i) To address the label distribution skew, we construct a local graph based on the local losses of clients and the JS divergence among clients. Based on the local graph, similar clients are selected for aggregation to make the global model highly consistent; (ii) To address the sample size skew, we use the sample size to select clients with more samples when designing the objective function. Experimental results demonstrated the accuracy of our proposed FBLG is higher than that of baseline methods. Especially with the increasing degrees of data skewness, the advantage of FBLG becomes more obvious. Meanwhile, both theoretical analysis and experimental results have successfully proven that our proposed FBLG can converge quickly. In the future, we will explore to address the problem of nonIID data consisting of more complex multiple skews.

Acknowledgements

The research was supported by National Natural Science Foundation of China (No. 62076002, 61402005, 62272403), and Natural Science Foundation of Anhui Province, China (No.2008085MF194).