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

2025-10-08

AI Summary

Existing clustering-based federated learning (CFL) methods require pre-specifying the number of clusters (K) and struggle to adapt to unknown client heterogeneity. To address this, we propose DPMM-CFL—a nonparametric Bayesian CFL framework built upon the Dirichlet process mixture model (DPMM) and variational inference, which automatically infers the optimal number of clusters without prior knowledge of (K). Our approach decouples client clustering from model training while jointly optimizing both components, integrating federated averaging for distributed model updates. Evaluated on Dirichlet- and label-skew-type non-IID data, DPMM-CFL achieves superior global accuracy and personalized performance trade-offs compared to fixed-(K) baselines. It provides a scalable, adaptive clustering paradigm for heterogeneous federated learning, eliminating reliance on manual cluster specification and enhancing robustness to unseen client data distributions.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Web Mining and Content Analysis: Normalization, clustering, classification, and summarization of Web text

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.

Problem

• Automatically determines cluster count in federated learning
• Handles non-IID client data through Bayesian clustering
• Simultaneously infers cluster assignments and trains models

Innovation

• Uses nonparametric Bayesian clustering for federated learning
• Automatically infers number of client clusters from data
• Combines federated updates with cluster inference jointly