Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DPMM-CFL: Clustered Federated Learning via Dirichlet Process Mixture Model Nonparametric Clustering

Mariona Jaramillo-Civill, Peng Wu, Pau Closas

TL;DR

The paper tackles the limitation of clustered federated learning (CFL) requiring a pre-specified number of clusters $K$ by introducing DPMM-CFL, which uses a Dirichlet Process prior to perform nonparametric clustering of clients. It jointly infers the number of clusters and client assignments while optimizing per-cluster federated objectives through an iterative, round-based process that couples federated updates with Bayesian clustering. The core contributions include (i) formulating CFL with a DP prior over cluster parameters, (ii) employing a split–merge MCMC strategy to sample client partitions under a CRP, and (iii) validating the method on Fashion-MNIST and CIFAR-10 under Dirichlet and class-split non-IID partitions, showing close alignment between inferred and ground-truth cluster counts or peak performance clusters and stable convergence. The approach enables automatic structure learning in CFL, reducing the need for hyperparameter sweeps and enhancing personalization in heterogeneous FL settings.

Abstract

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.

DPMM-CFL: Clustered Federated Learning via Dirichlet Process Mixture Model Nonparametric Clustering

TL;DR

The paper tackles the limitation of clustered federated learning (CFL) requiring a pre-specified number of clusters by introducing DPMM-CFL, which uses a Dirichlet Process prior to perform nonparametric clustering of clients. It jointly infers the number of clusters and client assignments while optimizing per-cluster federated objectives through an iterative, round-based process that couples federated updates with Bayesian clustering. The core contributions include (i) formulating CFL with a DP prior over cluster parameters, (ii) employing a split–merge MCMC strategy to sample client partitions under a CRP, and (iii) validating the method on Fashion-MNIST and CIFAR-10 under Dirichlet and class-split non-IID partitions, showing close alignment between inferred and ground-truth cluster counts or peak performance clusters and stable convergence. The approach enables automatic structure learning in CFL, reducing the need for hyperparameter sweeps and enhancing personalization in heterogeneous FL settings.

Abstract

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.

Paper Structure

This paper contains 8 sections, 14 equations, 2 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Clustered federated learning with Bayesian clustering
  3. Nonparametric Clustering of Federated Learning clients with DPMM
  4. Evaluation
  5. Datasets and Non-IID generation
  6. System Settings
  7. Results
  8. Conclusions

Figures (2)

  • Figure 1: Performance of FeSEM (fixed $K$) versus DPMM-CFL (unknown $K$) on Fashion-MNIST and CIFAR-10 under Dirichlet and class-split partitions. Solid curves with shaded regions show mean $\pm$1 standard deviation across seeds. Dotted horizontal lines indicate mean performance for the unknown $K$ case, and dotted vertical lines indicate the mean inferred number of clusters (with variability across seeds).
  • Figure 2: Empirical convergence of DPMM-CFL on CIFAR-10 with class-split (3 classes per cluster, 2 classes per client): accuracy and F1 score converge while the inferred $K_t$ stabilizes after early rounds.