Distributionally Robust Federated Learning with Client Drift Minimization

TL;DR

This paper introduces DRDM, a robust federated learning algorithm that minimizes client drift and improves worst-case performance through distributionally robust optimization, dynamic regularization, and efficient communication, with theoretical guarantees and extensive experiments.

Contribution

The paper proposes DRDM, a novel federated learning method combining DRO and dynamic regularization to enhance robustness and fairness in heterogeneous environments, reducing communication rounds and energy consumption.

Findings

DRDM improves worst-case test accuracy significantly.

DRDM reduces communication rounds compared to baselines.

Energy consumption analysis shows adaptive local updates are effective.

Abstract

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.