C2S: Class-aware client selection for effective aggregation in federated learning

Federated learning is proposed to train distributed data in a safe manner by avoiding to send data to server. The server maintains a global model and sends it to clients in each communication round, and then aggregates the updated local models to derive a new global model. Traditionally, the clients are randomly selected in each round and aggregation is based on weighted averaging. Researches show that the performance on IID data is satisfactory while significant accuracy drop can be observed for Non-IID data. In this paper, we explore the reasons and propose a novel aggregation approach for Non-IID data in federated learning. Specifically, we propose to group the clients according to classes of data they have, and select one set in each communication round. Local models from the same set are averaged as usual and the updated global model is sent to next group of clients for further training. In this way, the parameters are only averaged on similar clients and passed among different groups. Evaluation shows that the proposed scheme has advantages in terms of model accuracy and convergence speed with highly unbalanced data distribution and complex models.