Optimizing Federated Learning on Device Heterogeneity with A Sampling Strategy

Abstract

Federated learning (FL) is a novel machine learning that performs distributed training locally on devices and aggregating the local models into a global one. The limited network bandwidth and the tremendous amount of model data that need to be transported bring up expensive communication cost. Meanwhile, heterogeneity in the devices’ local datasets and computation power exerts a huge influence on the performance of FL. To address these issues, we provide an empirical and mathematical analysis of device heterogeneity on the performance of model convergence and quality, then propose a holistic design to efficiently sample devices. Furthermore, we design a dynamic strategy to further speed up convergence and propose the FedAgg algorithm to alleviate the deviation caused by device heterogeneity. With extensive experiments performed in PyTorch, we show that the number of communication rounds required in FL can be reduced by up to 52% on the MNIST dataset, 32% on CIFAR-10, and 28% on FashionMNIST in comparison to the Federated Averaging algorithm.