# A Fast and Flat Federated Learning Method via Weighted Momentum and Sharpness-Aware Minimization

arXiv: https://arxiv.org/abs/2511.22080 · HTML: https://mlsanity.com/paper/2511.22080/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 27 Nov 2025.

## Built on by (most cited first, 3 of 3)

- [FedSCAM (Federated Sharpness-Aware Minimization with Clustered Aggregation and Modulation): Scam-resistant SAM for Robust Federated Optimization in Heterogeneous Environments](https://mlsanity.com/paper/2601.00853/) (2025)
- [Beyond Fixed Rounds: Data-Free Early Stopping for Practical Federated Learning](https://mlsanity.com/paper/2601.22669/) (2026)
- [EarLog: Opportunistic Ear Physiology Sensing toward Personalized Hearing Risk Awareness](https://mlsanity.com/paper/2610.12501/) (2026)

## Abstract (arXiv)

In federated learning (FL), models must converge quickly under tight communication budgets while generalizing across non-IID client distributions. These twin requirements have naturally led to two widely used techniques: client/server momentum to accelerate progress, and 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: local-global curvature misalignment (local SAM directions need not reflect the global loss geometry) and 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 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 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 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/NeurlPSFedWMSAM.