Distributionally Robust Federated Learning with Client Drift Minimization
This work addresses critical challenges in federated learning for applications with heterogeneous data across clients, offering a robust and fair solution, though it is incremental as it builds on existing DRO and regularization methods.
The paper tackles the problem of unfair and inefficient model performance in federated learning due to heterogeneous client data by introducing DRDM, a novel algorithm that combines distributionally robust optimization with dynamic regularization to minimize client drift. The result is significantly improved worst-case test accuracy and reduced communication rounds compared to state-of-the-art baselines, with experiments on three benchmark datasets showing these gains.
Federated learning (FL) faces critical challenges, particularly in heterogeneous environments where non-independent and identically distributed data across clients can lead to unfair and inefficient model performance. In this work, we introduce \textit{DRDM}, a novel algorithm that addresses these issues by combining a distributionally robust optimization (DRO) framework with dynamic regularization to mitigate client drift. \textit{DRDM} frames the training as a min-max optimization problem aimed at maximizing performance for the worst-case client, thereby promoting robustness and fairness. This robust objective is optimized through an algorithm leveraging dynamic regularization and efficient local updates, which significantly reduces the required number of communication rounds. Moreover, we provide a theoretical convergence analysis for convex smooth objectives under partial participation. Extensive experiments on three benchmark datasets, covering various model architectures and data heterogeneity levels, demonstrate that \textit{DRDM} significantly improves worst-case test accuracy while requiring fewer communication rounds than existing state-of-the-art baselines. Furthermore, we analyze the impact of signal-to-noise ratio (SNR) and bandwidth on the energy consumption of participating clients, demonstrating that the number of local update steps can be adaptively selected to achieve a target worst-case test accuracy with minimal total energy cost across diverse communication environments.