Source-linked AI summary
Virtual Homogeneity Learning: Defending against Data Heterogeneity in Federated Learning
Zhenheng Tang, Yonggang Zhang, Shaohuai Shi, Xin He, Bo Han, Xiaowen Chu
TL;DR
Data heterogeneity can induce client drift and poor FL performance, while sharing natural data raises privacy concerns. VHL trains with an independent virtual dataset and calibrates natural and virtual features by label. The paper reports improved convergence and generalization, with less client drift.
Problem
Data heterogeneity can cause client drift and poor performance, while sharing natural data to reduce it raises privacy-leakage concerns.
Source: §1 · p. 1Method
VHL uses an independent virtual dataset and a conditional feature-mismatch penalty to calibrate natural and virtual data by label.
Source: §1 · p. 2, §3.2 · p. 4Results
Takeaways & Limitations
VHL offers a virtual-data perspective on mitigating FL data heterogeneity, with calibration associated with performance improvement and less client drift.
Source: §6 · p. 9Generating a sufficiently diverse virtual dataset for many classes may require more calculation, and training on virtual data adds computation cost.
Source: §5.5 · p. 9
Abstract
from arXiv · showhide
In federated learning (FL), model performance typically suffers from client drift induced by data heterogeneity, and mainstream works focus on correcting client drift. We propose a different approach named virtual homogeneity learning (VHL) to directly "rectify" the data heterogeneity. In particular, VHL conducts FL with a virtual homogeneous dataset crafted to satisfy two conditions: containing no private information and being separable. The virtual dataset can be generated from pure noise shared across clients, aiming to calibrate the features from the heterogeneous clients. Theoretically, we prove that VHL can achieve provable generalization performance on the natural distribution. Empirically, we demonstrate that VHL endows FL with drastically improved convergence speed and generalization performance. VHL is the first attempt towards using a virtual dataset to address data heterogeneity, offering new and effective means to FL.
1. Introduction
Data heterogeneity can drive client drift, unstable convergence, and poor performance, while sharing natural data raises privacy concerns. VHL instead uses an independent virtual dataset and calibrates natural and virtual features; experiments report improved convergence speed and generalization.
1. Introduction
- Data heterogeneity can cause client drift, unstable convergence, and poor model performance, while a benchmark found FedAvg can outperform variants in many settings.The difficulty of addressing drift across scenarios remains a fundamental FL challenge.Source: §1 · p. 1
- Sharing private data or private statistics can make client data more homogeneous, but it raises privacy-leakage concerns and differential privacy can degrade performance.These tradeoffs motivate the question of whether shared data containing no private information can defend against heterogeneity.Source: §1 · p. 1
- Virtual datasets risk distribution shift because they may differ from natural data, which can harm generalization.VHL calibrates features to address this risk; experiments with FedAvg, FedProx, SCAFFOLD, and FedNova on four datasets report improved generalization and convergence speed.Source: §1 · p. 2
- VHL addresses heterogeneity with an independent virtual dataset and calibrates natural-data features against virtual-data features, drawing on domain adaptation.The approach is proposed as a way to share data without sharing private natural data.Source: §1 · pp. 1–2
2. Preliminaries
Federated learning trains a shared model from client-held data through communication with a server, without collecting clients’ natural data centrally. Domain adaptation addresses distribution differences between training and test domains, including by matching label-conditioned distributions.
2.1. Federated Learning
- The FL objective combines client-local objectives with client weights, and training proceeds through server broadcasts, local updates, and model aggregation.The server communicates with selected clients each round; clients perform local updates and return parameters for aggregation.Source: §2.1 · pp. 2–3
2.2. Domain Adaptation
- Domain adaptation seeks to reduce source–target distribution discrepancy in feature space; for conditional shift, the discrepancy is conditioned on the label.The paper treats virtual data as the source domain and natural data as the target domain.Source: §2.2 · p. 3
3. Virtual Homogeneity Learning
VHL uses a virtual dataset to mitigate client data heterogeneity, calibrating natural and virtual feature distributions by label. Its analysis connects this mismatch penalty to bounded statistical margin under stated separability and label-distribution conditions.
3.1. Heterogeneity Mitigation
- VHL introduces a shared virtual dataset independent of clients’ natural data to reduce heterogeneity without exposing private information.The virtual dataset is intended to support performance on the natural distribution.Source: §3.1 · p. 3
3.2. Feature Calibration
- VHL treats virtual data as a source domain and natural data as a target domain, matching their feature distributions conditioned on labels.The mismatch penalty encourages client-natural and virtual features from the same class to have the same conditional feature distribution.Source: §3.2 · pp. 3–4
3.3. Analysis
- VHL’s feature-calibration rationale is to pull same-label natural features toward virtual features that are relatively close across clients.The paper presents this alignment as a route to more similar local-model features for each category.Source: §3.3 · p. 4
- The analysis defines statistical margin as a generalization measure and proves that the mismatch-penalty objective yields bounded margin when the natural and virtual distributions are separable and have identical label distributions.A larger statistical margin for the virtual distribution gives a tighter bound on the natural distribution’s statistical margin.Source: §3.3 · p. 4
3.4. Overview
- To apply VHL, the server distributes the virtual dataset and clients train using both natural and virtual data with the modified local objective.The training-data and objective changes can also be combined with algorithms such as FedProx and SCAFFOLD.Source: §3.4 · pp. 4–5
4. Related Work
Related work addresses client drift through local regularization or update and aggregation corrections, and explores data sharing and synthetic-data approaches. The paper highlights privacy risks in approaches that use information derived from raw data.
Mitigating Client Drift
- Client-drift methods include local regularization strategies and server-side or optimization-based update corrections.Examples include FedProx, SCAFFOLD, FedNova, and adaptive federated optimizers.Source: §4 · p. 5
Sharing Data
- Synthetic-data methods using GANs or adversarial learning may expose raw-data information, including intermediate features or statistics.Other approaches share statistical information, activations, or logits, which can reveal image profiles or high-level semantics and may permit reconstruction.Source: §4 · p. 5
5. Experimental Studies
Experiments compare VHL with federated-learning baselines across datasets and heterogeneous data partitions. They examine convergence, feature distributions, design choices, and computational and dataset-diversity constraints.
5.1. Experiment Setup
- Experiments use four datasets, standard FL baselines, and additional two-class and subset partitions to test VHL across heterogeneous settings.Datasets are partitioned using Latent Dirichlet Sampling at a = 0.1 and a = 0.05; the additional CIFAR-10 partitions use 10 clients.Source: §5.1 · p. 5
5.2. Experimental Results
- VHL gains more as heterogeneity increases and improves results across client-count settings, but does not work well on CIFAR-100 when E = 5.The reported larger gains occur at a = 0.05 than at a = 0.1; the authors suggest model capacity and drift in virtual-data feature representations as possible reasons for the CIFAR-100 result.Source: §5.1 · p. 5, §5.2 · p. 6
- VHL improves generalization and accelerates convergence across CIFAR-10, FMNIST, SVHN, and CIFAR-100 comparisons.The study compares best accuracy during training and communication rounds to reach FedAvg’s best accuracy.Source: §5.2 · p. 6, Table 1 · p. 7, Table 2 · p. 7
- After calibration, same-class features from different clients cluster closely, unlike the separated client representations observed with FedAvg.The t-SNE visualization uses test features from three clients at communication round 299.Source: §5.2 · pp. 6–7, Figure 2 · p. 6
5.3. Other Facets of VHL
- Virtual-dataset pretraining alone does not improve model performance, while naive virtual-data addition and random-feature calibration benefit FL; VFA trails VHL.The pretraining variant provides some speedup at the beginning of training.Source: §5.3 · pp. 7–8, Table 6 · p. 9
5.4. Ablation Study
- Pure-noise virtual data benefits FL, and increasing its sampling weight strengthens VHL’s effect.Using Tiny-ImageNet does not outperform pure noise; the authors suggest its data may be harder to separate. The sampling-weight study varies the virtual-data batch size.Source: §5.4 · pp. 8–9, Table 6 · p. 9
- VHL is not sensitive to calibration weight, performs best using ResNet-18’s fourth-layer features, and works across three tested model capacities.The authors suggest shallow features may not affect later feature shifts, while last-layer calibration may disrupt classification between natural and virtual data.Source: §5.4 · p. 9, Table 6 · p. 9
5.5. Limitations
- A virtual dataset needs as many classes as the natural data, which can make generating sufficient diversity computationally costly; VHL also adds training computation.The paper identifies VFA as a possible direction for reducing that computation cost.Source: §5.5 · p. 9
6. Conclusion
The paper concludes that federated learning can benefit from virtual data unrelated to natural data, with calibration on homogeneous virtual data improving performance and reducing client drift.
6. Conclusion
- Calibration on homogeneous virtual data lets FL attain significant performance improvement and much less client drift.The conclusion frames this as a virtual-data perspective on mitigating data heterogeneity.Source: §6 · p. 9
- The authors present VHL as a new viewpoint for FL, beyond its model-performance improvements.They hope future work will explore VHL further and use it to enhance FL or other machine-learning tasks.Source: §6 · p. 9
- FL can benefit remarkably from a virtual dataset containing no information related to the natural dataset.Source: §6 · p. 9
A.1. Proof of Theorem 3.2
The proof bounds the statistical margin of a model trained with the distribution-mismatch penalty, yielding bounded generalization performance. It does so by relating margins on the natural, virtual, and sampled virtual distributions.
A.1. Proof of Theorem 3.2
- The proof decomposes the natural-distribution margin into differences involving the virtual distribution and sampled virtual training set, plus the sampled-set margin.It uses an auxiliary optimal transport coupling between the natural and virtual conditional feature distributions.Source: §A.1 · p. 13
- Theorem A.2 states that the distribution-mismatch penalty in Eq. 4 elicits a model with bounded statistical margin when natural and virtual distributions share label distributions.The theorem assumes a neural network decomposed into feature extractor φ and classifier ρ.Source: §A.1 · p. 13
- The bound requires a separable virtual distribution and large statistical margins on both the sampled training set and the virtual distribution.The remaining margin term is bounded using the conditional-distribution distance in Eq. 4, leading to bounded generalization performance.Source: §A.1 · p. 13
B.1. Federated Learning
The section frames client drift as a consequence of heterogeneous client data and positions virtual data as an alternative way to address it. It distinguishes VHL from data-sharing approaches and summarizes its dataset, training, and implementation choices.
B.1. Federated Learning
- Heterogeneous client data is associated with damaged FL convergence and final performance, unstable local and global models, and client drift.The section describes local-model divergence as more severe than in distributed training with IID data.Source: §B.1 · p. 14
- VHL trains directly on virtual data that need not be meaningful or transferable, and the authors report it performs much better with noise than FedDF does with noise.They characterize VHL as introducing virtual datasets containing no private information; FedDF instead requires data with a distribution similar to private data, while VHL does not.Source: §B.1 · pp. 14–15
B.2. Virtual Dataset
- VHL connects FL with virtual datasets, a line of work that includes noise data, synthetic driving data, and virtual labels.The cited examples are Baradad et al. (2021), Huang et al. (2019), and Miyato et al. (2018).Source: §B.2 · p. 15
C.1. Hyper-parameters / C.2. Implementation
- Experiments use learning-rate decay of 0.992 per round, momentum-SGD, and batch size 128 for both natural and virtual data.Natural and virtual data are sampled together; the virtual-to-natural sample ratio per epoch is Bv/Bd, and virtual features are detached in the supervised contrastive loss.Source: §C.1 · p. 15, §C.2 · p. 16
C.3. Generating Virtual Dataset
The virtual dataset is generated from class-associated Gaussian noise using several generation methods. The resulting samples are described as distinguishable across labels, with Style-GAN samples within a class being the most diverse.
C.3. Generating Virtual Dataset
- All generation methods sample noise from class-specific Gaussian distributions with different means and a shared standard deviation.The server collects StyleGAN-v2 outputs and distributes them to clients at the beginning, avoiding extra client-side dataset-generation computation.Source: §C.3 · p. 16
- The simple CNN method uses ten differently initialized generators because one untrained CNN lacks output diversity; the upsampling method turns 8 × 8 noise into 32 × 32 color-block images.The paper says those color blocks can let client models learn basic vision features.Source: §C.3 · p. 17, Figure 6 · p. 18, Figure 7 · p. 18
- Virtual-data examples for CIFAR-10, FMNIST, and SVHN show 90 samples of the same class per figure, while the CIFAR-100 figure shows 10 of 100 classes.The displayed generated datasets are reported to have easily distinguishable images across labels, with Style-GAN producing the most diverse same-class images.Source: Figure 4 · p. 17, Figure 5 · p. 17, §C.3 · p. 17
D.1. Feature Shifts between Clients
The feature visualizations indicate that VHL clients reach consensus on virtual-feature expressions and that this is followed by more consistent natural-data expressions. Without VHL, client representations remain markedly different, including late in FedAvg training.
D.1. Feature Shifts between Clients
- VHL clients quickly reach consensus on same-class virtual features, and subsequent local training on those homogeneous features can benefit consistent natural-data expressions.The cited visualizations are the client-model feature plots in Figures 9 and 10.Source: §D.1 · p. 18, Figure 10 · p. 19
- Naive VHL, which trains only with virtual data, also benefits FL and yields faster consistent client-feature expressions than FedAvg.The comparison is described for the client-model visualizations in Figures 9 and 10.Source: §D.1 · p. 18, Figure 10 · p. 19
- Without VHL, t-SNE clusters features from each client together, indicating substantially different client models; this pattern persists for FedAvg at communication round 999.Figure 8 depicts features from different client models without VHL.Source: §D.1 · pp. 18, 20, Figure 8 · p. 19
D.2. More Facets of VHL
The appendix examines why virtual data can help VHL, how batch size and virtual-data choice affect it, and how its calculation cost compares with other algorithms.
D.2. More Facets of VHL
- Figures 16 and 17 show convergence comparisons for SVHN and CIFAR100, while Figure 18 presents convergence trends for loading pretrained models, VHL batch sizes, and training accuracy with different virtual data.Source: §D.2 · p. 20, Figure 16 · p. 21, Figure 17 · p. 21
- Common virtual data can alleviate class imbalance and improve feature extraction even when it is unrelated to clients’ raw data.Local training on severely imbalanced data biases classifiers toward dominant classes, hindering shallow-layer representations; the appendix says virtual-data vision information can help clients learn better feature extractors.Source: §D.2 · pp. 20–21
- Larger virtual-data batches improve VHL convergence and model performance, but add calculation cost.VHL pulls same-class data together, so batch size affects its effect; the authors identify the added computation as a potential target for future improvement.Source: §D.2 · p. 21
- Tiny-ImageNet as virtual data does not outperform virtual noise; the authors suspect it is harder to separate, producing inconsistent features and ineffective calibration.They examine virtual-data training accuracy in Figure 18(c) to assess this conjecture.Source: §D.2 · p. 21
D.3. Comparisons of Calculation Cost
- For most cases, VHL processes a similar or smaller number of samples to reach target test accuracy, but costs more on CIFAR-100.The authors relate the CIFAR-100 result to VHL's smaller performance gains there; Table 10 reports calculation cost relative to Ddata.Source: §D.3 · p. 22, Table 10 · p. 22
Related papers
This paper cites
- Analyzing and Improving the Image Quality of StyleGAN2019
- Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms2017
- Federated Optimization in Heterogeneous Networks2018
- Supervised Contrastive Learning2020
- Virtual Adversarial Training: A Regularization Method for Supervised and Semi-Supervised Learning2017
- Federated Learning with Non-IID Data2018