Source-linked AI summary
FedDC: Federated Learning with Non-IID Data via Local Drift Decoupling and Correction
Liang Gao, Huazhu Fu, Li Li, Yingwen Chen, Ming Xu, Cheng-Zhong Xu
TL;DR
Non-IID client data create local drift and inconsistent optimization in federated learning. FedDC learns auxiliary drift variables during lightweight local training to bridge local and global parameters, and experiments report faster convergence and better performance across image-classification settings. The method’s analysis assumes non-convex smooth losses with bounded dissimilarity, and its objective cannot directly update the global model because global data are unavailable.
Problem
Non-IID client distributions create inconsistent local objectives and optimization directions, producing drift between local models and the global model.
Method
FedDC learns auxiliary local drift variables, uses them to bridge local-global parameter gaps, and combines penalty and gradient-correction terms during local training.
Results
FedDC provides better performance and faster convergence across image-classification datasets, including robustness under heterogeneous data and partial client participation.
Takeaways & Limitations
Learning and bridging the parameter gap lets FedDC address local drift without hindering local models from fitting their local data distributions.
Takeaways & Limitations
The convergence guarantee assumes non-convex smooth local losses with bounded dissimilarity, and global parameters cannot be updated directly because global data are unavailable.
Abstract
from arXiv · showhide
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.
1. Introduction
Federated learning must reconcile decentralized private data with a shared global model, but non-IID client distributions create local drift and inconsistent optimization. FedDC addresses this by tracking and correcting local drift rather than forcing local and global models to coincide.
- Motivation: Non-IID client data creates inconsistent local objectives and optimization directions, producing client drift that slows convergence.The drift includes residual parameter drift from the previous round and gradient drift in the current round.
- Limitations of prior methods: Existing methods reduce local-update variance, but residual parameter deviation can accumulate and impair convergence and performance.FedProx may hinder movement toward the global stationary point, while Scaffold only approximately reduces gradient drift.
- FedDC: FedDC targets heterogeneous objectives by introducing auxiliary drift variables that track the gap between local and global model parameters.Its local objective combines a constraint penalty with a gradient correction term.
- FedDC: FedDC uses learned drift variables to decouple local and global models, reducing local drift’s impact on the global objective.The proposed modifications are lightweight and are applied during local training.
2. Related Work
Related work characterizes heterogeneous federated learning as a convergence challenge and develops methods that modify local objectives, gradients, regularization, or aggregation. FedDC belongs to the line of approaches addressing client drift and statistical heterogeneity.
- Federated learning under heterogeneity: FedAvg converges asymptotically for homogeneous clients, whereas heterogeneous clients can have different convergence behavior and face gradient divergence and biased optimization directions.A uniform global model is difficult to optimize when client data are non-IID.
- Local-objective methods: FedProx adds proximal regularization to keep each local model near the global model and improve stability under statistical heterogeneity.
- Implementation: The FedDC implementation provides publicly available code.
- Client-drift correction: Scaffold customizes client gradients, while FedDyn uses a dynamic device-specific regularizer to align local and global solutions.Other methods also optimize the server-side aggregation step.
3. Local Drift in Federated Learning
Local models trained on heterogeneous client data drift from the centralized global objective, and ignoring this drift can bias FedAvg’s aggregated model. With nonlinear model transformations, averaging parameters does not generally reproduce the centralized model’s output.
- Federated optimization: Federated learning seeks a global model from private client datasets, while FedAvg coordinates local optimization and server aggregation without sharing raw data.The server broadcasts updated global parameters as the starting point for the next round.
- Illustration: Figure 1 contrasts the centralized ideal parameter wc with the FedAvg parameter wf and the local parameters θ1 and θ2.The illustration uses a nonlinear sigmoid activation function f.
- Local drift: A local model trained on a client dataset can differ from the model trained directly on the global dataset, creating local drift.Highly skewed non-IID data significantly reduces FedAvg performance when this drift is ignored.
- Illustration: For nonlinear f, FedAvg’s parameter wf differs from the centralized parameter wc, and their outputs differ as well.Thus parameter averaging does not generally preserve the centralized output.
4. Proposed Method
FedDC modifies local training with drift variables, penalty and gradient-correction terms, then alternates client and server updates. Its analysis assumes non-convex smooth losses and bounded dissimilarity, while its discussion distinguishes learning the parameter gap from merely restricting local updates.
- 4. Proposed Method: FedDC learns a client-specific drift variable representing the gap between the local model and the global model, then uses it to correct local parameters.The method aims to improve convergence speed and robustness with lightweight local-training modifications.
- 4.1. Objectives in FedDC: Each client objective combines empirical loss, a penalty enforcing the drift relationship, and a gradient-correction term.The correction term reduces the variance of local gradients, while the penalty converts the constraint into an unconstrained objective.
- 4.2. Training Procedure: Each round initializes local models from the global parameters, performs K local updates, corrects local parameters with drift variables, and aggregates the corrected parameters at the server.The server uses a sample-weighted average to obtain the updated global model.
- 4.1. Objectives in FedDC: The global model cannot be updated directly through the drift constraint because global data are unavailable.
- 4.2. Updating Drift Variables: FedDC reuses local-model updates to update drift variables, avoiding an additional back-propagation process for the drift.This approximation is introduced to reduce computation.
- 4.2. Training Procedure: FedDC’s alternating optimization treats local and global parameter updates as client-side M-steps and server-side E-steps in an EM-style procedure.The alternating scheme handles variables updated on different devices.
- 4.3. Convergence: The convergence analysis assumes non-convex β-Lipschitz-smooth local losses, bounded dissimilarity, and suitable parameter conditions.Under these assumptions, the global empirical loss decreases according to the stated bound.
- 4.4. Discussion: FedDC learns and bridges the local-global parameter gap instead of restricting local models to remain close to the global model.This is intended to preserve local fitting while addressing parameter inconsistency.
5. Experiments
Experiments compare FedDC with established federated optimization methods across datasets, architectures, participation levels, and data distributions. FedDC consistently improves convergence speed and accuracy, including heterogeneous, unbalanced, partial-participation, massive-client, and pretrained-model settings.
- Experimental setup: Experiments evaluate convergence speed and model accuracy across multiple datasets, architectures, participation levels, and data distributions.The study uses six benchmark datasets, four network architectures, and comparisons with FedAvg, FedProx, Scaffold, and FedDyn.
- Convergence speed: 37 communication rounds reached 40% accuracy for FedDC versus over 1000 rounds for FedAvg on iid CIFAR100 with 100 fully participating clients.This corresponds to a FedDC convergence speedup of over 27.03× relative to FedAvg.
- Model accuracy: FedDC achieved the best test accuracy on CIFAR10 and CIFAR100 across evaluated settings, including 84.32% versus 79.14% for FedAvg on 0.3-Dirichlet CIFAR10.The comparison used 100 clients with full participation in that example.
- Robustness: On CIFAR100 with 100 fully participating clients, FedDC accuracy was 85.71% on iid data, 84.77% on D1 data, and 84.58% on D2 data.FedDC maintained its advantage over baselines as data heterogeneity increased.
- Robustness: FedDC remained effective with partial client participation, larger client populations, unbalanced data, and pretrained ResNet18 models on Tiny ImageNet.Its locally maintained drift variables preserve drift state when client training is intermittently interrupted.
6. Conclusion
FedDC addresses local drift from heterogeneous data by using learned local drift variables to bridge local and global models. Experiments report better performance, faster convergence, and robustness across participation and data settings.
- FedDC dynamically bridges local and global model parameters with a learned local drift variable.The method targets local drift caused by heterogeneous data.
- FedDC provides better performance and faster model convergence across various image classification datasets.
- FedDC remains robust and efficient under homogeneous or heterogeneous data and full or partial client participation.
A. Appendix: More Experiment Results
The experiments cover multiple real-world and synthetic image-classification datasets with several model architectures. The synthetic setup controls client-specific parameter and data heterogeneity through distributional parameters.
- Experiments evaluate FedDC on CIFAR10, CIFAR100, MNIST, fashion MNIST, Tiny ImageNet, EMNIST-L, and a Synthetic dataset.
- The study uses FCN, CNN, multi-class logistic, and ResNet18 models across the evaluated tasks.FCN is used for MNIST and EMNIST-L, CNN for CIFAR10 and CIFAR100, multi-class logistic for Synthetic, and ResNet18 for Tiny ImageNet.
- The Synthetic dataset generates client samples from client-specific linear classification parameters and labels defined by y = argmax(θ_ix + b_i).θ_i and b_i fit each client’s data distribution, while γ_1 controls their sampled variation.
A.2. Real World Dataset
The real-world evaluation uses several standard image datasets and corresponding network architectures. Data conditions include iid, Dirichlet-based non-iid, balanced, and unbalanced client distributions, with common hyperparameters across participation settings.
- Real-world datasets include MNIST, EMNIST-L, CIFAR10, fashion MNIST, Tiny ImageNet, and CIFAR100.EMNIST-L contains the first 10 character categories, while MNIST contains 10 handwritten-digit categories.
- The experiments use an FCN for MNIST and EMNIST-L and a CNN-based network for CIFAR10 and CIFAR100.
- Non-iid client label ratios follow a Dirichlet distribution, whereas iid clients receive equal numbers of independently identically distributed samples.
- Unbalanced settings assign client sample amounts according to a lognormal distribution with variance 0.3, compared with variance 0 for balanced settings.
- For true-world datasets, training uses batch size 50, five local epochs per round, initial learning rate 0.1, decay 0.998, and weight decay 0.001.
- The same dataset-specific hyperparameters are used across iid and non-iid data and full or 15% client participation settings.
A.3. Results
FedDC’s experiments examine hyperparameter sensitivity, component ablations, convergence, communication efficiency, and comparisons with adaptive federated optimizers. Across the reported settings, FedDC generally achieves strong or leading performance and faster convergence.
- Sensitivity of hyper-parameter in FedDC: All tested α values converge to a stationary point, while carefully selecting α improves model performance.α controls the weight of FedDC’s penalized term and is explored over [0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1].
- Ablation study: The ablation study separately evaluates the empirical loss, gradient correction term, and penalized term within FedDC’s training process.FedDC(lelglp) denotes the proposed combination of all three local-objective components.
- Convergence plots: FedDC is consistently the best method on the Synthetic dataset across three data settings and 15% or 100% client participation.
- Convergence plots: FedDC outperforms baselines and converges faster on fashion MNIST, CIFAR10, and CIFAR100, while achieving competitive performance on MNIST and EMNIST.
- The comparison of convergence speed: FedDC reaches target accuracy in fewer communication rounds than the comparison methods across iid and non-iid settings.Rounds exceeding 1000 indicate that a method did not reach the target within 1000 rounds.
- Comparison with other recent methods: FedDC performs best against FedAdam, FedYogi, and FedAdagrad in iid and two non-iid CIFAR100 settings with 15% participation.
A.4. Discussion
FedDC is evaluated across heterogeneous data, client participation, and client-scale settings, with experiments indicating faster convergence and robustness. Its algorithm uses local drift variables updated alongside local models before server aggregation.
- Higher non-iid severity reduces global-model accuracy and convergence speed compared with iid settings.
- The convergence experiments cover unbalanced data, massive-client settings, synthetic heterogeneity, and iid/non-iid image-classification datasets.
- FedDC is reported as robust across different client numbers, data distributions, and participation levels.
- FedDC samples active clients, trains local models and drift variables, then aggregates their sums to update the global model.
- FedDC’s procedure initializes local models from the global model and updates local drift using the difference between local parameters and the global model.
B.1. Discussion of FedDC.
FedDC addresses heterogeneous local optima by tracking parameter drift and combining a penalty with gradient correction. The analysis reports objective decrease, bounded gradient variance, and simultaneous convergence of local and global models.
- FedDC introduces auxiliary drift variables because clients’ local stationary points can differ arbitrarily under heterogeneous data.
- Convergence results: When local models reach local optima, FedDC’s global model also converges to a stationary point.
- Penalized Term: The penalized term tracks local-to-global parameter gaps and separates global-model training from client-local-model training.
- Gradient correction: The gradient correction term reduces gradient drift, with local-gradient variance bounded independently of local-objective dissimilarity.
- Convergence results: Under stated smoothness and dissimilarity assumptions, FedDC’s global objective decreases during training.
B.2. Detailed convergence Proof of FedDC.
The detailed proof establishes FedDC convergence under smoothness, bounded dissimilarity, and inexact-solution assumptions, including sampled-client settings. Its bounds rely on objective decrease proportional to the global gradient norm.
- FedDC’s local objective combines empirical loss with a quadratic penalty around the corrected parameter w − h_i.
- The convergence analysis represents the global parameter using averages of local model parameters and local drift variables.
- The proof assumes smooth local losses, bounded local dissimilarity, and γ-inexact local solutions.
- For non-convex objectives, the theorem states global-objective decrease under β-Lipschitz smoothness and bounded dissimilarity conditions.
- With an appropriately chosen α, the objective decrease is proportional to ||∇L(w^{t−1})||2.
- The partial-participation proof bounds deviations using expectations over the randomly selected active client set.
B.5. Convergence of FedDC in convex case
In the convex case, the analysis specializes the non-convex framework by setting β_d = 0 and characterizes convergence through decreases in the global objective.
- For suitable conditions, the gradient norm is stated to decrease proportionally during convergence.
- The convex analysis sets β_d = 0, so the adjusted parameter becomes α.
- FedDC’s convergence rate is characterized by the decrease in the global objective toward a threshold ε.
- The analysis defines Γ as the initial-to-optimal global-objective gap when deriving the communication-round requirement.