An Improved Federated Clustering Algorithm with Model-based Clustering

Harsh Vardhan, Avishek Ghosh, Arya Mazumdar

Abstract:

Federated learning (FL) is a distributed learning paradigm that allows multiple clients to collaboratively train a shared model via communications to a central server. However, optimal models of different clients often differ due to heterogeneity of data across clients.
In this paper, we address the dichotomy between heterogeneous models and simultaneous training in FL via a clustering structure among the clients. The clustering framework is one way to allow for high heterogeneity level between clients, while clients with similar data can still train a shared model. We define a new clustering framework for FL based on the (optimal) local models of the clients: two clients belong to the same cluster if their local models are close. We propose an algorithm, \emph{Successive Refine Federated Clustering Algorithm} (\texttt{SR-FCA}), that treats each client as a singleton cluster as an initialization, and then successively refine the cluster estimation via exploiting similarity with other clients. In any intermediate step, \texttt{SR-FCA} uses an {\em error-tolerant} federated learning algorithm within each cluster to exploit simultaneous training and to correct clustering errors. Unlike some prominent prior works \texttt{SR-FCA} does not require any \emph{good} initialization (or warm start), both in theory and practice. We show that with proper choice of learning rate, \texttt{SR-FCA} incurs arbitrarily small clustering error. Additionally, \texttt{SR-FCA} does not require the knowledge of the number of clusters apriori like some prior works. We validate the performance of \texttt{SR-FCA} on real-world FL datasets including FEMNIST and Shakespeare in non-convex problems and show the benefits of \texttt{SR-FCA} over several baselines.

Submission Length: Regular submission (no more than 12 pages of main content)

Code: https://github.com/harshv834/sr-fca

Assigned Action Editor: ~Ahmad_Beirami1

License: Creative Commons Attribution 4.0 International (CC BY 4.0)

Submission Number: 1838