[Paper Review] Data-Free Knowledge Distillation for Heterogeneous Federated Learning

Zhuangdi Zhu, Junyuan Hong+1 more|PubMed|May 20, 2021

Privacy-Preserving Technologies in Data66 citations

TL;DR

FeDGen introduces a data-free knowledge distillation framework for heterogeneous federated learning by learning a lightweight generator on the server to distill ensemble knowledge into local models without using proxy data.

ABSTRACT

Federated Learning (FL) is a decentralized machine-learning paradigm in which a global server iteratively aggregates the model parameters of local users without accessing their data. User <i>heterogeneity</i> has imposed significant challenges to FL, which can incur drifted global models that are slow to converge. <i>Knowledge Distillation</i> has recently emerged to tackle this issue, by refining the server model using aggregated knowledge from heterogeneous users, other than directly aggregating their model parameters. This approach, however, depends on a proxy dataset, making it impractical unless such 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. In this work, we propose a <i>data-free knowledge distillation</i> 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.

Motivation & Objective

• Address user heterogeneity in federated learning where proxy data for distillation is unavailable.
• Propose a data-free KD approach that leverages a generator learned from model predictions to augment local training.
• Enable knowledge distillation that directly guides local model updates, not just global model refinement.
• Provide a lightweight, communication-efficient framework suitable for privacy-sensitive FL scenarios.

Proposed method

• Learn a conditional generator G_w that, given a target label y, produces latent representations z that induce ensemble-consistent predictions from local user models.
• Optimize G_w by minimizing a loss that aligns ensemble predictions from user-prediction heads with target labels (Equation 4).
• Broadcast G_w to clients; clients sample z ~ G_w(z|y) to augment local training via an additional distillation term (Equation 5).
• Allow flexible parameter sharing by sharing only the prediction layer (θ^p_k) while keeping feature extractors local, reducing privacy and communication burdens.

Experimental results

Research questions

• RQ1 Can data-free knowledge distillation effectively mitigate heterogeneity in federated learning without access to proxy data?
• RQ2 Does distilling ensemble knowledge into locally trained models provide comparable or superior generalization with fewer communication rounds?
• RQ3 How does FeDGen perform under varying degrees of data non-iid distribution across clients?
• RQ4 What are the theoretical and empirical benefits of distilling knowledge into local models rather than solely refining a global model?

Key findings

• FeDGen consistently outperforms baselines (FedAvg, FedProx, FedEnsemble, FedDistill, FedDistill+, FedDFusion) across MNIST, CelebA, and EMNIST in heterogeneous FL settings.
• FeDGen remains robust across different levels of data heterogeneity (Dirichlet α varying) and active-user ratios.
• The learned generator G_w captures an aggregated latent distribution aligned with the global data distribution, improving local inductive bias and generalization.
• Distilling knowledge into local models via augmented latent samples provides advantages over methods that only refine the global model, especially under high heterogeneity.
• Sharing only the predictor layer (instead of full models) still yields significant gains, highlighting privacy and communication efficiency benefits.