FedWMSAM: Fast and Flat Federated Learning via Weighted Momentum and Sharpness-Aware Minimization
Tianle Li, Yongzhi Huang, Linshan Jiang, Chang Liu, Qipeng Xie, Wenfeng Du, Lu Wang, Kaishun Wu
Federated Learning (FL) has emerged as a distributed learning framework that allows multiple clients to train a global model while preserving data privacy collaboratively. However, FL in edge computing faces several unique challenges, including limited communication rounds, data heterogeneity, and client drift, which lead to slow convergence and poor model generalization. Meanwhile, Momentum and Sharpness-Aware Minimization (SAM) show their promotion on the convergence speed and flatness in FL, respectively.To seamlessly concatenate the Momentum and SAM while avoiding their flaws, we propose a novel approach called FedWMSAM to achieve \textit{fast and flat} FL. In detail, by introducing personalized momentum, we mitigate client inconsistency and accelerate convergence. Additionally, by dynamically adjusting the perturbation direction and leveraging momentum-based global flatness, we estimate global flat minima and use a dynamic weighting mechanism to balance the contributions of momentum and SAM adaptively throughout the training process. We conduct extensive experiments on multiple datasets and model architectures, and the results validate the effectiveness, adaptability, and robustness of our method, demonstrating its superiority in addressing the optimization challenges of Federated Learning. Our code is available at \url{https://anonymous.4open.science/r/NeurlPS_FedWMSAM-2AC5}.