Distributionally Robust Federated Learning with Client Drift Minimization

2025-05-21

AI Summary

Federated learning faces critical challenges under data heterogeneity, including severe client drift, poor worst-client performance, and insufficient fairness. To address these, we propose DRDM—a novel algorithm that integrates distributionally robust optimization (DRO) with adaptive dynamic regularization into a min-max optimization framework, explicitly optimizing for worst-client accuracy. We establish its convergence guarantee under partial client participation and convex smooth objectives. Furthermore, DRDM is the first method to jointly optimize communication rounds, worst-client accuracy, and client energy consumption, incorporating an SNR- and bandwidth-aware energy model. Evaluations across three benchmark datasets demonstrate substantial improvements: worst-client accuracy increases significantly, communication rounds decrease by up to 37%, local update steps adapt automatically to target accuracy, and total energy consumption is minimized across diverse communication environments.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Intelligent Robots: Learning & Optimization for ROB
• Search and Optimization: Distributed Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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 extit{DRDM}, a novel algorithm that addresses these issues by combining a distributionally robust optimization (DRO) framework with dynamic regularization to mitigate client drift. extit{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 extit{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.

Problem

• Addresses unfair model performance in federated learning due to data heterogeneity
• Minimizes client drift using robust optimization and dynamic regularization
• Improves worst-case accuracy with fewer communication rounds and energy cost

Innovation

• Combines DRO with dynamic regularization
• Min-max optimization for worst-case client
• Adaptive local updates reduce energy cost