Distributionally Robust Federated Learning for Mobile Edge Networks

Federated Learning (FL) revolutionizes data processing in mobile networks by enabling collaborative learning without data exchange. This not only reduces latency and enhances computational efficiency but also enables the system to adapt, learn and optimize the performance from the user’s context in real-time. Nevertheless, FL faces challenges in training and generalization due to statistical heterogeneity, stemming from the diverse data nature across varying user contexts. To address these challenges, we propose \(\textsf {WAFL}\), a robust FL framework grounded in Wasserstein distributionally robust optimization, aimed at enhancing model generalization against all adversarial distributions within a predefined Wasserstein ambiguity set. We approach \(\textsf {WAFL}\) by formulating it as an empirical surrogate risk minimization problem, which is then solved using a novel federated algorithm. Experimental results demonstrate that \(\textsf {WAFL}\) outperforms other robust FL baselines in non-i.i.d settings, showcasing superior generalization and robustness to significant distribution shifts.