# On the Generalization of Wasserstein Robust Federated Learning

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 50 references that are in the map.
- First posted on arXiv 3 Jun 2022.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Towards Evaluating the Robustness of Neural Networks](https://mlsanity.com/paper/1608.04644/) (2016)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Federated Learning: Strategies for Improving Communication Efficiency](https://mlsanity.com/paper/1610.05492/) (2016)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Adversarial Machine Learning at Scale](https://mlsanity.com/paper/1611.01236/) (2016)
- [Ensemble Adversarial Training: Attacks and Defenses](https://mlsanity.com/paper/1705.07204/) (2017)
- [The Limitations of Deep Learning in Adversarial Settings](https://mlsanity.com/paper/1511.07528/) (2015)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [Federated Multi-Task Learning](https://mlsanity.com/paper/1705.10467/) (2017)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [Personalized Federated Learning with Moreau Envelopes](https://mlsanity.com/paper/2006.08848/) (2020)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Adaptive Personalized Federated Learning](https://mlsanity.com/paper/2003.13461/) (2020)
- [FedPAQ: A Communication-Efficient Federated Learning Method with Periodic Averaging and Quantization](https://mlsanity.com/paper/1909.13014/) (2019)
- [Certifying Some Distributional Robustness with Principled Adversarial Training](https://mlsanity.com/paper/1710.10571/) (2017)
- [Three Approaches for Personalization with Applications to Federated Learning](https://mlsanity.com/paper/2002.10619/) (2020)
- [A Field Guide to Federated Optimization](https://mlsanity.com/paper/2107.06917/) (2021)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [Federated Multi-Task Learning under a Mixture of Distributions](https://mlsanity.com/paper/2108.10252/) (2021)
- [Distributionally Robust Stochastic Optimization with Wasserstein Distance](https://mlsanity.com/paper/1604.02199/) (2016)
- [Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learning](https://mlsanity.com/paper/1908.08729/) (2019)
- [Geometric Dataset Distances via Optimal Transport](https://mlsanity.com/paper/2002.02923/) (2020)
- [Distributed Mean Estimation with Limited Communication](https://mlsanity.com/paper/1611.00429/) (2016)

## Built on by (most cited first, 9 of 9)

- [Federated Domain Generalization: A Survey](https://mlsanity.com/paper/2306.01334/) (2023)
- [Label driven Knowledge Distillation for Federated Learning with non-IID Data](https://mlsanity.com/paper/2209.14520/) (2022)
- [Distributionally Robust Clustered Federated Learning: A Case Study in Healthcare](https://mlsanity.com/paper/2410.07039/) (2024)
- [Convergence of Agnostic Federated Averaging](https://mlsanity.com/paper/2507.10325/) (2025)
- [Distributionally Robust Federated Learning with Outlier Resilience](https://mlsanity.com/paper/2509.24462/) (2025)
- [Distributionally Robust Federated Learning: An ADMM Algorithm](https://mlsanity.com/paper/2503.18436/) (2025)
- [Data Valuation and Selection in a Federated Model Marketplace](https://mlsanity.com/paper/2509.18104/) (2025)
- [Distributionally Robust Federated Learning with Multi-Source Data](https://mlsanity.com/paper/2609.20501/) (2026)
- [FedeRage: Provably Convergent Agnostic Federated Learning under General Client Drift](https://mlsanity.com/paper/2609.21057/) (2026)

## Abstract (arXiv)

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.