FedDC: Federated Learning with Non-IID Data via Local Drift Decoupling and Correction (2203.11751v1)

Abstract: Federated learning (FL) allows multiple clients to collectively train a high-performance global model without sharing their private data. However, the key challenge in federated learning is that the clients have significant statistical heterogeneity among their local data distributions, which would cause inconsistent optimized local models on the client-side. To address this fundamental dilemma, we propose a novel federated learning algorithm with local drift decoupling and correction (FedDC). Our FedDC only introduces lightweight modifications in the local training phase, in which each client utilizes an auxiliary local drift variable to track the gap between the local model parameter and the global model parameters. The key idea of FedDC is to utilize this learned local drift variable to bridge the gap, i.e., conducting consistency in parameter-level. The experiment results and analysis demonstrate that FedDC yields expediting convergence and better performance on various image classification tasks, robust in partial participation settings, non-iid data, and heterogeneous clients.

An Overview of "FedDC: Federated Learning with Non-IID Data via Local Drift Decoupling and Correction"

The paper introduces FedDC, a federated learning (FL) algorithm designed to address statistical heterogeneity issues inherent in FL systems. Federated learning allows multiple clients to collaboratively train a global model while maintaining data privacy by keeping the data localized. However, the performance and convergence of federated models are often impacted by non-IID (Independent and Identically Distributed) data distributions across different clients. FedDC aims to mitigate these challenges by employing a local drift decoupling and correction technique.

Problem Statement and Motivations

In federated learning, data heterogeneity among clients leads to inconsistent optimal local models that collectively hinder the convergence of the global model. Traditional approaches like FedAvg struggle with this issue as they are not inherently designed to handle non-IID data distributions. As clients optimize their local models independently, variations in data distribution can result in significant parameter drift, thereby slowing convergence or leading to sub-optimal global models.

Methodology: Local Drift Decoupling and Correction

FedDC addresses the core issue of local parameter drift by introducing auxiliary drift variables that learn and track the gap between local and global model parameters. The methodological innovations in FedDC include:

• Local Drift Tracking: Each client maintains an auxiliary variable to track the local drift, allowing better alignment between local and global parameters.
• Parameter Correction: Leveraging the drift variables to adjust local updates before aggregation. This correction step aims to bridge the gap caused by statistical heterogeneity.
• Constrained Local Objectives: The local objective function incorporates a penalty term related to the drift, effectively balancing the local model updates with respect to the global model's parameter space.

The algorithm operates in rounds, where each client performs local updates by optimizing their models with respect to their drift-adjusted objectives. Following local training, clients send corrected updates to the server, where these are aggregated to update the global model.

Empirical Evaluation

The paper provides comprehensive empirical evaluations across various benchmarks, including MNIST, CIFAR10, and CIFAR100 datasets, employing scenarios with both IID and non-IID data distributions. These experiments demonstrate:

• Faster Convergence: FedDC consistently reaches target accuracy levels in fewer communication rounds compared to baselines like FedAvg, FedProx, and Scaffold.
• Improved Accuracy: The algorithm achieves better model performance across various settings, indicating its robustness to the statistical heterogeneity of client data.
• Scalability: FedDC performs well in both full and partial client participation scenarios, showcasing its applicability to practical federated learning deployments.

These results underscore the potential of FedDC in environments characterized by diverse data distributions, which is a common scenario in real-world federated deployments.

Implications and Future Directions

The introduction of FedDC represents a significant step towards enhancing federated learning systems' robustness to non-IID data distributions. By decoupling and correcting local drift, FedDC provides a methodological framework that can potentially be extended or combined with other optimization strategies to further enhance federated model performance.

Future research directions could explore the integration of adaptive optimizers within the FedDC framework or investigate the scalability of FedDC in massively distributed systems with limited computation resources. Additionally, extending this approach to encompass other federated learning challenges, such as communication efficiency and privacy-preserving mechanisms, could further strengthen the utility and adoption of federated learning models across industries.

In conclusion, FedDC offers a pragmatic solution to the pressing challenge of statistical heterogeneity in FL, enhancing convergence and model efficacy while respecting data privacy constraints.