FedCLS: Class-Aware Federated Learning in a Heterogeneous Environment

Federated learning is an approach of training the global model on the server by utilizing the personal data of the end users while data privacy is preserved. The users, referred to as clients, are responsible for performing local training using their respective datasets. Once trained, the clients forward their local models to the server, where the models are aggregated to update the global model. Practically, the datasets of clients have different classes of labels regardless of the number of samples. In other words, the data is non-independent and identically distributed (non-iid) among clients in terms of classes of labels, which creates heterogeneity among them. Hence, the local model weights updated by clients result in a broad variation due to heterogeneity among their local datasets. Thus, the process of aggregating the diversified local models of clients has a valuable impact on the performance of global training. When the server aggregates the local models by calculating the weighted average based solely on the number of samples available at the clients, the aggregation process may misguide the global training process. To address this issue, our paper proposes a novel reweighting method called FedCLS that performs based on the volume and variance of local datasets among clients. By taking into account the heterogeneity of data for aggregation in federated learning, the proposed method aims to achieve the minimum global point. The simulation results show that the proposed method achieves 28% performance improvement compared to the conventional federated learning methods.