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

📄 arXiv: 2504.04642v1 📥 PDF

作者: Hengrui Hu, Anai N. Kothari, Anjishnu Banerjee

分类: stat.ML, cs.LG

发布日期: 2025-04-06

💡 一句话要点

提出pFedKD-WCL算法，通过知识蒸馏和加权组合损失解决个性化联邦学习中的非独立同分布问题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 联邦学习 个性化联邦学习 知识蒸馏 非独立同分布 加权组合损失

📋 核心要点

  1. 联邦学习在非独立同分布数据下，模型容易产生漂移，泛化性能下降，这是当前方法面临的主要挑战。
  2. pFedKD-WCL算法的核心思想是利用全局模型作为教师模型，通过知识蒸馏指导本地模型的训练，实现个性化。
  3. 实验结果表明，pFedKD-WCL在MNIST和合成数据集上，相较于FedAvg等基线方法，在准确性和收敛速度上均有提升。

📝 摘要（中文）

联邦学习(FL)为分布式机器学习提供了一种保护隐私的框架，无需集中敏感数据即可实现跨多个客户端的协作模型训练。然而，统计异质性，即非独立同分布(non-IID)的客户端数据，带来了重大挑战，导致模型漂移和泛化能力差。本文提出了一种新算法pFedKD-WCL（Personalized Federated Knowledge Distillation with Weighted Combination Loss，基于加权组合损失的个性化联邦知识蒸馏），该算法集成了知识蒸馏和双层优化，以应对非IID挑战。pFedKD-WCL利用当前的全局模型作为教师来指导本地模型，从而有效地优化全局收敛和本地个性化。我们在MNIST数据集和一个具有非IID分区的合成数据集上评估了pFedKD-WCL，使用了多项逻辑回归和多层感知器模型。实验结果表明，在准确性和收敛速度方面，pFedKD-WCL优于最先进的算法，包括FedAvg、FedProx、Per-FedAvg和pFedMe。

🔬 方法详解

问题定义：联邦学习旨在保护用户隐私的前提下进行模型训练，但当客户端数据呈现非独立同分布(non-IID)时，传统的联邦学习算法（如FedAvg）容易导致模型在不同客户端上的性能差异较大，甚至出现模型漂移，无法有效泛化。现有方法难以兼顾全局收敛和局部个性化需求。

核心思路：pFedKD-WCL的核心思路是利用知识蒸馏，将全局模型的知识迁移到每个客户端的本地模型，从而在保证全局模型收敛的同时，实现客户端的个性化。通过加权组合损失函数，平衡全局模型知识和本地数据的影响，从而更好地适应非IID数据。

技术框架：pFedKD-WCL的整体框架如下：1) 服务器端维护一个全局模型。2) 每个客户端从服务器端接收全局模型。3) 客户端使用本地数据和全局模型进行知识蒸馏，训练本地模型。4) 客户端将本地模型的更新上传到服务器端。5) 服务器端聚合客户端的更新，更新全局模型。该过程迭代进行，直到模型收敛。

关键创新：pFedKD-WCL的关键创新在于将知识蒸馏与联邦学习相结合，并引入了加权组合损失函数。与传统的联邦学习算法相比，pFedKD-WCL能够更好地处理非IID数据，实现个性化建模。与现有的个性化联邦学习算法相比，pFedKD-WCL通过知识蒸馏，能够更有效地利用全局模型的知识，提高模型的泛化能力。

关键设计：pFedKD-WCL的关键设计包括：1) 知识蒸馏损失函数：使用全局模型的预测作为软标签，指导本地模型的训练。2) 加权组合损失函数：将知识蒸馏损失和本地数据损失进行加权组合，平衡全局知识和本地数据的影响。权重参数需要仔细调整，以获得最佳性能。3) 双层优化：外层优化全局模型，内层优化本地模型。这种双层优化结构能够有效地实现全局收敛和局部个性化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在MNIST数据集和合成数据集上，pFedKD-WCL算法在准确性和收敛速度方面均优于FedAvg、FedProx、Per-FedAvg和pFedMe等基线算法。例如，在非IID的MNIST数据集上，pFedKD-WCL的准确率比FedAvg提高了5%以上，并且收敛速度更快。

🎯 应用场景

pFedKD-WCL算法适用于各种需要保护用户隐私的分布式机器学习场景，例如：医疗健康领域，不同医院的数据分布可能存在差异；金融领域，不同银行的用户行为数据也可能存在差异；智能推荐系统，不同用户的偏好数据也可能存在差异。该算法能够有效提升模型在这些场景下的性能和泛化能力，具有重要的实际应用价值。

📄 摘要（原文）

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.