Ensemble Distillation for Robust Model Fusion in Federated Learning

In plain words

In federated learning, many devices each train a copy of a model on their own data, and a server combines the copies into one shared model. The usual way to combine them is to average their weights. That only works if every copy has the same shape. The authors instead let the server combine the devices' predictions. The server runs a short extra training phase in which a new model learns to imitate the average prediction of all the device models, on inputs that contain no labels. Because only predictions are averaged, the device models can differ in size or number precision. The raw training data never leaves the devices, and the extra work happens on the server, so the devices do not do more work. The authors report that the shared model learns in noticeably fewer rounds than standard averaging.

TL;DR

• FedDF replaces the weight-averaging step of FedAvg with knowledge distillation from an ensemble of the received client models §3 eq. 1.
• The server fusion uses unlabeled real data or generator samples as distillation inputs §3. Labels are not used in the distillation step §4.3.
• Client models may differ in architecture, size or precision §3 app. B, alg. 3.
• Author-reported: fewer communication rounds than the compared FL baselines on CIFAR-10/100, ImageNet-32, AG News and SST2 §4.2 §4.3.
• Cost: extra server forward passes of every sampled client model on each distillation batch, and a server-side validation split that uses labels §4.1 (see Pitfalls).

Problem

Parameter averaging (FedAvg and its variants) requires all client models to share one structure and size §1. Ensemble methods avoid that constraint by averaging predictions, but keeping every client model on the server and averaging logits at inference is impractical with many clients §1. Standard FedAvg also degrades as data becomes more heterogeneous across clients §4.2 tab. 1. Batch Normalization is a known source of quality loss under non-i.i.d. data, and replacing it with Group Normalization only partly helps §4.2 tab. 2. The authors aim for a fusion step that works with heterogeneous models, does not change client training, and keeps the communication pattern of FedAvg §3 app. A.

Method

Homogeneous case (main algorithm). Each round, the server samples a fraction C of the K clients alg. 1. Each sampled client trains locally for E epochs, starting from the current global model, and returns its weights [alg. 2] app. B. The server does not use these weights directly. Instead it first forms a data-size-weighted average of the returned models, and uses that average as the starting point of the fused model alg. 1. It then runs N distillation steps. In each step, a mini-batch of unlabeled or generated data is passed through every sampled client model. The softmax of the mean of those logits is the target distribution. The fused model is trained to match it with a KL divergence eq. 1 §3. The fused model is the global model for the next round alg. 1.

Update rule for the fused (student) model, as written in the paper eq. 1:

• Student parameters x_{t,j} = x_{t,j-1} − η · ∂KL( σ(mean of client logits), σ(student logits) ) / ∂x_{t,j-1}.
• σ is the softmax, η the step size. No temperature appears in the equation eq. 1.

Heterogeneous case. The server keeps one prototype model per distinct architecture, precision or size group app. B, alg. 3. Each prototype starts from the average of the sampled clients that share its prototype, weighted by their data sizes app. B, alg. 3. The distillation target for every prototype is the average logits of all sampled clients, across all prototypes §3 app. B, alg. 3. So knowledge is shared between architectures. Clients receive their prototype model in the next round, and no extra work is done on clients §3.

Distillation data. Each distillation batch can come from a different unlabeled domain, such as another image dataset, or from a frozen generator such as BigGAN, sampled before FL training starts §3 §4.1. The paper's footnote notes that GAN training is not part of the FL procedure §3.

Privacy. Clients send only model weights, as in FedAvg. Input data is not sent app. A. The authors name differential privacy for client models, and a decentralized variant where clients compute logits on shared random public data, as possible extensions, and leave both for future work §3.

Pseudocode of the homogeneous server round, following alg. 1, alg. 2 and eq. 1. Shapes are for CIFAR (32 × 32 RGB) with batch 128 §4.1:

# FedDF server round  [alg. 1] [eq. 1]
# Input:  global_model, clients, distill_data; C, N, eta
# Output: fused global model for the next round

def feddf_round(global_model, clients, distill_data):
    # Sample the participating clients S_t, a fraction C of all K  [alg. 1]
    participants = sample_fraction(clients, C)

    # Each participant trains a copy locally for E epochs  [alg. 2]
    local_models = [
        local_update(client, copy(global_model))
        for client in participants
    ]

    # Start fusion from the data-size-weighted parameter average  [alg. 1]
    student = weighted_average(local_models, sizes(participants))

    # Distill the averaged logits of all participants  [eq. 1]
    for step in range(N):
        # Unlabeled or generated batch; labels are never used  [§3] [§4.3]
        batch = sample_batch(distill_data)               # (128, 3, 32, 32)

        # Teacher targets: mean of all participants' logits  [eq. 1]
        teacher_logits = mean(
            [model(batch) for model in local_models]
        )                                                # (128, classes)
        student_logits = student(batch)                  # (128, classes)

        # KL between softmax of teacher and softmax of student  [eq. 1]
        loss = kl_divergence(
            softmax(teacher_logits), softmax(student_logits)
        )

        # Adam, lr 1e-3, cosine annealing  [§4.1]
        student = adam_step(student, loss, lr=1e-3)

        # Stop when validation accuracy has plateaued for 1e3 steps  [§4.1]
        if validation_plateaued(student, patience=1000):
            break

    # The fused student becomes the next global model  [alg. 1]
    return student

Hyperparameters

Values not stated in the paper are marked "not stated". Table entries cite the source of each value.

Results

All numbers are author-reported. They are mean ± std over three random seeds unless stated. Setup per table, as given in the paper.

CIFAR-10, ResNet-8, 20 clients, rounds needed to reach target accuracy tab. 1 (lower is better). Rows with 40 local epochs, C = 0.4 shown here:

The paper reports the same comparison for E = 1, 20, 40, C = 0.2, 0.4, 0.8 tab. 1. The centralized upper bound is 86% tab. 1. FedDF is best in every setting in that table, and its advantage grows with more local epochs §4.2.

CIFAR-10 and CIFAR-100, BN vs GN, 100 rounds, 40 local epochs, 20 clients, C = 0.4 tab. 2 (top-1 test accuracy, %):

VGG-9 without BN, CIFAR-10, 20 clients, C = 0.4, α = 1 tab. 3. Accuracy at round 100 (top-1, %): FedAvg with drop-worst 78.34 ± 1.42, FedProx 78.60 ± 1.91, FedDF 80.36 ± 0.63. Without drop-worst, FedAvg drops to 7.52 ± 4.29, and FedMA reaches 65.0 ± 0.02. FedMA was run only on VGG-9 because it does not support BN or residual connections §4.2 fn..

DistilBERT fine-tuning, AG News and SST2, 10 clients, α = 1 fig. 3. The paper reports faster convergence and higher accuracy than FedAvg over the 10 rounds. The numeric values are in the figure only and are not in the text, so they are not restated here.

Binarized ResNet-8, CIFAR-10, 20 clients, 40% sampled, 100 rounds tab. 4 (top-1, %):

Heterogeneous architectures (ResNet-20, ResNet-32, ShuffleNetV2) fig. 4: the paper reports that FedDF beats FedAvg in test accuracy each round, with much lower variance, on CIFAR-10, CIFAR-100 and ImageNet-32. The fused model is close to the ensemble accuracy on CIFAR-10. The gap grows to about 6% on ImageNet-32, which the authors attribute to the larger number of classes §4.3. No table values are given in the text for this experiment.

Ablations (appendix):

• Initialization matters: starting distillation from this round's average beats starting from the previous round's fused model. CIFAR-10, ResNet-8, 40 epochs: α = 1 gives 80.43 ± 0.37 vs 74.13 ± 0.91; α = 0.1 gives 71.84 ± 0.86 vs 62.94 ± 1.12 tab. 5 app. C.4.1.
• Local optimizer: Adam in both local training and fusion gives FedDF 83.32 (α = 1) and 72.58 (α = 0.1), against 78.13 and 62.53 for FedAvg tab. 6. The authors read this as showing that the gain is not just from Adam app. C.4.
• Fusion optimizer: SGD 76.68 / 57.33, Adam 80.27 / 71.52, SWAG 80.84 / 72.40 at α = 1 / α = 0.1 tab. 7. SWAG needs extra hyperparameters.
• Compatibility: FedAvg with a proximal term reaches 76.11 at α = 1, and FedDF with the same term reaches 80.56 tab. 8.
• Data source: fusion works on real unlabeled images and on generator samples, but pure random noise gives a clear drop §5 fig. 5.
• Fraction of distillation data: 1% of the ImageNet data is enough for reasonable fusion §5 fig. 6.
• Distillation steps: about 100 steps, roughly 5 local epochs of the CIFAR-100 partition, reach satisfactory performance §5 fig. 6.

Theory [thm. 1] app. D: an informal bound on the risk of the averaged ensemble on the global distribution. It has an empirical term over the pooled local data, plus terms for local-to-global distribution discrepancy and for the ideal joint error. It does not bound the fused model directly. The authors use it only to motivate the experiments in §5.

Paper vs code

No official code.

• The paper contains no code link §1 §4 app. C.2. The only URL in the source is a third-party VGG-9 implementation app. C.2.
• The arXiv abstract page lists no code, project page or data links.
• A GitHub search for the method name returned only unrelated or unverified forks, none owned by the authors or by epfml. None was used.
• The EPFL organization's repository search for distillation and federated terms returned no hits, and epfml/FedDF returned 404.

Pitfalls

• Two different averages. The fusion starting point is weighted by client data size alg. 1. The distillation target is an unweighted mean of all sampled client logits eq. 1 app. B, alg. 3. Use exactly these two weightings.
• Distillation data for CIFAR-100 and ImageNet is inconsistent. The main text and Table 2 use ImageNet-32 for CIFAR-100 §4.1 tab. 2. The Fig. 4 caption instead names STL-10 for CIFAR-100 and STL-10 for ImageNet, while its file names say CIFAR-100 for ImageNet and STL-10 for CIFAR-100. The text is also inconsistent: the caption repeats STL-10 twice. Treat the main text as the authoritative setting and mark the other as unclear.
• Labels are still used on the server. The paper says the fusion needs no labels, but validation accuracy drives early stopping §4.1, and drop-worst relies on server validation accuracy §4.2. The claim "no labels" applies only to the distillation loss. An implementation needs a labeled server validation set.
• Memory and compute scale with participants. Each distillation step runs every sampled client model forward on the batch eq. 1. The paper does not measure server memory or time for this, although it notes minor extra cost with about 100 steps §5.
• Unlabeled data must be domain-appropriate. Random noise causes an abrupt drop in fusion quality §5. The paper gives no rule for choosing a domain.
• Logit averaging assumes shared output classes. The paper averages logits across all participants, including prototypes of different architectures app. B. It does not state how this works when clients have different label sets. Assume a shared label space.
• Temperature and loss details. Eq. 1 uses plain softmax with no temperature. Whether the KL is summed or averaged over the batch is not stated.
• Local batch size, weight initialization and the GN group count are not stated. Implementers should check the appendix tables before assuming defaults.
• Early-stopping text. The text gives "1e{3}" steps and "1e{4}" total without a sign on the exponent. Read them as 1e3 and 1e4 §4.1.
• Figure and table numbers are source-derived. LaTeX was not compiled. The numbers in this card follow source order of the float environments. Check them against the PDF before citing them externally.
• Fig. 3 epoch ordering. The caption says the local-epoch setting is "10 and 1 respectively" for panels (a) AG News and (b) SST2, which is the reading used here. The text does not confirm it.
• Privacy is argued, not evaluated. The paper argues FedDF matches FedAvg in privacy cost, since only model weights leave the client. It gives no privacy analysis and notes that models can leak through memorization §3.
• Baselines differ in tuning. FedProx, FedAvgM and FedAvg were tuned per table tab. 1 caption. FedDF uses the learning rate tuned for FedAvg tab. 1 caption. Table 2 uses a constant learning rate and tunes the other hyperparameters tab. 2 caption.
• Scope of results. Results are on small models, 3 seeds, and CIFAR/ImageNet-32 at 32 pixels §4.1. Heterogeneous and quantized results are reported with fewer numbers in the text.

Lineage

• Builds on FedAvg [mcmahan2016communication] for the local-training and client-sampling loop alg. 1 §2.
• Builds on ensemble distillation [bucilua2006model, hinton2015distilling] §1 §2.
• Compared against: FedAvg, FedProx, FedAvgM (with momentum and a proximal term), FedMA for VGG-9 only §4.1. Also discussed but not numerically compared: FedMD and Cronus (personalization and robustness), FD (logit sharing), and contemporaneous work FedDistill and others §2 app. A.
• Superseded by: not known from content on this wiki.