[论文解读] Achieving Personalized Federated Learning with Sparse Local Models

Tiansheng Huang, Shiwei Liu+4 more|arXiv (Cornell University)|Jan 27, 2022

Privacy-Preserving Technologies in Data被引用 13

一句话总结

该论文提出FedSpa，一种个性化联邦学习框架，通过使用带有个性化掩码的稀疏本地模型，在保持低通信和计算成本的同时实现高精度。通过在整个训练过程中维持固定的稀疏度水平（稀疏到稀疏训练），FedSpa实现了高效且个性化的模型更新，同时以𝒪(1/√T)的速率收敛至局部极小值。

ABSTRACT

Federated learning (FL) is vulnerable to heterogeneously distributed data, since a common global model in FL may not adapt to the heterogeneous data distribution of each user. To counter this issue, personalized FL (PFL) was proposed to produce dedicated local models for each individual user. However, PFL is far from its maturity, because existing PFL solutions either demonstrate unsatisfactory generalization towards different model architectures or cost enormous extra computation and memory. In this work, we propose federated learning with personalized sparse mask (FedSpa), a novel PFL scheme that employs personalized sparse masks to customize sparse local models on the edge. Instead of training an intact (or dense) PFL model, FedSpa only maintains a fixed number of active parameters throughout training (aka sparse-to-sparse training), which enables users' models to achieve personalization with cheap communication, computation, and memory cost. We theoretically show that the iterates obtained by FedSpa converge to the local minimizer of the formulated SPFL problem at rate of $\mathcal{O}(\frac{1}{\sqrt{T}})$. Comprehensive experiments demonstrate that FedSpa significantly saves communication and computation costs, while simultaneously achieves higher model accuracy and faster convergence speed against several state-of-the-art PFL methods.

研究动机与目标

• 解决联邦学习中非独立同分布（Non-IID）数据分布带来的挑战，即单一全局模型难以在异构客户端上实现良好泛化。
• 克服现有个性化联邦学习（PFL）方法依赖密集模型或动态稀疏性所导致的高通信与计算开销。
• 开发一种在训练过程中保持恒定稀疏度的PFL解决方案，以降低边缘设备上的内存与计算成本。
• 在保留个性化与模型性能的前提下，实现稀疏子模型的高效聚合。
• 理论上分析收敛性，并通过实证验证其在准确率、速度与效率方面优于当前最先进PFL方法。

提出的方法

• 提出一种新颖的稀疏个性化联邦学习（SPFL）问题形式化，通过在全局模型上应用个性化稀疏掩码，生成客户端特定的稀疏子模型。
• 设计FedSpa作为PFL框架，其中每个客户端在整个训练过程中维护并训练一个具有固定稀疏度水平的稀疏模型。
• 采用动态稀疏训练（DST）机制，基于本地梯度与数据搜索新的个性化掩码，以确保模型个性化。
• 在掩码更新过程中应用基于梯度的权重恢复方法，选择梯度坐标中绝对值最大的αₜ个以恢复被剪枝的权重。
• 在服务器端聚合稀疏梯度更新与新掩码，通过标准平均法更新全局模型。
• 确保所有本地模型在整个训练过程中始终保持稀疏，从而最小化通信与计算开销。

实验结果

研究问题

• RQ1 联邦学习系统是否能在非独立同分布（Non-IID）设置下实现高个性化性能，同时保持低通信与计算成本？
• RQ2 与基于密集模型的PFL基线相比，使用个性化掩码的稀疏到稀疏训练是否能实现更快收敛与更高准确率？
• RQ3 基于梯度的权重恢复在掩码重训练过程中对提升模型性能有多有效？
• RQ4 通过FedSpa训练的全局模型是否能在数据异构性条件下达到与标准FedAvg相当的性能？
• RQ5 掩码初始化策略对FedSpa在异构数据环境下的最终性能有何影响？

主要发现

• 与密集PFL基线相比，FedSpa将通信开销减少了50%，因为每轮仅传输稀疏梯度与掩码。
• 在非独立同分布（Non-IID）设置下，FedSpa的测试准确率比FedAvg最高提升21.9%，比Ditto提升4.4%。
• FedSpa达到相同准确率水平所需的通信轮次比FedAvg减少76.2%，比Ditto减少38.1%。
• FedSpa中掩码搜索过程（DST）每轮仅占总运行时间的5%至20%，表明运行时开销极低。
• 与随机恢复相比，FedSpa中基于梯度的权重恢复（DST）能实现更快收敛与更高最终准确率。
• 在非独立同分布（Non-IID）设置下，FedSpa的全局模型性能显著劣于个性化模型，证实子网络在数据偏移下优于完整模型。