Contextual Client Selection for Efficient Federated Learning over Edge Devices

Federated learning (FL) has emerged as a prominent distributed learning paradigm, enabling collaborative training of neural network models across local devices with raw data stay local. However, FL systems often encounter significant challenges due to data heterogeneity. Specifically, the non-IID dataset in FL systems substantially slows down the convergence speed during training and adversely impacts the accuracy of the final model. In our paper, we introduce a novel client selection framework that judiciously leverages correlations across local datasets to accelerate training. Our framework first employs a lightweight locality-sensitive hashing algorithm to extract client features while respecting data privacy and incurring minimal overhead. We then design a novel Neural Contextual Combinatorial Bandit (NCCB) algorithm to establish relationships between client features and rewards, enabling intelligent selection of client combinations. We theoretically prove that our proposed NCCB has a bounded regret. Extensive experiments on real-world datasets further demonstrate that our framework surpasses state-of-the-art solutions, resulting in a 50% reduction in training time and a 17% increase in final model accuracy, closing to the performance in the ideal IID case.