CIC-FL: Enabling Class Imbalance-Aware Clustered Federated Learning over Shifted Distributions

Federated learning (FL) is a distributed training framework where decentralized clients collaboratively train a model. One challenge in FL is concept shift, i.e. that the conditional distributions of data in different clients are disagreeing. A natural solution is to group clients with similar conditional distributions into the same cluster. However, methods following this approach leverage features extracted in federated settings (e.g., model weights or gradients) which intrinsically reflect the joint distributions of clients. Considering the difference between conditional and joint distributions, they would fail in the presence of class imbalance (i.e. that the marginal distributions of different classes vary in a client’s data). Although adopting sampling techniques or cost-sensitive algorithms can alleviate class imbalance, they either skew the original conditional distributions or lead to privacy leakage. To address this challenge, we propose CIC-FL, a class imbalance-aware clustered federated learning method. CIC-FL iteratively bipartitions clients by leveraging a particular feature sensitive to concept shift but robust to class imbalance. In addition, CIC-FL is privacy-preserving and communication efficient. We test CIC-FL on benchmark datasets including Fashion-MNIST, CIFAR-10 and IMDB. The results show that CIC-FL outperforms state-of-the-art clustering methods in FL in the presence of class imbalance.