Lightweight Clustered Federated Learning via Feature Extraction

Abstract

Clustered federated learning (FL), which groups clients with similar data distributions for collaborative training, represents a pivotal technique within federated learning for effectively addressing the challenges posed by non-IID data on clients. Existing clustered FL algorithms typically endeavor to learn distribution similarities of clients iteratively or indirectly through representations like gradients and loss. This necessitates resource-intensive pre-training or multiple iterations to attain stable clusters, thereby incurring additional communication cost and computational overhead. To address the above issues, we propose lightweight Clustered Federated Learning via Feature Extraction (FECFL). FECFL adopts a simple but effective client representation, i.e., the features extracted from the clients’ data using identically initialized models without any pre-training, to perform efficient one-shot clustering. Moreover, client data distribution is often dynamic in practice. To tackle distribution shift, we embed a distribution monitoring mechanism in FECFL, enabling adaptive re-grouping for new distributions. Finally, we demonstrate the benefits of FECFL over the baselines by conducting experiments on various datasets and distributions.