FEDBS: Learning on Non-IID Data in Federated Learning using Batch Normalization

Federated learning (FL) is a well-established distributed machine-learning paradigm that enables training global models on massively distributed data i.e., training on multi-owner data. However, classic FL algorithms, such as Federated Averaging (FedAvg), generally underperform when faced with Non-Independent and Identically Distributed (Non-IID) data. Such a problem is aggravated for some hyperparametric methods such as optimizers, regularization, and normalization techniques. In this paper, we introduce FedBS, a new efficient strategy to handle global models having batch normalization layers, in the presence of Non-IID data. FedBS modifies FedAvg by introducing a new aggregation rule at the server-side, while also retaining full compatibility with Batch Normalization (BN). Through our evaluations, we have empirically proven that FedBS outperforms both classical FedAvg, as well as the state-of-the-art FedProx through a comprehensive set of experiments conducted on Cifar-10, Mnist, and Fashion-Mnist datasets under various Non-IID data settings. Furthermore, we observed that in some cases, FedBS can be 2× faster than other FL approaches, coupled with higher testing accuracy.