Lazy But Effective: Collaborative Personalized Federated Learning with Heterogeneous Data

TL;DR

This paper introduces pFedLIA, a personalized federated learning framework that clusters clients using influence approximation to improve model performance on heterogeneous, non-IID data across diverse applications.

Contribution

The paper proposes a novel, computationally efficient clustering method called Lazy Influence for personalized federated learning, enhancing performance on non-IID data.

Findings

Successfully recovers global model performance in non-IID settings

Matches Oracle clustering performance in experiments

Improves baseline methods by up to 17% on CIFAR100

Abstract

In Federated Learning, heterogeneity in client data distributions often means that a single global model does not have the best performance for individual clients. Consider for example training a next-word prediction model for keyboards: user-specific language patterns due to demographics (dialect, age, etc.), language proficiency, and writing style result in a highly non-IID dataset across clients. Other examples are medical images taken with different machines, or driving data from different vehicle types. To address this, we propose a simple yet effective personalized federated learning framework (pFedLIA) that utilizes a computationally efficient influence approximation, called `Lazy Influence', to cluster clients in a distributed manner before model aggregation. Within each cluster, data owners collaborate to jointly train a model that captures the specific data patterns of the clients. Our method has been shown to successfully recover the global model's performance drop due to the non-IID-ness in various synthetic and real-world settings, specifically a next-word prediction task on the Nordic languages as well as several benchmark tasks. It matches the performance of a hypothetical Oracle clustering, and significantly improves on existing baselines, e.g., an improvement of 17% on CIFAR100.