A unified Personalized Federated Learning framework ensuring Domain Generalization

Abstract

Personalized Federated Learning (pFL) allows for the development of customized models for personalized information from multiple distributed domains. In real-world scenarios, some testing data may originate from new target domains (unseen domains) outside of the federated network, resulting in another learning task called Federated Domain Generalization (FedDG). In this paper, we aim to tackle the new problem, named Personalized Federated Domain Generalization (pFedDG), which not only protects the personalization but also obtains a general model for unseen target domains. We observe that pFL and FedDG objectives can conflict, posing challenges in addressing both objectives simultaneously. To sufficiently moderate the conflict, we develop a unified framework, named Personalized Federated Decoupled Representation (pFedDR), which decouples the two objectives using two separate loss functions simultaneously and uses an integrated predictor to serve both two learning tasks. Specifically, the framework decouples domain-sensitive layers linked to different representations and design an entropy increase loss to encourage the separation of two representations to achieve the pFedDG. Extensive experiments show that our pFedDR method achieves state-of-the-art performance for both tasks while incurring almost no increase in communication cost. Code is available at https://github.com/CSU-YL/pFedDR.