Federated Learning convergence on IID features via optimized local model parameters

Federated learning enables collaborative training of global models without sharing local training data. In this paper, we propose a method to enhance the convergence of federated learning algorithm by modifying the training method at the client's side, such as we have deployed two different methods at local models, i) filter out batch normalization parameters and remove them from weight decay, and ii) zero initialize the last batch normalization in each residual branch (in case of resnet) so that the residual branch starts with zeros and each residual block behaves as identity. The following two methods were previously proposed for distributed training. In this paper, we propose to deploy the method in federated learning to enhance the convergence of the federated aggregation method. The proposed method is implemented at local model training and has shown significant improvements in the convergence, keeping the same rounds of federated learning. We have rigorously evaluated the efficacy of the proposed algorithm using a range of image classification models such as Resnet (-18, -50) and MobileNetv2 under the CIFAR-10 dataset. For all experiments, we have used an independent and identically distributed (IID) configuration for distributing the data among contributing clients and fedavg as a federated aggregation algorithm. We have deployed our code based on the open-source library 'Flower' for federated learning and communication. We have presented detailed results on various performance metrices such convergence under pretrained networks and without pretrained networks with different number of federated rounds, hyperparameter optimization, impact on different models to provide generability. The proposed algorithm outperforms state-of-the art algorithms in terms of convergence for federated learning by increasing up to 1.5% accuracy in limited federated rounds.