A Novel Neural Network-Based Federated Learning System for Imbalanced and Non-IID Data

Abstract With the growth of machine learning techniques, the privacy of users’ data has become a major concern. Most of the machine learning algorithms rely heavily on large amounts of data which may be collected from various sources. Collecting these data yet maintaining privacy policies has become one of the most challenging tasks for the researchers. To combat this issue, researchers have introduced federated learning, where a prediction model is learned by ensuring the privacy of clients’ data. However, the prevalent federated learning algorithms possess an accuracy and efficiency trade-off, especially for non-IID data. In this research, a centralized, neural network-based federated learning system is proposed. The centralized algorithm incorporates micro-level parallel processing inspired by the traditional mini-batch algorithm where the clients’ devices and the server handle the forward and backward propagation respectively. A semi-centralized version of the proposed algorithm is also devised. This algorithm takes advantage of edge computing to minimize the load from the central server, where clients handle both the forward and backward propagation while sacrificing the overall train time to some extent. The proposed systems are evaluated on five well-known benchmark datasets and achieve satisfactory performance in a reasonable time across various data distribution settings as compared to some existing benchmark algorithms.