Window-based Model Averaging Improves Generalization in Heterogeneous Federated Learning

TL;DR

WIMA enhances federated learning by aggregating models over a window of rounds, improving robustness and stability without extra communication, especially under data heterogeneity and distribution shifts.

Contribution

The paper introduces WIMA, a window-based model averaging method that improves federated learning robustness and generalization without additional communication overhead.

Findings

WIMA improves stability under data heterogeneity.

WIMA enhances robustness against distribution shifts.

WIMA integrates easily with existing algorithms.

Abstract

Federated Learning (FL) aims to learn a global model from distributed users while protecting their privacy. However, when data are distributed heterogeneously the learning process becomes noisy, unstable, and biased towards the last seen clients' data, slowing down convergence. To address these issues and improve the robustness and generalization capabilities of the global model, we propose WIMA (Window-based Model Averaging). WIMA aggregates global models from different rounds using a window-based approach, effectively capturing knowledge from multiple users and reducing the bias from the last ones. By adopting a windowed view on the rounds, WIMA can be applied from the initial stages of training. Importantly, our method introduces no additional communication or client-side computation overhead. Our experiments demonstrate the robustness of WIMA against distribution shifts and bad client sampling, resulting in smoother and more stable learning trends. Additionally, WIMA can be easily integrated with state-of-the-art algorithms. We extensively evaluate our approach on standard FL benchmarks, demonstrating its effectiveness.