Federated Learning with Class Imbalance Reduction
Miao Yang, Ximin Wang, Hongbin Zhu, Haifeng Wang, Hua Qian
Abstract
Federated learning (FL) is a promising technique that enables a large amount of edge computing devices to collaboratively train a global learning model. Due to the communication limitation, only a subset of devices can be engaged to train and transmit the trained model to centralized server for aggregation. Since the local data distribution varies among all devices, class imbalance problem arises along with the unfavorable client selection, resulting in a slow converge rate of the global model. In this paper, we design an estimation scheme to reveal the class distribution without the awareness of raw data. According to the estimation scheme, we propose a multi-arm bandit based algorithm that can select the client set with minimal class imbalance. The proposed algorithm can significantly improve the convergence performance of the global model. Simulation results demonstrate the effectiveness of the proposed algorithm.