Distributionally Robust Federated Learning with Client Drift Minimization

Abstract

Federated learning (FL) faces critical challenges, particularly in heterogeneous environments where non-independent and identically distributed (non-IID) data across clients can lead to unfair and inefficient model performance. We introduce DRDM, a novel algorithm that integrates distributionally robust optimization (DRO) with dynamic regularization to explicitly mitigate client drift. Compared to previous approaches that address robustness or drift separately, DRDM combines both aspects within a unified framework, dynamically aligning local updates with the global robust objective to improve convergence toward a worst-case optimal model while maintaining fairness across clients. The robust objective is optimized through efficient local updates, which significantly reduce the number of communication rounds. We provide a theoretical convergence analysis for convex smooth objectives under partial client participation and multiple local update steps. Experiments on three benchmark datasets, covering various model architectures and levels of data heterogeneity, show that DRDM consistently improves worst-case test accuracy while requiring fewer communication rounds than state-of-the-art baselines. Furthermore, we analyze the impact of signal-to-noise ratio (SNR) and bandwidth on energy consumption, demonstrating that adaptive selection of local updates can achieve a target worst-case accuracy with minimal total energy cost across diverse communication environments.