Lazy But Effective: Collaborative Personalized Federated Learning with Heterogeneous Data

TL;DR

The paper tackles the challenge of non-IID data across clients in Federated Learning by introducing pFedLIA, a clustering-based personalized FL framework that uses a Lazy Influence Approximation to cheaply estimate cross-client influence and form client clusters. The method supports both centralized and peer-to-peer clustering and is compatible with any FL aggregator, without requiring a priori knowledge of the number of clusters. Experiments on next-word prediction for Nordic languages and standard vision benchmarks show that pFedLIA matches Oracle clustering performance and outperforms strong baselines by up to 17%, while offering substantial computational efficiency gains. This approach enables personalized, cluster-specific models that improve per-client performance in heterogeneous FL settings and scales to decentralized training scenarios.

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.

Figures (4)

• Figure 1: Validation loss over time, when training a GPT-2 model for next word prediction. We compare using a single model trained on all client data using FedAvg (blue line), versus clustering clients in a peer-to-peer setting using the proposed pFedLIA and learning distinct models. pFedLIA automatically splits the clients into three clusters (green, orange, and pink lines, respectively -- which roughly correspond to Norwegian, Swedish, and Danish clients). FedAvg is clearly unable to learn in this setting, due to the data heterogeneity. In comparison, pFedLIA achieves loss comparable to centralized training, where all data is collected, and a single model is trained centrally (red line).
• Figure 2: Visual illustration of the proposed pFedLIA, from the perspective of client $i$ (black phone on the left). Client $i$ calculates a Lazy Influence Approximation score for every other client $j$ (phones on the top). To do so, every client $j$ performs a few local epochs to the joint model $M_0$ -- enough to get an estimate of the gradient (e.g., 5–20), no need to train the model fully -- using a batch of their training data. Subsequently, $j$ sends the updated partial model $M_{j}$ to client $i$. The Lazy Influence Approximation $\mathcal{I}(i,j)$ is simply the difference in the validation loss of models $M_0$ and $M_{j}$. Finally, we use these influence scores to cluster clients; in this case, the blue and orange phones will be in $i$'s cluster, while the green and purple will not.
• Figure 3: (\ref{['figure: LIA speedup']}) Computation time comparison between the proposed Lazy Influence Approximation (LIA), versus the exact influence (for three different convergence thresholds of varying precision). LIA achieves between 40–500 times speedup. (\ref{['figure: pFedLIA decentralized clustering illustration']}) pFedLIA decentralized clustering illustration. Consider a classification task (classes [0-7]). We have 32 clients, each assigned 2 out of the 8 classes. This figure depicts the Lazy Influence Approximation values of every client, as calculated by client 0 with data from classes [0, 1]. The red horizontal line represents the clustering frontier, using k-means (2 classes, representing beneficial or not beneficial clients for client 0).
• Figure 4: Training comparison of our proposed method and all baselines on CIFAR100. Note that we train the 'warm-up' model till round 20 (see Section \ref{['sec: Evaluation Setup']}), hence the jump for the proposed approach at this time-step.