PFed-DBA: Distribution Bias Aware Personalized Federated Learning for Data Heterogeneity

Personalized Federated Learning (PFL) aims to learn a custom model for each distributed client while benefiting from collaborative training in order to overcome the detrimental impact of data heterogeneity. Despite the promising benefits, the existing approaches often compromise the generalization performance of personalized models, as they solely focus on enhancing the personalization capability of models or merely aim to strike a balance between personalization and generalization. Indeed, increasing the personalization capability while preserving the strong generalization performance enabled by collaborative training remains a challenge for PFL, as the two objectives seem to compete with each other. To tackle this challenge, we investigate the relationship between model generalization and personalization under different degrees of heterogeneity. We find that besides the client-specific data distribution, the distribution bias between the unique data distribution of each client and that of the whole population is another critical factor that prominently impacts these two performances. Motivated by the above finding, we propose PFed-DBA, a novel PFL framework that effectively perceives this distribution bias to guide the training process. Concretely, we design the PFL models as a skip-connection network between a shared module for learning the shared representations delivering the common distribution of data across all clients and a personalized module for learning the personalized representations of the heterogeneous distribution bias. Then, we devise corresponding loss functions, aggregation strategy, and updating strategy in order to make the two modules intelligently complement each other. Moreover, we conduct extensive experiments to evaluate the effectiveness of PFed-DBA. The results show that PFed-DBA improves model accuracy to 12.34% at best compared with the state-of-the-art.