# FedSWA: Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Based Stochastic Controlled Weight Averaging

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

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 49 references that are in the map.
- First posted on arXiv 26 Jul 2025.

## Builds on (most cited first, 30 of 34)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [The Future of Digital Health with Federated Learning](https://mlsanity.com/paper/2003.08119/) (2020)
- [Cyclical Learning Rates for Training Neural Networks](https://mlsanity.com/paper/1506.01186/) (2015)
- [Convex Optimization: Algorithms and Complexity](https://mlsanity.com/paper/1405.4980/) (2014)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [Lookahead Optimizer: k steps forward, 1 step back](https://mlsanity.com/paper/1907.08610/) (2019)
- [Momentum-Based Variance Reduction in Non-Convex SGD](https://mlsanity.com/paper/1905.10018/) (2019)
- [FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Data](https://mlsanity.com/paper/2005.11418/) (2020)
- [Generalized Federated Learning via Sharpness Aware Minimization](https://mlsanity.com/paper/2206.02618/) (2022)
- [A Closer Look at Deep Learning Heuristics: Learning rate restarts, Warmup and Distillation](https://mlsanity.com/paper/1810.13243/) (2018)
- [Stochastic Controlled Averaging for Federated Learning with Communication Compression](https://mlsanity.com/paper/2308.08165/) (2023)
- [Improving Generalization in Federated Learning by Seeking Flat Minima](https://mlsanity.com/paper/2203.11834/) (2022)
- [How much does your data exploration overfit? Controlling bias via information usage](https://mlsanity.com/paper/1511.05219/) (2015)
- [When Do Flat Minima Optimizers Work?](https://mlsanity.com/paper/2202.00661/) (2022)
- [Differentially Private Secure Multi-Party Computation for Federated Learning in Financial Applications](https://mlsanity.com/paper/2010.05867/) (2020)
- [Stop Wasting My Time! Saving Days of ImageNet and BERT Training with Latest Weight Averaging](https://mlsanity.com/paper/2209.14981/) (2022)
- [FedSpeed: Larger Local Interval, Less Communication Round, and Higher Generalization Accuracy](https://mlsanity.com/paper/2302.10429/) (2023)
- [LoRA-FAIR: Federated LoRA Fine-Tuning with Aggregation and Initialization Refinement](https://mlsanity.com/paper/2411.14961/) (2024)
- [Stability of SGD: Tightness Analysis and Improved Bounds](https://mlsanity.com/paper/2102.05274/) (2021)
- [Taming Cross-Domain Representation Variance in Federated Prototype Learning with Heterogeneous Data Domains](https://mlsanity.com/paper/2403.09048/) (2024)
- [Locally Estimated Global Perturbations are Better than Local Perturbations for Federated Sharpness-aware Minimization](https://mlsanity.com/paper/2405.18890/) (2024)

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

- [HUD: Hierarchical Uncertainty-Aware Disambiguation Network for Composed Video Retrieval](https://mlsanity.com/paper/2512.02792/) (2025)
- [FedAdamW: A Communication-Efficient Optimizer with Convergence and Generalization Guarantees for Federated Large Models](https://mlsanity.com/paper/2510.27486/) (2025)
- [DP-FedPGN: Finding Global Flat Minima for Differentially Private Federated Learning via Penalizing Gradient Norm](https://mlsanity.com/paper/2510.27504/) (2025)
- [CompTrack: Information Bottleneck-Guided Low-Rank Dynamic Token Compression for Point Cloud Tracking](https://mlsanity.com/paper/2511.15580/) (2025)
- [FedNSAM:Consistency of Local and Global Flatness for Federated Learning](https://mlsanity.com/paper/2602.23827/) (2026)
- [FocusTrack: One-Stage Focus-and-Suppress Framework for 3D Point Cloud Object Tracking](https://mlsanity.com/paper/2602.24133/) (2026)
- [Adaptive LoRA Experts Allocation and Selection for Federated Fine-Tuning](https://mlsanity.com/paper/2509.15087/) (2025)
- [CATCH: A Modular Cross-domain Adaptive Template with Hook](https://mlsanity.com/paper/2510.26582/) (2025)
- [Rethinking LoRA for Privacy-Preserving Federated Learning in Large Models](https://mlsanity.com/paper/2602.19926/) (2026)
- [Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data](https://mlsanity.com/paper/2602.19271/) (2026)
- [FedTreeLoRA: Reconciling Statistical and Functional Heterogeneity in Federated LoRA Fine-Tuning](https://mlsanity.com/paper/2603.13282/) (2026)
- [Contrastive Prompt Clustering for Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/2508.17009/) (2025)
- [TOFA: Training-Free One-Shot Federated Adaptation for Vision-Language Models](https://mlsanity.com/paper/2511.16423/) (2025)
- [DP-FedAdamW: An Efficient Optimizer for Differentially Private Federated Large Models](https://mlsanity.com/paper/2602.19945/) (2026)
- [Deconstructing the Failure of Ideal Noise Correction: A Three-Pillar Diagnosis](https://mlsanity.com/paper/2603.12997/) (2026)
- [A Step Toward Federated Pretraining of Multimodal Large Language Models](https://mlsanity.com/paper/2603.26786/) (2026)
- [A Note on Stability for Orthogonalized Matrix Momentum with Client Sampling](https://mlsanity.com/paper/2606.01720/) (2026)
- [SeFoRA: Sketch-Aggregated Federated Low-Rank Adaptation with Heterogeneous Client Ranks](https://mlsanity.com/paper/2608.10144/) (2026)
- [FRIEREN: Federated Learning with Vision-Language Regularization for Segmentation](https://mlsanity.com/paper/2510.02114/) (2025)
- [FQ-PETR: Fully Quantized Position Embedding Transformation for Multi-View 3D Object Detection](https://mlsanity.com/paper/2511.09347/) (2025)
- [HiLoRA: Hierarchical Low-Rank Adaptation for Personalized Federated Learning](https://mlsanity.com/paper/2603.02785/) (2026)
- [Fusion and Alignment Enhancement with Large Language Models for Tail-item Sequential Recommendation](https://mlsanity.com/paper/2604.03688/) (2026)
- [CTP-FL: Common-Trajectory Gradient Prediction for Federated Learning](https://mlsanity.com/paper/2609.35130/) (2026)
- [FedMIX-P: Mixing Local and Global Preconditioners for Federated Vision and Language Model Training](https://mlsanity.com/paper/2610.01515/) (2026)
- [FedLore: Communication and Memory Efficient Federated Learning via Shared Gradient Low-Rank Projection](https://mlsanity.com/paper/2610.01620/) (2026)

## Abstract (arXiv)

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 textttFedSWA), 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 (textttFedMoSWA), which is designed to better align local and global models. Theoretically, we provide both convergence analysis and generalization bounds for textttFedSWA and textttFedMoSWA. We also prove that the optimization and generalization errors of textttFedMoSWA 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.