Federated Lottery: Private and Communication-Efficient Learning of Personalized Networks

Abstract

A promising approach to address privacy concerns, Federated learning (FL) enables distributed training of machine learning (ML) models where user data remains on edge devices and isn’t shared. However, classic FL paradigms struggle in real-world settings where they underperform on non- independent and identically distributed (non-iid) data, fail to generate personalized models, and face major bottlenecks in communication efficiency. Prior work has tackled each of these challenges individually, but finding one approach that simultaneously addresses all remains elusive. This work proposes FedLottery, a novel framework that learns personalized lottery ticket networks (LTNs). FedLottery presents a significant improvement from prior attempts at applying the Lottery Ticket Hypothesis (LTH) to FL by introducing new approaches of postpruning without rewinding and preserving batch normalization (BN) parameters during federated aggregation. This paper also presents Jump-Start as an extension of FedLottery that drastically reduces (halves) the number of communication rounds. FedLottery with Jump-Start significantly outperforms earlier approaches, achieving 9.95% higher personalized test accuracy than the best performing baseline on TinyImageNet data. Importantly, this is achieved while simultaneously reducing communication cost by up to 8.1X and pruning models to be 90% sparse. FedLottery maintains performance improvements when experiments are scaled to 100 clients running resource-constrained custom-CNN models that have a mere 31KB in memory footprint when 90% pruned.