Clust-PSI-PFL: A Population Stability Index Approach for Clustered Non-IID Personalized Federated Learning

TL;DR

The paper tackles the challenge of non-IID label-skew in federated learning by introducing Clust-PSI-PFL, a clustering-based personalized FL framework that uses Population Stability Index (PSI) to quantify cross-client distributional drift. It forms distributionally homogeneous client clusters via K-means++ on PSI features and trains a dedicated model per cluster, enabling more stable convergence and better generalization. Empirically, Clust-PSI-PFL achieves up to 18% higher global accuracy and a 37% improvement in client fairness over strong baselines across six datasets and two partition protocols, demonstrating robustness to severe non-IID conditions. The approach is lightweight and interpretable, leveraging PSI’s per-class breakdown ($PSI_{i,c}^L$) and the weighted metric $WPSI^L$, though privacy-preserving extensions (DP/MPC) and broader skew types remain future work.

Abstract

Federated learning (FL) supports privacy-preserving, decentralized machine learning (ML) model training by keeping data on client devices. However, non-independent and identically distributed (non-IID) data across clients biases updates and degrades performance. To alleviate these issues, we propose Clust-PSI-PFL, a clustering-based personalized FL framework that uses the Population Stability Index (PSI) to quantify the level of non-IID data. We compute a weighted PSI metric, $WPSI^L$, which we show to be more informative than common non-IID metrics (Hellinger, Jensen-Shannon, and Earth Mover's distance). Using PSI features, we form distributionally homogeneous groups of clients via K-means++; the number of optimal clusters is chosen by a systematic silhouette-based procedure, typically yielding few clusters with modest overhead. Across six datasets (tabular, image, and text modalities), two partition protocols (Dirichlet with parameter $α$ and Similarity with parameter S), and multiple client sizes, Clust-PSI-PFL delivers up to 18% higher global accuracy than state-of-the-art baselines and markedly improves client fairness by a relative improvement of 37% under severe non-IID data. These results establish PSI-guided clustering as a principled, lightweight mechanism for robust PFL under label skew.

Figures (8)

• Figure 1: (Left) Typical FL training process with FedAvg. (Right) Illustration of weight divergence for FL with non-IID data
• Figure 2: High-level Clust-PSI-PFL training process.
• Figure 3: Relationship between $WPSI^L$ and the level of non-IID data for the Dirichlet and Similarity partition protocols.
• Figure 4: Feature importance of models predicting the level of non-IID data for the Dirichlet (top) and Similarity (bottom) partition protocols.
• Figure 5: Silhouette score for K-means++ changing number of clusters ($\tau$) for the Dirichlet and Similarity partition protocols and the ACSIncome dataset.