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

TL;DR

This paper introduces Libra, a new framework for analyzing the generalization performance of federated learning algorithms, focusing on the trade-off between model stability and optimization, and providing insights for hyperparameter tuning.

Contribution

The paper proposes Libra, an innovative analysis framework that captures the trade-offs affecting generalization in federated learning, especially for non-convex neural networks.

Findings

Larger local steps or momentum improve convergence speed but reduce model stability.

Libra effectively predicts the impact of hyperparameters on generalization performance.

Experimental results validate the theoretical insights and guide hyperparameter tuning.

Abstract

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 \textit{convergence analysis} on gradient norms, focusing on how well a model fits training data, or through \textit{algorithmic 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 \textit{Libra}, 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.