Data-Free Knowledge Distillation for Heterogeneous Federated Learning

In plain words

Federated learning trains one model across many devices or clients without collecting their data. The usual method averages the clients' model weights. That works poorly when each client's data is skewed toward different classes, because the averaged model drifts and converges slowly.

Knowledge distillation is a common fix: the server combines the clients' predictions and teaches the global model from them. That needs an extra public dataset on the server, which is often not available. This paper removes that need. The server builds a small generator that produces feature vectors for each class, chosen so the clients' models agree on them. The generator is sent back to the clients, and each client trains on these generated examples alongside its own data. The generated examples carry what the other clients learned, so each client moves toward a shared view of the classes. The method only needs the last layer of each client model, so the full model does not have to be shared.

The authors report faster convergence and better accuracy than standard averaging and several distillation baselines on three image datasets.

TL;DR

FedGen replaces a server-side proxy dataset in federated distillation with a generator learned from the clients' prediction layers. The generator's samples regularize local training. Use it when clients have skewed label distributions and no shared public data. Results are author-reported on MNIST, EMNIST and CelebA with 20 users. Code is official and in PyTorch, but it departs from the paper in several places (see Paper vs code).

Problem

Averaging client models (FedAvg) can yield a poor global model when client data is non-iid, and local models drift toward their own skewed data §1. Knowledge distillation from the clients' ensemble improves the server model, but it needs an unlabeled proxy dataset on the server, and the knowledge only refines the global model, not the clients §1 §2.

The paper's setting is supervised multi-class classification with K users, each with its own domain (data plus labeling function) §2. The method assumes the clients share one labeling rule §2. The theory assumes binary classification with 0-1 loss app. §8.1. The experiments use MNIST and EMNIST, which the code configures with 10 and 26 labels model_config.py model_config.py, and CelebA, which is binary §6.1.

Method

FedGen §3 has two parts: knowledge extraction on the server §3.1 and knowledge distillation on the clients §3.2. A third part, sharing only the prediction layer, is an extension §3.3. Algorithm 1 gives the loop alg. 1.

Knowledge extraction (server) §3.1. The goal is a conditional distribution over the latent feature space that the ensemble of client models would label correctly. The paper first writes this as an optimization over the input space eq. 2, then moves it to a latent space because that is cheaper and leaks less about user data §3.1. The approximations are:

• label prior: estimated from the label counts the clients report §3.1;
• posterior: the average of the clients' log-probabilities §3.1.

The generator G_w(z | y) is trained with

\min_{w}\; J(w) = \mathbb{E}_{y \sim \tilde p(y)}\, \mathbb{E}_{z \sim G_w(z \mid y)}\Big[\, \ell\Big(\sigma\Big(\tfrac{1}{K}\textstyle\sum_{k=1}^{K} g(z;\theta_k^p)\Big), y\Big) \Big]

eq. 4, where g is the logit output of the predictor and sigma the activation §3.1. Only the predictor parameters theta_k^p of each client are needed §3.1. Noise eps ~ N(0, I) is concatenated to the generator input, a re-parameterization trick, so that z ~ G_w(. | y) is a function of (y, eps) §3.1.

Knowledge distillation (clients) §3.2. The server sends the generator to the clients. Each client adds a term to its local loss: the generator samples should be classified correctly by its own predictor eq. 5:

\min_{\theta_k}\; J(\theta_k) = \hat{L}_k(\theta_k) + \hat{\mathbb{E}}_{y \sim \tilde p(y),\, z \sim G_w(z \mid y)} \big[\ell(h(z;\theta_k^p), y)\big]

eq. 5, where L_k is the empirical risk on the client's data §3.2. Remark 1 shows that this is equivalent to minimizing a KL divergence between the generator's conditional distribution and the client's conditional distribution in latent space eq. 6 app. §8.2. The paper calls the generator samples an inductive bias for local learning §4.1.

Partial parameter sharing (extension) §3.3. Clients share only the last (prediction) layer and keep the feature extractor local. The server then averages only the predictor layers alg. 2 app. §9.3. The paper says this lowers communication and exposure of the feature extractor §3.3.

# FedGen training loop  [alg. 1]
# Input:  K users, global model theta, generator w,
#         generator step alpha, local step beta, T local steps
# Output: global model theta, generator w

def fedgen_round(users, theta, generator, label_prior):
    # Server samples active users uniformly at random  [alg. 1]
    active = sample_without_replacement(users, num_active)
    # Each active user starts from the broadcast model  [alg. 1]
    for user in active:
        theta_k = theta
        # Local training over T steps  [alg. 1]
        for step in range(T):
            # Draw a real batch and generator labels  [alg. 1]
            x, y = sample_batch(user.data, B)
            y_gen = sample_labels(label_prior, B)
            # Generate latent samples for the generator labels  [§3.1]
            z = generator(y_gen)
            # Local loss: empirical risk plus generator term  [eq. 5]
            loss = risk(theta_k, x, y) + gen_term(theta_k, z, y_gen)
            theta_k = theta_k - beta * grad(loss, theta_k)
        # Upload the local model and label counts  [alg. 1]
        send(theta_k, user.label_counts)
    # Server averages the uploaded models  [alg. 1]
    theta = mean(theta_k for theta_k in uploads)
    # Server does one generator update with Eq. (4)  [alg. 1] [eq. 4]
    w = generator_update(generator, active, alpha)

Pseudocode is in our words. The paper's Algorithm 1 is shown as a single generator update per round alg. 1, which the code does not do (see Paper vs code).

Hyperparameters

Values stated in the paper are cited to the paper. Values marked "code" are in the official repository and are not in the paper.

Results

Author-reported. Setup: 20 users, 200 rounds, T = 20, top-1 test accuracy (%), reported as mean ± spread (the paper does not define the ±; see Pitfalls) §6.1 tab. 1. Baselines are FedAvg, FedProx, FedEnsemble, FedDistill, FedDistill-FL (the authors' stronger variant that also shares the model), and FedDFusion (a data-based distillation method that uses a proxy dataset) §6.1. The FedGen column is the proposed method.

Table 1 (main results, T = 20 unless stated) tab. 1

Bold marks the best mean in each row, as in the paper tab. 1. FedGen is not best in every row: FedDFusion leads at EMNIST alpha = 0.05 and FedEnsemble at CelebA r = 5/25 tab. 1.

Table 9 (appendix, extra rows) tab. 9: MNIST alpha = 10: FedAvg 94.23±0.13, FedDistill-FL 95.04±0.21, FedDFusion 93.36±0.45, FedGen 95.79±0.10. EMNIST T = 40: alpha = 0.05: FedGen 67.64±0.75 (FedDFusion 70.46±1.16 is higher); alpha = 0.1: FedGen 70.90±0.49 (FedDFusion 71.55±0.43 is higher); alpha = 1: FedGen 78.92±0.73; alpha = 10: FedGen 79.29±0.53 (FedDistill-FL 78.37±0.76, FedEnsemble 77.92±0.62).

Table 4 (last layer only shared, MNIST) tab. 4: alpha = 0.05: FedAvg 59.67±0.76, FedDistill-FL 58.83±0.62, FedDFusion 59.62±0.84, FedGen 63.60±0.63. alpha = 0.1: 58.39±0.74, 56.25±0.98, 58.38±0.81, 65.42±0.29. alpha = 1: 74.49±0.57, 74.24±0.60, 74.51±0.55, 79.72±0.52. alpha = 10: 86.35±0.60, 86.89±0.26, 86.28±0.69, 87.92±0.46. FedDFusion was modified to upload whole models and fine-tune with proxy data in this comparison §6.4.

Table 2 (generator architecture, EMNIST, alpha = 0.1) tab. 2: FedAvg 66.22±2.58. FedGen with (noise, hidden) = (64, 256): 71.61±0.25; (32, 256): 72.09±0.46; (32, 128): 72.43±0.57; (16, 128): 72.01±0.76; (32, 64): 70.98±0.85.

Table 3 (generator samples, local training time, EMNIST alpha = 0.1) tab. 3: FedAvg 47.66±1.68 ms; FedGen with B_G = 8, 16, 32, 64, 128: 57.20±2.22, 57.39±2.21, 58.17±2.24, 58.91±2.29, 60.06±2.32 ms. The accuracy row of this table is commented out in the source and is not reported.

Figures (qualitative). Learning curves show FedGen reaching a given accuracy in fewer rounds than the baselines fig. 6. Its gain is largest at high label heterogeneity fig. 5. In the 2-D toy setting, one user's accuracy rises from 81.2% to 98.4% after distillation, against 93.2% for parameter averaging fig. 2 (see Pitfalls for a conflicting number in the appendix).

Paper vs code

Repository: https://github.com/zhuangdizhu/FedGen at commit dd5807c5625404e99035d97e2ef8c712c2e007a6. The paper's abstract footnote links it, and the repository README says it "accompanies the paper". The owner account (zhuangdizhu) matches the first author's name. Code license: the repository has no LICENSE file and nothing in it names a license, so the code is treated as unlicensed. Nothing is quoted from it; citations are for description only.

Each difference below is cited to both paper and code.

1. Generator noise. Paper: eps ~ N(0, I) §3.1. Code: torch.rand, which is uniform on [0, 1) generator.py.
2. Teacher loss. Paper: loss of the averaged client logits eq. 4. Code: the sum over active users of a per-user NLL, weighted by label share serverpFedGen.py. The averaged logits are only used for the student term serverpFedGen.py, which is off in every config (ensemble_beta = 0) model_config.py model_config.py model_config.py.
3. Generator updates per round. Paper: one gradient step per round alg. 1. Code: 10 epochs of 5 steps, i.e. 50 Adam steps per round serverpFedGen.py serverpFedGen.py model_config.py serverbase.py.
4. Generator learning rate. Paper: 1e-4 tab. 8. Code: 3e-4 (MNIST, CelebA), 1e-4 (EMNIST) model_config.py model_config.py model_config.py. The code also decays it by 0.98 per round serverpFedGen.py; the paper does not mention a decay.
5. Label prior for sampling. Paper: y ~ p~(y) from reported label counts §3.1. Code: y drawn uniformly over labels with at least one sample among the active users serverpFedGen.py serverpFedGen.py.
6. Model averaging. Paper: uniform mean over active users alg. 1. Code: weighted by each user's training sample count serverbase.py.
7. Client loss. Paper: Eq. (5) has only the cross-entropy term on generator samples eq. 5. Code adds a second term: KL between the client's prediction on real data and a detached copy of its own prediction on generated samples for the batch's labels, scaled by beta userpFedGen.py. The cross-entropy term is scaled by gen_batch_size / batch_size userpFedGen.py. Neither appears in Eq. (5).
8. Schedules. Client alpha and beta decay as 0.98^round with a floor of 1e-4 userpFedGen.py userpFedGen.py. Regularization is skipped in round 0 serverpFedGen.py. Early stop after 20 local epochs is hard-coded serverpFedGen.py. None of this is in the paper.
9. Generator architecture. Paper: two MLP layers, no normalization stated §6.1 app. §9.2. Code: one hidden Linear, BatchNorm and ReLU, then a linear representation layer generator.py.
10. Diversity loss. Paper: "idea of diversity loss" from Mao et al., no formula §6.1. Code: exp(mean(-D_noise * D_latent)), where D_noise is the pairwise L2 distance between noise vectors and D_latent the pairwise L1 distance between latent outputs, computed over the batch generator.py.
11. Partial sharing. Paper: separate Algorithm 2 alg. 2. Code: a partial mode, switched on when the algorithm name contains "partial" serverbase.py. I did not run this path.
12. Dataset ratio. Paper: at most 50% of training data for all sets §6.1. README: 10% for EMNIST README.md.
13. Generator batch size. Paper: §6.3 says B_G = B by default, while Table 8 says 128 §6.3 tab. 8. Code: server batch 128 (ensemble_batch_size) model_config.py, client gen_batch_size 32 by default main.py.

The baselines are implemented in the same repository README.md.

Pitfalls

• Pre-set theory assumptions. Theorem 1 and Corollary 1 assume binary classification and 0-1 loss app. §8.1, and Corollary 1 needs the augmented distribution's divergence to be bounded by a small epsilon §4.3. Neither is checked in the experiments, and the paper does not test the bound §4.3 app. §8.4.
• Conflicting prototype numbers. The main-text Figure 2 caption gives one user 81.2% -> 98.4% and the oracle 98.6% fig. 2. Appendix Table 5 gives user 3 81.2 -> 98.2 and the oracle 98.4 tab. 5. Use Table 5 for per-user numbers and treat the caption as approximate.
• Mean ± std definition. The paper does not say how the ± is computed or over which seeds; learning curves average 3 seeds fig. 6 and the code runs --times 3 README.md.
• Dirichlet split details. The paper cites prior work for the Dirichlet split §6.1. The per-label split and the minimum samples per user depend on the generator script (--min_sample default 10) generate_niid_dirichlet.py. Regenerate the partition from the code, not from the paper's text.
• Unstated generator choices. Weights alpha and beta, the diversity weight eta, weight decay and the generator's schedule are absent from the paper. The code's values (see Hyperparameters) are the only source. Tuning them changes the outcome; Table 2 shows about 1.5 points between the best and worst generator sizes tested tab. 2.
• Cost of the generator. Each client step makes extra generator calls and extra predictor passes on generated samples userpFedGen.py. The authors' local-time measurement is 47.66 ms for FedAvg vs 57.20–60.06 ms for FedGen tab. 3. The generator is downloaded every round; the authors call it lightweight §6.2 but do not report its size in bytes.
• Privacy. The paper says the latent generator is more compact and "can alleviate certain privacy-related concerns" §3.1. It gives no privacy analysis, so do not claim a privacy guarantee from this card.
• Single-round statistics. The prototype is one round with three users app. §9.1, which the paper uses to illustrate the mechanism, not to measure it.
• Appendix numbering. The source has no \appendix command, so the appendix sections are numbered continuously (§8, §9) in the compiled paper. Check against the PDF before quoting section numbers from the supplement.

Lineage

• Builds on: FedAvg (McMahan et al., 2017) for averaging §2; ensemble distillation with a proxy dataset (Lin et al., 2020; FedDFusion here) §1 §6.1; FedDistill (Jeong et al., 2018; arXiv:2011.02367 as cited in the repository README) for label-wise logit sharing §6.1 README.md.
• Also cited: the re-parameterization technique (Kingma and Welling, 2013) §3.1; the diversity loss idea (Mao et al., 2019) §6.1.
• Compared against: FedAvg, FedProx, FedEnsemble, FedDistill, FedDistill-FL, FedDFusion §6.1.
• Superseded by: not stated in this card.