Towards Provably Personalized Federated Learning via Threshold-Clustering of Similar Clients

Mariel Werner, Lie He, Sai Praneeth Karimireddy, Michael Jordan, Martin Jaggi

Keywords: Federated Learning, Collaborative Learning, Personalized Learning, Clustering

TL;DR: We develop simple clustering-based algorithms that achieve personalization in federated learning and that have optimal convergence guarantees.

Abstract: Clustering clients with similar objectives together and learning a model per cluster is an intuitive and interpretable approach to personalization in federated learning (PFL). However, doing so with provable and optimal guarantees has remained an open challenge. In this work, we formalize personalized federated learning as a stochastic optimization problem where the stochastic gradients on a client may correspond to one of $K$ distributions. In such a setting, we show that using i) a simple thresholding based clustering algorithm, and ii) local client momentum obtains optimal convergence guarantees. In fact, our rates asymptotically match those obtained if we knew the true underlying clustering of the clients. Further, we extend our algorithm to the decentralized setting where each node performs clustering using itself as the center.

Is Student: Yes