[Paper Review] FedMD: Heterogenous Federated Learning via Model Distillation

Daliang Li, Junpu Wang|arXiv (Cornell University)|Oct 8, 2019

Privacy-Preserving Technologies in DataComputer Science17 references480 citations

TL;DR

FedMD enables federated learning where participants use independently designed models by translating knowledge via distillation on a public dataset, yielding gains over isolated training and approaching pooled-data performance.

ABSTRACT

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.

Motivation & Objective

• Motivate federated learning in settings where participants deploy their own model architectures.
• Propose a framework that allows model heterogeneity without sharing private data or architectures.
• Leverage transfer learning and knowledge distillation to enable cross-model collaboration.
• Evaluate FedMD on standard datasets to demonstrate performance gains over isolated training.

Proposed method

• Use a public dataset as a common communication basis.
• Each party pre-trains its own model on public data then on private data (transfer learning).
• Models share class scores on public data; a central server averages these scores to form a consensus.
• Each participant updates its model to align with the consensus on the public data (distillation).
• Repeat the digest and revisit steps, with occasional sub-sampling of the public data for efficiency.
• Optionally weight participants differently when forming the consensus.

Experimental results

Research questions

• RQ1 Can heterogeneous models collaborate in federated learning without sharing data or architectures?
• RQ2 How can knowledge be translated between diverse models to improve each participant’s performance?
• RQ3 What performance gains are achievable relative to isolated training and pooled-data upper bounds?
• RQ4 How does the framework perform under i.i.d. and non-i.i.d. data distributions across participants?

Key findings

• Across MNIST/FEMNIST and CIFAR10/CIFAR100, FedMD yields significant improvements over isolated transfer learning.
• With 10 participants, final test accuracy gains average around 20% over non-collaborative baselines.
• Performance approaches the pooled-private-data upper bound, differing by only a few percent.
• Preliminary results show typical pre-collaboration accuracies around 99% on MNIST and 76% on CIFAR10.
• FedMD remains effective across both i.i.d. and non-i.i.d. settings, with heterogeneity in model architectures.
• In some cases, certain simple models can achieve competitive or superior performance within the FedMD framework.