Hierarchical Sinkhorn Distance Clustering Federated Learning for Non-IID Data

Federated learning, as an emerging distributed neural network training paradigm in edge computing, faces a key challenge of data heterogeneity over clients, which significantly impacts global model performance. While clustering federated learning, serving as a promising solution, has shown improvements in model accuracy, its effectiveness is limited by the insufficient utilization of local data distribution statistics. To address this issue, this paper proposes a novel federated learning framework named HS-CFA (Hierarchical Sinkhorn Distance-Based Clustering Federated Algorithm), which is designed to optimize the aggregation weights of local models under the heterogeneous data environments, thereby enhancing the performance of the global model. It utilizes the entropy-regularized optimal transport cost to capture the characteristics of local data distribution and dynamically adjusts the aggregation weights of local models with a hierarchical clustering strategy. Specially, HS-CFA employs the Sinkhorn distance as a lightweight optimal transport cost metric to measure distributional dissimilarities across clients. Furthermore, it adopts a hierarchical two-layer clustering mechanism, combining density-based spatial clustering and average aggregation during the server training phase, facilitating the dynamic and adaptive adjustment of local model aggregation weights. Experimental results on multiple benchmark datasets demonstrate that HS-CFA significantly enhances the accuracy and robustness of the global model in scenarios with highly heterogeneous distribution settings.