[Paper Review] Distilled One-Shot Federated Learning

Yanlin Zhou, George Pu+3 more|arXiv (Cornell University)|Sep 17, 2020

Privacy-Preserving Technologies in Data42 references73 citations

TL;DR

DOSFL distills each client’s private data into synthetic examples in a single round, sending only the distilled data to the server to train a global model with up to 1000x less communication while achieving 93%–99% of centralized performance.

ABSTRACT

Current federated learning algorithms take tens of communication rounds transmitting unwieldy model weights under ideal circumstances and hundreds when data is poorly distributed. Inspired by recent work on dataset distillation and distributed one-shot learning, we propose Distilled One-Shot Federated Learning (DOSFL) to significantly reduce the communication cost while achieving comparable performance. In just one round, each client distills their private dataset, sends the synthetic data (e.g. images or sentences) to the server, and collectively trains a global model. The distilled data look like noise and are only useful to the specific model weights, i.e., become useless after the model updates. With this weight-less and gradient-less design, the total communication cost of DOSFL is up to three orders of magnitude less than FedAvg while preserving between 93% to 99% performance of a centralized counterpart. Afterwards, clients could switch to traditional methods such as FedAvg to finetune the last few percent to fit personalized local models with local datasets. Through comprehensive experiments, we show the accuracy and communication performance of DOSFL on both vision and language tasks with different models including CNN, LSTM, Transformer, etc. We demonstrate that an eavesdropping attacker cannot properly train a good model using the leaked distilled data, without knowing the initial model weights. DOSFL serves as an inexpensive method to quickly converge on a performant pre-trained model with less than 0.1% communication cost of traditional methods.

Motivation & Objective

• Reduce communication rounds and data transmitted in federated learning without sacrificing accuracy.
• Leverage dataset distillation to create synthetic data that trains a global model in one round.
• Address non-IID data challenges with soft labels, soft resets, and random masking.
• Demonstrate the applicability of DOSFL across vision and language tasks and model types.

Proposed method

• Server initializes a global model θ0 and broadcasts it to clients.
• Each client distills its private data into a small synthetic dataset with associated labels and learning rates.
• Distilled data from all clients are merged on the server and used to update the global model through multiple gradient steps on the distilled sequence.
• Soft labels are used to improve robustness on non-IID data.
• Two techniques—soft reset and random masking—are introduced to mitigate interference from non-IID distilled data.
• The final model is distributed back to clients for optional fine-tuning.

Experimental results

Research questions

• RQ1 Can federated learning be achieved with orders-of-magnitude less communication by transmitting distilled data instead of model weights or gradients?
• RQ2 How do soft labels, soft resets, and random masking affect performance on IID and non-IID data across vision and language tasks?
• RQ3 To what extent can a single round of distilled data training approximate centralized training across diverse model types?
• RQ4 Is the distilled data approach resilient to eavesdropping and initial-weight uncertainty?

Key findings

• DOSFL achieves up to ~1000x reduction in communication compared with FedAvg in the reported settings.
• On IID data, DOSFL preserves 93%–99% of centralized training performance across tasks.
• On non-IID data, soft resets yield the largest gains among the proposed techniques, improving robustness considerably.
• DOSFL supports multiple model types (CNNs, LSTMs, Transformers) and tasks (vision and language) with comparable accuracy to centralized baselines.
• An eavesdropper cannot reproduce the global model using leaked distilled data without knowledge of the initial server weights, indicating a privacy advantage over traditional FL.