Rethinking the initialization of Momentum in Federated Learning with Heterogeneous Data

TL;DR

This paper introduces Reversed Momentum Federated Learning (RMFL), a novel initialization method that assigns exponentially decayed weights to gradients over time, aiming to better handle data heterogeneity in federated learning.

Contribution

The paper proposes a new momentum calculation method for federated learning that reverses traditional weighting, improving performance under data heterogeneity.

Findings

RMFL outperforms traditional methods on benchmark datasets.

RMFL shows robustness across different heterogeneity levels.

Experimental results validate the effectiveness of reversed momentum in federated settings.

Abstract

Data Heterogeneity is a major challenge of Federated Learning performance. Recently, momentum based optimization techniques have beed proved to be effective in mitigating the heterogeneity issue. Along with the model updates, the momentum updates are transmitted to the server side and aggregated. Therefore, the local training initialized with a global momentum is guided by the global history of the gradients. However, we spot a problem in the traditional cumulation of the momentum which is suboptimal in the Federated Learning systems. The momentum used to weight less on the historical gradients and more on the recent gradients. This however, will engage more biased local gradients in the end of the local training. In this work, we propose a new way to calculate the estimated momentum used in local initialization. The proposed method is named as Reversed Momentum Federated Learning (RMFL). The key idea is to assign exponentially decayed weights to the gradients with the time going forward, which is on the contrary to the traditional momentum cumulation. The effectiveness of RMFL is evaluated on three popular benchmark datasets with different heterogeneity levels.