FedConceptEM: Robust Federated Learning Under Diverse Distribution Shifts

Abstract

Federated Learning (FL) is a machine learning paradigm that protects privacy by keeping client data on edge devices. However, optimizing FL in practice can be challenging due to the diversity and heterogeneity of the learning system. Recent research efforts have aimed to improve the optimization of FL with distribution shifts, but it is still an open problem how to train FL models when multiple types of distribution shifts, i.e., feature distribution skew, label distribution skew, and concept shift occur simultaneously. To address this challenge, we propose a novel algorithm framework, FedConceptEM, for handling diverse distribution shifts in FL. FedConceptEM automatically assigns clients with concept shifts to different models, avoiding the performance drop caused by these shifts. At the same time, clients without concept shifts, even with feature or label skew, are assigned to the same model, improving the robustness of the trained models. Extensive experiments demonstrate that FedConceptEM outperforms other state-of-the-art cluster-based FL methods by a significant margin.