Source-linked AI summary
Communication-Efficient On-Device Machine Learning: Federated Distillation and Augmentation under Non-IID Private Data
Eunjeong Jeong, Seungeun Oh, Hyesung Kim, Jihong Park, Mehdi Bennis, Seong-Lyun Kim
TL;DR
On-device ML must limit communication while learning from private, non-IID data, since FL communication scales with model size and non-IID data can reduce accuracy. The paper proposes FD, which exchanges aggregated model outputs, and FAug, which uses a shared GAN to restore IID-like local data. Empirically, FD with FAug achieves 95-98% test accuracy with around 26x lower communication cost than FL.
Problem
FL incurs model-size-proportional communication overhead, while non-IID private data can substantially reduce prediction accuracy and raw data exchange raises privacy concerns.
Method
FD exchanges per-label averaged logit vectors instead of model parameters, while FAug trains a conditional GAN to generate missing local data samples.
Results
FD with FAug achieves 95-98% test accuracy with around 26x lower aggregate communication cost than FL.
Takeaways & Limitations
FD and FAug achieve comparably high accuracy with much smaller communication overhead than FL.
Takeaways & Limitations
FAug’s privacy guarantee could be further improved using differential privacy, and FD performance could be improved through weighted local-logit averaging.
Abstract
from arXiv · showhide
On-device machine learning (ML) enables the training process to exploit a massive amount of user-generated private data samples. To enjoy this benefit, inter-device communication overhead should be minimized. With this end, we propose federated distillation (FD), a distributed model training algorithm whose communication payload size is much smaller than a benchmark scheme, federated learning (FL), particularly when the model size is large. Moreover, user-generated data samples are likely to become non-IID across devices, which commonly degrades the performance compared to the case with an IID dataset. To cope with this, we propose federated augmentation (FAug), where each device collectively trains a generative model, and thereby augments its local data towards yielding an IID dataset. Empirical studies demonstrate that FD with FAug yields around 26x less communication overhead while achieving 95-98% test accuracy compared to FL.
1 Introduction
On-device ML can use private user-generated data without directly sharing it, but FL faces model-size communication costs and accuracy losses under non-IID data. The paper proposes FD for communication efficiency and FAug to make local data more IID.
- On-device ML exploits private user-generated data while preserving privacy by exchanging local model parameters rather than raw data.
- Under FL, communication overhead scales with model size, limiting the use of large models.
- Non-IID data reduces FL prediction accuracy by up to 11% on MNIST and 51% on CIFAR-10 compared with IID data.
- FD uses output-based distributed knowledge distillation whose communication payload depends on output dimension rather than model size.
- FAug uses a collectively trained GAN to augment local data toward an IID training dataset, balancing privacy leakage and communication overhead.
2 Federated distillation
FD adapts online knowledge distillation to reduce communication by exchanging per-label average outputs instead of model parameters or sample-level outputs. Devices use global per-label averages as teacher outputs during local training.
- FD exchanges model outputs rather than model parameters, enabling large-sized local models while retaining periodic communication.
- The distillation regularizer uses cross entropy between the student output and the teacher output during local optimization.
- Standard co-distillation is communication-intensive because teacher and student outputs must correspond to identical training samples.
- FD stores per-label mean logit vectors locally, averages them at the server, and redistributes global-average vectors to devices.
- Each device selects the global-average logit vector matching the current sample’s ground-truth label as its teacher output.
3 Federated augmentation
FAug addresses non-IID local data by training a shared conditional GAN from sparse seed samples and downloading its generator to devices. It also quantifies privacy leakage arising from target-label disclosure and shared generation capabilities.
- FAug locally generates missing data samples with a generative model instead of exchanging complete datasets between devices.
- Devices upload seed samples for lacking target labels, while the server oversamples them to train a conditional GAN.
- Downloading the trained GAN generator lets each device replenish target labels until its training dataset becomes IID.
- Devices upload redundant non-target-label samples to reduce server-side privacy leakage, at the cost of extra uplink communication.
- Shared generators can reveal other devices’ target labels, creating inter-device privacy leakage.
- Inter-device privacy leakage is minimized when the denominator reaches the number of entire labels, achievable with sufficiently many devices.
4 Evaluation
The evaluation on non-IID MNIST compares FD and FL, with and without FAug, on test accuracy, communication cost, and privacy leakage. FAug substantially improves FD accuracy while the combined approach retains much lower communication cost than FL.
- Dataset and setup: The evaluation uses a non-IID MNIST dataset with 2,000 samples allocated per device, after reducing each target label to 5 samples.Target labels are selected randomly, and approximately 97.5% of their samples are removed.
- Communication cost: FD exchanges 100 logits per global iteration, whereas FL exchanges 1,199,648 CNN model parameters in both uplink and downlink.Each FD logit comprises 10 logit vectors with 10 elements each.
- Accuracy and communication: 26x smaller communication cost is achieved by FD with FAug than by FL, while test accuracy reaches 95-98%.Without FAug, FD achieves 77-90% test accuracy; FAug raises it by 7-22%.
- Accuracy and communication: FAug increases FD test accuracy by 7-22% and FL test accuracy by 0.8-2.7% relative to training without FAug.The comparison is against the corresponding non-IID cases without FAug.
- Privacy and per-label accuracy: Around 2x higher test accuracy is obtained by FD with FAug than by standalone training across device and redundant-label counts.For target label ‘2’, standalone accuracy is 3.585%, while FAug with FD or FL reaches 73.44% or 92.19%, respectively.
- Privacy and per-label accuracy: Inter-device privacy leakage decreases with more devices or redundant labels, while device-server privacy leakage decreases with redundant labels and increases with target labels.Figure 2 reports these trends alongside accuracy under FD with FAug.
5 Concluding remarks
The paper concludes that FD and FAug provide communication-efficient training and augmentation for on-device ML. It reports comparable accuracy with much smaller communication overhead than FL, while identifying several directions for improving accuracy and privacy.
- Conclusion: FD and FAug are introduced as communication-efficient training and data augmentation algorithms, respectively.FAug addresses data augmentation, while FD addresses distributed training communication.
- Conclusion: Comparable high accuracy is achieved with much smaller communication overhead than FL.This is the paper’s reported empirical conclusion.
- Future work: Future work proposes weighted local-logit averaging, hybrid FD-FL communication, and differential-privacy noise for uploaded seed samples.The proposed combinations aim to improve FD performance, balance communication efficiency and accuracy, and ameliorate FAug privacy guarantees.