[Paper Review] LotteryFL: Personalized and Communication-Efficient Federated Learning with Lottery Ticket Hypothesis on Non-IID Datasets

Ang Li, Jingwei Sun+5 more|arXiv (Cornell University)|Aug 7, 2020

Privacy-Preserving Technologies in Data25 references73 citations

TL;DR

LotteryFL personalizes federated learning by training and exchanging lottery ticket subnetworks (LTNs) per client, achieving better personalization with lower communication cost on non-IID data.

ABSTRACT

Federated learning is a popular distributed machine learning paradigm with enhanced privacy. Its primary goal is learning a global model that offers good performance for the participants as many as possible. The technology is rapidly advancing with many unsolved challenges, among which statistical heterogeneity (i.e., non-IID) and communication efficiency are two critical ones that hinder the development of federated learning. In this work, we propose LotteryFL -- a personalized and communication-efficient federated learning framework via exploiting the Lottery Ticket hypothesis. In LotteryFL, each client learns a lottery ticket network (i.e., a subnetwork of the base model) by applying the Lottery Ticket hypothesis, and only these lottery networks will be communicated between the server and clients. Rather than learning a shared global model in classic federated learning, each client learns a personalized model via LotteryFL; the communication cost can be significantly reduced due to the compact size of lottery networks. To support the training and evaluation of our framework, we construct non-IID datasets based on MNIST, CIFAR-10 and EMNIST by taking feature distribution skew, label distribution skew and quantity skew into consideration. Experiments on these non-IID datasets demonstrate that LotteryFL significantly outperforms existing solutions in terms of personalization and communication cost.

Motivation & Objective

• Address statistical heterogeneity (non-IID data) in federated learning by enabling per-client personalization.
• Reduce communication overhead in FL by transmitting compact lottery ticket subnetworks instead of full models.
• Leverage the Lottery Ticket hypothesis to identify sparse, trainable subnetworks tailored to each client’s data.
• Provide non-IID datasets and a metric to quantify data heterogeneity to facilitate FL research under non-IID conditions.

Proposed method

• Each client learns a Lottery Ticket Network (LTN) by pruning the base model using local data.
• Only LTNs are communicated between server and clients; the server aggregates LTNs via FedAvg.
• LTNs are re-initialized from a common baseline and re-pruned over rounds to adapt to local data.
• Client validation governs pruning and mask updating to produce a personalized mask for the next round.
• The server updates the global base model by aggregating the LTNs received from clients.
• A new non-IID data generation protocol and the Client-Wise Non-IID Index (CNI) metric are introduced.

Experimental results

Research questions

• RQ1 Does LotteryFL achieve better personalization on non-IID data compared to standard FL baselines?
• RQ2 Can communication costs be substantially reduced by exchanging LTNs instead of full models?
• RQ3 How does non-IID data degree (as measured by CNI) affect LotteryFL performance?
• RQ4 What is the trade-off between pruning rate and model performance in a personalized FL setting?

Key findings

• LotteryFL achieves superior personalization and lower communication costs across MNIST, CIFAR-10, and EMNIST non-IID datasets compared to FedAvg, LG-FedAvg, and Standalone baselines.
• LTNs are highly sparse and personalized, with a large fraction of parameters unique to individual clients as pruning increases.
• Increasing the target pruning rate generally increases the share of personalized parameters per layer, while maintaining or improving accuracy.
• More participating clients per round improves personalization for all methods, with LotteryFL maintaining the best performance and lowest communication cost.
• Higher data per class on each client (more samples) accelerates convergence, enabling even more aggressive pruning and further communication savings.
• The authors propose and validate the Client-Wise Non-IID Index (CNI) as a simple, fixed-encoder-based measure of data heterogeneity across clients.