针对联邦学习在异构数据下的泛化能力差问题,提出新算法提升模型稳定性。
FedSWA: Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Based Stochastic Controlled Weight Averaging
- 基于动量的随机加权平均策略,寻找更平坦的损失极小值。
- 在CIFAR10/100和Tiny ImageNet上优于FedSAM等基线方法。
- 适合数据异构严重的实际联邦学习场景,代码开源可复现。
针对联邦学习(FL)算法如FedSAM的泛化能力在真实应用中至关重要。本文重新审视了联邦学习中的泛化问题,研究数据异构性对泛化的影响。发现当数据高度异构时,FedSAM的表现通常劣于FedAvg,因此提出一种新的联邦学习算法FedSWA,采用随机权重平均(Stochastic Weight Averaging),旨在高异构数据下寻找更平坦的极小值点。进一步设计了基于动量的受控权重平均算法FedMoSWA,以更好对齐局部与全局模型。理论上,提供了FedSWA和FedMoSWA的收敛性分析及泛化界,并证明其优化与泛化误差小于包括FedSAM在内的基准方法。实验在CIFAR10/100和Tiny ImageNet上验证了所提算法的优越性。开源代码已发布于https://github.com/junkangLiu0/FedSWA。
原文摘要 · Abstract (English)
For federated learning (FL) algorithms such as FedSAM, their generalization capability is crucial for real-word applications. In this paper, we revisit the generalization problem in FL and investigate the impact of data heterogeneity on FL generalization. We find that FedSAM usually performs worse than FedAvg in the case of highly heterogeneous data, and thus propose a novel and effective federated learning algorithm with Stochastic Weight Averaging (called \texttt{FedSWA}), which aims to find flatter minima in the setting of highly heterogeneous data. Moreover, we introduce a new momentum-based stochastic controlled weight averaging FL algorithm (\texttt{FedMoSWA}), which is designed to better align local and global models. Theoretically, we provide both convergence analysis and generalization bounds for \texttt{FedSWA} and \texttt{FedMoSWA}. We also prove that the optimization and generalization errors of \texttt{FedMoSWA} are smaller than those of their counterparts, including FedSAM and its variants. Empirically, experimental results on CIFAR10/100 and Tiny ImageNet demonstrate the superiority of the proposed algorithms compared to their counterparts. Open source code at: https://github.com/junkangLiu0/FedSWA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。