Robust Federated Learning: A Heterogeneity Index Based Clustering Approach

Federated learning (FL) has made possible the collaborative training of machine learning models between aggregation server and clients without sharing their private data. With the massive volume of heterogeneous data from various clients, the server faces challenges such as data unbalance, data corruption, and/or data irrelevancy. As a result, the FL setting is exposed to numerous security risks that lead to performance deterioration of learning effectiveness. To tackle the issue, in this paper we propose the Heterogeneity Index Based Clustering (HIC) approach, which enables the dynamic categorization of clients into clusters. Particularly, the model weights are dynamically clustered based on their heterogeneity level using an affinity propagation method. The HIC approach uses a simple, but effective way of scaling data heterogeneity and dynamic clustering to create a resilient learning system against backdoor attacks that outperforms the existing works on FL robustness. Our experimental results demonstrate that the clustering client’s weight based on their heterogeneity level decreases data unbalance and reduces attack success rate, increasing model performance, and encouraging clients’ contribution in FL.