A Cluster-Based Client Selection Model for Federated Learning With Heterogeneous Clients

Client selection can result in more homogeneous data and improve the generalization performance of the global model. In the actual training environment, the heterogeneity of client devices and the non-IID characteristics of data seriously reduce the performance of the model. Therefore, this paper proposes a federated learning cluster selection model to alleviate client drift. To address the issue of unstable computing and communication resources among participants, we propose an algorithm that adaptively adjusts the task load based on the linear upper confidence bound algorithm (LinUCB). The algorithm mines the probability distribution of resources from historical training data,whilch helps alleviate the problem of dropout caused by device overload. To solve the high computational complexity problem caused by high-dimensional gradients, a method for extracting the updated direction of a local model is proposed based on updated weights. The updated direction is jointly represented by the class distribution calculated from the cross-entropy loss of the model on a balanced dataset, and the updated weights that are reduced in dimensionality. Finally, we propose a client clustering selection algorithm based on the updated direction to address the problem of client drift in collaborative training with heterogeneous data. The algorithm prioritizes selecting clients whose update directions are close to that of the global model, while simultaneously selecting clients with tilted update directions with equal preference. This makes the data distribution more homogeneous and alleviates the problem of client drift. Experimental results demonstrate that the proposed algorithm for adaptive task allocation can reduce the number of dropouts under resource-limited and unstable scenarios. The client selection algorithm shows good performance on three public datasets, MNIST, CIFAR-10 and AG_NEWS, outperforming six baseline methods across varying levels of data heterogeneity.Our code is public athttps://github.com/Wannery/MBUT-CS.