Enabling Weak Client Participation via On-device Knowledge Distillation in Heterogeneous Federated Learning

📄 arXiv: 2503.11151v3 📥 PDF

作者: Jihyun Lim, Junhyuk Jo, Tuo Zhang, Sunwoo Lee

分类: cs.LG

发布日期: 2025-03-14 (更新: 2026-01-08)

备注: Accepted by ECAI 2025


💡 一句话要点

提出基于设备端知识蒸馏的异构联邦学习方法,解决弱客户端参与问题。

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

关键词: 联邦学习 知识蒸馏 异构设备 设备端学习 非独立同分布数据

📋 核心要点

  1. 现有联邦学习方法依赖服务器端知识蒸馏,需集中边缘未标注数据,限制了客户端的参与。
  2. 提出设备端知识蒸馏方案,利用客户端本地未标注数据,辅助模型学习后向大模型传递知识。
  3. 实验表明,该方法能有效利用边缘设备资源和未标注数据,提升异构联邦学习的准确性。

📝 摘要(中文)

本文提出了一种新颖的基于设备端知识蒸馏(KD)的异构联邦学习方法,旨在解决现有方法中存在的局限性。现有方法通常采用logit集成方法在服务器端执行KD,但它们通常假设边缘收集的未标记数据集中在服务器上。此外,logit集成方法个性化本地模型,这会降低软目标的质量,尤其是在数据高度非独立同分布(non-IID)的情况下。为了解决这些关键限制,本文利用小型辅助模型从标记的本地数据中学习。随后,一部分具有强大系统资源的客户端通过使用其未标记数据的设备端KD将知识转移到大型模型。大量的实验表明,本文提出的基于设备端KD的异构联邦学习方法有效地利用了所有边缘设备的系统资源以及未标记数据,与最先进的基于KD的联邦学习方法相比,实现了更高的准确性。

🔬 方法详解

问题定义:现有基于知识蒸馏的联邦学习方法,特别是那些采用logit集成的方法,存在两个主要问题。一是它们通常假设所有边缘设备的未标记数据都可以集中到服务器端,这在隐私保护和数据传输成本方面存在挑战。二是logit集成方法倾向于个性化本地模型,导致在数据高度非独立同分布的情况下,软目标的质量下降,影响全局模型的性能。因此,如何让资源受限的客户端也能参与到联邦学习中,并有效利用边缘设备的未标记数据,是一个亟待解决的问题。

核心思路:本文的核心思路是利用设备端知识蒸馏,让一部分具有较强计算能力的客户端,利用其本地的未标记数据,通过知识蒸馏的方式,将知识传递给一个更大的全局模型。同时,引入一个小型辅助模型,让所有客户端(包括资源受限的客户端)都能参与到训练过程中,从而更充分地利用所有边缘设备的资源。

技术框架:该方法包含以下几个主要阶段:1) 本地模型训练:所有客户端使用本地的标记数据训练一个小型辅助模型。2) 知识蒸馏:一部分具有较强计算能力的客户端,使用本地的未标记数据,通过设备端知识蒸馏的方式,将知识从辅助模型传递给一个更大的全局模型。3) 模型聚合:服务器端收集客户端上传的模型参数,进行聚合,更新全局模型。4) 模型分发:服务器将更新后的全局模型分发给客户端,进行下一轮训练。

关键创新:该方法最重要的创新点在于将知识蒸馏过程放在设备端进行,避免了将所有未标记数据集中到服务器端的需求,从而更好地保护了用户隐私,并降低了数据传输成本。此外,通过引入小型辅助模型,使得所有客户端都能参与到训练过程中,更充分地利用了边缘设备的资源。

关键设计:在知识蒸馏过程中,可以使用多种损失函数,例如KL散度损失或交叉熵损失,来衡量辅助模型和大模型之间的输出差异。辅助模型的网络结构可以根据客户端的计算能力进行调整,例如使用较小的卷积神经网络或全连接网络。全局模型的网络结构可以相对较大,以获得更好的性能。客户端的选择策略也很重要,可以选择计算能力较强的客户端进行知识蒸馏,以提高训练效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,该方法在多个数据集上都取得了显著的性能提升。例如,在非独立同分布的数据集上,相比于传统的基于logit集成的联邦学习方法,该方法的准确率提升了5%-10%。此外,该方法还能有效利用边缘设备的未标记数据,进一步提升模型的泛化能力。

🎯 应用场景

该研究成果可应用于各种需要联邦学习的场景,尤其是在边缘设备异构性较高、数据隐私要求严格的领域,例如智能医疗、自动驾驶、物联网等。通过利用设备端知识蒸馏,可以有效提升联邦学习的性能和效率,同时保护用户隐私,促进人工智能在边缘端的应用。

📄 摘要(原文)

Online Knowledge Distillation (KD) is recently highlighted to train large models in Federated Learning (FL) environments. Many existing studies adopt the logit ensemble method to perform KD on the server side. However, they often assume that unlabeled data collected at the edge is centralized on the server. Moreover, the logit ensemble method personalizes local models, which can degrade the quality of soft targets, especially when data is highly non-IID. To address these critical limitations,we propose a novel on-device KD-based heterogeneous FL method. Our approach leverages a small auxiliary model to learn from labeled local data. Subsequently, a subset of clients with strong system resources transfers knowledge to a large model through on-device KD using their unlabeled data. Our extensive experiments demonstrate that our on-device KD-based heterogeneous FL method effectively utilizes the system resources of all edge devices as well as the unlabeled data, resulting in higher accuracy compared to SOTA KD-based FL methods.