Improving Accuracy and Convergence in Group-Based Federated Learning on Non-IID Data

Federated learning (FL) enables a large number of edge devices to learn a shared model without data sharing collaboratively. However, the imbalanced data distribution among users poses challenges to the convergence performance of FL. Group-based FL is a novel framework to improve FL performance, which appropriately groups users and allows localized aggregations within the group before a global aggregation. Nevertheless, most existing Group-based FL methods are K-means-based approaches that need to explicitly specify the number of groups, which may severely reduce the efficacy and optimality of the proposed solutions. In this paper, we propose a grouping mechanism called Auto-Group, which can automatically group users without specifying the number of groups. Specifically, various grouping strategies with different numbers of groups are generated with our mechanism. In particular, equipped with an optimized Genetic Algorithm, Auto-Group ensures that the data distribution of each group is similar to the global distribution, further reducing the communication delay. We conduct extensive experiments in various settings to evaluate Auto-Group. Experimental results show that, compared with the baselines, our mechanism can significantly improve the model accuracy while accelerating the training speed.

I. Introduction

With the increase of mobile and Internet of Things devices [1], [2], an unprecedented amount of data is constantly generated by users, which accelerates the development of machine learning [3], [4]. However, due to data privacy and security concerns [5], [6], it is impractical and often unnecessary for traditional centralized machine learning [7], [8] to collect a large amount of data generated by the users and train a model on a central server. Recently, federated learning is proposed in [9]–[11] to mitigate these issues by training a shared model without access to the raw data of the users. Specifically, in the prevalent Federated Averaging algorithm [9], users train models with their local data and only upload the local model updates to the central server for model aggregation, thereby protecting the privacy of user data.