A Fast and Flat Federated Learning Method via Weighted Momentum and Sharpness-Aware Minimization
Tianle Li, Yongzhi Huang, Linshan Jiang, Chang Liu, Qipeng Xie, Wenfeng Du, Lu Wang, Kaishun Wu
TL;DR
This paper tackles the challenge of achieving fast convergence with flat minima in non-IID federated learning by integrating momentum with Sharpness-Aware Minimization (SAM). It identifies two failure modes—local-global curvature misalignment and momentum-echo—and introduces FedWMSAM, which uses a momentum-guided global perturbation and a cosine-similarity based adaptive coupling to realize an early-momentum, late-SAM schedule. A non-IID convergence bound is derived that explicitly accounts for perturbation-induced variance, showcasing robustness to data heterogeneity and system factors. Empirically, FedWMSAM demonstrates strong performance across multiple datasets and heterogeneity settings, achieving target accuracies in fewer rounds with near-FedAvg per-round cost, thereby delivering fast-and-flat optimization in realistic FL scenarios.
Abstract
In federated learning (FL), models must \emph{converge quickly} under tight communication budgets while \emph{generalizing} across non-IID client distributions. These twin requirements have naturally led to two widely used techniques: client/server \emph{momentum} to accelerate progress, and \emph{sharpness-aware minimization} (SAM) to prefer flat solutions. However, simply combining momentum and SAM leaves two structural issues unresolved in non-IID FL. We identify and formalize two failure modes: \emph{local-global curvature misalignment} (local SAM directions need not reflect the global loss geometry) and \emph{momentum-echo oscillation} (late-stage instability caused by accumulated momentum). To our knowledge, these failure modes have not been jointly articulated and addressed in the FL literature. We propose \textbf{FedWMSAM} to address both failure modes. First, we construct a momentum-guided global perturbation from server-aggregated momentum to align clients' SAM directions with the global descent geometry, enabling a \emph{single-backprop} SAM approximation that preserves efficiency. Second, we couple momentum and SAM via a cosine-similarity adaptive rule, yielding an early-momentum, late-SAM two-phase training schedule. We provide a non-IID convergence bound that \emph{explicitly models the perturbation-induced variance} $σ_ρ^2=σ^2+(Lρ)^2$ and its dependence on $(S, K, R, N)$ on the theory side. 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 https://github.com/Huang-Yongzhi/NeurlPS_FedWMSAM.