Toward Universal Personalization in Federated Learning via Collaborative Foundation Generative Models

Personalized federated learning (PFL) enhances the performance of customized client models through collaborative training without compromising data privacy and ownership. Some previous PFL methods rely on rich prior knowledge about the types of data heterogeneity (such as class imbalance or feature skew), which greatly limits their application ranges. In this paper, we study the Universal Personalization in Federated Learning (UniPFL), the problem that has no prior knowledge about the types of data heterogeneity. In real-world PFL scenarios, UniPFL is potential because the data distributions of clients are usually heterogeneous and unknown to the server, where quantity imbalance, class imbalance, feature skew, or hybrid heterogeneity are possible contingencies. To address UniPFL, we propose FedFD, a novel framework with local data augmentation and global concept fusion, which is based on the recent advances in the foundation generative models (e.g., diffusion models, BLIP-2). On the client side, FedFD utilizes a diffusion model to assist local training by generating augmented data samples, and is then efficiently fine-tuned to be personalized. On the server side, we customize the aggregation strategies based on model similarities to learn both personalized models and diverse feature concepts. Extensive experiments show that FedFD reaches the state-of-the-art on (1) CIFAR-10 and CIFAR-100 for class imbalance; (2) DomainNet and Office-10 for feature skew, and (3) hybrid heterogeneity with both class and feature shifts.