CFARMS: A clustered federated learning framework with recursive model selection

Abstract

Federated learning, originally devised to train a single global model over diverse client populations while upholding data privacy, faces challenges due to variations in data distributions among clients. The learning and generalization performance of the global model tends to be suboptimal. To address this, we propose CFARMS (Clustered Federated Learning Algorithm with Recursive Model Selection), a novel framework that effectively enhances model learning within clusters. CFARMS employs an iterative process to cluster clients, primarily leveraging gradients from the loss function of each client's local training. Moreover, in instances of a tie, CFARMS also considers clients’ local model losses to inform clustering decisions. In contrast to existing frameworks, CFARMS produces highly personalized cluster-specific models, mitigating challenges in global model generalization. Additionally, it significantly reduces communication costs within the federated learning network by recursively refining the number of local models per client and updating the cluster models periodically. During training, CFARMS uniquely facilitates decentralized search for the most optimal model for each client. Extensive experimental evaluations on benchmark image and tabular datasets with non-convex models demonstrate the proposed framework's superior performance. It achieved higher prediction accuracy and a remarkable reduction in communication bottleneck, exceeding 80 % and 49 % respectively, compared to state-of-the-art federated learning frameworks, including IFCA, FedProx and FedAvg. It is worth mentioning that the proposed framework exhibits faster convergence and effective scalability to a higher number of models and clients.