On the Generalization of Wasserstein Robust Federated Learning

TL;DR

This paper introduces WAFL, a Wasserstein distributionally robust optimization method for federated learning, which improves generalization and robustness against distribution shifts and unseen domains, with theoretical guarantees and empirical validation.

Contribution

We propose WAFL, a novel Wasserstein-based robust federated learning framework with convergence guarantees and enhanced generalization to unseen distributions.

Findings

WAFL outperforms FedAvg in non-i.i.d. settings.

WAFL is more robust under distribution shifts.

WAFL generalizes well to unseen target domains.

Abstract

In federated learning, participating clients typically possess non-i.i.d. data, posing a significant challenge to generalization to unseen distributions. To address this, we propose a Wasserstein distributionally robust optimization scheme called WAFL. Leveraging its duality, we frame WAFL as an empirical surrogate risk minimization problem, and solve it using a local SGD-based algorithm with convergence guarantees. We show that the robustness of WAFL is more general than related approaches, and the generalization bound is robust to all adversarial distributions inside the Wasserstein ball (ambiguity set). Since the center location and radius of the Wasserstein ball can be suitably modified, WAFL shows its applicability not only in robustness but also in domain adaptation. Through empirical evaluation, we demonstrate that WAFL generalizes better than the vanilla FedAvg in non-i.i.d. settings, and is more robust than other related methods in distribution shift settings. Further, using benchmark datasets we show that WAFL is capable of generalizing to unseen target domains.