Robust-PFedproto: robust federated prototype learning based on personalized layers

Federated learning (FL), a distributed machine learning framework, was recognized for retaining training data on remote clients. However, two critical challenges were identified. First, heterogeneous data distributions were commonly observed across clients, which significantly degraded overall training efficiency. Second, the central server could’t access authentic remote client data, allowing adversarial clients to upload malicious model updates, thus negatively affecting global training accuracy and efficiency. To address data heterogeneity, the PFedproto framework was proposed, incorporating personalized layers based on prototype learning. These layers, positioned before each client’s decision layer, were optimized to enhance client model adaptation to local datasets and improve localized task prediction accuracy. Building on this framework, a three-stage defense scheme named RobustPFedproto was developed to strengthen the PFedproto framework’s robustness against data poisoning attacks. Experiments were conducted on four real-world image classification datasets. Results demonstrate that the PFedproto framework effectively mitigates data heterogeneity impacts. Additionally, the Robust-PFedproto scheme not only efficiently handles data heterogeneity but also shows strong robustness against data poisoning attacks.