Fast converging Federated Learning with Non-IID Data

Abstract

With the advancement of device capabilities, Internet of Things (IoT) devices can employ built-in hardware to perform machine learning (ML) tasks, extending their horizons in many promising directions. In traditional ML, data are sent to a server for training. However, this approach raises user privacy concerns. On the other hand, transferring user data to a cloud-centric environment results in increased latency. A decentralized ML technique, Federated learning (FL), has been proposed to enable devices to train locally on personal data and then send the data to a server for model aggregation. In these models, malicious devices, or devices with a minor contribution to a global model, increase communication rounds and resource usage. Likewise, heterogeneous data, such as non-independent and identically distributed (Non-IID), may decrease accuracy of the FL model. This paper proposes a mechanism to quantify device contributions based on weight divergence. We propose an outlier-removal approach which identifies irrelevant device updates. Client selection probabilities are computed using a Bayesian model. To obtain a global model, we employ a novel merging algorithm utilizing weight shifting values to ensure convergence towards more accurate predictions. A simulation using the MNIST dataset employing both non-iid and iid devices, distributed on 10 Jetson Nano devices, shows that our approach converges faster, significantly reduces communication cost, and improves accuracy.