Data Distribution Profiling for Clustered Federated Learning in Extreme NIID Settings

The homogeneity of the training data is a critical determinant of the performance, reliability, and stability of any AI model. Federated learning (FL) provides a promising framework for developing global AI models in a distributed heterogeneous environment, where individual data nodes can maintain data privacy while contributing to a shared model. One of the main challenges that arises from the heterogeneity introduced by FL is the presence of non-independently and identically distributed (NIID) datasets, which often results in a global model training degradation and large average performance variance on each local dataset. Clustered federated learning (CFL) has been introduced to mitigate this issue by grouping nodes into clusters with jointly trainable data distributions. CFL is a post-processing approach serving personalized FL, however; very few works utilize CFL as a preprocessing method for further generalized global training. This paper addresses this gap by leveraging CFL as a preprocessing step, focusing on enhancing the node profile based on data distribution to enable efficient clustering, thereby promoting a more stable and generalized FL model under the challenge of NIID. The proposed approach utilizes a disentangled variational auto encoder (β- VAE) to produce more generalizable nodes’ embeddings with the most uncorrelated features. The statistics embedding vectors (mean and three quantiles) are utilized to serve as the lightweight clustering object. This paper presents extreme simulation models of pathological nodes using the MNIST dataset. These models compare the proposed node profiling and clustering approach against previous methods, evaluating global model convergence rate, binary cross-entropy, and silhouette score. The proposed method demonstrates faster convergence rate by about 18%, reduced binary cross-entropy by 26.4%, and more effective clustering.