Source-linked AI summary
FedMD: Heterogenous Federated Learning via Model Distillation
Daliang Li, Junpu Wang
TL;DR
Existing federated learning frameworks do not support participants independently designing fully heterogeneous models, despite heterogeneous data and privacy-sensitive settings. FedMD combines transfer learning and knowledge distillation to enable black-box collaboration, producing significant gains across tested tasks and performance usually within a few percent of pooled-data training.
Problem
Existing federated learning frameworks require centralized control over local model design, limiting full model independence for heterogeneous participants.
Method
FedMD uses transfer learning and knowledge distillation to translate knowledge through output class scores while leaving participants free to design their model architectures.
Results
Significant performance gains were observed across tested local models, usually reaching levels only a few percent below pooled private-data performance.
Takeaways & Limitations
FedMD demonstrates federated collaboration among independently designed models across various tasks and datasets without centralized architectural control.
Takeaways & Limitations
The framework assumes each participant has a very small labeled private dataset and access to a large public dataset.
Abstract
from arXiv · showhide
Federated learning enables the creation of a powerful centralized model without compromising data privacy of multiple participants. While successful, it does not incorporate the case where each participant independently designs its own model. Due to intellectual property concerns and heterogeneous nature of tasks and data, this is a widespread requirement in applications of federated learning to areas such as health care and AI as a service. In this work, we use transfer learning and knowledge distillation to develop a universal framework that enables federated learning when each agent owns not only their private data, but also uniquely designed models. We test our framework on the MNIST/FEMNIST dataset and the CIFAR10/CIFAR100 dataset and observe fast improvement across all participating models. With 10 distinct participants, the final test accuracy of each model on average receives a 20% gain on top of what's possible without collaboration and is only a few percent lower than the performance each model would have obtained if all private datasets were pooled and made directly available for all participants.
1 Introduction
Federated learning protects private data but traditionally assumes a shared model architecture, leaving independently designed black-box models unsupported. FedMD addresses this gap with knowledge distillation and transfer learning, enabling collaboration across heterogeneous local models.
- Motivation: Federated learning can train centralized models while keeping sensitive user data on device, but large datasets remain difficult to curate in settings such as health care.Hospitals may possess only dozens of annotated MRI images for a disease, making privacy-preserving collaboration valuable.
- Problem: Model heterogeneity extends federated learning’s challenges beyond differing bandwidth, computation, and data distributions.The paper focuses specifically on differences among local model architectures.
- Problem: Participants may need unique black-box models because of distinct specifications, privacy concerns, intellectual property, or client-specific tasks.The paper highlights health-care collaborations and AI-as-a-service providers as examples.
- Related work: Existing approaches address statistical heterogeneity with customization but require centralized control over local model design, leaving full model independence unresolved.The paper identifies full model independence as a distinct research direction.
- Approach: FedMD uses knowledge distillation as a translation protocol, with models communicating through outputs on public data rather than sharing private data or architectures.Transfer learning first trains each model on public data and then its private data; public-data class scores support model-agnostic communication.
- Contribution: FedMD enables independently designed models through limited black-box access and reports significant local-model gains on FEMNIST and CIFAR10/CIFAR100.The framework’s communication module translates knowledge between participants.
2 Methods
FedMD formulates collaboration among participants with small private datasets, a shared public dataset, and independently designed models. It combines transfer learning with repeated score communication, server aggregation, distillation, and private-data revisiting.
- Problem definition: FedMD targets m participants whose small labeled private datasets may differ in distribution, while all participants can access a large public dataset.Each participant independently designs a potentially different architecture and need not share hyperparameters.
- Problem definition: The framework seeks to improve each local model beyond individual training on the public dataset and its own private data.This defines the collaboration objective.
- Transfer learning: Transfer learning trains each independently designed model to convergence on the public dataset and then on its own private dataset before collaboration.This establishes the baseline against which later improvements are compared.
- Aggregation and distillation: The server averages participants’ outputs into a consensus, which each party downloads and uses to train its model on public data.Consensus distribution lets participants exchange knowledge without explicitly sharing private data or model architectures.
- Local adaptation: After digesting the consensus, each participant briefly retrains its model on private data before the next communication round.This alternates shared public-data distillation with local private-data adaptation.
- Communication: During communication, each participant computes class scores on public examples and sends them to a central server.The public dataset is reused as the communication basis through knowledge distillation.
3 Results
FedMD is evaluated with heterogeneous convolutional models across FEMNIST and CIFAR10/CIFAR100 settings, including both i.i.d. and non-i.i.d. private-data conditions. The results compare collaborative learning with individual and pooled-data baselines.
- FEMNIST environment: FedMD is tested with MNIST public data and FEMNIST private data under both i.i.d. and non-i.i.d. conditions.The non-i.i.d. setting trains each participant on letters from one writer but tests classification across all writers.
- CIFAR environment: The CIFAR environment uses CIFAR10 as public data and CIFAR100 as private data, with a non-i.i.d. task requiring superclass recognition from one subclass per superclass.Participants must use information communicated by others to generalize beyond the subclasses seen during training.
- Experimental setup: 10 participants design unique convolutional networks differing in channels and layers, then undergo public-data pretraining, private-data training, and collaborative training.Pre-collaboration models typically reach around 99% MNIST and 76% CIFAR10 test accuracy on the public datasets.
- Baselines: FedMD’s learning curves are evaluated against a dashed individual-data baseline and a dash-dot pooled-private-data performance reference.The dashed baseline follows full transfer learning on public data plus each model’s own small private dataset.
4 Discussion and conclusion
The paper presents FedMD as a knowledge-distillation framework for federated learning with independently designed models and reports successful testing across varied tasks and datasets. It identifies richer communication and more extreme heterogeneity as future directions.
- Conclusion: FedMD enables federated learning for independently designed models using knowledge distillation across various tasks and datasets.The paper frames this as the central contribution and reports that the framework works across its tested settings.
- Future work: Future work will investigate feature transformations, emergent communication protocols, NLP and reinforcement-learning applications, and larger discrepancies in data, model capacity, and local tasks.These directions extend the communication module and the scope of heterogeneity considered.
Supplementary Material
The supplementary material documents the model architectures, datasets, implementation details, and performance comparisons used to evaluate FedMD. Results are compared with models trained using pooled private data.
- The supplementary material provides additional details about the models, datasets, algorithm, and results.
- Communication uses public-data logits, a randomly selected subset of 5000 public samples per round, and usually equal consensus weights.The implementation notes state that logits are aligned without softmax, 5000 samples speed communication without sacrificing performance, and equal weights are typical except for two weaker CIFAR models.
- The results compare federated performance with test accuracy obtained when all participants' private data are pooled.Table 4 reports performance for models trained with pooled private data.