系统异构性是联邦学习中的一个重要挑战,尤其是在处理多样化客户端设备群体时。客户端通常具有显著不同的计算能力(CPU功率、可用内存)、能源预算(特别是移动设备)和网络带宽。一个协同训练的大型复杂模型可能整体表现良好,但在资源受限的客户端上可能完全无法使用。部署单一、庞大的全局模型通常是不切实际的。
本节讨论专门设计用于根据联邦学习过程中各个客户端能力调整模型复杂度的方法:模型剪枝和模型适应。这些方法让更多种类的设备参与,并能使联邦系统更有效率和更个性化。
联邦模型剪枝
模型剪枝旨在通过移除冗余或次要参数 (parameter)来减小深度学习 (deep learning)模型的尺寸和计算成本。在联邦环境中,剪枝可以策略性地应用以适应设备限制。
什么是剪枝?
其主要目的是识别并去除对模型性能贡献最小的模型参数(权重 (weight)、神经元、滤波器甚至层)。这会带来更小的模型占用空间、更低的内存使用、更短的推理 (inference)延迟,并可能降低能耗。挑战在于去除参数而不显著降低准确率。
剪枝主要有两类:
-
非结构化剪枝: 这涉及根据特定标准(通常是其大小,小数值权重被认为不那么重要)移除单个权重。所得的权重矩阵变得稀疏,包含大量零值。
- 优点: 可以在准确率损失最小的情况下实现高压缩比。
- 缺点: 稀疏模型通常需要专用库(例如,用于稀疏矩阵乘法)或硬件支持才能实现实际加速。标准硬件可能不会从非结构化稀疏性中获得太多好处。稀疏更新的通信也可能很复杂。
-
结构化剪枝: 此方法移除模型的整个结构元素,例如卷积层中的滤波器、通道,甚至整个层。
- 优点: 得到更小、密集的模型,可以在标准硬件上高效运行,无需专用库。更易于实现和部署。
- 缺点: 粒度不如非结构化剪枝细,可能导致在相同参数减少量下准确率下降更多。需要仔细选择要移除的结构。
在联邦学习中应用剪枝
将剪枝集成到联邦学习流程中需要决定剪枝发生的时间和地点:
- 服务器端剪枝: 中央服务器在将聚合的全局模型广播给客户端进行下一轮训练之前对其进行剪枝。这种方式管理更简单,但对所有客户端应用统一的剪枝级别,无法适应各个设备差异。对于性能强的设备,剪枝可能过于激进;而对于受限设备,剪枝可能不足。
- 客户端剪枝: 每个客户端接收(可能密集的)全局模型,并在开始本地训练之前,根据自身的资源限制在本地对其进行剪枝。或者,客户端可以在将计算出的模型更新发回服务器之前对其进行剪枝。这允许细粒度的适应,但引入了复杂性:
- 服务器如何聚合来自不同结构模型(不同稀疏掩码或剪枝架构)的更新?只平均剩余的共享参数是一种方法,但可能导致信息损失。
- 剪枝掩码如何确定和通信?客户端可能需要报告它们的能力或期望的剪枝级别。
- 联邦剪枝算法: 更完善的方法涉及协调。例如,客户端可以根据本地数据提出要剪枝的参数,服务器聚合这些提议以做出全局剪枝决定。有些技术试图迭代地学习跨客户端的共享剪枝掩码。
示例:基于幅度的滤波器剪枝(结构化)
CNN中一种常见的结构化剪枝技术是滤波器剪枝。重要性分数较低的滤波器被移除。一个简单的度量滤波器Fi重要性的分数可以是其权重的L1范数:
Importance(Fi)=∑∣w∣其中 w∈Fi
在联邦学习情境中:
- 服务器发送当前全局模型Wg。
- 客户端k根据Wg或经过一些本地训练后计算滤波器重要性。
- 客户端k根据设备限制确定其目标剪枝比例。
- 客户端k通过移除重要性最低的滤波器来创建剪枝模型Wk′。
- 客户端k在本地训练Wk′。
- 客户端k将更新Δk′=Wk′−Wg′(其中Wg′是使用与Wk′相同掩码的Wg的剪枝版本)发送回服务器。
- 服务器需要一个策略来聚合这些可能结构不同的更新。一个常见方法是只平均该轮中没有被任何参与客户端剪枝的滤波器/权重的更新,或使用更高级的聚合方法来处理异构结构。
模型适应策略
剪枝、模型适应技术提供了将模型复杂度与设备能力匹配的替代方法。
模型缩放与子网络提取
与训练单一固定模型不同,这里的理念是训练一个灵活的“超网络”,它包含许多大小和复杂度不同的子网络。客户端可以提取并微调 (fine-tuning)一个符合其资源预算的子网络。
- 一次性网络 (OFA): 一个代表性例子是训练一个单一的大型网络,使得可以从中提取出多样化子网络(通过改变深度、宽度、核大小获得),并且无需重新训练即可表现良好。
- 在联邦学习中: 服务器可以协调OFA超网络的训练。在每一轮中,客户端根据它们的资源采样一个子网络,在该子网络上执行本地训练,并将与采样架构对应的更新发回。服务器将这些更新聚合回超网络参数 (parameter)中。这使得不同能力的客户端能有效参与训练共享的、可伸缩的模型表示。聚合需要仔细地将子网络更新映射回超网络权重 (weight)。
联邦知识蒸馏 (knowledge distillation) (FKD)
知识蒸馏涉及训练一个更小的“学生”模型来模仿更大、预训练 (pre-training)的“教师”模型的行为。在联邦学习中,全局模型可以充当教师模型。
- 流程:
- 服务器训练一个可能较大、高准确率的全局模型Wglobal(教师模型)。
- 服务器将Wglobal(或者仅其架构和推理 (inference)所需的参数)发送给客户端。
- 每个客户端k定义一个适合其设备的、更小的“学生”模型架构Wstudent,k。
- 在本地训练期间,客户端k不仅在其本地数据Dk上训练其Wstudent,k,还使用教师模型Wglobal在相同本地数据上的输出(例如 logits)作为正则化 (regularization)目标。损失函数 (loss function)可能如下所示:
Lk=αLCE(Wstudent,k,Dk)+(1−α)LKD(Wglobal(x),Wstudent,k(x))对于 x∈Dk
其中LCE是标准交叉熵损失,LKD是蒸馏损失(例如,教师模型和学生模型输出分布之间的Kullback-Leibler散度)。
- 优点: 允许客户端训练适合其设备限制的模型,同时受益于强大的全局模型中编码的知识。
- 挑战: 需要客户端使用教师模型进行推理,这对于某些设备来说可能仍然过大(尽管只需要推理,不需要反向传播 (backpropagation))。共享 logits 也可能涉及隐私问题,尽管通常认为不如梯度敏感。聚合通常涉及根据客户端贡献更新全局教师模型,这可能是间接的(例如,客户端发送它们改进的学生模型,服务器用它们来改进教师模型)。
自适应层/计算跳过
客户端可以在执行过程中基于可用资源或实时限制,动态跳过模型的某些部分。
- 层丢弃: 训练期间随机丢弃层可作为正则化手段。在推理过程中,客户端可以确定性地丢弃特定层(例如,靠后的、更复杂的层)以满足延迟预算。
- 提前退出: 设计带有中间分类器的模型。如果做出了有信心的预测或达到资源限制,客户端可以在这些分类器处提前退出计算。
- 在联邦学习中: 全局模型可以设计为可丢弃层或具有提前退出机制。客户端根据其限制在本地决定使用哪些层或在何处退出。聚合需要处理来自使用不同计算路径的客户端的更新。
将适应方法集成到联邦学习系统中
成功实施剪枝和适应方法需要系统级别的支持:
- 能力报告: 客户端需要一个机制向服务器报告其资源限制(例如,内存、CPU FLOPS、目标延迟)。这可以在初始握手期间或定期发生。
- 服务器协调: 服务器必须管理可能异构的模型配置。这可能涉及:
- 维护全局模型的不同剪枝版本。
- 管理超网络的参数 (parameter)并映射子网络更新。
- 协调知识蒸馏 (knowledge distillation)流程。
- 异构聚合: 核心联邦学习挑战被放大。如果客户端拥有结构不同的模型或更新,简单的平均可能不起作用。方法包括:
- 只平均共享参数。
- 根据客户端模型的“大小”或贡献进行加权平均。
- 使用基于知识的聚合,其中聚合的是客户端输出或模型表示,而非原始参数。
权衡与考量
实施这些方法需要权衡多个因素:
-
准确率与效率: 几乎总存在权衡。激进剪枝或使用非常小的适应模型会减少资源使用,但可能降低能达到的最高准确率。可视化这种权衡会有所帮助。
通过剪枝或适应减小模型尺寸时,准确率通常会下降,需要根据应用需求和设备限制仔细平衡。
-
系统复杂度: 与使用同构模型的标准FedAvg相比,这些方法显著增加了联邦学习系统设计、实现和调试的复杂性。
-
通信开销: 尽管目标通常是设备上的更小模型,但有些技术可能引入新的通信开销(例如,发送能力配置文件、交换知识蒸馏 (knowledge distillation)目标、通信剪枝掩码或架构选择)。
-
公平性: 确保严重受限的设备客户端不被落下。它们适应后的模型仍应提供有用的功能,其贡献应适当地整合到全局学习过程中。
“通过仔细选择和实施模型剪枝和适应策略,联邦学习系统可以更能兼容多样化硬件,从而在以系统异构性为特征的场景中带来更实用和高效的部署。”