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.
翻译:聚类联邦学习(CFL)通过将客户端聚类并为每个聚类训练一个模型,从而在全局模型与完全个性化模型之间取得平衡,以提升在非独立同分布(non-IID)客户端异构性下的性能。然而,大多数CFL方法需要预先固定聚类数量K,这在潜在结构未知的情况下是不切实际的。我们提出了DPMM-CFL,一种在聚类参数分布上施加狄利克雷过程(DP)先验的CFL算法。这使得非参数贝叶斯推断能够联合推断聚类数量和客户端分配,同时优化每个聚类的联邦学习目标。由此产生的方法在每一轮中,联邦更新与聚类推断是耦合的,如本文所述。该算法在狄利克雷分布和类别分割的非IID划分下的基准数据集上得到了验证。