A Novel Algorithm for Personalized Federated Learning: Knowledge Distillation with Weighted Combination Loss

2025-04-06

AI Summary

To address model drift and degraded generalization caused by non-IID data in federated learning, this paper proposes pFedKD-WCL—a novel algorithm integrating knowledge distillation with bi-level optimization to jointly achieve global convergence and local personalization. Its key contribution is the first incorporation of a weighted composite loss into the federated knowledge distillation framework, dynamically balancing the teacher-guided loss (ensuring global consistency) and the local fitting loss (enhancing client-specific adaptation). The method employs multi-layer perceptrons and logistic regression, and is empirically validated on MNIST and synthetic non-IID datasets. Results demonstrate that pFedKD-WCL significantly outperforms FedAvg, FedProx, Per-FedAvg, and pFedMe, achieving state-of-the-art performance in both test accuracy and convergence speed.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Non-convex Optimization
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

Federated learning (FL) offers a privacy-preserving framework for distributed machine learning, enabling collaborative model training across diverse clients without centralizing sensitive data. However, statistical heterogeneity, characterized by non-independent and identically distributed (non-IID) client data, poses significant challenges, leading to model drift and poor generalization. This paper proposes a novel algorithm, pFedKD-WCL (Personalized Federated Knowledge Distillation with Weighted Combination Loss), which integrates knowledge distillation with bi-level optimization to address non-IID challenges. pFedKD-WCL leverages the current global model as a teacher to guide local models, optimizing both global convergence and local personalization efficiently. We evaluate pFedKD-WCL on the MNIST dataset and a synthetic dataset with non-IID partitioning, using multinomial logistic regression and multilayer perceptron models. Experimental results demonstrate that pFedKD-WCL outperforms state-of-the-art algorithms, including FedAvg, FedProx, Per-FedAvg, and pFedMe, in terms of accuracy and convergence speed.

Problem

• Addresses non-IID data challenges in federated learning
• Improves model accuracy and convergence speed
• Combines knowledge distillation with bi-level optimization

Innovation

• Knowledge distillation with weighted loss
• Bi-level optimization for non-IID data
• Global teacher model guides local training