pFedSD & Personalized Edge Intelligence via Federated Self-Knowledge Distillation
# Personalized Edge Intelligence via Federated Self-Knowledge Distillation
摘要——联邦学习(Federated Learning,FL)是一种新兴的边缘计算方法,通过多个设备协作训练机器学习模型,旨在解决传统集中式训练中的带宽限制、系统异构性和隐私问题。然而,现有的联邦学习方法主要关注为所有设备学习一个共享的全局模型,这对于不同设备来说可能并不理想。当每个边缘设备拥有其特定的数据分布或任务时,这种情况会变得更加严重。在本文中,我们研究了个性化联邦学习,其目标是训练在每个客户端上表现优异的模型。我们发现,每轮通信中的初始化会导致历史个性化知识的遗忘。基于这一观察,我们提出了一种通过自知识蒸馏实现的全新个性化联邦学习(PFL)框架,称为 pFedSD。通过允许客户端将先前个性化模型的知识提炼到当前的本地模型中,pFedSD 加速了最新初始化客户端对个性化知识的回忆过程。此外,自知识蒸馏在特征空间中为数据提供了不同的视角,实现了本地模型的隐式集成。在多种数据集和设置上的广泛实验表明,pFedSD 的效果和鲁棒性显著优于现有方法。
关键词——边缘计算,个性化联邦学习,知识蒸馏
# 1 引言
联邦学习(Federated Learning,FL)[1] 是一种在异构边缘计算环境中高效的分布式学习范式。当本地数据因隐私、存储或通信限制无法上传到云端时,FL 显得尤为重要。例如,物联网设备可能在网络边缘生成大量数据,而医疗应用可能因隐私问题禁止共享私人数据。在 FL 设置中,多个边缘设备协作训练一个全局模型,而无需传输本地数据。作为一种通信高效且保护隐私的学习方案,FL 已经展现了其在实际边缘计算应用中的潜力[2][3][4],如医疗[5]、目标检测[6] 和自然语言处理[7]。
然而,当设备间的数据分布不同或每个边缘设备的任务各异时,学习单一的全局模型可能会失败[8][9][10]。例如,客户端数据可能来源于不同的设备或地理位置,其数据可能是非独立同分布(non-IID)的[11][12]。在这种数据异质的场景下,单一全局模型的性能可能甚至比不进行协作而直接本地训练的模型还要差。此外,在多任务应用中,每个边缘设备可能有不同的最优方向。以移动设备上的下一词预测任务为例[13]。给定文本“我喜欢玩”,下一个词的建议应因人而异。传统联邦学习方案中共享的全局模型具有统一的训练目标,无法提供个性化的建议。个性化在许多领域同样重要,如医疗。不同的个人、医院或地区通常有不同的生活方式、活动模式和其他健康相关特征[14]。因此,为每个客户端获取个性化模型是必要的,以在利用共同信息的同时保持其特定知识。
为了解决上述异质性挑战,提出了个性化联邦学习(Personalized Federated Learning,PFL)[8][9][10],以从客户端特定的视角处理 FL。正如图 1 所示,当前的 FL 可以分为两类:传统联邦学习和个性化联邦学习。传统联邦学习,如 FedAvg [1] 和 FedProx [12],旨在学习一个强大且通用的共享模型以找到全局最优解。相比之下,个性化联邦学习旨在允许每个客户端训练一个在特定任务上表现良好的个性化模型。与缺乏数据的仅本地训练不同,PFL 中的个性化模型可以通过协作训练受益于共享知识。
个性化 FL 不仅需要适应本地数据分布的个性化,还需要通过协作挖掘通用知识的泛化能力。PFL 的关键挑战在于如何在共享知识与本地任务特定知识之间取得平衡。为了研究个性化与泛化之间的关系,我们对比了全局模型与本地模型的性能。正如图 2 所示,全局模型的个性化性能明显低于本地模型(图 2 的详细描述见第 2.2 节)。在 FL 中,本地模型在每轮通信开始时用最新的全局模型初始化。这种初始化会丢弃历史的本地知识,导致本地模型性能的严重下降。客户端可能需要在新一轮通信中从头开始训练本地模型。当数据分布更异质或参与率更低时,这一现象会更加严重。因此,我们假设 FL 中的初始化破坏了个性化与泛化之间的平衡。
基于我们的观察,我们提出了一种通过本地自知识蒸馏实现的全新个性化联邦学习框架,称为 pFedSD。具体而言,pFedSD 让客户端保留其本地训练的模型,并将其视为每轮结束时的个性化模型。在下一轮中,本地模型以接收到的全局模型进行初始化,并在先前个性化模型的指导下进行训练。通过知识蒸馏(Knowledge Distillation,KD),先前的个性化模型将历史个性化知识传递到当前本地模型。因此,pFedSD 避免了历史本地知识的遗忘,并在个性化与泛化之间实现了平衡。此外,先前的本地模型在特征空间中为本地数据提供了额外的视角。因此,我们可以将自知识蒸馏视为与 KD 一起执行的隐式集成[15]。
在多种数据集和设置上的实验结果表明,我们的方法显著提高了个性化性能和收敛性。本文的主要贡献如下:
- 我们首次探索了 PFL 中的个性化知识遗忘现象。每轮训练中,客户端在用聚合模型初始化后会遗忘历史个性化知识(第 2.2 节)。
- 我们提出了一种通过自知识蒸馏的有效个性化 FL 框架,用于传递历史个性化知识,并在个性化与泛化之间实现更好的平衡(第 3 节)。
- 我们进一步分析了个性化与泛化之间的权衡,以及自知识蒸馏对个性化联邦学习的益处(第 3.3 节)。
- 我们在多种实际数据集和设置上进行了广泛实验。与当前最先进的 FL 方法相比,我们的方法在个性化本地模型训练中展现了卓越的效率和鲁棒性(第 4 节)。
# 2. 预备知识
# 2.1 问题定义
本节正式定义个性化联邦学习(PFL)问题。目标是在联邦学习系统中为一组客户端协同训练个性化模型。本文将监督式
传统联邦学习系统的目标是找到一个全局模型
其中
经典的 FedAvg 算法在多个通信轮次中迭代进行本地训练和全局聚合:客户端本地优化:
然后在服务器端执行全局模型聚合
由于客户端之间的数据分布非独立同分布(non-IID),即
个性化联邦学习的目标是在客户端
# 2.2 个性化知识遗忘现象的观察
初始化后遗忘个性化知识
在通信轮次
如图 2 所示,最新下载的全局模型的个性化测试精度显著低于上一轮本地模型的性能。在每一轮的开始,参与客户端直接用全局模型替换本地模型,导致性能出现大幅下降。在更现实的设置下(如较小的
结果是,客户端
尽管 FedAvg 算法通过共享通用知识在客户端间提供了较强的泛化能力,但在一定程度上损害了针对特定客户端的个性化性能,特别是在数据异质性较高的场景中。个性化联邦学习的核心在于寻求泛化与个性化之间的平衡。简单地用全局模型初始化本地模型虽然增强了泛化能力,但丢失了之前的个性化信息。基于上述经验观察,我们假设防止个性化知识遗忘可以改善个性化联邦学习中的性能和收敛速度。
# 3 个性化联邦自我知识蒸馏(pFedSD)
本节首先介绍我们提出的 pFedSD 的核心思想,并分析我们提出的方法与现有算法的兼容性。接着,我们详细说明了框架的系统设计。pFedSD 在边缘计算环境中的工作流程如图 3 所示。
图 3. pFedSD 工作流程示意图训练过程按轮次进行,直到满足终止条件。工作流程包括以下 5 个步骤:
- 服务器选择一个客户端子集
,并广播全局模型 ; - 激活客户端
使用接收到的全局模型初始化本地模型 ,然后执行本地训练,并使用自我知识蒸馏; - 客户端将训练好的本地模型
存储为下轮的教师模型 ; - 客户端将更新后的本地模型
发送回服务器; - 服务器聚合所有接收到的本地模型,得到新的全局模型
。
FL 过程基于客户端-服务器框架。请注意,我们将边缘设备视为客户端。
# 3.1 提出的 pFedSD
受到第 2.2 节分析的启发,我们提出了个性化联邦自我知识蒸馏(pFedSD)。为了保留历史个性化知识,我们将个性化模型
为此,客户端
其中超参数
其中
与其他算法的兼容性由于 pFedSD 侧重于本地侧,它与近年来的个性化联邦学习(PFL)方法是正交的。在这里,我们将本地模型
算法 1. pFedSD 算法:服务器端
输入: 参与比例
输出: 本地个性化模型
过程
ServerExecute
初始化模型
循环
(通信轮次): 采样大小为
的客户端子集 (通过设备采样器) 将全局模型
发送给 中的所有客户端(通过通信管理器) 并行地
对每个客户端
: 结束并行循环
聚合得到新的全局模型:$w^{t+1} = \sum_{k \in S^t} \frac{w_k^{t+1}}{|S^t|} $(通过聚合器)
结束循环
结束过程
算法 2. pFedSD 算法:客户端端
输入: 学习率
输出: 本地个性化模型
过程
接收服务器发送的全局模型(通过通信管理器)
初始化本地模型(通过模型管理器):
计算本地迭代次数
for
(本地训练轮次): for
(本地迭代): 更新本地模型:
(优化公式 (4),通过训练优化器) end for
end for
使用
更新个性化模型 (通过模型存储器) 返回本地模型
(通过通信管理器) 结束过程
过程概述
- 服务器端:
- 服务器随机初始化全局模型,并在多个轮次中与选定的客户端进行迭代通信。
- 每轮,服务器将全局模型发送给一部分活动客户端,客户端进行本地训练并将更新后的模型返回给服务器。
- 服务器聚合接收到的本地模型来更新全局模型。
- 客户端:
- 客户端接收全局模型,并在其私有数据集上进行本地训练。
- 每个客户端使用自己的训练数据和来自个性化模型的知识蒸馏相结合进行更新。
- 训练结束后,客户端存储个性化模型并返回更新后的本地模型给服务器。
该过程涉及服务器与客户端之间的分布式协作,服务器负责整体模型聚合,而客户端则通过自我知识蒸馏进行本地更新,从而增强个性化学习。
# 3.2 系统设计
# 系统架构
我们在图 4 中提供了系统设计和架构的高层次示意图。它由一个云服务器和大量边缘设备组成。联邦学习系统采用模块化设计,将训练、通信和存储解耦。这样的设计具有以下优点:
- 用户可以灵活扩展到新的数据集、模型和算法。
- 系统支持模块复用,并为不同联邦学习算法的公平比较提供了可能。
# 云服务器的组成模块
- 协调器(Coordinator)
- 云服务器的顶层操作器,负责管理整个联邦学习流程。
- 指挥服务器上的其他组件,管理一组边缘设备,并在这些设备间协调联邦学习任务。
- 所有设备、模块和事件都必须在协调器中注册。
- 设备采样器(Device Sampler)
- 选择每轮通信中参与的边缘设备子集,支持部分参与的联邦学习。
- 简单的客户端选择策略是均匀采样。该模块可以扩展为更复杂的选择策略,以提高联邦学习的质量。
- 聚合器(Aggregator)
- 负责聚合接收到的所有模型,以生成全局模型。
- 聚合器的输入是客户端的模型或模型更新,输出是全局模型。
- 当服务器收到所有参与客户端的模型,或者达到预定义的时间限制时,会触发该模块。
- 默认的聚合方法是简单平均模型参数,此外还可以采用许多论文提出的增强型聚合策略。
- 通信管理器(Communication Manager)
- 负责服务器和客户端之间的交互。
- 包括向客户端发送工作请求消息、交换模型参数及其他联邦学习算法的辅助信息。
- 本模块基于 MPI 和 Gloo2 实现。
# 边缘设备的组成模块
- 控制器(Controller)
- 客户端的核心组件,控制本地过程并指挥边缘端的其他模块。
- 系统监控器(System Monitor)
- 获取当前设备的系统统计信息,例如电量和负载情况。
- 系统信息会报告给控制器,并用于决定是否接受来自服务器的工作请求。
- 模型管理器(Model Manager)
- 准备本地模型。
- 在本地训练开始时,用全局模型初始化本地模型;在本地训练结束时,存储本地模型作为个性化模型。
- 数据管理器(Data Manager)
- 加载本地训练和测试数据集。
- 训练优化器(Training Optimizer)
- 负责客户端的本地训练。
- 可以根据不同的联邦学习算法进行专门设计,例如在本地目标中添加一个邻近项。
- pFedSD 的主要改动集中在此模块中。
# 系统开销
与其他个性化联邦学习方法不同,pFedSD 能够在不显著增加系统开销的情况下很好地运行,与普通联邦学习相比开销差异不大。
# 软预测的获取方式
- 存储过去预测的结果
- 不需要重复前向传播,能够节省计算成本,但需要更多的磁盘空间来存储过去的预测结果。
- 对于本地数据不断更新(例如流数据)的情况,这种方式可能不适用。
- 保留过去的模型
- 可能需要更多的 GPU 内存来计算软预测,但灵活性更高。
# 适配无状态的边缘设备
在某些高可用性系统中,边缘设备需要无状态以便于迁移。
- 在这种情况下,可以将个性化模型存储在服务器端。
- 在模型下载阶段,将个性化模型与全局模型一起传输即可满足故障容错和负载均衡的需求。
# 3.3 pFedSD 分析
为更深入理解我们提出的 pFedSD 的效率,我们讨论以下两个问题:
- 解释个性化联邦学习算法中的个性化与泛化之间的权衡。
- 说明为什么自知识蒸馏在我们的方法中能够起作用。
# 3.3.1 泛化与个性化的权衡
如上文所分析,每轮开始时的初始化会导致对过去个性化知识的遗忘,进而扰乱了泛化与个性化之间的平衡。
- 个性化的极端是本地独立的算法,无需协作。
- 泛化的极端是传统联邦学习算法,如 FedAvg。
个性化联邦学习可以视为本地独立算法与传统联邦学习之间的一种中间状态,其个性化表现优于两者。当前的个性化联邦学习算法均试图在个性化与泛化之间找到更好的平衡。
# 初始化的作用
初始化在个性化联邦学习中发挥了重要作用,影响了本地有限训练步数中的个性化优化过程:
- 文献 [9]、[20]、[21] 提出,用上一个本地模型
初始化当前本地模型 。这种初始化方式使本地模型具有很强的个性化但较弱的泛化能力。为此,这些方法通过正则化项在全局模型和本地模型之间传递公共知识,缓解对稀缺本地数据的过拟合。 - 当正则化权重为 0 时,这些方法退化为本地独立算法。
- 当正则化权重为 1 时,则退化为传统联邦学习。
- 文献 [18]、[22]、[23] 将模型划分为个性化参数和全局参数两部分。这些方法用两部分混合后的结果初始化本地模型,无需额外的正则化项进行训练。
# 我们的方法
- 我们的方法使用全局模型(如 FedAvg)初始化本地模型,并在当前本地模型与先前个性化模型之间添加正则化项进行训练。
- 现有方法更倾向于本地独立算法,优先考虑个性化。我们的方法则更接近 FedAvg,优先考虑泛化。
# 3.3.2 为什么自知识蒸馏起作用?
为了帮助客户端回忆历史个性化知识,我们采用自知识蒸馏,将先前个性化模型中的知识传递到最新初始化的本地模型中。
# 知识迁移的其他选择
例如,可以使用
# 隐式地结合集成学习与知识蒸馏
自知识蒸馏的优势部分源于隐式的集成学习与知识蒸馏:
- 最近的研究 [15] 表明,许多真实世界的数据集具有“多视角”结构。每类数据可能包含多个视角特征,但并非所有数据都具有完整的视角特征。
- 例如,一张汽车照片可以通过识别挡风玻璃、车轮或车灯正确分类。然而,从侧面拍摄的汽车照片可能仅包含车轮特征。
- 包含多个特征的数据称为多视角数据;其余数据为单视角数据。
- 多视角数据只需部分特征即可正确分类,不再提供梯度。神经网络随后记忆剩余数据,而不学习新特征。因此,单个网络在训练过程中只能学习部分视角特征。
# FL 场景中的多视角学习
在联邦学习中,本地模型
- 这一过程可以视为“首先对
和 进行集成学习,然后将集成的知识蒸馏给 ”。
因此,与其他联邦学习方法相比,pFedSD 的个性化模型能够覆盖本地数据的更多视角特征。
# 性能改进
基于上述机制,pFedSD 不仅提高了个性化模型的性能,还实现了更快的收敛速度和更高的训练稳定性。
# 4 实验
# 4.1 实验设置
我们假设了两种不同的联邦学习 (FL) 场景:
个客户端,参与率 ; 个客户端,参与率 。
分别运行
数据集与任务
我们在不同的图像分类任务上评估所提出的 FL 框架,包括 FashionMNIST (FMNIST) [24] 和 CIFAR-10/100 [25]。表 1 列出了数据集的详细信息。
对于每个数据集,我们采用两种不同的非独立同分布 (non-IID) 数据设置:
- 病理性 non-IID [1]:每个客户端最多只持有
类样本。即 表示每个客户端可以拥有的最大类别数。当 减小时,数据异构性增加。 - Dirichlet non-IID [26]:使用 Dirichlet 分布
创建异构的客户端数据。 的值控制标签分布的倾斜程度。较小的 会生成更多 non-IID 数据分区,因为它使客户端 的标签分布 更偏向。我们按照 [27] 使用 Dirichlet 分布策略平衡分区。
图 5 展示了 CIFAR-10 数据集在 20 个客户端中的样本分布情况。所有数据集按照 80%-20% 的比例随机划分为训练集和测试集。
模型结构
参考先前的评估设置 [9], [28], [29],我们为 FashionMNIST 数据集使用一个简单的卷积神经网络 (CNN) [28],为 CIFAR 数据集使用一个类似于 [1] 的 5 层 CNN。
基线方法
我们将 pFedSD 与当前最先进的 FL 方法进行比较,这些方法可分为以下两类:
- 非个性化 FL 方法
- FedAvg [1]:作为联邦学习的开创性工作,直接平均本地模型的参数来学习全局模型。
- FedProx [12]:在本地目标中加入了一个近端项,防止客户端的更新偏离全局模型太多,从而有效缓解传统 FL 中的 non-IID 问题。
- 个性化 FL 方法
- FedPer [22]:将神经网络分为共享部分(body)和个性化部分(head)。共享部分为全局模型,个性化部分针对客户端优化。
- LG-FedAvg [28]:也采用参数分离方法,但个性化客户端学习的是局部特征表示,而全局模型聚合的是全局分类头。
- pFedMe [9]:使用 Moreau 包络作为客户端的正则化损失函数,控制个性化模型与全局模型之间的距离。
- FedFomo [29]:通过对每个客户端的最优模型组合进行一阶近似计算权重,客户端需要维护一个额外的验证集来衡量其他客户端对本地目标的适应性。
超参数设置
我们遵循相关工作的通用超参数设置:
- 学习率设为 0.01;
- 优化器为 SGD,权重衰减系数
,动量 0.9; - 批量大小为
。
我们在推荐的参数下运行所有基线方法:
- 对于 FedProx,近端项
从 中选择,最佳值分别为 0.001(FashionMNIST)、0.01(CIFAR-10)和 0.001(CIFAR-100)。 - 对于 FedPer,FashionMNIST 的个性化部分为最后两层,CIFAR 的个性化部分为最后的线性层。
- 对于 LG-FedAvg,所有数据集的共享头为最后两层。
- 对于 pFedSD,正则化参数从
中选择。 - 对于 FedFomo,下载的模型数
, -贪婪参数 ,每轮衰减 0.05。 - 对于我们的方法,权重参数从
中选择,温度 从 中调优。
实现细节
我们使用 PyTorch 实现所有基线方法,模拟服务器和一组客户端,采用多进程方式并使用 MPI 作为通信后端。所有实验均在配备四块 V100 GPU 的深度学习服务器上运行。
表1. 数据集统计信息
数据集 | 样本数量 | 类别数量 | 任务 |
---|---|---|---|
Fashion-MNIST | 70,000 | 10 | 图像分类 |
CIFAR-10 | 60,000 | 10 | 图像分类 |
CIFAR-100 | 60,000 | 100 | 图像分类 |
# 5 相关工作
# 5.1 个性化联邦学习
个性化联邦学习(Personalized FL)的主要动机是帮助客户端在其本地数据分布上获得优于单独本地模型的性能。大多数近期研究仍保留了传统联邦学习(FL)范式,服务器上维护一个全局模型,客户端从该全局模型中定制个性化模型。这些方法可以分为以下四类:
- 局部微调(local fine-tuning)[16], [34]
- 正则化方法(regularization)[9], [20], [21], [35]
- 模型插值(model interpolation)[36]
- 元学习(meta-learning)[37], [38]
- 模型混合/参数解耦(model mixture/parameter decoupling)[6], [22], [23], [28], [30], [36], [39]
上述方法利用单一全局模型传递所有客户端间的群体知识,但单一全局模型在某些数据分布下可能丧失泛化能力。为解决这一问题,最近的研究尝试在服务器或客户端端保留多个全局模型以实现更细粒度的泛化。
基于聚类的FL方法(Cluster-based FL)试图将客户端分组,每组对应一个全局模型。例如,FedAMP[40] 在服务器上为每个客户端维护个性化的云模型,增强相似客户端间的协作。FedFomo[29] 允许客户端下载其他客户端的模型,并通过本地验证数据计算个性化加权组合。
然而,这些方法仅在客户端本地数据分布明显聚类时表现良好。此外,FedFomo需要额外的验证数据以代表目标分布,这在实际应用中并不现实。与上述方法不同,本文提出的方法在通用设置中更加灵活且稳健,并且不需要额外数据支持。
# 5.2 联邦学习中的知识蒸馏
知识蒸馏(Knowledge Distillation,KD)最早由[41]提出,并在多个领域表现出显著成功。通过让学生模型模仿教师模型的输出,知识可以从教师传递到学生模型。传统KD通过强大的教师模型为学生网络生成软目标。近期研究提出了自蒸馏(self-distillation),即学生模型通过模仿结构相同的教师模型输出,可以超越教师模型的泛化能力[42], [43]。
值得注意的是,本文方法受自知识蒸馏启发,但并未严格遵循其范式。KD及其变体[44], [45]已被证明是联邦学习中的有效技术。大多数方法[7], [27]采用KD将客户端知识传递到服务器,即本地模型为教师,服务器上的全局模型为学生。这种方式能提升全局模型性能并加速收敛,但对代理数据的依赖较强,而获取或生成这种数据并不容易。
此外,[46], [47], [48]通过交换模型输出而非参数,实现了联合蒸馏,以减少通信成本并解决模型异构性问题。其中,[46]公开客户端私有数据的logits,而[47], [48]假设客户端间有公共数据集。FedGen[49]在数据无关的情况下,通过让服务器学习生成器整合客户端信息并传递给客户端以监督本地训练,使KD在FL中成为可能。FML[50]基于深度互学习(DML)允许客户端与全局模型相互训练。FedGKT[51]结合了FL与分裂学习,通过KD实现了边缘和服务器间的双向知识传递。
最近的FedLSD[52]在非IID环境下应用了自蒸馏以优化全局模型。然而,FedLSD在本地训练期间从全局模型蒸馏知识,其更像FedProx[12]的软版本。
简而言之,大多数现有方法假设公共或代理数据,需要仔细考虑甚至提前了解客户端私有数据。而本文方法并未放宽标准FL系统假设,也不引入额外数据或本地信息。更重要的是,该方法充分利用现有资源,且能轻松集成到主流个性化FL系统中。
# 6 结论
本文提出了一种简单且有效的个性化联邦学习框架——pFedSD,旨在解决边缘场景中的统计异质性问题。通过实验观察,我们发现初始化会引发个性化知识遗忘现象。
本文方法通过自蒸馏充分利用历史个性化模型,缓解遗忘问题并在个性化和泛化间取得更优平衡。实验结果表明,pFedSD在真实数据集上的性能和鲁棒性均优于现有方法。