CO-PFL: Contribution-Oriented Personalized Federated Learning for Heterogeneous Networks

TL;DR

CO-PFL addresses aggregation bias in personalized federated learning under data heterogeneity by introducing COWA, a contribution-aware aggregation module that uses gradient discrepancy and prediction deviation to weight client updates. It is complemented by PWPM for dynamic parameter-wise personalization and MAMO for stable optimization across mixed shared/personalized parameters. Extensive experiments on CIFAR10, CIFAR10C, CINIC10, and Mini-ImageNet show that CO-PFL achieves superior personalization accuracy, robustness to distribution shifts, and faster convergence compared with state-of-the-art methods. The framework offers a practical pathway to scalable, privacy-preserving, highly personalized models in heterogeneous federated networks.

Abstract

Personalized federated learning (PFL) addresses a critical challenge of collaboratively training customized models for clients with heterogeneous and scarce local data. Conventional federated learning, which relies on a single consensus model, proves inadequate under such data heterogeneity. Its standard aggregation method of weighting client updates heuristically or by data volume, operates under an equal-contribution assumption, failing to account for the actual utility and reliability of each client's update. This often results in suboptimal personalization and aggregation bias. To overcome these limitations, we introduce Contribution-Oriented PFL (CO-PFL), a novel algorithm that dynamically estimates each client's contribution for global aggregation. CO-PFL performs a joint assessment by analyzing both gradient direction discrepancies and prediction deviations, leveraging information from gradient and data subspaces. This dual-subspace analysis provides a principled and discriminative aggregation weight for each client, emphasizing high-quality updates. Furthermore, to bolster personalization adaptability and optimization stability, CO-PFL cohesively integrates a parameter-wise personalization mechanism with mask-aware momentum optimization. Our approach effectively mitigates aggregation bias, strengthens global coordination, and enhances local performance by facilitating the construction of tailored submodels with stable updates. Extensive experiments on four benchmark datasets (CIFAR10, CIFAR10C, CINIC10, and Mini-ImageNet) confirm that CO-PFL consistently surpasses state-of-the-art methods in in personalization accuracy, robustness, scalability and convergence stability.

Figures (6)

• Figure 1: Conventional federated learning averages client updates uniformly, disregarding differences in data quality and model contribution, which motivates the development of contribution-weighted aggregation strategies.
• Figure 2: Overall architecture of the proposed CO-PFL framework, illustrating the client–server collaboration process and the interaction among the PWPM, MAMO, and COWA modules.
• Figure 3: The test accuracy under few client settings converges with communication rounds in CO-PFL and benchmark.
• Figure 4: The train loss under few client settings converges with communication rounds in CO-PFL and benchmark.
• Figure 5: Performance of CO-PFL on CIFAR10 with varying personalization rates $p$ and budgets $\gamma$.

Theorems & Definitions (1)

• Remark 1