FedReMa: Improving Personalized Federated Learning via Leveraging the Most Relevant Clients

TL;DR

FedReMa is proposed, an efficient PFL algorithm that can tackle class-imbalance by utilizing an adaptive inter-client co-learning approach to identify and harness different clients' expertise on different data classes throughout various phases of the training process, and employing distinct aggregation methods for clients' feature extractors and classifiers.

Abstract

Federated Learning (FL) is a distributed machine learning paradigm that achieves a globally robust model through decentralized computation and periodic model synthesis, primarily focusing on the global model's accuracy over aggregated datasets of all participating clients. Personalized Federated Learning (PFL) instead tailors exclusive models for each client, aiming to enhance the accuracy of clients' individual models on specific local data distributions. Despite of their wide adoption, existing FL and PFL works have yet to comprehensively address the class-imbalance issue, one of the most critical challenges within the realm of data heterogeneity in PFL and FL research. In this paper, we propose FedReMa, an efficient PFL algorithm that can tackle class-imbalance by 1) utilizing an adaptive inter-client co-learning approach to identify and harness different clients' expertise on different data classes throughout various phases of the training process, and 2) employing distinct aggregation methods for clients' feature extractors and classifiers, with the choices informed by the different roles and implications of these model components. Specifically, driven by our experimental findings on inter-client similarity dynamics, we develop critical co-learning period (CCP), wherein we introduce a module named maximum difference segmentation (MDS) to assess and manage task relevance by analyzing the similarities between clients' logits of their classifiers. Outside the CCP, we employ an additional scheme for model aggregation that utilizes historical records of each client's most relevant peers to further enhance the personalization stability. We demonstrate the superiority of our FedReMa in extensive experiments.

Figures (3)

• Figure 1: Workflow of FedReMa. (1) Local training. (2) Upload model to server. (3) Aggregate global feature extractors. (4) Determine whether the current period is a critical learning period, if so, go to (5), otherwise go to (8). (5) Measure the task relevance $\mathcal{S}$ using similarity of their logits, then perform MDS to filter out irrelevant clients. At the same time, the server will record the client's historical dependencies $\mathcal{G}$. (6) Update CCP status. (7) Aggregate personalized classifiers by $\mathcal{A}_k^{(t)}$. (8) Aggregate personalized classifiers based on historical dependencies $\mathcal{G}$. (9) Server communicates personalized models to all clients.
• Figure 2: The heatmap of $\mathbf{S}^{(t)}$ on different communication round $t$ (FedAvg)
• Figure 3: The test accuracy while training process.