DPMM-CFL: Clustered Federated Learning via Dirichlet Process Mixture Model Nonparametric Clustering
This addresses the impracticality of pre-specifying cluster counts in federated learning for heterogeneous clients, though it is an incremental improvement over existing CFL methods.
The paper tackles the problem of clustered federated learning (CFL) requiring a fixed number of clusters by proposing DPMM-CFL, which uses a Dirichlet Process prior to infer the number of clusters and client assignments nonparametrically, and it is validated on benchmark datasets under non-IID partitions.
Clustered Federated Learning (CFL) improves performance under non-IID client heterogeneity by clustering clients and training one model per cluster, thereby balancing between a global model and fully personalized models. However, most CFL methods require the number of clusters K to be fixed a priori, which is impractical when the latent structure is unknown. We propose DPMM-CFL, a CFL algorithm that places a Dirichlet Process (DP) prior over the distribution of cluster parameters. This enables nonparametric Bayesian inference to jointly infer both the number of clusters and client assignments, while optimizing per-cluster federated objectives. This results in a method where, at each round, federated updates and cluster inferences are coupled, as presented in this paper. The algorithm is validated on benchmark datasets under Dirichlet and class-split non-IID partitions.