PaperlayerView exact revision on arXiv

Source-linked AI summary

Data-Free Knowledge Distillation for Heterogeneous Federated Learning

Zhuangdi Zhu, Junyuan Hong, Jiayu Zhou

arXiv:2105.10056v2cs.LGcs.DC

TL;DR

Heterogeneous federated learning can make parameter averaging ineffective, while knowledge-distillation approaches require proxy data and do not fully guide local learning. FEDGEN learns a lightweight generator from user prediction rules without external data, broadcasts it to regulate local training, and improves generalization with fewer communication rounds than the state-of-the-art. The method is subject to limitations when its generated data are poorly relevant to the original task.

  • Problem

    User heterogeneity can drift global models, while federated knowledge distillation depends on a carefully selected proxy dataset and may not directly guide local model learning.

  • Method

    FEDGEN learns a lightweight generator from local prediction rules without external data and broadcasts it to users as an inductive bias for local training.

  • Results

    FEDGEN delivers better generalization performance using fewer communication rounds than the state-of-the-art.

  • Takeaways & Limitations

    Data-free generator-based distillation can extract ensemble knowledge and directly regulate local models under non-iid data distributions.

  • Takeaways & Limitations

    When the generated data distribution is overly distant from the original FL task, it may negatively affect generalization performance.

Abstract

from arXiv · show

Federated Learning (FL) is a decentralized machine-learning paradigm, in which a global server iteratively averages the model parameters of local users without accessing their data. User heterogeneity has imposed significant challenges to FL, which can incur drifted global models that are slow to converge. Knowledge Distillation has recently emerged to tackle this issue, by refining the server model using aggregated knowledge from heterogeneous users, other than directly averaging their model parameters. This approach, however, depends on a proxy dataset, making it impractical unless such a prerequisite is satisfied. Moreover, the ensemble knowledge is not fully utilized to guide local model learning, which may in turn affect the quality of the aggregated model. Inspired by the prior art, we propose a data-free knowledge distillation} approach to address heterogeneous FL, where the server learns a lightweight generator to ensemble user information in a data-free manner, which is then broadcasted to users, regulating local training using the learned knowledge as an inductive bias. Empirical studies powered by theoretical implications show that, our approach facilitates FL with better generalization performance using fewer communication rounds, compared with the state-of-the-art.

1. Introduction

Federated learning struggles with heterogeneous user data and models, while existing distillation methods require proxy data and mainly refine the global model. FEDGEN addresses these limitations through data-free generator-based distillation that guides local training and improves performance with fewer communication rounds.

  • Challenges: Non-iid user data and permutation-invariant neural networks can produce divergent local optima, making element-wise model averaging ineffective.These sources of heterogeneity may prevent the averaged parameters from forming an ideal global model.
  • Limitations of prior work: Proxy-data knowledge distillation aggregates local-model knowledge into the global model, but depends on carefully selected unlabeled data unavailable in many applications.Prior approaches also primarily refine the global model rather than directly guiding local updates.
  • FEDGEN: FEDGEN learns a lightweight conditional generator from local prediction rules and broadcasts it so users can distill peer knowledge during training.The generator produces latent-space feature representations consistent with ensemble predictions without observing user data.
  • FEDGEN: FEDGEN extracts knowledge without external data and directly regulates local model updates with an inductive bias under non-iid distributions.It requires only local prediction layers for knowledge extraction, supporting settings where sharing full model parameters is impractical.
  • Results: FEDGEN achieves better generalization performance with fewer communication rounds than the state-of-the-art.The paper attributes this conclusion to extensive empirical studies supported by theoretical elaborations.

2. Notations and Preliminaries

This section formalizes supervised federated classification and reviews knowledge distillation as a teacher–student approach. It identifies proxy-dataset dependence as a central limitation and introduces the generator-based, data-free procedure used by FEDGEN.

  • Problem setting: The paper models supervised FL as multi-class classification over domains with a shared labeling function and loss, using feature extractors and predictors.Each model consists of a feature extractor mapping inputs to a lower-dimensional latent space and a predictor mapping latent features to label probabilities.
  • Problem setting: Federated learning seeks global parameters minimizing risk across user tasks, with the global dataset distributed among local domains.The formal setting assumes users share the same labeling rules and loss function while their data distributions may differ.
  • Knowledge distillation: Knowledge distillation trains a lightweight student from one or more teacher models, typically matching their outputs on a proxy dataset using KL divergence.In federated learning, local models serve as teachers whose information is aggregated into the global student model.
  • Knowledge distillation: Proxy-dataset dependence is a primary limitation because the dataset requires delicate selection and strongly affects distillation performance.FEDGEN is presented as a data-free alternative to this prerequisite.
  • FEDGEN procedure: FEDGEN learns a server-side generator from client information and sends it to users for local knowledge distillation.The overview describes this process as aggregating client information without observing client data and using the generator to adjust local feature-distribution interpretations.

3. FEDGEN: Data-Free Federated Distillation via Generative Learning

FEDGEN performs data-free federated distillation by learning a lightweight conditional generator from user-model predictions and broadcasting it to guide local training. It operates in a compact feature space, can use only predictor modules, and supports partial parameter sharing.

  • Knowledge extraction: FEDGEN recovers knowledge in a latent space smaller than the input space, reducing computation and alleviating certain privacy-related concerns.Direct optimization over high-dimensional input space can be computationally prohibitive and may expose user-data profiles.
  • Knowledge extraction: The generator produces label-conditioned feature representations that induce ideal predictions from the ensemble of user models.Noise is introduced to obtain diversified outputs from the generator.
  • Knowledge distillation: Broadcasted generator samples augment local training, altering each user objective to encourage ideal predictions over the feature space.These augmented samples provide an inductive bias for local models.
  • Knowledge extraction: FEDGEN learns a conditional generator from user-model prediction rules, without requiring access to external data.The generator models consensual knowledge from the ensemble of user predictions.
  • Flexible parameter sharing: FEDGEN can share only local prediction layers while keeping feature extractors localized, improving communication efficiency and reducing data-leakage vulnerability.Empirical study reportedly finds benefits even without sharing feature-extraction modules.

4. FEDGEN Analysis

FEDGEN distills ensemble knowledge into local models as an inductive bias and aligns their feature distributions with a learned global distribution. The analysis connects this augmentation to improved generalization, while experiments illustrate gains over parameter averaging.

  • Knowledge Distillation for Inductive Bias: FEDGEN’s generator learns an aggregated feature distribution that can approximate a global data distribution even when individual local models are biased.In the visualization, generator samples gradually coincide with the ground-truth conditional distribution.
  • Knowledge Distillation for Inductive Bias: Users sampling generator outputs can observe beyond their local training data and adjust decision boundaries toward ensemble knowledge.The generator therefore supplies an inductive bias to users with limited data.
  • Knowledge Distillation for Distribution Matching: Regulating local models with generator samples can minimize conditional KL-divergence between generator-derived and model-derived feature distributions.The practical objective uses empirical samples drawn from the estimated label prior and generator.
  • Knowledge Distillation for Improved Generalization: 98.4% accuracy follows distillation from 81.2% for one user, compared with 93.2% for parameter averaging and 98.6% for an oracle model.These values are reported for the illustrated prototype.
  • Knowledge Distillation for Improved Generalization: Large user heterogeneity increases feature-distribution divergence and undermines global-model quality, whereas more empirical samples favor generalization.These are the two key implications derived from the federated-learning generalization bound.
  • Knowledge Distillation for Improved Generalization: Augmented data can relax sample-related numerical constraints and reduce discrepancy between local and global feature distributions.The analysis treats generator samples as augmented data from an aggregated distribution with small deviation from the global induced distribution.

5. Related Work

Prior FL knowledge-distillation methods address heterogeneity but often depend on proxy data or exchange limited user-model information. FEDGEN instead uses a lightweight generative model to ensemble multiple users’ knowledge over the latent space.

  • Federated learning research addresses challenges including heterogeneity, privacy, communication efficiency, and convergence analysis.
  • Knowledge distillation compresses knowledge from one or more teacher models into a student, and conventional approaches typically rely on a proxy dataset.
  • Federated knowledge-distillation methods commonly depend on data, including unlabeled proxy datasets or user-shared batch-averaged data.
  • FEDGEN differs by extracting knowledge from user-model parameters and learning a lightweight latent-space generative model, while FEDDISTILL shares logit-vector statistics.

6. Experiments

Experiments compare FEDGEN with parameter-averaging, proximal, ensemble, and distillation baselines across three datasets and heterogeneous training settings. FEDGEN consistently improves performance and learning efficiency, especially under stronger heterogeneity, while remaining robust to architectural and communication changes.

  • Setup: Experiments compare FEDGEN with FEDAVG, FEDPROX, FEDENSEMBLE, FEDDFUSION, and FEDDISTILL across several federated-learning strategies.FEDDFUSION uses unlabeled proxy samples, whereas FEDDISTILL shares label-wise average logit vectors.
  • Setup: Experiments use MNIST, EMNIST, and CELEBA for digit, character, and celebrity-smile classification tasks.
  • Setup: The default configuration uses 200 communication rounds, 20 users, a 50% active-user ratio, 20 local steps, and batch size 32.
  • Setup: For MNIST and EMNIST, smaller Dirichlet α values represent higher heterogeneity; CELEBA uses naturally non-iid data with additional grouping-based heterogeneity.
  • Performance Overview: FEDGEN outperforms the other baselines by a considerable margin and remains robust across different levels of user heterogeneity.
  • Performance Overview: FEDGEN’s advantage becomes more notable at high heterogeneity, while FEDDFUSION’s gain diminishes as heterogeneity is reduced.
  • Performance Overview: FEDDISTILL suffers non-negligible performance drops, and FEDDISTILL+ performs worse than FEDAVG as α decreases.
  • Learning Efficiency: FEDGEN has the most rapid learning curves, outperforms other baselines, and requires fewer communication rounds to reach high performance.

7. Conclusions

The paper concludes that FEDGEN enables data-free knowledge distillation for heterogeneous federated learning. Experiments and theoretical implications associate the approach with better generalization and fewer communication rounds than state-of-the-art methods.

  • FEDGEN enables efficient knowledge distillation for heterogeneous federated learning without requiring external data.
  • Empirical experiments guided by theoretical implications report better generalization performance using fewer communication rounds than the state-of-the-art.

8. Theoretical Derivations

The theoretical analysis bounds federated generalization using local risk, distribution divergence, oracle performance, and sample complexity. It further shows that augmented data can improve the bound when its induced distribution remains sufficiently close to the global distribution, but can hurt when that discrepancy is too large.

  • Setup: The analysis assumes binary classification with a 0-1 bounded loss and represents domains through a shared feature mapping R.The induced distributions over the latent feature space support the subsequent divergence-based bounds.
  • Generalization Bounds: Theorem 1 bounds aggregated-model performance by local empirical risk, feature-space distribution divergence, oracle performance, and numerical constraints involving sample count and VC-dimension.The divergence is induced by the symmetric-difference hypothesis space H△H.
  • Augmented Data: Augmented data affects the global-performance bound through its divergence from the global distribution, which is therefore a key quality criterion.The corollary explicitly isolates 2dH△H(˜DA, ˜D) as the relevant augmentation term.
  • Augmented Data: When augmented-data divergence is small, augmentation reduces discrepancy and enriches empirical samples, benefiting local-user generalization.The theoretical argument compares augmented and local distributions after feature mapping.
  • Augmented Data: When augmented-data divergence is moderate, the augmented distribution can still tighten the global-performance upper bound; when it is over-large, it may negatively affect generalization.The negative effect occurs when the augmented distribution is not relevant to the original FL task.

9. Extended Experiments

The extended experiments examine implementation choices, prototype distillation, partial parameter sharing, and performance across heterogeneous MNIST, CELEBA, and EMNIST settings. They report improved user-model generalization after distillation and evaluate learning curves under varied heterogeneity and communication conditions.

  • Implementation Considerations: FEDGEN weights user models using label-availability statistics to account for differences in predictive ability across labels.The resulting weight matrix adjusts the generator objective.
  • Implementation Considerations: The weighted objective can further mitigate negative ensemble effects when a teacher lacks training samples for particular labels.This is the reported motivation for weighting the generator objective.
  • Implementation Considerations: Stochastic generative learning uses noise and a label to infer feature representations, with diversity loss increasing generator-output diversity.The generator samples z from Gw(·|y) using Gaussian auxiliary noise.
  • Prototype Results: In the one-round prototype experiment, distilled knowledge notably improves user-model generalization performance on the global dataset.The experiment visualizes local distributions and decision boundaries before and after distillation and summarizes accuracy in Table 5.
  • Partial Parameter Sharing: The partial-sharing variant broadcasts and aggregates only the predictor layer while keeping feature-extraction layers localized.Algorithm 2 describes this setting and its iterative server-user updates.
  • Extended Evaluation: Experiments evaluate learning curves on MNIST, CELEBA, and EMNIST across data-heterogeneity and communication settings summarized in Table 9.For MNIST and EMNIST, smaller Dirichlet α denotes higher heterogeneity; CELEBA varies the active-user ratio.
Loading 2105.10056v2…