On the Convergence of Clustered Federated Learning

TL;DR

This paper introduces a unified bi-level optimization framework for clustered federated learning, providing convergence analysis and a new algorithm, WeCFL, which effectively handles non-IID data by clustering clients.

Contribution

It formulates clustered FL into a bi-level optimization framework, offers a convergence proof considering client clusterability, and proposes the WeCFL algorithm for improved performance.

Findings

Theoretical convergence guarantees under clusterability assumptions.

Empirical results show WeCFL outperforms existing methods in non-IID settings.

The framework unifies various clustered FL approaches.

Abstract

Knowledge sharing and model personalization are essential components to tackle the non-IID challenge in federated learning (FL). Most existing FL methods focus on two extremes: 1) to learn a shared model to serve all clients with non-IID data, and 2) to learn personalized models for each client, namely personalized FL. There is a trade-off solution, namely clustered FL or cluster-wise personalized FL, which aims to cluster similar clients into one cluster, and then learn a shared model for all clients within a cluster. This paper is to revisit the research of clustered FL by formulating them into a bi-level optimization framework that could unify existing methods. We propose a new theoretical analysis framework to prove the convergence by considering the clusterability among clients. In addition, we embody this framework in an algorithm, named Weighted Clustered Federated Learning (WeCFL). Empirical analysis verifies the theoretical results and demonstrates the effectiveness of the proposed WeCFL under the proposed cluster-wise non-IID settings.