PaperlayerView exact revision on arXiv

Source-linked AI summary

Ensemble Distillation for Robust Model Fusion in Federated Learning

Tao Lin, Lingjing Kong, Sebastian U. Stich, Martin Jaggi

arXiv:2006.07242v3cs.LGstat.ML

TL;DR

Federated learning needs aggregation methods that handle communication constraints and heterogeneous client models without requiring identical architectures. The paper proposes FedDF, which distills client-model outputs on unlabeled or generated data into a server model. Across diverse vision and NLP tasks, FedDF converges faster and requires fewer communication rounds, while model exchange still leaves potential privacy leakage.

  • Problem

    Parameter-averaging FL methods require client models to have the same size and structure, limiting aggregation under client heterogeneity and resource diversity.

  • Method

    FedDF performs server-side ensemble distillation, using client-model logits on unlabeled or artificially generated data to train a central student model.

  • Results

    Across CIFAR-10/100, ImageNet, AG News, and SST2, FedDF trains the server model faster and requires fewer communication rounds than existing FL techniques.

  • Takeaways & Limitations

    FedDF supports model fusion across heterogeneous architectures, data settings, and numerical precision, broadening collaborative federated learning.

  • Takeaways & Limitations

    The simplest FedDF framework exchanges models and therefore remains vulnerable to potential privacy leakage from memorization; additional protection mechanisms are left for future work.

Abstract

from arXiv · show

Federated Learning (FL) is a machine learning setting where many devices collaboratively train a machine learning model while keeping the training data decentralized. In most of the current training schemes the central model is refined by averaging the parameters of the server model and the updated parameters from the client side. However, directly averaging model parameters is only possible if all models have the same structure and size, which could be a restrictive constraint in many scenarios. In this work we investigate more powerful and more flexible aggregation schemes for FL. Specifically, we propose ensemble distillation for model fusion, i.e. training the central classifier through unlabeled data on the outputs of the models from the clients. This knowledge distillation technique mitigates privacy risk and cost to the same extent as the baseline FL algorithms, but allows flexible aggregation over heterogeneous client models that can differ e.g. in size, numerical precision or structure. We show in extensive empirical experiments on various CV/NLP datasets (CIFAR-10/100, ImageNet, AG News, SST2) and settings (heterogeneous models/data) that the server model can be trained much faster, requiring fewer communication rounds than any existing FL technique so far.

1 Introduction

Federated learning must address communication costs and client heterogeneity, while parameter averaging restricts participating models to identical structures. FedDF uses ensemble distillation to support heterogeneous fusion and reports faster training across diverse tasks and settings.

  • Motivation: FL clients preserve private datasets while collaborating on a centralized model, but communication delays and client heterogeneity remain key challenges.Heterogeneity includes non-i.i.d. data and differences in hardware and computing resources.
  • Limitations of parameter averaging: FEDAVG and related methods directly average client parameters, requiring all participating models to share the same size and structure.Direct parameter averaging therefore constrains model architecture choices.
  • Limitations of parameter averaging: A toy 3-class MLP example shows that the averaged model produces much more blurred decision boundaries than the individually trained local models.The figure compares local models, aggregated models, and the global data distribution.
  • FedDF: FedDF applies ensemble distillation for robust model fusion, aggregating knowledge from heterogeneous client models using unlabeled or artificially generated data.The approach is designed to address quality loss associated with Batch Normalization and knowledge barriers across heterogeneous models.
  • Empirical findings: Across CIFAR-10/100, ImageNet, AG News, and SST2, FedDF trains the server model faster and requires fewer communication rounds than existing FL techniques.The experiments cover heterogeneous models and/or data settings.
  • Analysis: The paper analyzes when FedDF can outperform FEDAVG and which factors influence its performance.These analyses are linked to the intrinsic limitation of parameter-averaging approaches.

2 Related Work

Related work covers parameter averaging and alignment methods, ensemble and knowledge distillation, and earlier federated distillation approaches. FedDF differs by targeting multiple-round neural FL across diverse architectures and tasks.

  • Federated learning: FEDAVG performs weighted parameter averaging after local SGD updates, while alignment methods first match neurons layer-wise before averaging.The weights in FEDAVG are proportional to each client’s local data size.
  • Knowledge distillation: Knowledge distillation trains a student to approximate teacher logits, while ensemble methods average logits or extract knowledge from teacher features.These approaches provide the background for ensemble distillation.
  • Comparison with close FL work: One-shot fusion uses unlabeled data for an SVM objective, whereas FedDF considers multiple-round scenarios across diverse neural architectures and tasks.The comparison distinguishes the scope of the federated fusion settings.
  • Comparison with close FL work: FD synchronizes accumulated per-label logits for local distillation and experiences roughly 15% quality drop on MNIST relative to FEDAVG.FedDF is described as achieving stronger performance and reducing rounds to target accuracy on challenging tasks.
  • Comparison with close FL work: The paper notes contemporary related work and defers its discussion to Appendix A.No further characterization is provided in the supplied passage.

3 Ensemble Distillation for Robust Model Fusion

FedDF adds server-side ensemble distillation to federated training: client models provide teacher outputs on unlabeled or generated data, and a server student is optimized from those outputs. The framework extends to heterogeneous architectures and addresses privacy-related deployment considerations.

  • Federated procedure: Unlike FEDAVG, FedDF performs N server-side model-fusion iterations after initializing the fused model from the averaged client models.The resulting server model is returned after these additional iterations.
  • Federated procedure: FedDF runs for T communication rounds, selecting a random C-fraction subset of clients and collecting their locally updated models.Client updates are computed in parallel before server-side fusion begins.
  • Ensemble distillation: The server samples mini-batches from an unlabeled dataset or generator and updates the server student using the ensemble of received client models.This procedure uses AVGLOGITS to train the student from client predictions.
  • Ensemble distillation: In homogeneous FL, the server distills the ensemble of client teacher models into one server student model using their logits on unlabeled mini-batches.The client models are evaluated with a forward pass before student training.
  • Ensemble distillation: The distillation update uses KL divergence between softmax outputs, with η denoting the stepsize.The supplied prose explicitly defines KL, σ, and η in the update rule.
  • Heterogeneous models: FedDF extends to p heterogeneous model prototype groups differing in architecture, structure, or numerical precision, with each group receiving knowledge from logits averaged across received models.Fusion occurs on the server, and activated clients receive their corresponding fused prototype model.
  • Distillation data: Unlabeled datasets from other domains and synthetic generator data can provide distillation inputs without using labeled training-domain data.Synthetic data may alleviate acquisition and storage limitations of real unlabeled datasets.
  • Privacy considerations: The basic framework still exchanges models between server and clients, so model memorization can create potential privacy leakage.Differential privacy and decentralized or hierarchical logit fusion are identified as possible protections, with further exploration left for future work.

4 Experiments

Experiments evaluate FedDF across heterogeneous data, architectures, precision levels, normalization choices, and CV/NLP tasks. Across these settings, FedDF converges faster and is more robust than competing federated learning methods, while distillation quality depends on local training and data conditions.

  • Experimental setup: FedDF is evaluated on CIFAR-10/100, ImageNet, AG News, and SST2 using ResNet, VGG, ShuffleNetV2, and DistilBERT models.The experiments cover heterogeneous models and/or data, including federated fine-tuning for NLP tasks.
  • Common federated learning settings: FedDF consistently outperforms FEDAVG across client fractions and non-i.i.d. degrees when local training is reasonably sufficient.The reported setting includes more than 40 local epochs per communication round.
  • Common federated learning settings: FedDF requires significantly fewer communication rounds than other state-of-the-art methods to reach designated target accuracies in the evaluated CV scenarios.The benefits become more pronounced with more local training epochs.
  • Common federated learning settings: Reducing α from 1 to 0.1 nearly triples the rounds required by FEDAVG, FEDPROX, and FEDAVGM, while FedDF needs less than twice as many.This comparison indicates improved robustness to increasing non-i.i.d. data heterogeneity.
  • Normalization and precision: FedDF remains effective for BN and low-precision settings, outperforming GN-based baselines and matching or exceeding FEDAVG for binarized ResNet-8.This avoids extra Group Normalization tuning and supports model fusion for 1-bit models.
  • NLP tasks: FedDF achieves significantly faster convergence than FEDAVG and consistently outperforms it on federated DistilBERT fine-tuning for AG News and SST2.The experiments use non-i.i.d. local data with α = 1.
  • Heterogeneous systems: In heterogeneous architectures, FedDF dominates FEDAVG in test accuracy with less variance, but the fused-to-ensemble gap reaches around 6% on ImageNet.The gap is negligible on CIFAR-10 and widens when the training dataset has many more classes than the distillation dataset.

5 Understanding FedDF

FedDF combines ensemble predictions with knowledge distillation using out-of-domain data, and its behavior depends on distribution alignment, data choice, and distillation effort.

  • Core components: FedDF combines ensembling with knowledge distillation through out-of-domain data.The section studies ensemble performance and how distillation-data properties affect the distilled student.
  • Generalization: FedDF’s ensemble performance depends on discrepancies between local and global distributions, while distillation-to-global distribution shift also affects the distilled model.Theorem 5.1 relates ensemble performance to local/global discrepancy and identifies distillation/global shift as another determinant.
  • Distillation data: FedDF remains consistent across realistic distillation-data sources, but performance declines abruptly when data come from a dramatically different manifold such as random noise.Synthetic data generated by a pretrained GAN does not incur noticeable quality loss.
  • Distillation data: Distillation-data diversity generally has limited impact, with optimal performance when the two domains have similar numbers of classes.The number of classes is used as a proxy for distribution shift, and labels are not used during distillation.
  • Distillation data: Even 1% of the distillation data can provide reasonably good fusion performance.This corresponds to approximately 48% of the local training dataset in the reported setting.
  • Distillation effort: 100 distillation steps, corresponding to five local CIFAR-100 epochs, suffice for satisfactory performance and introduce minor time-wise expense.A moderate number of steps approaches optimal fusion performance.

Broader Impact

The paper situates FedDF within privacy-preserving and resource-efficient federated learning, emphasizing its ability to fuse heterogeneous models where parameter averaging cannot be applied.

  • Broader impact: A practical distillation technique for heterogeneous models and low-resource clients is framed as supporting broader collaborative, privacy-preserving, and efficient decentralized learning.The broader-impact claim connects heterogeneity support with lower resource requirements in distributed training.
  • Related work: The paper identifies FedMD and Cronus as prior approaches using neural-network logits for federated learning, but distinguishes their personalized and poisoning-robust objectives from FedDF’s model-fusion objective.FedDistill is described as contemporaneous and empirically compared with FedDF.
  • Model heterogeneity: Parameter-averaging methods require client models with identical size and structure, whereas FedDF aggregates knowledge from heterogeneous models.FedDF uses server-side prototypes for groups sharing architecture, size, and precision characteristics.

C Additional Experimental Setup and Evaluations

The appendix gives a detailed toy-example illustration of FEDAVG’s limitation using decision boundaries for a three-class, three-layer MLP.

  • Toy example: Figure 8 illustrates FEDAVG’s limitation through a toy three-class classification task with a three-layer MLP.The figure displays decision boundaries as probabilities over RGB channels.
  • Toy example: The toy example uses 60 uniformly sampled distillation points from the range (−3, 3).The figure compares individually trained local models with aggregated models.

C.2 Detailed Experiment Setup

The appendix details experimental conventions, FEDMA reproduction settings, low-bit FedDF training, synthetic non-i.i.d. data generation, and dataset-partition visualizations.

  • General setup: The experiments use tuned learning rates of 0.1 for ResNet-like architectures, 0.05 for VGG, and 1e−5 for DistilBERT.Corresponding FEDPROX and FEDAVGM hyperparameters are also tuned.
  • FEDMA setup: FEDMA uses layer-wise local training, freezing already matched layers while progressively fusing layers from bottom to top.The reported reproduction uses five local epochs per layer, or 45 epochs per complete model update.
  • Low-bit models: For binarized ResNet-8, FedDF simulates local training with a straight-through estimator or error-feedback and receives binarized models from activated clients.The setup targets low-bit quantized client models.
  • Data partitioning: Client non-i.i.d. distributions are synthesized by drawing class-probability vectors from a Dirichlet distribution parameterized by αp.The vector p represents a prior class distribution, while α controls the concentration of the draw.
  • Data partitioning: The appendix visualizes CIFAR-10 and CIFAR-100 partitions across α values for 20 clients and AG News and SST-2 partitions at α=1 for 10 clients.Dot sizes represent sample counts in the class-allocation figures.
  • FEDAVG setup: FEDAVG experiments use a constant learning rate because no benefit from local learning-rate decay was observed.The appendix also reports learning curves for different normalization techniques.

C.4.1 Ablation Study

The ablations examine initialization, optimizer choices, local-training quality, normalization, and compatibility with other federated training schemes. Results identify uniformly averaged initialization as important, show FedDF’s gain is robust to local-training choices, and find compatibility with proximal regularization.

  • Normalization and heterogeneity: FedDF addresses Batch Normalization quality loss in homogeneous federated learning and supports heterogeneous client models and data.The framework is designed for flexible model fusion across architectures and settings.
  • Model initialization: FedDF performs better when ensemble distillation starts from the uniformly averaged model of the current communication round rather than the previous fused model.This comparison is evaluated on CIFAR-10 with ResNet-8 over 100 communication rounds.
  • Local training quality: FedDF’s performance gain remains robust to data heterogeneity and is orthogonal to learning-rate and Adam effects in local training and model fusion.The study compares Adam-based local training and fusion across heterogeneous-data settings.
  • Distillation optimizer: SWAG slightly improves ensemble-distillation performance, while Adam achieves similar performance with trivial tuning overhead.The comparison considers alternative optimization schemes for server-side distillation.
  • Compatibility: FedDF retains a significant performance gain over FEDAVG when combined with a local proximal regularizer.The proximal regularizer reduces local-model diversity, yet the reported FedDF advantage remains.

C.4.2 Comparison with FEDAVG

This comparison studies FedDF and FEDAVG across local-training budgets, data fractions, and non-i.i.d. degrees on CIFAR-10. The appendix provides detailed learning curves for both methods and several heterogeneity settings.

  • Comparison setup: FedDF and FEDAVG are compared across different local training epochs, data fractions, and local-data non-i.i.d. degrees.The comparison uses CIFAR-10 with ResNet-8 and reports top-1 test accuracy over three seeds.

D Details on Generalization Bounds

The generalization analysis bounds the ensemble risk using global empirical risk and terms measuring discrepancies between local and global distributions. It also identifies test-domain shift as a factor affecting the later distilled model.

  • Setup: The analysis assumes each client has a local distribution D_i over Ξ := X × Y, while D denotes the server’s global distribution.Each local empirical distribution uses an equal amount m of data, giving equal contribution to the global empirical distribution.
  • Assumptions: The theoretical result studies binary classification with a hypothesis class H and convex absolute-error loss.The classifier maps X to {0,1}, and the loss is defined as ℓ(h(x), y) = |ŷ − y|.
  • Bound interpretation: The ensemble-risk bound compares the ensemble of local models with a virtual ERM trained on the global empirical distribution.The bound depends on hypothesis-class complexity and local-to-global distribution discrepancies.
  • Test-domain implications: The later distilled model cannot exceed the ensemble’s performance upper bound on the test domain, and distillation-to-test distribution shift affects knowledge transfer quality.The analysis also notes that uniform averaging is used because the server lacks further local information needed for distribution-derived weights.

D.1 Proof for Generalization Bounds

The proof derives the generalization bound by combining domain-adaptation inequalities, empirical-risk bounds, convexity, and uniform convergence. The resulting expression aggregates local risks and complexity and discrepancy terms across sources.

  • Uniform convergence: Uniform convergence bounds empirical local risk using the growth function τ_H and confidence parameter δ.Sauer’s lemma provides a polynomial bound on τ_H when the hypothesis class has finite VC dimension.
  • Proof strategy: The proof begins from the risk of the ensemble model and derives a sequence of upper bounds.The argument proceeds by comparing ensemble risk with local and global-domain quantities.
  • Aggregation: Convexity and Jensen’s inequality connect the ensemble risk with the risks of individual local models.The proof then aggregates the resulting inequalities over K sources.
  • Domain adaptation: Domain-adaptation theory transfers bounds between the global distribution D and each client distribution D_k.The transfer introduces a domain-discrepancy term and a joint-error term λ_k.
Loading 2006.07242v3…