FedProb: An Aggregation Method Based on Feature Probability Distribution for Federated Learning on Non-IID Data

Federated learning (FL) has been used to protect data contributors’ privacy by allowing training at clients and then feeding back machine learning models to servers for aggregation. Conventional methods of FL aggregation methods average model weights to produce a fused global model. However, in real-world applications in cyber-space systems, which often have heterogeneous Internet of Things data configuration and collection, FL encounters obstacles with non-independent and identically distributed (Non-IID) data. The main problem is the aggregated global models deviating from the optimal model trained on centralized servers. According to recent research, most Non-IID FL aggregation methods attempt to direct the movement of gradients to the optimal one using differentiation from trained models. In this paper, we propose a framework for using feature probability distribution in aggregation calculation. The proposed aggregation algorithm shows robustness on different Non-IID datasets and outperforms state-of-the-art methods in various FL experiments.