Tackling Data Heterogeneity in Federated Learning through Knowledge Distillation with Inequitable Aggregation

TL;DR

This paper introduces KDIA, a novel federated learning strategy that effectively handles client data heterogeneity and limited client participation by leveraging knowledge distillation and weighted aggregation, leading to improved model accuracy.

Contribution

The paper proposes KDIA, a new method combining knowledge distillation with inequitable aggregation to address client heterogeneity and partial participation in federated learning.

Findings

KDIA outperforms existing methods in accuracy and convergence speed.

The approach is especially effective under severe data heterogeneity.

Experimental results on CIFAR datasets validate the method's robustness.

Abstract

Federated learning aims to train a global model in a distributed environment that is close to the performance of centralized training. However, issues such as client label skew, data quantity skew, and other heterogeneity problems severely degrade the model's performance. Most existing methods overlook the scenario where only a small portion of clients participate in training within a large-scale client setting, whereas our experiments show that this scenario presents a more challenging federated learning task. Therefore, we propose a Knowledge Distillation with teacher-student Inequitable Aggregation (KDIA) strategy tailored to address the federated learning setting mentioned above, which can effectively leverage knowledge from all clients. In KDIA, the student model is the average aggregation of the participating clients, while the teacher model is formed by a weighted aggregation of all clients based on three frequencies: participation intervals, participation counts, and data volume proportions. During local training, self-knowledge distillation is performed. Additionally, we utilize a generator trained on the server to generate approximately independent and identically distributed (IID) data features locally for auxiliary training. We conduct extensive experiments on the CIFAR-10/100/CINIC-10 datasets and various heterogeneous settings to evaluate KDIA. The results show that KDIA can achieve better accuracy with fewer rounds of training, and the improvement is more significant under severe heterogeneity.