# Understanding Generalization of Federated Learning: the Trade-off between Model Stability and Optimization

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 56 of its 76 references that are in the map.
- First posted on arXiv 25 Nov 2024.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [TinyBERT: Distilling BERT for Natural Language Understanding](https://mlsanity.com/paper/1909.10351/) (2019)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/1901.08584/) (2019)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [A Unified Theory of Decentralized SGD with Changing Topology and Local Updates](https://mlsanity.com/paper/2003.10422/) (2020)
- [Federated Learning for Internet of Things: A Comprehensive Survey](https://mlsanity.com/paper/2104.07914/) (2021)
- [SGD on Neural Networks Learns Functions of Increasing Complexity](https://mlsanity.com/paper/1905.11604/) (2019)
- [Reasoning About Generalization via Conditional Mutual Information](https://mlsanity.com/paper/2001.09122/) (2020)
- [Generalized Federated Learning via Sharpness Aware Minimization](https://mlsanity.com/paper/2206.02618/) (2022)
- [Federated Learning for Open Banking](https://mlsanity.com/paper/2108.10749/) (2021)
- [Non-convex Optimization for Machine Learning](https://mlsanity.com/paper/1712.07897/) (2017)
- [Rademacher Complexity for Adversarially Robust Generalization](https://mlsanity.com/paper/1810.11914/) (2018)
- [Towards Theoretically Understanding Why SGD Generalizes Better Than ADAM in Deep Learning](https://mlsanity.com/paper/2010.05627/) (2020)
- [DiLoCo: Distributed Low-Communication Training of Language Models](https://mlsanity.com/paper/2311.08105/) (2023)
- [Fine-Grained Analysis of Stability and Generalization for Stochastic Gradient Descent](https://mlsanity.com/paper/2006.08157/) (2020)
- [Robust Federated Learning: The Case of Affine Distribution Shifts](https://mlsanity.com/paper/2006.08907/) (2020)
- [Stability and Generalization of Learning Algorithms that Converge to Global Optima](https://mlsanity.com/paper/1710.08402/) (2017)
- [Federated Learning for Generalization, Robustness, Fairness: A Survey and Benchmark](https://mlsanity.com/paper/2311.06750/) (2023)

## Abstract (arXiv)

Federated Learning (FL) is a distributed learning approach that trains machine learning models across multiple devices while keeping their local data private. However, FL often faces challenges due to data heterogeneity, leading to inconsistent local optima among clients. These inconsistencies can cause unfavorable convergence behavior and generalization performance degradation. Existing studies often describe this issue through textitconvergence analysis on gradient norms, focusing on how well a model fits training data, or through textitalgorithmic stability, which examines the generalization gap. However, neither approach precisely captures the generalization performance of FL algorithms, especially for non-convex neural network training. In response, this paper introduces an innovative generalization dynamics analysis framework, namely textitLibra, for algorithm-dependent excess risk minimization, highlighting the trade-offs between model stability and gradient norms. We present Libra towards a standard federated optimization framework and its variants using server momentum. Through this framework, we show that larger local steps or momentum accelerate convergence of gradient norms, while worsening model stability, yielding better excess risk. Experimental results on standard FL settings prove the insights of our theories. These insights can guide hyperparameter tuning and future algorithm design to achieve stronger generalization.